diff --git a/.compatibility b/.compatibility index c8ac4083d2a2..32da32be5521 100644 --- a/.compatibility +++ b/.compatibility @@ -1,3 +1,3 @@ 1.12.0-11.3.0 -1.11.0-11.3.0 -1.10.1-11.3.0 +1.13.0-11.6.0 +2.0.0-11.7.0 diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 000000000000..b065e6eb9b77 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,4 @@ +[run] +concurrency = multiprocessing +parallel = true +sigterm = true diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index 673b1274c94b..b310fcfefc15 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -8,4 +8,4 @@ contact_links: about: This issue tracker is not for technical support. Please use WeChat, and ask the community for help. - name: 😊 Advanced question - GitHub Discussions url: https://github.com/hpcaitech/ColossalAI/discussions - about: Use GitHub Discussions for advanced and unanswered technical questions, requiring a maintainer's answer. \ No newline at end of file + about: Use GitHub Discussions for advanced and unanswered technical questions, requiring a maintainer's answer. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index d05bc25f6f41..f12c41b52e6f 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -22,7 +22,7 @@ body: If applicable, add screenshots to help explain your problem. **Suggest a potential alternative/fix** Tell us how we could improve this project. - **Optional: Affiliation** + **Optional: Affiliation** Institution/email information helps better analyze and evaluate users to improve the project. Welcome to establish in-depth cooperation. placeholder: | A clear and concise description of your idea. diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 9634b84b8ff8..3fad7e36f14c 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -14,7 +14,7 @@ - [Compatibility Test on Dispatch](#compatibility-test-on-dispatch) - [Release](#release) - [User Friendliness](#user-friendliness) - - [Commmunity](#commmunity) + - [Community](#community) - [Configuration](#configuration) - [Progress Log](#progress-log) @@ -30,7 +30,7 @@ In the section below, we will dive into the details of different workflows avail Refer to this [documentation](https://docs.github.com/en/actions/managing-workflow-runs/manually-running-a-workflow) on how to manually trigger a workflow. I will provide the details of each workflow below. -**A PR which changes the `version.txt` is considered as a release PR in the following coontext.** +**A PR which changes the `version.txt` is considered as a release PR in the following context.** ### Code Style Check @@ -43,10 +43,18 @@ I will provide the details of each workflow below. | Workflow Name | File name | Description | | ---------------------- | -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------- | -| `Build on PR` | `build_on_pr.yml` | This workflow is triggered when the label `Run build and Test` is assigned to a PR. It will run all the unit tests in the repository with 4 GPUs. | +| `Build on PR` | `build_on_pr.yml` | This workflow is triggered when a PR changes essential files and a branch is created/deleted. It will run all the unit tests in the repository with 4 GPUs. | | `Build on Schedule` | `build_on_schedule.yml` | This workflow will run the unit tests everyday with 8 GPUs. The result is sent to Lark. | | `Report test coverage` | `report_test_coverage.yml` | This PR will put up a comment to report the test coverage results when `Build` is done. | +To reduce the average time of the unit test on PR, `Build on PR` workflow manages testmon cache. + +1. When creating a new branch, it copies `cache/main/.testmondata*` to `cache//`. +2. When creating a new PR or change the base branch of a PR, it copies `cache//.testmondata*` to `cache/_pull//`. +3. When running unit tests for each PR, it restores testmon cache from `cache/_pull//`. After the test, it stores the cache back to `cache/_pull//`. +4. When a PR is closed, if it's merged, it copies `cache/_pull//.testmondata*` to `cache//`. Otherwise, it just removes `cache/_pull/`. +5. When a branch is deleted, it removes `cache/`. + ### Example Test | Workflow Name | File name | Description | @@ -58,15 +66,15 @@ I will provide the details of each workflow below. #### Example Test on Dispatch This workflow is triggered by manually dispatching the workflow. It has the following input parameters: -- `example_directory`: the example directory to test. Multiple directories are supported and must be separated b$$y comma. For example, language/gpt, images/vit. Simply input language or simply gpt does not work. +- `example_directory`: the example directory to test. Multiple directories are supported and must be separated by comma. For example, language/gpt, images/vit. Simply input language or simply gpt does not work. ### Compatibility Test | Workflow Name | File name | Description | | -------------------------------- | ------------------------------------ | -------------------------------------------------------------------------------------------------------------------- | -| `Compatibility Test on PR` | `compatibility_test_on_pr.yml` | Check Colossal-AI's compatiblity when `version.txt` is changed in a PR. | -| `Compatibility Test on Schedule` | `compatibility_test_on_schedule.yml` | This workflow will check the compatiblity of Colossal-AI against PyTorch specified in `.compatibility` every Sunday. | -| `Compatiblity Test on Dispatch` | `compatibility_test_on_dispatch.yml` | Test PyTorch Compatibility manually. | +| `Compatibility Test on PR` | `compatibility_test_on_pr.yml` | Check Colossal-AI's compatibility when `version.txt` is changed in a PR. | +| `Compatibility Test on Schedule` | `compatibility_test_on_schedule.yml` | This workflow will check the compatibility of Colossal-AI against PyTorch specified in `.compatibility` every Sunday. | +| `Compatibility Test on Dispatch` | `compatibility_test_on_dispatch.yml` | Test PyTorch Compatibility manually. | #### Compatibility Test on Dispatch @@ -74,7 +82,7 @@ This workflow is triggered by manually dispatching the workflow. It has the foll - `torch version`:torch version to test against, multiple versions are supported but must be separated by comma. The default is value is all, which will test all available torch versions listed in this [repository](https://github.com/hpcaitech/public_assets/tree/main/colossalai/torch_build/torch_wheels). - `cuda version`: cuda versions to test against, multiple versions are supported but must be separated by comma. The CUDA versions must be present in our [DockerHub repository](https://hub.docker.com/r/hpcaitech/cuda-conda). -> It only test the compatiblity of the main branch +> It only test the compatibility of the main branch ### Release @@ -97,7 +105,7 @@ This workflow is triggered by manually dispatching the workflow. It has the foll | `Synchronize submodule` | `submodule.yml` | This workflow will check if any git submodule is updated. If so, it will create a PR to update the submodule pointers. | | `Close inactive issues` | `close_inactive.yml` | This workflow will close issues which are stale for 14 days. | -### Commmunity +### Community | Workflow Name | File name | Description | | -------------------------------------------- | -------------------------------- | -------------------------------------------------------------------------------- | @@ -113,7 +121,7 @@ This `.compatibility` file is to tell GitHub Actions which PyTorch and CUDA vers 2. `.cuda_ext.json` -This file controls which CUDA versions will be checked against CUDA extenson built. You can add a new entry according to the json schema below to check the AOT build of PyTorch extensions before release. +This file controls which CUDA versions will be checked against CUDA extension built. You can add a new entry according to the json schema below to check the AOT build of PyTorch extensions before release. ```json { @@ -144,7 +152,7 @@ This file controls which CUDA versions will be checked against CUDA extenson bui - [x] check on PR - [x] regular check - [x] manual dispatch -- [x] compatiblity check +- [x] compatibility check - [x] check on PR - [x] manual dispatch - [x] auto test when release diff --git a/.github/workflows/build_on_pr.yml b/.github/workflows/build_on_pr.yml index f595e677394a..380c8e9f882c 100644 --- a/.github/workflows/build_on_pr.yml +++ b/.github/workflows/build_on_pr.yml @@ -2,22 +2,93 @@ name: Build on PR on: pull_request: - types: [synchronize, labeled] + types: [synchronize, opened, reopened, ready_for_review, closed, edited] + branches: + - "main" + - "develop" + - "feature/**" + paths: + - ".github/workflows/build_on_pr.yml" # run command & env variables change + - "colossalai/**" # source code change + - "!colossalai/**.md" # ignore doc change + - "op_builder/**" # cuda extension change + - "!op_builder/**.md" # ignore doc change + - "requirements/**" # requirements change + - "tests/**" # test change + - "!tests/**.md" # ignore doc change + - "pytest.ini" # test config change + - "setup.py" # install command change + create: + delete: jobs: + prepare_cache: + name: Prepare testmon cache + if: | + github.event_name == 'create' && + github.event.ref_type == 'branch' && + github.event.repository.full_name == 'hpcaitech/ColossalAI' + runs-on: [self-hosted, gpu] + container: + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 + options: --rm + timeout-minutes: 5 + defaults: + run: + shell: bash + steps: + - name: Copy testmon cache + run: | # branch name may contain slash, we need to replace it with space + export REF_BRANCH=$(echo ${{ github.event.ref }} | sed "s/\// /") + if [ -d /github/home/testmon_cache/${MAIN_BRANCH} ]; then + cp -p -r /github/home/testmon_cache/${MAIN_BRANCH} "/github/home/testmon_cache/${REF_BRANCH}" + fi + env: + MAIN_BRANCH: ${{ github.event.master_branch }} + + prepare_cache_for_pr: + name: Prepare testmon cache for PR + if: | + github.event_name == 'pull_request' && + (github.event.action == 'opened' || github.event.action == 'reopened' || (github.event.action == 'edited' && github.event.changes.base != null)) && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' + runs-on: [self-hosted, gpu] + container: + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 + options: --rm + timeout-minutes: 5 + defaults: + run: + shell: bash + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false + steps: + - name: Copy testmon cache + run: | # branch name may contain slash, we need to replace it with space + export BASE=$(echo ${{ github.event.pull_request.base.ref }} | sed "s/\// /") + if [ -d "/github/home/testmon_cache/${BASE}" ] && [ ! -z "$(ls -A "/github/home/testmon_cache/${BASE}")" ]; then + mkdir -p /github/home/testmon_cache/_pull/${PR_NUMBER} && cp -p -r "/github/home/testmon_cache/${BASE}"/.testmondata* /github/home/testmon_cache/_pull/${PR_NUMBER} + fi + env: + PR_NUMBER: ${{ github.event.number }} + detect: name: Detect file change if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && - contains( github.event.pull_request.labels.*.name, 'Run Build and Test') + github.event_name == 'pull_request' && + (github.event.action == 'synchronize' || github.event.action == 'opened' || github.event.action == 'reopened' || github.event.action == 'ready_for_review') && + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' outputs: changedExtenisonFiles: ${{ steps.find-extension-change.outputs.all_changed_files }} anyExtensionFileChanged: ${{ steps.find-extension-change.outputs.any_changed }} changedLibraryFiles: ${{ steps.find-lib-change.outputs.all_changed_files }} anyLibraryFileChanged: ${{ steps.find-lib-change.outputs.any_changed }} runs-on: ubuntu-latest + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v2 with: @@ -27,10 +98,10 @@ jobs: - name: Locate base commit id: locate-base-sha run: | - curBranch=$(git rev-parse --abbrev-ref HEAD) - commonCommit=$(git merge-base origin/main $curBranch) - echo $commonCommit - echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT + curBranch=$(git rev-parse --abbrev-ref HEAD) + commonCommit=$(git merge-base origin/main $curBranch) + echo $commonCommit + echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT - name: Find the changed extension-related files id: find-extension-change @@ -63,18 +134,21 @@ jobs: echo "$file was changed" done - build: name: Build and Test Colossal-AI needs: detect + if: needs.detect.outputs.anyLibraryFileChanged == 'true' runs-on: [self-hosted, gpu] container: - image: hpcaitech/pytorch-cuda:1.11.0-11.3.0 + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10 - timeout-minutes: 40 + timeout-minutes: 60 defaults: run: shell: bash + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - name: Checkout TensorNVMe uses: actions/checkout@v2 @@ -85,7 +159,9 @@ jobs: - name: Restore TensorNVMe Cache run: | - [ ! -z "$(ls -A /github/home/tensornvme_cache/)" ] && cp -p -r /github/home/tensornvme_cache/* /__w/ColossalAI/ColossalAI/TensorNVMe + if [ -d /github/home/tensornvme_cache ] && [ ! -z "$(ls -A /github/home/tensornvme_cache/)" ]; then + cp -p -r /github/home/tensornvme_cache/* /__w/ColossalAI/ColossalAI/TensorNVMe + fi - name: Install TensorNVMe run: | @@ -108,10 +184,11 @@ jobs: if: needs.detect.outputs.anyExtensionFileChanged != 'true' run: | # -p flag is required to preserve the file timestamp to avoid ninja rebuild - [ ! -z "$(ls -A /github/home/cuda_ext_cache/)" ] && cp -p -r /github/home/cuda_ext_cache/* /__w/ColossalAI/ColossalAI/ + if [ -d /github/home/cuda_ext_cache ] && [ ! -z "$(ls -A /github/home/cuda_ext_cache/)" ]; then + cp -p -r /github/home/cuda_ext_cache/* /__w/ColossalAI/ColossalAI/ + fi - name: Install Colossal-AI - if: needs.detect.outputs.anyLibraryFileChanged == 'true' run: | CUDA_EXT=1 pip install -v -e . pip install -r requirements/requirements-test.txt @@ -121,15 +198,29 @@ jobs: # -p flag is required to preserve the file timestamp to avoid ninja rebuild cp -p -r /__w/ColossalAI/ColossalAI/build /github/home/cuda_ext_cache/ + - name: Restore Testmon Cache + run: | + if [ -d /github/home/testmon_cache/_pull/${PR_NUMBER} ] && [ ! -z "$(ls -A /github/home/testmon_cache/_pull/${PR_NUMBER})" ]; then + cp -p -r /github/home/testmon_cache/_pull/${PR_NUMBER}/.testmondata* /__w/ColossalAI/ColossalAI/ + fi + env: + PR_NUMBER: ${{ github.event.number }} + - name: Execute Unit Testing - if: needs.detect.outputs.anyLibraryFileChanged == 'true' run: | - PYTHONPATH=$PWD pytest --cov=. --cov-report xml tests/ + CURL_CA_BUNDLE="" PYTHONPATH=$PWD pytest --testmon --testmon-cov=. --durations=10 tests/ env: DATA: /data/scratch/cifar-10 NCCL_SHM_DISABLE: 1 LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 + - name: Store Testmon Cache + run: | + mkdir -p /github/home/testmon_cache/_pull/${PR_NUMBER} + cp -p -r /__w/ColossalAI/ColossalAI/.testmondata* /github/home/testmon_cache/_pull/${PR_NUMBER}/ + env: + PR_NUMBER: ${{ github.event.number }} + - name: Collate artifact env: PR_NUMBER: ${{ github.event.number }} @@ -141,7 +232,7 @@ jobs: echo $PR_NUMBER > ./report/pr_number # generate coverage.xml if any - if [ "$anyLibraryFileChanged" == "true" ]; then + if [ "$anyLibraryFileChanged" == "true" ] && [ -e .coverage ]; then allFiles="" for file in $changedLibraryFiles; do if [ "$allFiles" == "" ]; then @@ -166,3 +257,54 @@ jobs: with: name: report path: report/ + + store_cache: + name: Store testmon cache for PR + if: | + github.event_name == 'pull_request' && + github.event.action == 'closed' && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' + runs-on: [self-hosted, gpu] + container: + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 + options: --rm + timeout-minutes: 5 + defaults: + run: + shell: bash + steps: + - name: Store testmon cache if possible + if: github.event.pull_request.merged == true + run: | # branch name may contain slash, we need to replace it with space + export BASE=$(echo ${{ github.event.pull_request.base.ref }} | sed "s/\// /") + if [ -d /github/home/testmon_cache/_pull/${PR_NUMBER} ] && [ ! -z "$(ls -A /github/home/testmon_cache/_pull/${PR_NUMBER})" ]; then + cp -p -r /github/home/testmon_cache/_pull/${PR_NUMBER}/.testmondata* "/github/home/testmon_cache/${BASE}/" + fi + env: + PR_NUMBER: ${{ github.event.pull_request.number }} + + - name: Remove testmon cache + run: | + rm -rf /github/home/testmon_cache/_pull/${PR_NUMBER} + env: + PR_NUMBER: ${{ github.event.pull_request.number }} + + remove_cache: + name: Remove testmon cache + if: | + github.event_name == 'delete' && + github.event.ref_type == 'branch' && + github.event.repository.full_name == 'hpcaitech/ColossalAI' + runs-on: [self-hosted, gpu] + container: + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 + options: --rm + timeout-minutes: 5 + defaults: + run: + shell: bash + steps: + - name: Remove testmon cache + run: | # branch name may contain slash, we need to replace it with space + export BASE=$(echo ${{ github.event.ref }} | sed "s/\// /") + rm -rf "/github/home/testmon_cache/${BASE}" diff --git a/.github/workflows/build_on_schedule.yml b/.github/workflows/build_on_schedule.yml index 6afdf581e6ca..03b47e6cb5b6 100644 --- a/.github/workflows/build_on_schedule.yml +++ b/.github/workflows/build_on_schedule.yml @@ -3,7 +3,7 @@ name: Build on Schedule on: schedule: # run at 00:00 of every Sunday - - cron: '0 0 * * *' + - cron: "0 0 * * *" workflow_dispatch: jobs: @@ -12,7 +12,7 @@ jobs: if: github.repository == 'hpcaitech/ColossalAI' runs-on: [self-hosted, 8-gpu] container: - image: hpcaitech/pytorch-cuda:1.11.0-11.3.0 + image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10 timeout-minutes: 40 steps: @@ -60,7 +60,7 @@ jobs: - name: Unit Testing if: steps.check-avai.outputs.avai == 'true' run: | - PYTHONPATH=$PWD pytest tests + PYTHONPATH=$PWD pytest --durations=0 tests env: DATA: /data/scratch/cifar-10 LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 diff --git a/.github/workflows/compatiblity_test_on_dispatch.yml b/.github/workflows/compatiblity_test_on_dispatch.yml index 717cf729b3f3..3dcc4dfd182a 100644 --- a/.github/workflows/compatiblity_test_on_dispatch.yml +++ b/.github/workflows/compatiblity_test_on_dispatch.yml @@ -19,26 +19,26 @@ jobs: outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} steps: - - id: set-matrix - env: - TORCH_VERSIONS: ${{ inputs.torch_version }} - CUDA_VERSIONS: ${{ inputs.cuda_version }} - run: | - IFS=',' - DOCKER_IMAGE=() + - id: set-matrix + env: + TORCH_VERSIONS: ${{ inputs.torch_version }} + CUDA_VERSIONS: ${{ inputs.cuda_version }} + run: | + IFS=',' + DOCKER_IMAGE=() - for tv in $TORCH_VERSIONS - do - for cv in $CUDA_VERSIONS - do - DOCKER_IMAGE+=("\"hpcaitech/pytorch-cuda:${tv}-${cv}\"") - done - done + for tv in $TORCH_VERSIONS + do + for cv in $CUDA_VERSIONS + do + DOCKER_IMAGE+=("\"hpcaitech/pytorch-cuda:${tv}-${cv}\"") + done + done - container=$( IFS=',' ; echo "${DOCKER_IMAGE[*]}" ) - container="[${container}]" - echo "$container" - echo "::set-output name=matrix::{\"container\":$(echo "$container")}" + container=$( IFS=',' ; echo "${DOCKER_IMAGE[*]}" ) + container="[${container}]" + echo "$container" + echo "::set-output name=matrix::{\"container\":$(echo "$container")}" build: name: Test for PyTorch Compatibility @@ -70,6 +70,17 @@ jobs: - uses: actions/checkout@v2 with: ssh-key: ${{ secrets.SSH_KEY_FOR_CI }} + - name: Download cub for CUDA 10.2 + run: | + CUDA_VERSION=$(cat $CUDA_HOME/version.txt | grep "CUDA Version" | awk '{print $NF}' | cut -d. -f1,2) + + # check if it is CUDA 10.2 + # download cub + if [ "$CUDA_VERSION" = "10.2" ]; then + wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip + unzip 1.8.0.zip + cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ + fi - name: Install Colossal-AI run: | pip install -r requirements/requirements.txt diff --git a/.github/workflows/compatiblity_test_on_pr.yml b/.github/workflows/compatiblity_test_on_pr.yml index 2fca67b820a1..5098b8e364d0 100644 --- a/.github/workflows/compatiblity_test_on_pr.yml +++ b/.github/workflows/compatiblity_test_on_pr.yml @@ -3,8 +3,8 @@ name: Compatibility Test on PR on: pull_request: paths: - - 'version.txt' - - '.compatibility' + - "version.txt" + - ".compatibility" jobs: matrix_preparation: @@ -12,6 +12,9 @@ jobs: runs-on: ubuntu-latest outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v3 - id: set-matrix @@ -40,6 +43,9 @@ jobs: image: ${{ matrix.container }} options: --gpus all --rm -v /data/scratch/cifar-10:/data/scratch/cifar-10 timeout-minutes: 120 + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - name: Install dependencies run: | @@ -58,6 +64,18 @@ jobs: - uses: actions/checkout@v2 with: ssh-key: ${{ secrets.SSH_KEY_FOR_CI }} + - name: Download cub for CUDA 10.2 + run: | + CUDA_VERSION=$(cat $CUDA_HOME/version.txt | grep "CUDA Version" | awk '{print $NF}' | cut -d. -f1,2) + + # check if it is CUDA 10.2 + # download cub + if [ "$CUDA_VERSION" = "10.2" ]; then + wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip + unzip 1.8.0.zip + cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ + fi + - name: Install Colossal-AI run: | pip install -v --no-cache-dir . diff --git a/.github/workflows/doc_build_after_merge.yml b/.github/workflows/doc_build_on_schedule_after_release.yml similarity index 69% rename from .github/workflows/doc_build_after_merge.yml rename to .github/workflows/doc_build_on_schedule_after_release.yml index ede04b336620..62dfdc67257c 100644 --- a/.github/workflows/doc_build_after_merge.yml +++ b/.github/workflows/doc_build_on_schedule_after_release.yml @@ -1,18 +1,16 @@ -name: Build Documentation After Merge +name: Build Documentation On Schedule & After Release on: workflow_dispatch: - pull_request: - paths: - - 'version.txt' - - 'docs/**' - types: - - closed + schedule: + - cron: "0 12 * * *" # build doc every day at 8pm Singapore time (12pm UTC time) + release: + types: [published] jobs: build-doc: name: Trigger Documentation Build Workflow - if: ( github.event_name == 'workflow_dispatch' || github.event.pull_request.merged == true ) && github.repository == 'hpcaitech/ColossalAI' + if: github.repository == 'hpcaitech/ColossalAI' runs-on: ubuntu-latest steps: - name: trigger workflow in ColossalAI-Documentation diff --git a/.github/workflows/doc_check_on_pr.yml b/.github/workflows/doc_check_on_pr.yml index 2022c957fba8..848991bd3a82 100644 --- a/.github/workflows/doc_check_on_pr.yml +++ b/.github/workflows/doc_check_on_pr.yml @@ -2,57 +2,68 @@ name: Check Documentation on PR on: pull_request: + branches: + - "main" + - "develop" + - "feature/**" paths: - - 'docs/**' + - "docs/**" jobs: check-i18n: name: Check docs in diff languages if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' runs-on: ubuntu-latest + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v2 - uses: actions/setup-python@v2 with: - python-version: '3.8.14' + python-version: "3.8.14" - run: python .github/workflows/scripts/check_doc_i18n.py -d docs/source check-doc-build: name: Test if the docs can be built if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' runs-on: ubuntu-latest + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v2 with: - path: './ColossalAI' + path: "./ColossalAI" fetch-depth: 0 - uses: actions/checkout@v2 with: - path: './ColossalAI-Documentation' - repository: 'hpcaitech/ColossalAI-Documentation' + path: "./ColossalAI-Documentation" + repository: "hpcaitech/ColossalAI-Documentation" - uses: actions/setup-python@v2 with: - python-version: '3.8.14' + python-version: "3.8.14" # we use the versions in the main branch as the guide for versions to display # checkout will give your merged branch # therefore, we need to make the merged branch as the main branch + # there is no main branch, so it's safe to checkout the main branch from the merged branch + # docer will rebase the remote main branch to the merged branch, so we have to config user - name: Make the merged branch main run: | cd ColossalAI - curBranch=$(git rev-parse --abbrev-ref HEAD) - git checkout main - git merge $curBranch # fast-forward master up to the merge + git checkout -b main + git branch -u origin/main + git config user.name 'github-actions' + git config user.email 'github-actions@github.com' - name: Build docs run: | diff --git a/.github/workflows/doc_test_on_pr.yml b/.github/workflows/doc_test_on_pr.yml index a083362a7f0f..2a07a2297bfb 100644 --- a/.github/workflows/doc_test_on_pr.yml +++ b/.github/workflows/doc_test_on_pr.yml @@ -1,21 +1,27 @@ name: Test Documentation on PR on: pull_request: + branches: + - "main" + - "develop" + - "feature/**" # any change in the examples folder will trigger check for the corresponding example. paths: - - 'docs/source/**.md' + - "docs/source/**.md" jobs: # This is for changed example files detect and output a matrix containing all the corresponding directory name. detect-changed-doc: if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' runs-on: ubuntu-latest outputs: any_changed: ${{ steps.changed-files.outputs.any_changed }} changed_files: ${{ steps.changed-files.outputs.all_changed_files }} + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false name: Detect changed example files steps: - uses: actions/checkout@v3 @@ -26,10 +32,10 @@ jobs: - name: Locate base commit id: locate-base-sha run: | - curBranch=$(git rev-parse --abbrev-ref HEAD) - commonCommit=$(git merge-base origin/main $curBranch) - echo $commonCommit - echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT + curBranch=$(git rev-parse --abbrev-ref HEAD) + commonCommit=$(git merge-base origin/main $curBranch) + echo $commonCommit + echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT - name: Get all changed example files id: changed-files @@ -43,10 +49,9 @@ jobs: check-changed-doc: # Add this condition to avoid executing this job if the trigger event is workflow_dispatch. if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' && - needs.detect-changed-doc.outputs.any_changed == 'true' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' && + needs.detect-changed-doc.outputs.any_changed == 'true' name: Test the changed Doc needs: detect-changed-doc runs-on: [self-hosted, gpu] @@ -57,12 +62,15 @@ jobs: defaults: run: shell: bash + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - name: Checkout ColossalAI-Documentation uses: actions/checkout@v2 with: - path: './ColossalAI-Documentation' - repository: 'hpcaitech/ColossalAI-Documentation' + path: "./ColossalAI-Documentation" + repository: "hpcaitech/ColossalAI-Documentation" - name: Install Docer run: | @@ -71,7 +79,7 @@ jobs: - name: Checkout ColossalAI uses: actions/checkout@v3 - + - name: Install Doc Test Requirements run: | source activate pytorch @@ -86,7 +94,7 @@ jobs: - name: Test the Doc run: | source activate pytorch - for file in ${{ steps.changed-files.outputs.all_changed_files }}; do + for file in ${{ needs.detect-changed-doc.outputs.changed_files }}; do echo "Testing $file now..." docer test -p $file done diff --git a/.github/workflows/example_check_on_pr.yml b/.github/workflows/example_check_on_pr.yml index b22664ee47cc..ee456c25f2b5 100644 --- a/.github/workflows/example_check_on_pr.yml +++ b/.github/workflows/example_check_on_pr.yml @@ -1,22 +1,28 @@ name: Test Example on PR on: pull_request: + branches: + - "main" + - "develop" + - "feature/**" # any change in the examples folder will trigger check for the corresponding example. paths: - - 'examples/**' + - "examples/**" jobs: # This is for changed example files detect and output a matrix containing all the corresponding directory name. detect-changed-example: if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' runs-on: ubuntu-latest outputs: matrix: ${{ steps.setup-matrix.outputs.matrix }} anyChanged: ${{ steps.setup-matrix.outputs.anyChanged }} name: Detect changed example files + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v3 with: @@ -26,10 +32,10 @@ jobs: - name: Locate base commit id: locate-base-sha run: | - curBranch=$(git rev-parse --abbrev-ref HEAD) - commonCommit=$(git merge-base origin/main $curBranch) - echo $commonCommit - echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT + curBranch=$(git rev-parse --abbrev-ref HEAD) + commonCommit=$(git merge-base origin/main $curBranch) + echo $commonCommit + echo "baseSHA=$commonCommit" >> $GITHUB_OUTPUT - name: Get all changed example files id: changed-files @@ -61,10 +67,9 @@ jobs: check-changed-example: # Add this condition to avoid executing this job if the trigger event is workflow_dispatch. if: | - github.event.pull_request.draft == false && - github.base_ref == 'main' && - github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' && - needs.detect-changed-example.outputs.anyChanged == 'true' + github.event.pull_request.draft == false && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' && github.event_name == 'pull_request' && + needs.detect-changed-example.outputs.anyChanged == 'true' name: Test the changed example needs: detect-changed-example runs-on: [self-hosted, gpu] @@ -75,6 +80,9 @@ jobs: image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 options: --gpus all --rm -v /data/scratch/examples-data:/data/ timeout-minutes: 10 + concurrency: + group: ${{ github.head_ref }} + cancel-in-progress: false steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/post_commit.yml b/.github/workflows/post_commit.yml index bf93eabbf43f..1bbc0d2f5c34 100644 --- a/.github/workflows/post_commit.yml +++ b/.github/workflows/post_commit.yml @@ -82,7 +82,7 @@ jobs: # create pull request - name: Create Pull Request - if: steps.commit.outputs.status == 'success' + if: steps.commit.outcome == 'success' id: cpr uses: peter-evans/create-pull-request@v4 with: @@ -90,7 +90,7 @@ jobs: title: "[format] applied code formatting on changed files in PR ${{ github.event.pull_request.number }}" - name: Enable Auto-merge for the New PR - if: steps.commit.outputs.status == 'success' + if: steps.commit.outcome == 'success' uses: peter-evans/enable-pull-request-automerge@v2 with: pull-request-number: ${{ steps.cpr.outputs.pull-request-number }} diff --git a/.github/workflows/release_docker_after_merge.yml b/.github/workflows/release_docker_after_publish.yml similarity index 81% rename from .github/workflows/release_docker_after_merge.yml rename to .github/workflows/release_docker_after_publish.yml index 607c19b05472..6c8df9730b0d 100644 --- a/.github/workflows/release_docker_after_merge.yml +++ b/.github/workflows/release_docker_after_publish.yml @@ -1,17 +1,14 @@ -name: Publish Docker Image to DockerHub after Merge +name: Publish Docker Image to DockerHub after Publish on: workflow_dispatch: - pull_request: - paths: - - 'version.txt' - types: - - closed + release: + types: [published] jobs: release: name: Publish Docker Image to DockerHub - if: ( github.event_name == 'workflow_dispatch' || github.event.pull_request.merged == true ) && github.repository == 'hpcaitech/ColossalAI' + if: github.repository == 'hpcaitech/ColossalAI' runs-on: [self-hosted, gpu] container: image: "hpcaitech/docker-in-docker:latest" @@ -26,8 +23,11 @@ jobs: run: | version=$(cat version.txt) tag=hpcaitech/colossalai:$version - docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 -t $tag ./docker + latest=hpcaitech/colossalai:latest + docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 --build-arg VERSION=v${version} -t $tag ./docker + docker tag $tag $latest echo "tag=${tag}" >> $GITHUB_OUTPUT + echo "latest=${latest}" >> $GITHUB_OUTPUT - name: Log in to Docker Hub uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9 @@ -39,6 +39,7 @@ jobs: id: docker-push run: | docker push ${{ steps.build.outputs.tag }} + docker push ${{ steps.build.outputs.latest }} notify: name: Notify Lark via webhook @@ -50,7 +51,7 @@ jobs: - uses: actions/setup-python@v2 with: - python-version: '3.8.14' + python-version: "3.8.14" - name: Install requests run: pip install requests diff --git a/.github/workflows/report_test_coverage.yml b/.github/workflows/report_test_coverage.yml index bbada74e6850..c9dc541b8a33 100644 --- a/.github/workflows/report_test_coverage.yml +++ b/.github/workflows/report_test_coverage.yml @@ -9,8 +9,9 @@ on: jobs: report-test-coverage: runs-on: ubuntu-latest + if: ${{ github.event.workflow_run.conclusion == 'success' }} steps: - - name: 'Download artifact' + - name: "Download artifact" uses: actions/github-script@v6 with: script: | @@ -31,7 +32,7 @@ jobs: let fs = require('fs'); fs.writeFileSync(`${process.env.GITHUB_WORKSPACE}/report.zip`, Buffer.from(download.data)); - - name: 'Unzip artifact' + - name: "Unzip artifact" id: unzip run: | unzip report.zip @@ -58,7 +59,7 @@ jobs: echo "" >> coverage_report.txt mv coverage_report.txt coverage.txt - - name: 'Comment on PR' + - name: "Comment on PR" if: steps.unzip.outputs.hasReport == 'true' uses: actions/github-script@v6 with: diff --git a/.github/workflows/run_chatgpt_examples.yml b/.github/workflows/run_chatgpt_examples.yml index 51bb9d074644..510f6b6f0985 100644 --- a/.github/workflows/run_chatgpt_examples.yml +++ b/.github/workflows/run_chatgpt_examples.yml @@ -4,19 +4,22 @@ on: pull_request: types: [synchronize, opened, reopened] paths: - - 'applications/ChatGPT/chatgpt/**' - - 'applications/ChatGPT/requirements.txt' - - 'applications/ChatGPT/setup.py' - - 'applications/ChatGPT/examples/**' - + - "applications/Chat/coati/**" + - "applications/Chat/requirements.txt" + - "applications/Chat/setup.py" + - "applications/Chat/examples/**" jobs: tests: name: Run ChatGPT examples + if: | + github.event.pull_request.draft == false && + github.base_ref == 'main' && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' runs-on: [self-hosted, gpu] container: image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 - options: --gpus all --rm -v /data/scratch/chatgpt:/data/scratch/chatgpt + options: --gpus all --rm -v /data/scratch/github_actions/chat:/data/scratch/github_actions/chat --shm-size=10.24gb timeout-minutes: 30 defaults: run: @@ -27,17 +30,23 @@ jobs: - name: Install ColossalAI and ChatGPT run: | - pip install -v . - cd applications/ChatGPT + pip install -e . + cd applications/Chat pip install -v . pip install -r examples/requirements.txt + - name: Install Transformers + run: | + pip install transformers==4.30.2 + - name: Execute Examples run: | - cd applications/ChatGPT + cd applications/Chat rm -rf ~/.cache/colossalai ./examples/test_ci.sh env: NCCL_SHM_DISABLE: 1 MAX_JOBS: 8 - PROMPT_PATH: /data/scratch/chatgpt/prompts.csv + SFT_DATASET: /data/scratch/github_actions/chat/data.json + PROMPT_PATH: /data/scratch/github_actions/chat/prompts_en.jsonl + PRETRAIN_DATASET: /data/scratch/github_actions/chat/alpaca_data.json diff --git a/.github/workflows/run_chatgpt_unit_tests.yml b/.github/workflows/run_chatgpt_unit_tests.yml index 4e539bfe06fd..47c80fc9a9fe 100644 --- a/.github/workflows/run_chatgpt_unit_tests.yml +++ b/.github/workflows/run_chatgpt_unit_tests.yml @@ -4,16 +4,20 @@ on: pull_request: types: [synchronize, opened, reopened] paths: - - 'applications/ChatGPT/chatgpt/**' - - 'applications/ChatGPT/requirements.txt' - - 'applications/ChatGPT/setup.py' - - 'applications/ChatGPT/requirements-test.txt' - - 'applications/ChatGPT/tests/**' - - 'applications/ChatGPT/pytest.ini' + - 'applications/Chat/coati/**' + - 'applications/Chat/requirements.txt' + - 'applications/Chat/setup.py' + - 'applications/Chat/requirements-test.txt' + - 'applications/Chat/tests/**' + - 'applications/Chat/pytest.ini' jobs: tests: name: Run ChatGPT unit tests + if: | + github.event.pull_request.draft == false && + github.base_ref == 'main' && + github.event.pull_request.base.repo.full_name == 'hpcaitech/ColossalAI' runs-on: [self-hosted, gpu] container: image: hpcaitech/pytorch-cuda:1.12.0-11.3.0 @@ -28,14 +32,14 @@ jobs: - name: Install ColossalAI and ChatGPT run: | - pip install -v . - cd applications/ChatGPT + pip install -e . + cd applications/Chat pip install -v . pip install -r requirements-test.txt - name: Execute Unit Testing run: | - cd applications/ChatGPT + cd applications/Chat rm -rf ~/.cache/colossalai pytest tests/ env: diff --git a/.github/workflows/scripts/example_checks/check_dispatch_inputs.py b/.github/workflows/scripts/example_checks/check_dispatch_inputs.py index 04d2063ec5fc..5bec96187e0c 100644 --- a/.github/workflows/scripts/example_checks/check_dispatch_inputs.py +++ b/.github/workflows/scripts/example_checks/check_dispatch_inputs.py @@ -1,27 +1,27 @@ -import argparse -import os - - -def check_inputs(input_list): - for path in input_list: - real_path = os.path.join('examples', path) - if not os.path.exists(real_path): - return False - return True - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('-f', '--fileNameList', type=str, help="List of file names") - args = parser.parse_args() - name_list = args.fileNameList.split(",") - is_correct = check_inputs(name_list) - - if is_correct: - print('success') - else: - print('failure') - - -if __name__ == '__main__': - main() +import argparse +import os + + +def check_inputs(input_list): + for path in input_list: + real_path = os.path.join('examples', path) + if not os.path.exists(real_path): + return False + return True + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('-f', '--fileNameList', type=str, help="List of file names") + args = parser.parse_args() + name_list = args.fileNameList.split(",") + is_correct = check_inputs(name_list) + + if is_correct: + print('success') + else: + print('failure') + + +if __name__ == '__main__': + main() diff --git a/.github/workflows/scripts/example_checks/check_example_weekly.py b/.github/workflows/scripts/example_checks/check_example_weekly.py index 941e90901f3d..83eff644e315 100644 --- a/.github/workflows/scripts/example_checks/check_example_weekly.py +++ b/.github/workflows/scripts/example_checks/check_example_weekly.py @@ -1,37 +1,37 @@ -import os - - -def show_files(path, all_files): - # Traverse all the folder/file in current directory - file_list = os.listdir(path) - # Determine the element is folder or file. If file, pass it into list, if folder, recurse. - for file_name in file_list: - # Get the abs directory using os.path.join() and store into cur_path. - cur_path = os.path.join(path, file_name) - # Determine whether folder - if os.path.isdir(cur_path): - show_files(cur_path, all_files) - else: - all_files.append(cur_path) - return all_files - - -def join(input_list, sep=None): - return (sep or ' ').join(input_list) - - -def main(): - contents = show_files('examples/', []) - all_loc = [] - for file_loc in contents: - split_loc = file_loc.split('/') - # must have two sub-folder levels after examples folder, such as examples/images/vit is acceptable, examples/images/README.md is not, examples/requirements.txt is not. - if len(split_loc) >= 4: - re_loc = '/'.join(split_loc[1:3]) - if re_loc not in all_loc: - all_loc.append(re_loc) - print(all_loc) - - -if __name__ == '__main__': - main() +import os + + +def show_files(path, all_files): + # Traverse all the folder/file in current directory + file_list = os.listdir(path) + # Determine the element is folder or file. If file, pass it into list, if folder, recurse. + for file_name in file_list: + # Get the abs directory using os.path.join() and store into cur_path. + cur_path = os.path.join(path, file_name) + # Determine whether folder + if os.path.isdir(cur_path): + show_files(cur_path, all_files) + else: + all_files.append(cur_path) + return all_files + + +def join(input_list, sep=None): + return (sep or ' ').join(input_list) + + +def main(): + contents = show_files('examples/', []) + all_loc = [] + for file_loc in contents: + split_loc = file_loc.split('/') + # must have two sub-folder levels after examples folder, such as examples/images/vit is acceptable, examples/images/README.md is not, examples/requirements.txt is not. + if len(split_loc) >= 4: + re_loc = '/'.join(split_loc[1:3]) + if re_loc not in all_loc: + all_loc.append(re_loc) + print(all_loc) + + +if __name__ == '__main__': + main() diff --git a/.github/workflows/scripts/example_checks/detect_changed_example.py b/.github/workflows/scripts/example_checks/detect_changed_example.py index df4fd67368fc..c69d95a552e9 100644 --- a/.github/workflows/scripts/example_checks/detect_changed_example.py +++ b/.github/workflows/scripts/example_checks/detect_changed_example.py @@ -1,24 +1,24 @@ -import argparse - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('-f', '--fileNameList', type=str, help="The list of changed files") - args = parser.parse_args() - name_list = args.fileNameList.split(":") - folder_need_check = set() - for loc in name_list: - # Find only the sub-sub-folder of 'example' folder - # the examples folder structure is like - # - examples - # - area - # - application - # - file - if loc.split("/")[0] == "examples" and len(loc.split("/")) >= 4: - folder_need_check.add('/'.join(loc.split("/")[1:3])) - # Output the result using print. Then the shell can get the values. - print(list(folder_need_check)) - - -if __name__ == '__main__': - main() +import argparse + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('-f', '--fileNameList', type=str, help="The list of changed files") + args = parser.parse_args() + name_list = args.fileNameList.split(":") + folder_need_check = set() + for loc in name_list: + # Find only the sub-sub-folder of 'example' folder + # the examples folder structure is like + # - examples + # - area + # - application + # - file + if loc.split("/")[0] == "examples" and len(loc.split("/")) >= 4: + folder_need_check.add('/'.join(loc.split("/")[1:3])) + # Output the result using print. Then the shell can get the values. + print(list(folder_need_check)) + + +if __name__ == '__main__': + main() diff --git a/.github/workflows/scripts/generate_leaderboard_and_send_to_lark.py b/.github/workflows/scripts/generate_leaderboard_and_send_to_lark.py index 16b8957c1d88..2884e38dd3dd 100644 --- a/.github/workflows/scripts/generate_leaderboard_and_send_to_lark.py +++ b/.github/workflows/scripts/generate_leaderboard_and_send_to_lark.py @@ -1,5 +1,4 @@ import os -from dataclasses import dataclass from datetime import datetime, timedelta from typing import Any, Dict, List @@ -10,8 +9,7 @@ from requests_toolbelt import MultipartEncoder -@dataclass -class Contributor: +class Counter(dict): """ Dataclass for a github contributor. @@ -19,8 +17,40 @@ class Contributor: name (str): name of the contributor num_commits_this_week (int): number of commits made within one week """ - name: str - num_commits_this_week: int + + def record(self, item: str): + if item in self: + self[item] += 1 + else: + self[item] = 1 + + def to_sorted_list(self): + data = [(key, value) for key, value in self.items()] + data.sort(key=lambda x: x[1], reverse=True) + return data + + +def get_utc_time_one_week_ago(): + """ + Get the UTC time one week ago. + """ + now = datetime.utcnow() + start_datetime = now - timedelta(days=7) + return start_datetime + + +def datetime2str(dt): + """ + Convert datetime to string in the format of YYYY-MM-DDTHH:MM:SSZ + """ + return dt.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def str2datetime(string): + """ + Convert string in the format of YYYY-MM-DDTHH:MM:SSZ to datetime + """ + return datetime.strptime(string, "%Y-%m-%dT%H:%M:%SZ") def plot_bar_chart(x: List[Any], y: List[Any], xlabel: str, ylabel: str, title: str, output_path: str) -> None: @@ -36,9 +66,30 @@ def plot_bar_chart(x: List[Any], y: List[Any], xlabel: str, ylabel: str, title: plt.savefig(output_path, dpi=1200) -def get_issue_pull_request_comments(github_token: str, since: str) -> Dict[str, int]: +def get_organization_repositories(github_token, organization_name) -> List[str]: + """ + Retrieve the public repositories under the organization. """ - Retrive the issue/PR comments made by our members in the last 7 days. + url = f"https://api.github.com/orgs/{organization_name}/repos?type=public" + + # prepare header + headers = { + 'Authorization': f'Bearer {github_token}', + 'Accept': 'application/vnd.github+json', + 'X-GitHub-Api-Version': '2022-11-28' + } + + res = requests.get(url, headers=headers).json() + repo_list = [] + + for item in res: + repo_list.append(item['name']) + return repo_list + + +def get_issue_pull_request_comments(github_token: str, org_name: str, repo_name: str, since: str) -> Dict[str, int]: + """ + Retrieve the issue/PR comments made by our members in the last 7 days. Args: github_token (str): GitHub access token for API calls @@ -56,7 +107,7 @@ def get_issue_pull_request_comments(github_token: str, since: str) -> Dict[str, # do pagination to the API page = 1 while True: - comment_api = f'https://api.github.com/repos/hpcaitech/ColossalAI/issues/comments?since={since}&page={page}' + comment_api = f'https://api.github.com/repos/{org_name}/{repo_name}/issues/comments?since={since}&page={page}' comment_response = requests.get(comment_api, headers=headers).json() if len(comment_response) == 0: @@ -70,7 +121,7 @@ def get_issue_pull_request_comments(github_token: str, since: str) -> Dict[str, continue issue_id = item['issue_url'].split('/')[-1] - issue_api = f'https://api.github.com/repos/hpcaitech/ColossalAI/issues/{issue_id}' + issue_api = f'https://api.github.com/repos/{org_name}/{repo_name}/issues/{issue_id}' issue_response = requests.get(issue_api, headers=headers).json() issue_author_relationship = issue_response['author_association'] @@ -87,9 +138,9 @@ def get_issue_pull_request_comments(github_token: str, since: str) -> Dict[str, return user_engagement_count -def get_discussion_comments(github_token, since) -> Dict[str, int]: +def get_discussion_comments(github_token: str, org_name: str, repo_name: str, since: str) -> Dict[str, int]: """ - Retrive the discussion comments made by our members in the last 7 days. + Retrieve the discussion comments made by our members in the last 7 days. This is only available via the GitHub GraphQL API. Args: @@ -105,7 +156,7 @@ def _generate_discussion_query(num, cursor: str = None): offset_str = f", after: \"{cursor}\"" query = f""" {{ - repository(owner: "hpcaitech", name: "ColossalAI"){{ + repository(owner: "{org_name}", name: "{repo_name}"){{ discussions(first: {num} {offset_str}){{ edges {{ cursor @@ -134,7 +185,7 @@ def _generate_comment_reply_count_for_discussion(discussion_number, num, cursor: offset_str = f", before: \"{cursor}\"" query = f""" {{ - repository(owner: "hpcaitech", name: "ColossalAI"){{ + repository(owner: "{org_name}", name: "{repo_name}"){{ discussion(number: {discussion_number}){{ title comments(last: {num} {offset_str}){{ @@ -191,10 +242,10 @@ def _call_graphql_api(query): for edge in edges: # print the discussion title discussion = edge['node'] + discussion_updated_at = str2datetime(discussion['updatedAt']) - discussion_updated_at = datetime.strptime(discussion['updatedAt'], "%Y-%m-%dT%H:%M:%SZ") # check if the updatedAt is within the last 7 days - # if yes, add it to dicussion_numbers + # if yes, add it to discussion_numbers if discussion_updated_at > since: if discussion['authorAssociation'] != 'MEMBER': discussion_numbers.append(discussion['number']) @@ -207,14 +258,14 @@ def _call_graphql_api(query): # update cursor cursor = edges[-1]['cursor'] - # get the dicussion comments and replies made by our member + # get the discussion comments and replies made by our member user_engagement_count = {} - for dicussion_number in discussion_numbers: + for discussion_number in discussion_numbers: cursor = None num_per_request = 10 while True: - query = _generate_comment_reply_count_for_discussion(dicussion_number, num_per_request, cursor) + query = _generate_comment_reply_count_for_discussion(discussion_number, num_per_request, cursor) data = _call_graphql_api(query) # get the comments @@ -249,7 +300,8 @@ def _call_graphql_api(query): reply = reply_edge['node'] if reply['authorAssociation'] == 'MEMBER': # check if the updatedAt is within the last 7 days - # if yes, add it to dicussion_numbers + # if yes, add it to discussion_numbers + reply_updated_at = datetime.strptime(reply['updatedAt'], "%Y-%m-%dT%H:%M:%SZ") if reply_updated_at > since: member_name = reply['author']['login'] @@ -260,7 +312,7 @@ def _call_graphql_api(query): return user_engagement_count -def generate_user_engagement_leaderboard_image(github_token: str, output_path: str) -> bool: +def generate_user_engagement_leaderboard_image(github_token: str, org_name: str, repo_list: List[str], output_path: str) -> bool: """ Generate the user engagement leaderboard image for stats within the last 7 days @@ -270,23 +322,29 @@ def generate_user_engagement_leaderboard_image(github_token: str, output_path: s """ # request to the Github API to get the users who have replied the most in the last 7 days - now = datetime.utcnow() - start_datetime = now - timedelta(days=7) - start_datetime_str = start_datetime.strftime("%Y-%m-%dT%H:%M:%SZ") + start_datetime = get_utc_time_one_week_ago() + start_datetime_str = datetime2str(start_datetime) # get the issue/PR comments and discussion comment count - issue_pr_engagement_count = get_issue_pull_request_comments(github_token=github_token, since=start_datetime_str) - discussion_engagement_count = get_discussion_comments(github_token=github_token, since=start_datetime) total_engagement_count = {} - # update the total engagement count - total_engagement_count.update(issue_pr_engagement_count) - for name, count in discussion_engagement_count.items(): - if name in total_engagement_count: - total_engagement_count[name] += count - else: - total_engagement_count[name] = count + def _update_count(counter): + for name, count in counter.items(): + if name in total_engagement_count: + total_engagement_count[name] += count + else: + total_engagement_count[name] = count + + for repo_name in repo_list: + print(f"Fetching user engagement count for {repo_name}/{repo_name}") + issue_pr_engagement_count = get_issue_pull_request_comments(github_token=github_token, org_name=org_name, repo_name=repo_name, since=start_datetime_str) + discussion_engagement_count = get_discussion_comments(github_token=github_token, org_name=org_name, repo_name=repo_name, since=start_datetime) + + # update the total engagement count + _update_count(issue_pr_engagement_count) + _update_count(discussion_engagement_count) + # prepare the data for plotting x = [] y = [] @@ -302,9 +360,6 @@ def generate_user_engagement_leaderboard_image(github_token: str, output_path: s x.append(count) y.append(name) - # use Shanghai time to display on the image - start_datetime_str = datetime.now(pytz.timezone('Asia/Shanghai')).strftime("%Y-%m-%dT%H:%M:%SZ") - # plot the leaderboard xlabel = f"Number of Comments made (since {start_datetime_str})" ylabel = "Member" @@ -315,7 +370,7 @@ def generate_user_engagement_leaderboard_image(github_token: str, output_path: s return False -def generate_contributor_leaderboard_image(github_token, output_path) -> bool: +def generate_contributor_leaderboard_image(github_token, org_name, repo_list, output_path) -> bool: """ Generate the contributor leaderboard image for stats within the last 7 days @@ -324,54 +379,81 @@ def generate_contributor_leaderboard_image(github_token, output_path) -> bool: output_path (str): the path to save the image """ # request to the Github API to get the users who have contributed in the last 7 days - URL = 'https://api.github.com/repos/hpcaitech/ColossalAI/stats/contributors' headers = { 'Authorization': f'Bearer {github_token}', 'Accept': 'application/vnd.github+json', 'X-GitHub-Api-Version': '2022-11-28' } - while True: - response = requests.get(URL, headers=headers).json() + counter = Counter() + start_datetime = get_utc_time_one_week_ago() - if len(response) != 0: - # sometimes the Github API returns empty response for unknown reason - # request again if the response is empty - break + def _get_url(org_name, repo_name, page): + return f'https://api.github.com/repos/{org_name}/{repo_name}/pulls?per_page=50&page={page}&state=closed' + + def _iterate_by_page(org_name, repo_name): + page = 1 + stop = False + + while not stop: + print(f"Fetching pull request data for {org_name}/{repo_name} - page{page}") + url = _get_url(org_name, repo_name, page) - contributor_list = [] + while True: + response = requests.get(url, headers=headers).json() - # get number of commits for each contributor - start_timestamp = None - for item in response: - num_commits_this_week = item['weeks'][-1]['c'] - name = item['author']['login'] - contributor = Contributor(name=name, num_commits_this_week=num_commits_this_week) - contributor_list.append(contributor) + if isinstance(response, list): + # sometimes the Github API returns nothing + # request again if the response is not a list + break + print("Empty response, request again...") - # update start_timestamp - start_timestamp = item['weeks'][-1]['w'] + if len(response) == 0: + # if the response is empty, stop + stop = True + break + + # count the pull request and author from response + for pr_data in response: + merged_at = pr_data['merged_at'] + author = pr_data['user']['login'] + + if merged_at is None: + continue + + merge_datetime = str2datetime(merged_at) + + if merge_datetime < start_datetime: + # if we found a pull request that is merged before the start_datetime + # we stop + stop = True + break + else: + # record the author1 + counter.record(author) + + # next page + page += 1 + + for repo_name in repo_list: + _iterate_by_page(org_name, repo_name) # convert unix timestamp to Beijing datetime - start_datetime = datetime.fromtimestamp(start_timestamp, tz=pytz.timezone('Asia/Shanghai')) - start_datetime_str = start_datetime.strftime("%Y-%m-%dT%H:%M:%SZ") + bj_start_datetime = datetime.fromtimestamp(start_datetime.timestamp(), tz=pytz.timezone('Asia/Shanghai')) + bj_start_datetime_str = datetime2str(bj_start_datetime) - # sort by number of commits - contributor_list.sort(key=lambda x: x.num_commits_this_week, reverse=True) + contribution_list = counter.to_sorted_list() # remove contributors who has zero commits - contributor_list = [x for x in contributor_list if x.num_commits_this_week > 0] - - # prepare the data for plotting - x = [x.num_commits_this_week for x in contributor_list] - y = [x.name for x in contributor_list] + author_list = [x[0] for x in contribution_list] + num_commit_list = [x[1] for x in contribution_list] # plot - if len(x) > 0: - xlabel = f"Number of Commits (since {start_datetime_str})" + if len(author_list) > 0: + xlabel = f"Number of Pull Requests (since {bj_start_datetime_str})" ylabel = "Contributor" title = 'Active Contributor Leaderboard' - plot_bar_chart(x, y, xlabel=xlabel, ylabel=ylabel, title=title, output_path=output_path) + plot_bar_chart(num_commit_list, author_list, xlabel=xlabel, ylabel=ylabel, title=title, output_path=output_path) return True else: return False @@ -438,10 +520,14 @@ def send_message_to_lark(message: str, webhook_url: str): GITHUB_TOKEN = os.environ['GITHUB_TOKEN'] CONTRIBUTOR_IMAGE_PATH = 'contributor_leaderboard.png' USER_ENGAGEMENT_IMAGE_PATH = 'engagement_leaderboard.png' + ORG_NAME = "hpcaitech" + + # get all open source repositories + REPO_LIST = get_organization_repositories(GITHUB_TOKEN, ORG_NAME) # generate images - contrib_success = generate_contributor_leaderboard_image(GITHUB_TOKEN, CONTRIBUTOR_IMAGE_PATH) - engagement_success = generate_user_engagement_leaderboard_image(GITHUB_TOKEN, USER_ENGAGEMENT_IMAGE_PATH) + contrib_success = generate_contributor_leaderboard_image(GITHUB_TOKEN, ORG_NAME, REPO_LIST, CONTRIBUTOR_IMAGE_PATH) + engagement_success = generate_user_engagement_leaderboard_image(GITHUB_TOKEN, ORG_NAME, REPO_LIST, USER_ENGAGEMENT_IMAGE_PATH) # upload images APP_ID = os.environ['LARK_APP_ID'] @@ -457,8 +543,8 @@ def send_message_to_lark(message: str, webhook_url: str): 2. 用户互动榜单 注： -- 开发贡献者测评标准为：本周由公司成员提交的commit次数 -- 用户互动榜单测评标准为：本周由公司成员在非成员创建的issue/PR/discussion中回复的次数 +- 开发贡献者测评标准为：本周由公司成员与社区在所有开源仓库提交的Pull Request次数 +- 用户互动榜单测评标准为：本周由公司成员在非成员在所有开源仓库创建的issue/PR/discussion中回复的次数 """ send_message_to_lark(message, LARK_WEBHOOK_URL) @@ -467,7 +553,7 @@ def send_message_to_lark(message: str, webhook_url: str): if contrib_success: send_image_to_lark(contributor_image_key, LARK_WEBHOOK_URL) else: - send_message_to_lark("本周没有成员贡献commit，无榜单图片生成。", LARK_WEBHOOK_URL) + send_message_to_lark("本周没有成员贡献PR，无榜单图片生成。", LARK_WEBHOOK_URL) # send user engagement image to lark if engagement_success: diff --git a/.gitignore b/.gitignore index bf74a753894f..81113fa99dd5 100644 --- a/.gitignore +++ b/.gitignore @@ -155,3 +155,7 @@ colossalai/version.py # ignore coverage test file coverage.lcov coverage.xml + +# ignore testmon and coverage files +.coverage +.testmondata* diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 00abcf650158..a3dc020f74e9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -30,6 +30,12 @@ pip install -e . ### Unit Tests We use [PyTest](https://docs.pytest.org/en/latest/) to execute tests. You can install pytest by `pip install pytest`. As some of the tests require initialization of the distributed backend, GPUs are needed to execute these tests. +To set up the environment for unit testing, first change your current directory to the root directory of your local ColossalAI repository, then run +```bash +pip install -r requirements/requirements-test.txt +``` +If you encounter an error telling "Could not find a version that satisfies the requirement fbgemm-gpu==0.2.0", please downgrade your python version to 3.8 or 3.9 and try again. + If you only want to run CPU tests, you can run ```bash @@ -138,4 +144,4 @@ You can now create a pull request on the GitHub webpage of your repository. The Do write clearly the description of your pull request and [link the pull request to your target issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue). This will automatically close the issue when the pull request is approved. -In case of code conflict, you should rebase your branch and resolve the conflicts manually. \ No newline at end of file +In case of code conflict, you should rebase your branch and resolve the conflicts manually. diff --git a/LICENSE b/LICENSE index 394791da2771..c7a5bb16880e 100644 --- a/LICENSE +++ b/LICENSE @@ -326,3 +326,73 @@ Copyright 2021- HPC-AI Technology Inc. All rights reserved. CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + ---------------- LICENSE FOR Flash Attention ---------------- + + BSD 3-Clause License + + Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + ---------------- LICENSE FOR Facebook xFormers ---------------- + + From xFormers: + + Copyright (c) Facebook, Inc. and its affiliates + + + === + + BSD 3-Clause License + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + + 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE + LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md index 3b55649b44bb..21670e1e59fb 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [![logo](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/colossal-ai_logo_vertical.png)](https://www.colossalai.org/) - Colossal-AI: Making large AI models cheaper, faster and more accessible + Colossal-AI: Making large AI models cheaper, faster, and more accessible

Paper | Documentation | @@ -20,13 +20,16 @@ [![WeChat badge](https://img.shields.io/badge/微信-加入-green?logo=wechat&)](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png) - | [English](README.md) | [中文](README-zh-Hans.md) | + | [English](README.md) | [中文](docs/README-zh-Hans.md) | ## Latest News +* [2023/07] [65B Model Pretraining Accelerated by 38%, Best Practices for Building LLaMA-Like Base Models Open-Source](https://www.hpc-ai.tech/blog/large-model-pretraining) +* [2023/03] [ColossalChat: An Open-Source Solution for Cloning ChatGPT With a Complete RLHF Pipeline](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +* [2023/03] [Intel and Colossal-AI Partner to Deliver Cost-Efficient Open-Source Solution for Protein Folding Structure Prediction](https://www.hpc-ai.tech/blog/intel-habana) * [2023/03] [AWS and Google Fund Colossal-AI with Startup Cloud Programs](https://www.hpc-ai.tech/blog/aws-and-google-fund-colossal-ai-with-startup-cloud-programs) -* [2023/02] [Open source solution replicates ChatGPT training process! Ready to go with only 1.6GB GPU memory](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) +* [2023/02] [Open Source Solution Replicates ChatGPT Training Process! Ready to go with only 1.6GB GPU Memory](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) * [2023/01] [Hardware Savings Up to 46 Times for AIGC and Automatic Parallelism](https://medium.com/pytorch/latest-colossal-ai-boasts-novel-automatic-parallelism-and-offers-savings-up-to-46x-for-stable-1453b48f3f02) * [2022/11] [Diffusion Pretraining and Hardware Fine-Tuning Can Be Almost 7X Cheaper](https://www.hpc-ai.tech/blog/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper) * [2022/10] [Use a Laptop to Analyze 90% of Proteins, With a Single-GPU Inference Sequence Exceeding 10,000](https://www.hpc-ai.tech/blog/use-a-laptop-to-analyze-90-of-proteins-with-a-single-gpu-inference-sequence-exceeding) @@ -36,9 +39,18 @@ @@ -113,21 +117,122 @@ distributed training and inference in a few lines. - [PatrickStar](https://arxiv.org/abs/2108.05818) - Friendly Usage - - Parallelism based on configuration file + - Parallelism based on the configuration file - Inference - [Energon-AI](https://github.com/hpcaitech/EnergonAI) +

(back to top)

+ +## Colossal-AI in the Real World + +### ColossalChat + + + +[ColossalChat](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat): An open-source solution for cloning [ChatGPT](https://openai.com/blog/chatgpt/) with a complete RLHF pipeline. +[[code]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat) +[[blog]](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +[[demo]](https://www.youtube.com/watch?v=HcTiHzApHm0) +[[tutorial]](https://www.youtube.com/watch?v=-qFBZFmOJfg) + +

+ +

+ +- Up to 10 times faster for RLHF PPO Stage3 Training + +

+ +

+ +- Up to 7.73 times faster for single server training and 1.42 times faster for single-GPU inference + +

+ +

+ +- Up to 10.3x growth in model capacity on one GPU +- A mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU) + +

+ +

+ +- Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPU +- Keep at a sufficiently high running speed + +

(back to top)

+ + +### AIGC +Acceleration of AIGC (AI-Generated Content) models such as [Stable Diffusion v1](https://github.com/CompVis/stable-diffusion) and [Stable Diffusion v2](https://github.com/Stability-AI/stablediffusion). +

+ +

+ +- [Training](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): Reduce Stable Diffusion memory consumption by up to 5.6x and hardware cost by up to 46x (from A100 to RTX3060). + +

+ +

+ +- [DreamBooth Fine-tuning](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/dreambooth): Personalize your model using just 3-5 images of the desired subject. + +

+ +

+ +- [Inference](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): Reduce inference GPU memory consumption by 2.5x. + + +

(back to top)

+ +### Biomedicine +Acceleration of [AlphaFold Protein Structure](https://alphafold.ebi.ac.uk/) + +

+ +

+ +- [FastFold](https://github.com/hpcaitech/FastFold): Accelerating training and inference on GPU Clusters, faster data processing, inference sequence containing more than 10000 residues. + +

+ +

+ +- [FastFold with Intel](https://github.com/hpcaitech/FastFold): 3x inference acceleration and 39% cost reduce. + +

+ +

+ +- [xTrimoMultimer](https://github.com/biomap-research/xTrimoMultimer): accelerating structure prediction of protein monomers and multimer by 11x. + +

(back to top)

## Parallel Training Demo +### LLaMA +

+ +

+ +- 65-billion-parameter large model pretraining accelerated by 38% +[[code]](https://github.com/hpcaitech/ColossalAI/tree/example/llama/examples/language/llama) +[[blog]](https://www.hpc-ai.tech/blog/large-model-pretraining) + ### GPT-3

-- Save 50% GPU resources, and 10.7% acceleration +- Save 50% GPU resources and 10.7% acceleration ### GPT-2 @@ -149,7 +254,7 @@ distributed training and inference in a few lines. ### OPT -- [Open Pretrained Transformer (OPT)](https://github.com/facebookresearch/metaseq), a 175-Billion parameter AI language model released by Meta, which stimulates AI programmers to perform various downstream tasks and application deployments because public pretrained model weights. +- [Open Pretrained Transformer (OPT)](https://github.com/facebookresearch/metaseq), a 175-Billion parameter AI language model released by Meta, which stimulates AI programmers to perform various downstream tasks and application deployments because of public pre-trained model weights. - 45% speedup fine-tuning OPT at low cost in lines. [[Example]](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/opt) [[Online Serving]](https://colossalai.org/docs/advanced_tutorials/opt_service) Please visit our [documentation](https://www.colossalai.org/) and [examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples) for more details. @@ -211,79 +316,6 @@ Please visit our [documentation](https://www.colossalai.org/) and [examples](htt - [BLOOM](https://github.com/hpcaitech/EnergonAI/tree/main/examples/bloom): Reduce hardware deployment costs of 176-billion-parameter BLOOM by more than 10 times. -

(back to top)

- -## Colossal-AI in the Real World -### ChatGPT -A low-cost [ChatGPT](https://openai.com/blog/chatgpt/) equivalent implementation process. [[code]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/ChatGPT) [[blog]](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) -

- -

- -- Up to 7.73 times faster for single server training and 1.42 times faster for single-GPU inference - -

- -

- -- Up to 10.3x growth in model capacity on one GPU -- A mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU) - -

- -

- -- Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPU -- Keep in a sufficiently high running speed - -

(back to top)

- - -### AIGC -Acceleration of AIGC (AI-Generated Content) models such as [Stable Diffusion v1](https://github.com/CompVis/stable-diffusion) and [Stable Diffusion v2](https://github.com/Stability-AI/stablediffusion). -

- -

- -- [Training](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): Reduce Stable Diffusion memory consumption by up to 5.6x and hardware cost by up to 46x (from A100 to RTX3060). - -

- -

- -- [DreamBooth Fine-tuning](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/dreambooth): Personalize your model using just 3-5 images of the desired subject. - -

- -

- -- [Inference](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): Reduce inference GPU memory consumption by 2.5x. - - -

(back to top)

- -### Biomedicine -Acceleration of [AlphaFold Protein Structure](https://alphafold.ebi.ac.uk/) - -

- -

- -- [FastFold](https://github.com/hpcaitech/FastFold): Accelerating training and inference on GPU Clusters, faster data processing, inference sequence containing more than 10000 residues. - -

- -

- -- [FastFold with Intel](https://github.com/hpcaitech/FastFold): 3x inference acceleration and 39% cost reduce. - -

- -

- -- [xTrimoMultimer](https://github.com/biomap-research/xTrimoMultimer): accelerating structure prediction of protein monomers and multimer by 11x. - -

(back to top)

## Installation @@ -292,8 +324,10 @@ Requirements: - PyTorch >= 1.11 (PyTorch 2.x in progress) - Python >= 3.7 - CUDA >= 11.0 +- [NVIDIA GPU Compute Capability](https://developer.nvidia.com/cuda-gpus) >= 7.0 (V100/RTX20 and higher) +- Linux OS -If you encounter any problem about installation, you may want to raise an [issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose) in this repository. +If you encounter any problem with installation, you may want to raise an [issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose) in this repository. ### Install from PyPI @@ -311,9 +345,9 @@ However, if you want to build the PyTorch extensions during installation, you ca CUDA_EXT=1 pip install colossalai ``` -**Otherwise, CUDA kernels will be built during runtime when you actually need it.** +**Otherwise, CUDA kernels will be built during runtime when you actually need them.** -We also keep release the nightly version to PyPI on a weekly basis. This allows you to access the unreleased features and bug fixes in the main branch. +We also keep releasing the nightly version to PyPI every week. This allows you to access the unreleased features and bug fixes in the main branch. Installation can be made via ```bash @@ -322,7 +356,7 @@ pip install colossalai-nightly ### Download From Source -> The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. :) +> The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problems. :) ```shell git clone https://github.com/hpcaitech/ColossalAI.git @@ -339,6 +373,22 @@ If you want to install and enable CUDA kernel fusion (compulsory installation wh CUDA_EXT=1 pip install . ``` +For Users with CUDA 10.2, you can still build ColossalAI from source. However, you need to manually download the cub library and copy it to the corresponding directory. + +```bash +# clone the repository +git clone https://github.com/hpcaitech/ColossalAI.git +cd ColossalAI + +# download the cub library +wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip +unzip 1.8.0.zip +cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ + +# install +CUDA_EXT=1 pip install . +``` +

(back to top)

## Use Docker @@ -375,7 +425,7 @@ Join the Colossal-AI community on [Forum](https://github.com/hpcaitech/ColossalA [Slack](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w), and [WeChat(微信)](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png "qrcode") to share your suggestions, feedback, and questions with our engineering team. -## Invitation to open-source contribution +## Contributing Referring to the successful attempts of [BLOOM](https://bigscience.huggingface.co/) and [Stable Diffusion](https://en.wikipedia.org/wiki/Stable_Diffusion), any and all developers and partners with computing powers, datasets, models are welcome to join and build the Colossal-AI community, making efforts towards the era of big AI models! You may contact us or participate in the following ways: @@ -385,9 +435,10 @@ You may contact us or participate in the following ways: Thanks so much to all of our amazing contributors! - + + + -*The order of contributor avatars is randomly shuffled.*

(back to top)

@@ -399,7 +450,7 @@ We leverage the power of [GitHub Actions](https://github.com/features/actions) t ## Cite Us -This project is inspired by some related projects (some by our team and some by other organizations). We would like to credit these amazing projects as listed in the [Reference List](./REFERENCE.md). +This project is inspired by some related projects (some by our team and some by other organizations). We would like to credit these amazing projects as listed in the [Reference List](./docs/REFERENCE.md). To cite this project, you can use the following BibTeX citation. @@ -412,6 +463,6 @@ To cite this project, you can use the following BibTeX citation. } ``` -Colossal-AI has been accepted as official tutorials by top conference [SC](https://sc22.supercomputing.org/), [AAAI](https://aaai.org/Conferences/AAAI-23/), [PPoPP](https://ppopp23.sigplan.org/), [CVPR](https://cvpr2023.thecvf.com/), [ISC](https://www.isc-hpc.com/), etc. +Colossal-AI has been accepted as official tutorial by top conferences [SC](https://sc22.supercomputing.org/), [AAAI](https://aaai.org/Conferences/AAAI-23/), [PPoPP](https://ppopp23.sigplan.org/), [CVPR](https://cvpr2023.thecvf.com/), [ISC](https://www.isc-hpc.com/), etc.

(back to top)

diff --git a/applications/ChatGPT/.gitignore b/applications/Chat/.gitignore similarity index 96% rename from applications/ChatGPT/.gitignore rename to applications/Chat/.gitignore index 40f3f6debeee..2b9b4f345d0f 100644 --- a/applications/ChatGPT/.gitignore +++ b/applications/Chat/.gitignore @@ -142,5 +142,7 @@ docs/.build # pytorch checkpoint *.pt -# ignore version.py generated by setup.py -colossalai/version.py +# wandb log +example/wandb/ + +examples/awesome-chatgpt-prompts/ \ No newline at end of file diff --git a/applications/ChatGPT/LICENSE b/applications/Chat/LICENSE similarity index 100% rename from applications/ChatGPT/LICENSE rename to applications/Chat/LICENSE diff --git a/applications/Chat/README.md b/applications/Chat/README.md new file mode 100644 index 000000000000..162528cee414 --- /dev/null +++ b/applications/Chat/README.md @@ -0,0 +1,461 @@ +

+ +
+ ColossalChat +

+ + +## Table of Contents + +- [Table of Contents](#table-of-contents) +- [What is ColossalChat and Coati ?](#what-is-colossalchat-and-coati-) +- [Online demo](#online-demo) +- [Install](#install) + - [Install the environment](#install-the-environment) + - [Install the Transformers](#install-the-transformers) +- [How to use?](#how-to-use) + - [Supervised datasets collection](#supervised-datasets-collection) + - [RLHF Training Stage1 - Supervised instructs tuning](#RLHF-training-stage1---supervised-instructs-tuning) + - [RLHF Training Stage2 - Training reward model](#RLHF-training-stage2---training-reward-model) + - [RLHF Training Stage3 - Training model with reinforcement learning by human feedback](#RLHF-training-stage3---training-model-with-reinforcement-learning-by-human-feedback) + - [Inference Quantization and Serving - After Training](#inference-quantization-and-serving---after-training) +- [Coati7B examples](#coati7b-examples) + - [Generation](#generation) + - [Open QA](#open-qa) + - [Limitation for LLaMA-finetuned models](#limitation) + - [Limitation of dataset](#limitation) +- [FAQ](#faq) + - [How to save/load checkpoint](#faq) + - [How to train with limited resources](#faq) +- [The Plan](#the-plan) + - [Real-time progress](#real-time-progress) +- [Invitation to open-source contribution](#invitation-to-open-source-contribution) +- [Quick Preview](#quick-preview) +- [Authors](#authors) +- [Citations](#citations) +- [Licenses](#licenses) +--- +## What is ColossalChat and Coati ? + +[ColossalChat](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat) is the project to implement LLM with RLHF, powered by the [Colossal-AI](https://github.com/hpcaitech/ColossalAI) project. + +Coati stands for `ColossalAI Talking Intelligence`. It is the name for the module implemented in this project and is also the name of the large language model developed by the ColossalChat project. + +The Coati package provides a unified large language model framework that has implemented the following functions +- Supports comprehensive large-model training acceleration capabilities for ColossalAI, without requiring knowledge of complex distributed training algorithms +- Supervised datasets collection +- Supervised instructions fine-tuning +- Training reward model +- Reinforcement learning with human feedback +- Quantization inference +- Fast model deploying +- Perfectly integrated with the Hugging Face ecosystem, a high degree of model customization + +
+

+ +

+ + Image source: https://openai.com/blog/chatgpt +
+ +**As Colossal-AI is undergoing some major updates, this project will be actively maintained to stay in line with the Colossal-AI project.** + + +More details can be found in the latest news. +* [2023/03] [ColossalChat: An Open-Source Solution for Cloning ChatGPT With a Complete RLHF Pipeline](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +* [2023/02] [Open Source Solution Replicates ChatGPT Training Process! Ready to go with only 1.6GB GPU Memory](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) + +## Online demo + + +[ColossalChat](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat): An open-source solution for cloning [ChatGPT](https://openai.com/blog/chatgpt/) with a complete RLHF pipeline. +[[code]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat) +[[blog]](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +[[demo]](https://www.youtube.com/watch?v=HcTiHzApHm0) +[[tutorial]](https://www.youtube.com/watch?v=-qFBZFmOJfg) + +

+ +

+ +> DeepSpeedChat performance comes from its blog on 2023 April 12, ColossalChat performance can be reproduced on an AWS p4d.24xlarge node with 8 A100-40G GPUs with the following command: torchrun --standalone --nproc_per_node 8 benchmark_opt_lora_dummy.py --num_collect_steps 1 --use_kernels --strategy colossalai_zero2 --experience_batch_size 64 --train_batch_size 32 + +## Install + +### Install the environment + +```shell +conda create -n coati +conda activate coati +git clone https://github.com/hpcaitech/ColossalAI.git +cd ColossalAI/applications/Chat +pip install . +``` + +### Install the Transformers + +```shell +pip install transformers==4.30.2 +``` + +## How to use? + +### Supervised datasets collection + +we collected 104K bilingual datasets of Chinese and English, and you can find the datasets in this repo +[InstructionWild](https://github.com/XueFuzhao/InstructionWild) + +Here is how we collected the data +

+ +

+ +### RLHF Training Stage1 - Supervised instructs tuning + +Stage1 is supervised instructs fine-tuning, which uses the datasets mentioned earlier to fine-tune the model. + +You can run the `examples/train_sft.sh` to start a supervised instructs fine-tuning. +[[Stage1 tutorial video]](https://www.youtube.com/watch?v=-qFBZFmOJfg) + +### RLHF Training Stage2 - Training reward model + +Stage2 trains a reward model, which obtains corresponding scores by manually ranking different outputs for the same prompt and supervises the training of the reward model + +You can run the `examples/train_rm.sh` to start a reward model training. +[[Stage2 tutorial video]](https://www.youtube.com/watch?v=gMx2CApKhuo) + +### RLHF Training Stage3 - Training model with reinforcement learning by human feedback + +Stage3 uses reinforcement learning algorithm, which is the most complex part of the training process: + +

+ +

+ +You can run the `examples/train_prompts.sh` to start training PPO with human feedback. +[[Stage3 tutorial video]](https://www.youtube.com/watch?v=Z8wwSHxPL9g) + +For more details, see [`examples/`](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples). + +### Inference Quantization and Serving - After Training + +We provide an online inference server and a benchmark. We aim to run inference on single GPU, so quantization is essential when using large models. + +We support 8-bit quantization (RTN), 4-bit quantization (GPTQ), and FP16 inference. You can +Online inference server scripts can help you deploy your own services. + +For more details, see [`inference/`](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/inference). + +## Coati7B examples + +### Generation + +
E-mail + +![phd](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/Phd.png) +
+ +
coding + +![sort](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/quick_sort.png) + +
+ +
regex + +![regex](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/regex.png) + +
+ +
Tex + +![tex](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/tex.png) + +
+ +
writing + +![writing](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/writing.png) + +
+ +
Table + +![Table](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/table.png) + +
+ +### Open QA +
Game + +![Game](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/game.png) + +
+ +
Travel + +![Travel](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/travel.png) + +
+ +
Physical + +![Physical](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/physical.png) + +
+ +
Chemical + +![Chemical](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/chemical.png) + +
+ +
Economy + +![Economy](https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chat/economy.png) + +
+ +You can find more examples in this [repo](https://github.com/XueFuzhao/InstructionWild/blob/main/comparison.md). + +### Limitation +
Limitation for LLaMA-finetuned models +- Both Alpaca and ColossalChat are based on LLaMA. It is hard to compensate for the missing knowledge in the pre-training stage. +- Lack of counting ability: Cannot count the number of items in a list. +- Lack of Logics (reasoning and calculation) +- Tend to repeat the last sentence (fail to produce the end token). +- Poor multilingual results: LLaMA is mainly trained on English datasets (Generation performs better than QA). +
+ +
Limitation of dataset +- Lack of summarization ability: No such instructions in finetune datasets. +- Lack of multi-turn chat: No such instructions in finetune datasets +- Lack of self-recognition: No such instructions in finetune datasets +- Lack of Safety: + - When the input contains fake facts, the model makes up false facts and explanations. + - Cannot abide by OpenAI's policy: When generating prompts from OpenAI API, it always abides by its policy. So no violation case is in the datasets. +
+ +## FAQ + +
How to save/load checkpoint + +We have integrated the Transformers save and load pipeline, allowing users to freely call Hugging Face's language models and save them in the HF format. + +``` +from coati.models.llama import LlamaLM +from coati.trainer import SFTTrainer + +model = LlamaLM(pretrained=args.pretrain) +tokenizer = AutoTokenizer.from_pretrained(args.pretrain) + +(model, optim) = strategy.prepare((model, optim)) +trainer = SFTTrainer(model=model, + strategy=strategy, + optim=optim, + train_dataloader=train_dataloader, + eval_dataloader=eval_dataloader, + batch_size=args.batch_size, + max_epochs=args.max_epochs, + accumulation_steps = args.accumulation_steps +) + +trainer.fit() +# this saves in pytorch format +strategy.save_model(model, args.save_path, only_rank0=True) + +# this saves in HF format. ColossalAI strategy with stage-3 doesn't support this method +strategy.save_pretrained(model, args.save_path, only_rank0=True, tokenizer=tokenizer) +``` + +
+ +
How to train with limited resources + +Here are some examples that can allow you to train a 7B model on a single or multiple consumer-grade GPUs. + +If you only have a single 24G GPU, you can use the following script. `batch_size`, `lora_rank` and `grad_checkpoint` are the most important parameters to successfully train the model. +``` +torchrun --standalone --nproc_per_node=1 train_sft.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy ddp \ + --log_interval 10 \ + --save_path /path/to/Coati-7B \ + --dataset /path/to/data.json \ + --batch_size 1 \ + --accumulation_steps 8 \ + --lr 2e-5 \ + --max_datasets_size 512 \ + --max_epochs 1 \ + --lora_rank 16 \ + --grad_checkpoint +``` + +`colossalai_gemini` strategy can enable a single 24G GPU to train the whole model without using LoRA if you have sufficient CPU memory. You can use the following script. +``` +torchrun --standalone --nproc_per_node=1 train_sft.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_gemini \ + --log_interval 10 \ + --save_path /path/to/Coati-7B \ + --dataset /path/to/data.json \ + --batch_size 1 \ + --accumulation_steps 8 \ + --lr 2e-5 \ + --max_datasets_size 512 \ + --max_epochs 1 \ + --grad_checkpoint +``` + +If you have 4x32 GB GPUs, you can even train the whole 7B model using our `colossalai_zero2_cpu` strategy! The script is given as follows. +``` +torchrun --standalone --nproc_per_node=4 train_sft.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_zero2_cpu \ + --log_interval 10 \ + --save_path /path/to/Coati-7B \ + --dataset /path/to/data.json \ + --batch_size 1 \ + --accumulation_steps 8 \ + --lr 2e-5 \ + --max_datasets_size 512 \ + --max_epochs 1 \ + --grad_checkpoint +``` +
+ + +## The Plan + +- [x] implement PPO fine-tuning +- [x] implement training reward model +- [x] support LoRA +- [x] support inference +- [x] support llama from [facebook](https://github.com/facebookresearch/llama) +- [x] implement PPO-ptx fine-tuning +- [ ] integrate with Ray +- [ ] support more RL paradigms, like Implicit Language Q-Learning (ILQL), +- [ ] support chain-of-thought by [langchain](https://github.com/hwchase17/langchain) + +### Real-time progress +You will find our progress in github project broad + +[Coati](https://github.com/orgs/hpcaitech/projects/17/views/1) + +## Invitation to open-source contribution +Referring to the successful attempts of [BLOOM](https://bigscience.huggingface.co/) and [Stable Diffusion](https://en.wikipedia.org/wiki/Stable_Diffusion), any and all developers and partners with computing powers, datasets, models are welcome to join and build the Colossal-AI community, making efforts towards the era of big AI models from the starting point of replicating ChatGPT! + +You may contact us or participate in the following ways: +1. [Leaving a Star ⭐](https://github.com/hpcaitech/ColossalAI/stargazers) to show your like and support. Thanks! +2. Posting an [issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose), or submitting a PR on GitHub follow the guideline in [Contributing](https://github.com/hpcaitech/ColossalAI/blob/main/CONTRIBUTING.md). +3. Join the Colossal-AI community on +[Slack](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w), +and [WeChat(微信)](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png "qrcode") to share your ideas. +4. Send your official proposal to email contact@hpcaitech.com + +Thanks so much to all of our amazing contributors! + +## Quick Preview + + +- An open-source low cost solution for cloning [ChatGPT](https://openai.com/blog/chatgpt/) with a complete RLHF pipeline. [[demo]](https://chat.colossalai.org) + +

+ +

+ +- Up to 7.73 times faster for single server training and 1.42 times faster for single-GPU inference + +

+ +

+ +- Up to 10.3x growth in model capacity on one GPU +- A mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU) + +

+ +

+ +- Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPU +- Keep in a sufficiently high running speed + +| Model Pair | Alpaca-7B ⚔ Coati-7B | Coati-7B ⚔ Alpaca-7B | +| :-----------: | :------------------: | :------------------: | +| Better Cases | 38 ⚔ **41** | **45** ⚔ 33 | +| Win Rate | 48% ⚔ **52%** | **58%** ⚔ 42% | +| Average Score | 7.06 ⚔ **7.13** | **7.31** ⚔ 6.82 | +- Our Coati-7B model performs better than Alpaca-7B when using GPT-4 to evaluate model performance. The Coati-7B model we evaluate is an old version we trained a few weeks ago and the new version is around the corner. + +## Authors + +Coati is developed by ColossalAI Team: +- [Fazzie](https://fazzie-key.cool/about/index.html) +- [FrankLeeeee](https://github.com/FrankLeeeee) +- [BlueRum](https://github.com/ht-zhou) +- [ver217](https://github.com/ver217) +- [ofey404](https://github.com/ofey404) + +The Phd student from [(HPC-AI) Lab](https://ai.comp.nus.edu.sg/) also contributed a lot to this project. +- [Zangwei Zheng](https://github.com/zhengzangw) +- [Xue Fuzhao](https://github.com/XueFuzhao) + +## Citations + +```bibtex +@article{Hu2021LoRALA, + title = {LoRA: Low-Rank Adaptation of Large Language Models}, + author = {Edward J. Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Weizhu Chen}, + journal = {ArXiv}, + year = {2021}, + volume = {abs/2106.09685} +} + +@article{ouyang2022training, + title={Training language models to follow instructions with human feedback}, + author={Ouyang, Long and Wu, Jeff and Jiang, Xu and Almeida, Diogo and Wainwright, Carroll L and Mishkin, Pamela and Zhang, Chong and Agarwal, Sandhini and Slama, Katarina and Ray, Alex and others}, + journal={arXiv preprint arXiv:2203.02155}, + year={2022} +} + +@article{touvron2023llama, + title={LLaMA: Open and Efficient Foundation Language Models}, + author={Touvron, Hugo and Lavril, Thibaut and Izacard, Gautier and Martinet, Xavier and Lachaux, Marie-Anne and Lacroix, Timoth{\'e}e and Rozi{\`e}re, Baptiste and Goyal, Naman and Hambro, Eric and Azhar, Faisal and Rodriguez, Aurelien and Joulin, Armand and Grave, Edouard and Lample, Guillaume}, + journal={arXiv preprint arXiv:2302.13971}, + year={2023} +} + +@misc{alpaca, + author = {Rohan Taori and Ishaan Gulrajani and Tianyi Zhang and Yann Dubois and Xuechen Li and Carlos Guestrin and Percy Liang and Tatsunori B. Hashimoto }, + title = {Stanford Alpaca: An Instruction-following LLaMA model}, + year = {2023}, + publisher = {GitHub}, + journal = {GitHub repository}, + howpublished = {\url{https://github.com/tatsu-lab/stanford_alpaca}}, +} + +@misc{instructionwild, + author = {Fuzhao Xue and Zangwei Zheng and Yang You }, + title = {Instruction in the Wild: A User-based Instruction Dataset}, + year = {2023}, + publisher = {GitHub}, + journal = {GitHub repository}, + howpublished = {\url{https://github.com/XueFuzhao/InstructionWild}}, +} +``` + +## Licenses + +Coati is licensed under the [Apache 2.0 License](LICENSE). diff --git a/applications/Chat/benchmarks/README.md b/applications/Chat/benchmarks/README.md new file mode 100644 index 000000000000..bc8ad8ba9816 --- /dev/null +++ b/applications/Chat/benchmarks/README.md @@ -0,0 +1,35 @@ +# Benchmarks + +## Benchmark OPT with LoRA on dummy prompt data + +We provide various OPT models (string in parentheses is the corresponding model name used in this script): + +- OPT-125M (125m) +- OPT-350M (350m) +- OPT-700M (700m) +- OPT-1.3B (1.3b) +- OPT-2.7B (2.7b) +- OPT-3.5B (3.5b) +- OPT-5.5B (5.5b) +- OPT-6.7B (6.7b) +- OPT-10B (10b) +- OPT-13B (13b) + +We also provide various training strategies: + +- ddp: torch DDP +- colossalai_gemini: ColossalAI GeminiDDP with `placement_policy="cuda"`, like zero3 +- colossalai_gemini_cpu: ColossalAI GeminiDDP with `placement_policy="cpu"`, like zero3-offload +- colossalai_zero2: ColossalAI zero2 +- colossalai_zero2_cpu: ColossalAI zero2-offload +- colossalai_zero1: ColossalAI zero1 +- colossalai_zero1_cpu: ColossalAI zero1-offload + +We only support `torchrun` to launch now. E.g. + +```shell +# run OPT-125M with no lora (lora_rank=0) on single-node single-GPU with min batch size +torchrun --standalone --nproc_per_node 1 benchmark_opt_lora_dummy.py --model 125m --critic_model 125m --strategy ddp --experience_batch_size 1 --train_batch_size 1 --lora_rank 0 +# run Actor (OPT-1.3B) and Critic (OPT-350M) with lora_rank=4 on single-node 4-GPU +torchrun --standalone --nproc_per_node 4 benchmark_opt_lora_dummy.py --model 1.3b --critic_model 350m --strategy colossalai_zero2 --lora_rank 4 +``` diff --git a/applications/ChatGPT/benchmarks/benchmark_opt_lora_dummy.py b/applications/Chat/benchmarks/benchmark_opt_lora_dummy.py similarity index 70% rename from applications/ChatGPT/benchmarks/benchmark_opt_lora_dummy.py rename to applications/Chat/benchmarks/benchmark_opt_lora_dummy.py index 207edbca94b5..90471ed727b0 100644 --- a/applications/ChatGPT/benchmarks/benchmark_opt_lora_dummy.py +++ b/applications/Chat/benchmarks/benchmark_opt_lora_dummy.py @@ -4,12 +4,13 @@ import torch import torch.distributed as dist import torch.nn as nn -from chatgpt.models.base import RewardModel -from chatgpt.models.opt import OPTActor, OPTCritic -from chatgpt.trainer import PPOTrainer -from chatgpt.trainer.callbacks import PerformanceEvaluator -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy +from coati.models.base import RewardModel +from coati.models.opt import OPTActor, OPTCritic +from coati.trainer import PPOTrainer +from coati.trainer.callbacks import PerformanceEvaluator +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy, Strategy from torch.optim import Adam +from torch.utils.data import DataLoader from transformers import AutoTokenizer from transformers.models.opt.configuration_opt import OPTConfig @@ -18,7 +19,7 @@ def get_model_numel(model: nn.Module, strategy: Strategy) -> int: numel = sum(p.numel() for p in model.parameters()) - if isinstance(strategy, ColossalAIStrategy) and strategy.stage == 3 and strategy.shard_init: + if isinstance(strategy, GeminiStrategy) and strategy.shard_init: numel *= dist.get_world_size() return numel @@ -75,30 +76,35 @@ def main(args): if args.strategy == 'ddp': strategy = DDPStrategy() elif args.strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) + strategy = GeminiStrategy(placement_policy='cuda', initial_scale=2**5) elif args.strategy == 'colossalai_gemini_cpu': - strategy = ColossalAIStrategy(stage=3, placement_policy='cpu', initial_scale=2**5) + strategy = GeminiStrategy(placement_policy='cpu', initial_scale=2**5) elif args.strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') elif args.strategy == 'colossalai_zero2_cpu': - strategy = ColossalAIStrategy(stage=2, placement_policy='cpu') + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cpu') elif args.strategy == 'colossalai_zero1': - strategy = ColossalAIStrategy(stage=1, placement_policy='cuda') + strategy = LowLevelZeroStrategy(stage=1, placement_policy='cuda') elif args.strategy == 'colossalai_zero1_cpu': - strategy = ColossalAIStrategy(stage=1, placement_policy='cpu') + strategy = LowLevelZeroStrategy(stage=1, placement_policy='cpu') else: raise ValueError(f'Unsupported strategy "{args.strategy}"') torch.cuda.set_per_process_memory_fraction(args.cuda_mem_frac) model_config = get_gpt_config(args.model) - + critic_config = get_gpt_config(args.critic_model) with strategy.model_init_context(): actor = OPTActor(config=model_config, lora_rank=args.lora_rank).cuda() - critic = OPTCritic(config=model_config, lora_rank=args.lora_rank).cuda() + critic = OPTCritic(config=critic_config, lora_rank=args.lora_rank).cuda() + + initial_model = deepcopy(actor).cuda().half() + reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda().half() - initial_model = deepcopy(actor).cuda() - reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() + if args.use_kernels: + from coati.kernels import convert_to_xformer_model + actor, critic, initial_model, reward_model = map(convert_to_xformer_model, + (actor, critic, initial_model, reward_model)) actor_numel = get_model_numel(actor, strategy) critic_numel = get_model_numel(critic, strategy) @@ -127,8 +133,13 @@ def main(args): tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m') tokenizer.pad_token = tokenizer.eos_token - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare( - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model) + (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim)) + + random_prompts = torch.randint(tokenizer.vocab_size, (1000, 256), device=torch.cuda.current_device()) + dataloader = DataLoader(random_prompts, + batch_size=args.experience_batch_size, + shuffle=True, + collate_fn=preprocess_batch) trainer = PPOTrainer(strategy, actor, @@ -137,23 +148,23 @@ def main(args): initial_model, actor_optim, critic_optim, - max_epochs=args.max_epochs, + ptx_coef=0, train_batch_size=args.train_batch_size, - experience_batch_size=args.experience_batch_size, - tokenizer=preprocess_batch, + offload_inference_models=args.offload_inference_models, max_length=512, do_sample=True, temperature=1.0, top_k=50, + use_cache=True, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, callbacks=[performance_evaluator]) - random_prompts = torch.randint(tokenizer.vocab_size, (1000, 400), device=torch.cuda.current_device()) - trainer.fit(random_prompts, + trainer.fit(prompt_dataloader=dataloader, + pretrain_dataloader=None, num_episodes=args.num_episodes, - max_timesteps=args.max_timesteps, - update_timesteps=args.update_timesteps) + num_update_steps=args.num_update_steps, + num_collect_steps=args.num_collect_steps) print_rank_0(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB') @@ -161,6 +172,7 @@ def main(args): if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model', default='125m') + parser.add_argument('--critic_model', default='125m') parser.add_argument('--strategy', choices=[ 'ddp', 'colossalai_gemini', 'colossalai_gemini_cpu', 'colossalai_zero2', @@ -168,12 +180,13 @@ def main(args): ], default='ddp') parser.add_argument('--num_episodes', type=int, default=3) - parser.add_argument('--max_timesteps', type=int, default=8) - parser.add_argument('--update_timesteps', type=int, default=8) - parser.add_argument('--max_epochs', type=int, default=3) + parser.add_argument('--num_collect_steps', type=int, default=8) + parser.add_argument('--num_update_steps', type=int, default=1) parser.add_argument('--train_batch_size', type=int, default=8) parser.add_argument('--experience_batch_size', type=int, default=8) - parser.add_argument('--lora_rank', type=int, default=4) + parser.add_argument('--lora_rank', type=int, default=0) parser.add_argument('--cuda_mem_frac', type=float, default=1.0) + parser.add_argument('--offload_inference_models', action='store_true', default=False) + parser.add_argument('--use_kernels', action='store_true', default=False) args = parser.parse_args() main(args) diff --git a/applications/Chat/benchmarks/ray/1mmt_dummy.py b/applications/Chat/benchmarks/ray/1mmt_dummy.py new file mode 100644 index 000000000000..7fc990448805 --- /dev/null +++ b/applications/Chat/benchmarks/ray/1mmt_dummy.py @@ -0,0 +1,178 @@ +import argparse +import os +import socket +from functools import partial + +import ray +import torch +from coati.quant import llama_load_quant, low_resource_init +from coati.ray.detached_trainer_ppo import DetachedPPOTrainer +from coati.ray.experience_maker_holder import ExperienceMakerHolder +from coati.ray.utils import ( + get_actor_from_args, + get_critic_from_args, + get_receivers_per_sender, + get_reward_model_from_args, + get_strategy_from_args, +) +from torch.utils.data import DataLoader +from transformers import AutoConfig, AutoTokenizer +from transformers.modeling_utils import no_init_weights + + +def get_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + return s.getsockname()[1] + + +def get_local_ip(): + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + s.connect(('8.8.8.8', 80)) + return s.getsockname()[0] + + +def main(args): + master_addr = str(get_local_ip()) + # trainer_env_info + trainer_port = str(get_free_port()) + env_info_trainers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_trainers), + 'master_port': trainer_port, + 'master_addr': master_addr + } for rank in range(args.num_trainers)] + + # maker_env_info + maker_port = str(get_free_port()) + env_info_maker = { + 'local_rank': '0', + 'rank': '0', + 'world_size': '1', + 'master_port': maker_port, + 'master_addr': master_addr + } + + # configure tokenizer + tokenizer = AutoTokenizer.from_pretrained(args.pretrain) + tokenizer.pad_token = tokenizer.eos_token + + def model_fn(): + actor_cfg = AutoConfig.from_pretrained(args.pretrain) + critic_cfg = AutoConfig.from_pretrained(args.critic_pretrain) + actor = get_actor_from_args(args.model, config=actor_cfg).requires_grad_(False).half().cuda() + critic = get_critic_from_args(args.critic_model, config=critic_cfg).requires_grad_(False).half().cuda() + reward_model = get_reward_model_from_args(args.critic_model, + config=critic_cfg).requires_grad_(False).half().cuda() + if args.initial_model_quant_ckpt is not None and args.model == 'llama': + # quantize initial model + with low_resource_init(), no_init_weights(): + initial_model = get_actor_from_args(args.model, config=actor_cfg) + initial_model.model = llama_load_quant(initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, + args.quant_group_size).cuda().requires_grad_(False) + else: + initial_model = get_actor_from_args(args.model, config=actor_cfg).requires_grad_(False).half().cuda() + return actor, critic, reward_model, initial_model + + # configure Experience Maker + experience_holder_ref = ExperienceMakerHolder.options(name="maker0", num_gpus=1, max_concurrency=2).remote( + detached_trainer_name_list=[f'trainer{i}' for i in range(args.num_trainers)], + strategy_fn=partial(get_strategy_from_args, args.maker_strategy), + model_fn=model_fn, + env_info=env_info_maker, + kl_coef=0.1, + debug=args.debug, + # sync_models_from_trainers=True, + # generation kwargs: + max_length=512, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + eval_performance=True, + use_cache=True, + ) + + def trainer_model_fn(): + actor = get_actor_from_args(args.model, config=AutoConfig.from_pretrained(args.pretrain)).half().cuda() + critic = get_critic_from_args(args.critic_model, + config=AutoConfig.from_pretrained(args.critic_pretrain)).half().cuda() + return actor, critic + + # configure Trainer + trainer_refs = [ + DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote( + experience_maker_holder_name_list=[ + f'maker{x}' for x in get_receivers_per_sender(i, args.num_trainers, 1, allow_idle_sender=True) + ], + strategy_fn=partial(get_strategy_from_args, args.trainer_strategy), + model_fn=trainer_model_fn, + env_info=env_info_trainer, + train_batch_size=args.train_batch_size, + buffer_limit=16, + eval_performance=True, + debug=args.debug, + ) for i, env_info_trainer in enumerate(env_info_trainers) + ] + + dataset_size = args.experience_batch_size * 4 + + def data_gen_fn(): + input_ids = torch.randint(tokenizer.vocab_size, (256,), device=torch.cuda.current_device()) + attn_mask = torch.ones_like(input_ids) + return {'input_ids': input_ids, 'attention_mask': attn_mask} + + def build_dataloader(size): + dataset = [data_gen_fn() for _ in range(size)] + dataloader = DataLoader(dataset, batch_size=args.experience_batch_size) + return dataloader + + # uncomment this function if sync_models_from_trainers is True + # ray.get([ + # trainer_ref.sync_models_to_remote_makers.remote() + # for trainer_ref in trainer_refs + # ]) + + wait_tasks = [] + + wait_tasks.append( + experience_holder_ref.workingloop.remote(partial(build_dataloader, dataset_size), + num_steps=args.experience_steps)) + + total_steps = args.experience_batch_size * args.experience_steps // (args.num_trainers * args.train_batch_size) + for trainer_ref in trainer_refs: + wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs)) + + ray.get(wait_tasks) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num_trainers', type=int, default=1) + parser.add_argument('--trainer_strategy', + choices=[ + 'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu', + 'colossalai_zero2_cpu' + ], + default='ddp') + parser.add_argument('--maker_strategy', choices=['naive'], default='naive') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--critic_model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--critic_pretrain', type=str, default=None) + parser.add_argument('--experience_steps', type=int, default=4) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--train_epochs', type=int, default=1) + parser.add_argument('--update_steps', type=int, default=2) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + + parser.add_argument('--initial_model_quant_ckpt', type=str, default=None) + parser.add_argument('--quant_bits', type=int, default=4) + parser.add_argument('--quant_group_size', type=int, default=128) + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)}) + main(args) diff --git a/applications/Chat/benchmarks/ray/mmmt_dummy.py b/applications/Chat/benchmarks/ray/mmmt_dummy.py new file mode 100644 index 000000000000..ca1df22070fc --- /dev/null +++ b/applications/Chat/benchmarks/ray/mmmt_dummy.py @@ -0,0 +1,189 @@ +import argparse +import os +import socket +from functools import partial + +import ray +import torch +from coati.quant import llama_load_quant, low_resource_init +from coati.ray.detached_trainer_ppo import DetachedPPOTrainer +from coati.ray.experience_maker_holder import ExperienceMakerHolder +from coati.ray.utils import ( + get_actor_from_args, + get_critic_from_args, + get_receivers_per_sender, + get_reward_model_from_args, + get_strategy_from_args, +) +from torch.utils.data import DataLoader +from transformers import AutoConfig, AutoTokenizer +from transformers.modeling_utils import no_init_weights + + +def get_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + return s.getsockname()[1] + + +def get_local_ip(): + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + s.connect(('8.8.8.8', 80)) + return s.getsockname()[0] + + +def main(args): + master_addr = str(get_local_ip()) + # trainer_env_info + trainer_port = str(get_free_port()) + env_info_trainers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_trainers), + 'master_port': trainer_port, + 'master_addr': master_addr + } for rank in range(args.num_trainers)] + + # maker_env_info + maker_port = str(get_free_port()) + env_info_makers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_makers), + 'master_port': maker_port, + 'master_addr': master_addr + } for rank in range(args.num_makers)] + + # configure tokenizer + tokenizer = AutoTokenizer.from_pretrained(args.pretrain) + tokenizer.pad_token = tokenizer.eos_token + + def model_fn(): + actor_cfg = AutoConfig.from_pretrained(args.pretrain) + critic_cfg = AutoConfig.from_pretrained(args.critic_pretrain) + actor = get_actor_from_args(args.model, config=actor_cfg).requires_grad_(False).half().cuda() + critic = get_critic_from_args(args.critic_model, config=critic_cfg).requires_grad_(False).half().cuda() + reward_model = get_reward_model_from_args(args.critic_model, + config=critic_cfg).requires_grad_(False).half().cuda() + if args.initial_model_quant_ckpt is not None and args.model == 'llama': + # quantize initial model + with low_resource_init(), no_init_weights(): + initial_model = get_actor_from_args(args.model, config=actor_cfg) + initial_model.model = llama_load_quant(initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, + args.quant_group_size).cuda().requires_grad_(False) + else: + initial_model = get_actor_from_args(args.model, config=actor_cfg).requires_grad_(False).half().cuda() + return actor, critic, reward_model, initial_model + + # configure Experience Maker + experience_holder_refs = [ + ExperienceMakerHolder.options(name=f"maker{i}", num_gpus=1, max_concurrency=2).remote( + detached_trainer_name_list=[ + f'trainer{x}' + for x in get_receivers_per_sender(i, args.num_makers, args.num_trainers, allow_idle_sender=False) + ], + strategy_fn=partial(get_strategy_from_args, args.maker_strategy), + model_fn=model_fn, + env_info=env_info_maker, + kl_coef=0.1, + debug=args.debug, + # sync_models_from_trainers=True, + # generation kwargs: + max_length=512, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + eval_performance=True, + use_cache=True, + ) + for i, env_info_maker in enumerate(env_info_makers) + ] + + def trainer_model_fn(): + actor = get_actor_from_args(args.model, config=AutoConfig.from_pretrained(args.pretrain)).half().cuda() + critic = get_critic_from_args(args.critic_model, + config=AutoConfig.from_pretrained(args.critic_pretrain)).half().cuda() + return actor, critic + + # configure Trainer + trainer_refs = [ + DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote( + experience_maker_holder_name_list=[ + f"maker{x}" + for x in get_receivers_per_sender(i, args.num_trainers, args.num_makers, allow_idle_sender=True) + ], + strategy_fn=partial(get_strategy_from_args, args.trainer_strategy), + model_fn=trainer_model_fn, + env_info=env_info_trainer, + train_batch_size=args.train_batch_size, + buffer_limit=16, + eval_performance=True, + debug=args.debug, + ) + for i, env_info_trainer in enumerate(env_info_trainers) + ] + + dataset_size = args.experience_batch_size * 4 + + def data_gen_fn(): + input_ids = torch.randint(tokenizer.vocab_size, (256,), device=torch.cuda.current_device()) + attn_mask = torch.ones_like(input_ids) + return {'input_ids': input_ids, 'attention_mask': attn_mask} + + def build_dataloader(size): + dataset = [data_gen_fn() for _ in range(size)] + dataloader = DataLoader(dataset, batch_size=args.experience_batch_size) + return dataloader + + # uncomment this function if sync_models_from_trainers is True + # ray.get([ + # trainer_ref.sync_models_to_remote_makers.remote() + # for trainer_ref in trainer_refs + # ]) + + wait_tasks = [] + + for experience_holder_ref in experience_holder_refs: + wait_tasks.append( + experience_holder_ref.workingloop.remote(partial(build_dataloader, dataset_size), + num_steps=args.experience_steps)) + + total_steps = args.experience_batch_size * args.experience_steps * \ + args.num_makers // (args.num_trainers * args.train_batch_size) + for trainer_ref in trainer_refs: + wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs)) + + ray.get(wait_tasks) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num_makers', type=int, default=1) + parser.add_argument('--num_trainers', type=int, default=1) + parser.add_argument('--trainer_strategy', + choices=[ + 'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu', + 'colossalai_zero2_cpu' + ], + default='ddp') + parser.add_argument('--maker_strategy', choices=['naive'], default='naive') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--critic_model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--critic_pretrain', type=str, default=None) + parser.add_argument('--experience_steps', type=int, default=4) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--train_epochs', type=int, default=1) + parser.add_argument('--update_steps', type=int, default=2) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + + parser.add_argument('--initial_model_quant_ckpt', type=str, default=None) + parser.add_argument('--quant_bits', type=int, default=4) + parser.add_argument('--quant_group_size', type=int, default=128) + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)}) + main(args) diff --git a/applications/ChatGPT/chatgpt/__init__.py b/applications/Chat/coati/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/__init__.py rename to applications/Chat/coati/__init__.py diff --git a/applications/Chat/coati/dataset/__init__.py b/applications/Chat/coati/dataset/__init__.py new file mode 100644 index 000000000000..f650668e90b0 --- /dev/null +++ b/applications/Chat/coati/dataset/__init__.py @@ -0,0 +1,9 @@ +from .prompt_dataset import PromptDataset +from .reward_dataset import HhRlhfDataset, RmStaticDataset +from .sft_dataset import DataCollatorForSupervisedDataset, SFTDataset, SupervisedDataset +from .utils import is_rank_0 + +__all__ = [ + 'RmStaticDataset', 'HhRlhfDataset', 'is_rank_0', 'SFTDataset', 'SupervisedDataset', + 'DataCollatorForSupervisedDataset', 'PromptDataset' +] diff --git a/applications/Chat/coati/dataset/prompt_dataset.py b/applications/Chat/coati/dataset/prompt_dataset.py new file mode 100644 index 000000000000..0bdcbbc5928e --- /dev/null +++ b/applications/Chat/coati/dataset/prompt_dataset.py @@ -0,0 +1,51 @@ +import copy +import random +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Callable, Dict, Sequence + +import torch +import torch.distributed as dist +import transformers +from torch.utils.data import Dataset +from tqdm import tqdm + +from colossalai.logging import get_dist_logger + +from .utils import is_rank_0, jload + +logger = get_dist_logger() + + +class PromptDataset(Dataset): + """Dataset for supervised fine-tuning.""" + + def __init__(self, + data_path: str, + tokenizer: transformers.PreTrainedTokenizer, + max_datasets_size: int = None, + max_length: int = 96): + super(PromptDataset, self).__init__() + self.keyed_prompt = defaultdict(list) + logger.info("Loading data...") + list_data_dict = jload(data_path) + logger.info(f"Loaded {len(list_data_dict)} examples.") + + if max_datasets_size is not None: + logger.info(f"Limiting dataset to {max_datasets_size} examples.") + list_data_dict = list_data_dict[:max_datasets_size] + + instructions = [data_dict["instruction"] for data_dict in list_data_dict] + tokens = tokenizer(instructions, + return_tensors='pt', + max_length=max_length, + padding='max_length', + truncation=True) + for k, tensor in tokens.items(): + self.keyed_prompt[k] = tensor.to(torch.cuda.current_device()).unbind() + + def __len__(self): + return len(self.keyed_prompt["input_ids"]) + + def __getitem__(self, i) -> Dict[str, torch.Tensor]: + return {k: v[i] for k, v in self.keyed_prompt.items()} diff --git a/applications/Chat/coati/dataset/reward_dataset.py b/applications/Chat/coati/dataset/reward_dataset.py new file mode 100644 index 000000000000..5dacf7e81464 --- /dev/null +++ b/applications/Chat/coati/dataset/reward_dataset.py @@ -0,0 +1,112 @@ +from typing import Callable + +from torch.utils.data import Dataset +from tqdm import tqdm + +from .utils import is_rank_0 + + +# Dahoas/rm-static +class RmStaticDataset(Dataset): + """ + Dataset for reward model + + Args: + dataset: dataset for reward model + tokenizer: tokenizer for reward model + max_length: max length of input + special_token: special token at the end of sentence + """ + + def __init__(self, dataset, tokenizer: Callable, max_length: int, special_token=None) -> None: + super().__init__() + self.chosen = [] + self.reject = [] + if special_token is None: + self.end_token = tokenizer.eos_token + else: + self.end_token = special_token + for data in tqdm(dataset, disable=not is_rank_0()): + prompt = data['prompt'] + + chosen = prompt + data['chosen'] + self.end_token + chosen_token = tokenizer(chosen, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.chosen.append({ + "input_ids": chosen_token['input_ids'], + "attention_mask": chosen_token['attention_mask'] + }) + + reject = prompt + data['rejected'] + self.end_token + reject_token = tokenizer(reject, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.reject.append({ + "input_ids": reject_token['input_ids'], + "attention_mask": reject_token['attention_mask'] + }) + + def __len__(self): + length = len(self.chosen) + return length + + def __getitem__(self, idx): + return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][ + "input_ids"], self.reject[idx]["attention_mask"] + + +# Anthropic/hh-rlhf +class HhRlhfDataset(Dataset): + """ + Dataset for reward model + + Args: + dataset: dataset for reward model + tokenizer: tokenizer for reward model + max_length: max length of input + special_token: special token at the end of sentence + """ + + def __init__(self, dataset, tokenizer: Callable, max_length: int, special_token=None) -> None: + super().__init__() + self.chosen = [] + self.reject = [] + if special_token is None: + self.end_token = tokenizer.eos_token + else: + self.end_token = special_token + for data in tqdm(dataset, disable=not is_rank_0()): + chosen = data['chosen'] + self.end_token + chosen_token = tokenizer(chosen, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.chosen.append({ + "input_ids": chosen_token['input_ids'], + "attention_mask": chosen_token['attention_mask'] + }) + + reject = data['rejected'] + self.end_token + reject_token = tokenizer(reject, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.reject.append({ + "input_ids": reject_token['input_ids'], + "attention_mask": reject_token['attention_mask'] + }) + + def __len__(self): + length = len(self.chosen) + return length + + def __getitem__(self, idx): + return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][ + "input_ids"], self.reject[idx]["attention_mask"] diff --git a/applications/Chat/coati/dataset/sft_dataset.py b/applications/Chat/coati/dataset/sft_dataset.py new file mode 100644 index 000000000000..3702d00cc609 --- /dev/null +++ b/applications/Chat/coati/dataset/sft_dataset.py @@ -0,0 +1,166 @@ +# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import random +from dataclasses import dataclass, field +from typing import Callable, Dict, Sequence + +import torch +import torch.distributed as dist +import transformers +from torch.utils.data import Dataset +from tqdm import tqdm + +from colossalai.logging import get_dist_logger + +from .utils import is_rank_0, jload + +logger = get_dist_logger() + +IGNORE_INDEX = -100 +PROMPT_DICT = { + "prompt_input": + ("Below is an instruction that describes a task, paired with an input that provides further context. " + "Write a response that appropriately completes the request.\n\n" + "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"), + "prompt_no_input": ("Below is an instruction that describes a task. " + "Write a response that appropriately completes the request.\n\n" + "### Instruction:\n{instruction}\n\n### Response:"), +} + + +class SFTDataset(Dataset): + """ + Dataset for sft model + + Args: + dataset: dataset for supervised model + tokenizer: tokenizer for supervised model + max_length: max length of input + """ + + def __init__(self, dataset, tokenizer: Callable, max_length: int = 512) -> None: + super().__init__() + self.input_ids = [] + + for data in tqdm(dataset, disable=not is_rank_0()): + prompt = data['prompt'] + data['completion'] + tokenizer.eos_token + prompt_token = tokenizer(prompt, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + + self.input_ids.append(prompt_token['input_ids'][0]) + self.labels = copy.deepcopy(self.input_ids) + + def __len__(self): + length = len(self.input_ids) + return length + + def __getitem__(self, idx): + return dict(input_ids=self.input_ids[idx], labels=self.labels[idx]) + + +def _tokenize_fn(strings: Sequence[str], + tokenizer: transformers.PreTrainedTokenizer, + max_length: int + ) -> Dict[str, torch.Tensor]: + """Tokenize a list of strings.""" + tokenized_list = tokenizer( + strings, return_tensors="pt", padding="longest", + max_length=max_length, truncation=True + ) + input_ids = labels = tokenized_list["input_ids"] + input_ids_lens = labels_lens = \ + tokenized_list["input_ids"].ne(tokenizer.pad_token_id).sum(dim=-1) + return dict( + input_ids=input_ids, + labels=labels, + input_ids_lens=input_ids_lens, + labels_lens=labels_lens, + ) + + +def preprocess( + sources: Sequence[str], + targets: Sequence[str], + tokenizer: transformers.PreTrainedTokenizer, + max_length: int, +) -> Dict: + """Preprocess the data by tokenizing.""" + examples = [s + t for s, t in zip(sources, targets)] + examples_tokenized, sources_tokenized = [ + _tokenize_fn(strings, tokenizer, max_length) + for strings in (examples, sources) + ] + input_ids = examples_tokenized["input_ids"] + labels = copy.deepcopy(input_ids) + for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]): + label[:source_len] = IGNORE_INDEX + return dict(input_ids=input_ids, labels=labels) + + +class SupervisedDataset(Dataset): + """Dataset for supervised fine-tuning.""" + + def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer, max_datasets_size: int = None, max_length: int = 512): + super(SupervisedDataset, self).__init__() + logger.info("Loading data...") + list_data_dict = jload(data_path) + logger.info(f"Loaded {len(list_data_dict)} examples.") + + if max_datasets_size is not None: + logger.info(f"Limiting dataset to {max_datasets_size} examples.") + list_data_dict = list_data_dict[:max_datasets_size] + + logger.info("Formatting inputs...") + prompt_input, prompt_no_input = PROMPT_DICT["prompt_input"], PROMPT_DICT["prompt_no_input"] + sources = [ + prompt_input.format_map(example) if example.get("input", "") != "" else prompt_no_input.format_map(example) + for example in list_data_dict + ] + targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict] + + logger.info("Tokenizing inputs... This may take some time...") + data_dict = preprocess(sources, targets, tokenizer, max_length) + + self.input_ids = data_dict["input_ids"] + self.labels = data_dict["labels"] + + def __len__(self): + return len(self.input_ids) + + def __getitem__(self, i) -> Dict[str, torch.Tensor]: + return dict(input_ids=self.input_ids[i], labels=self.labels[i]) + + +@dataclass +class DataCollatorForSupervisedDataset(object): + """Collate examples for supervised fine-tuning.""" + + tokenizer: transformers.PreTrainedTokenizer + + def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: + input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels")) + input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, + batch_first=True, + padding_value=self.tokenizer.pad_token_id) + labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX) + return dict( + input_ids=input_ids, + labels=labels, + attention_mask=input_ids.ne(self.tokenizer.pad_token_id), + ) diff --git a/applications/Chat/coati/dataset/utils.py b/applications/Chat/coati/dataset/utils.py new file mode 100644 index 000000000000..f37fce67a7c6 --- /dev/null +++ b/applications/Chat/coati/dataset/utils.py @@ -0,0 +1,22 @@ +import io +import json + +import torch.distributed as dist + + +def is_rank_0() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +def _make_r_io_base(f, mode: str): + if not isinstance(f, io.IOBase): + f = open(f, mode=mode) + return f + + +def jload(f, mode="r"): + """Load a .json file into a dictionary.""" + f = _make_r_io_base(f, mode) + jdict = json.load(f) + f.close() + return jdict diff --git a/applications/ChatGPT/chatgpt/experience_maker/__init__.py b/applications/Chat/coati/experience_maker/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/experience_maker/__init__.py rename to applications/Chat/coati/experience_maker/__init__.py diff --git a/applications/ChatGPT/chatgpt/experience_maker/base.py b/applications/Chat/coati/experience_maker/base.py similarity index 97% rename from applications/ChatGPT/chatgpt/experience_maker/base.py rename to applications/Chat/coati/experience_maker/base.py index f3640fc1e496..ff75852576c8 100644 --- a/applications/ChatGPT/chatgpt/experience_maker/base.py +++ b/applications/Chat/coati/experience_maker/base.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn -from chatgpt.models.base import Actor +from coati.models.base import Actor @dataclass @@ -18,7 +18,7 @@ class Experience: action_log_probs: (B, A) values: (B) reward: (B) - advatanges: (B) + advantages: (B) attention_mask: (B, S) action_mask: (B, A) diff --git a/applications/ChatGPT/chatgpt/experience_maker/naive.py b/applications/Chat/coati/experience_maker/naive.py similarity index 64% rename from applications/ChatGPT/chatgpt/experience_maker/naive.py rename to applications/Chat/coati/experience_maker/naive.py index 64835cfa1918..e5bb029e63d0 100644 --- a/applications/ChatGPT/chatgpt/experience_maker/naive.py +++ b/applications/Chat/coati/experience_maker/naive.py @@ -1,5 +1,6 @@ import torch -from chatgpt.models.utils import compute_reward, normalize +from coati.models.generation import generate_with_actor +from coati.models.utils import calc_action_log_probs, compute_reward, normalize from .base import Experience, ExperienceMaker @@ -16,16 +17,18 @@ def make_experience(self, input_ids: torch.Tensor, **generate_kwargs) -> Experie self.initial_model.eval() self.reward_model.eval() - sequences, attention_mask, action_mask = self.actor.generate(input_ids, + sequences, attention_mask, action_mask = generate_with_actor(self.actor, + input_ids, return_action_mask=True, **generate_kwargs) num_actions = action_mask.size(1) - action_log_probs = self.actor(sequences, num_actions, attention_mask) - base_action_log_probs = self.initial_model(sequences, num_actions, attention_mask) + actor_output = self.actor(sequences, attention_mask) + action_log_probs = calc_action_log_probs(actor_output, sequences, num_actions) + base_model_output = self.initial_model(sequences, attention_mask) + base_action_log_probs = calc_action_log_probs(base_model_output, sequences, num_actions) value = self.critic(sequences, action_mask, attention_mask) r = self.reward_model(sequences, attention_mask) - reward = compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_mask=action_mask) advantage = reward - value diff --git a/applications/Chat/coati/kernels/__init__.py b/applications/Chat/coati/kernels/__init__.py new file mode 100644 index 000000000000..230eedf7ecba --- /dev/null +++ b/applications/Chat/coati/kernels/__init__.py @@ -0,0 +1,6 @@ +from .wrapper import convert_to_xformer_model, recover_from_xformer_model + +__all__ = [ + 'convert_to_xformer_model', + 'recover_from_xformer_model', +] diff --git a/applications/Chat/coati/kernels/opt_attn.py b/applications/Chat/coati/kernels/opt_attn.py new file mode 100644 index 000000000000..e99f9c2247d1 --- /dev/null +++ b/applications/Chat/coati/kernels/opt_attn.py @@ -0,0 +1,87 @@ +from typing import Optional, Tuple + +import torch +import xformers.ops as xops +from torch import Tensor +from transformers.models.opt.modeling_opt import OPTAttention + + +# This is modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py +class XOPTAttention(OPTAttention): + # def _shape(self, tensor: Tensor, seq_len: int, bsz: int): + # return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).contiguous() + + def forward( + self, + hidden_states: Tensor, + key_value_states: Optional[Tensor] = None, + past_key_value: Optional[Tensor] = None, + attention_mask: Optional[Tensor] = None, + layer_head_mask: Optional[Tensor] = None, + output_attentions: bool = False, + ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor]]]: + if not self.training: + return super().forward(hidden_states, key_value_states, past_key_value, attention_mask, layer_head_mask, + output_attentions) + """Input shape: Batch x Time x Channel""" + assert layer_head_mask is None, 'Xformers attention does not support layer_head_mask' + assert not output_attentions, 'Xformers attention does not support output_attentions' + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + bsz, tgt_len, _ = hidden_states.size() + + # get query proj + query_states = self.q_proj(hidden_states) + # get key, value proj + if is_cross_attention and past_key_value is not None: + # reuse k,v, cross_attentions + key_states = past_key_value[0] + value_states = past_key_value[1] + elif is_cross_attention: + # cross_attentions + key_states = self._shape(self.k_proj(key_value_states), -1, bsz) + value_states = self._shape(self.v_proj(key_value_states), -1, bsz) + elif past_key_value is not None: + # reuse k, v, self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) + else: + # self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + + if self.is_decoder: + # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_states, value_states) + + query_states = self._shape(query_states, tgt_len, bsz).transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = xops.memory_efficient_attention(query_states, + key_states, + value_states, + attn_bias=xops.LowerTriangularMask(), + p=self.dropout if self.training else 0.0, + scale=self.scaling) + + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + attn_weights_reshaped = None + + return attn_output, attn_weights_reshaped, past_key_value diff --git a/applications/Chat/coati/kernels/wrapper.py b/applications/Chat/coati/kernels/wrapper.py new file mode 100644 index 000000000000..c55bda600230 --- /dev/null +++ b/applications/Chat/coati/kernels/wrapper.py @@ -0,0 +1,18 @@ +import torch.nn as nn +from transformers.models.opt.modeling_opt import OPTAttention + +from .opt_attn import XOPTAttention + + +def convert_to_xformer_model(model: nn.Module) -> nn.Module: + for module in model.modules(): + if isinstance(module, OPTAttention): + module.__class__ = XOPTAttention + return model + + +def recover_from_xformer_model(model: nn.Module) -> nn.Module: + for module in model.modules(): + if isinstance(module, XOPTAttention): + module.__class__ = OPTAttention + return model diff --git a/applications/Chat/coati/models/__init__.py b/applications/Chat/coati/models/__init__.py new file mode 100644 index 000000000000..709bc5ac0948 --- /dev/null +++ b/applications/Chat/coati/models/__init__.py @@ -0,0 +1,8 @@ +from .base import Actor, Critic, RewardModel +from .lora import LoRAModule, convert_to_lora_module +from .loss import LogExpLoss, LogSigLoss, PolicyLoss, PPOPtxActorLoss, ValueLoss + +__all__ = [ + 'Actor', 'Critic', 'RewardModel', 'PolicyLoss', 'ValueLoss', 'PPOPtxActorLoss', 'LogSigLoss', 'LogExpLoss', + 'LoRAModule', 'convert_to_lora_module' +] diff --git a/applications/Chat/coati/models/base/__init__.py b/applications/Chat/coati/models/base/__init__.py new file mode 100644 index 000000000000..c5f748a0c85a --- /dev/null +++ b/applications/Chat/coati/models/base/__init__.py @@ -0,0 +1,26 @@ +from typing import Union + +import torch.nn as nn + +from .actor import Actor +from .critic import Critic +from .reward_model import RewardModel + + +def get_base_model(model: Union[Actor, Critic, RewardModel]) -> nn.Module: + """Get the base model of our wrapper classes. + For Actor, Critic and RewardModel, return ``model.model``, + it's usually a ``transformers.PreTrainedModel``. + + Args: + model (nn.Module): model to get base model from + + Returns: + nn.Module: the base model + """ + assert isinstance(model, (Actor, Critic, RewardModel)), \ + f'Expect Actor, Critic or RewardModel, got {type(model)}, use unwrap_model first.' + return model.model + + +__all__ = ['Actor', 'Critic', 'RewardModel', 'get_base_model'] diff --git a/applications/Chat/coati/models/base/actor.py b/applications/Chat/coati/models/base/actor.py new file mode 100644 index 000000000000..2034d5cc81d4 --- /dev/null +++ b/applications/Chat/coati/models/base/actor.py @@ -0,0 +1,36 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from ..lora import LoRAModule + + +class Actor(LoRAModule): + """ + Actor model base class. + + Args: + model (nn.Module): Actor Model. + lora_rank (int): LoRA rank. + lora_train_bias (str): LoRA bias training mode. + """ + + def __init__(self, model: nn.Module, lora_rank: int = 0, lora_train_bias: str = 'none') -> None: + super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias) + self.model = model + self.convert_to_lora() + + def forward(self, + input_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + **model_kwargs, # HACK: `generate` method may pass more kwargs + ) -> torch.Tensor: + """Returns model output. + """ + output = self.model( + input_ids, + attention_mask=attention_mask, + **model_kwargs + ) + return output diff --git a/applications/ChatGPT/chatgpt/models/base/critic.py b/applications/Chat/coati/models/base/critic.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/base/critic.py rename to applications/Chat/coati/models/base/critic.py diff --git a/applications/ChatGPT/chatgpt/models/base/reward_model.py b/applications/Chat/coati/models/base/reward_model.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/base/reward_model.py rename to applications/Chat/coati/models/base/reward_model.py diff --git a/applications/ChatGPT/chatgpt/models/bloom/__init__.py b/applications/Chat/coati/models/bloom/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/bloom/__init__.py rename to applications/Chat/coati/models/bloom/__init__.py diff --git a/applications/ChatGPT/chatgpt/models/bloom/bloom_actor.py b/applications/Chat/coati/models/bloom/bloom_actor.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/bloom/bloom_actor.py rename to applications/Chat/coati/models/bloom/bloom_actor.py diff --git a/applications/ChatGPT/chatgpt/models/bloom/bloom_critic.py b/applications/Chat/coati/models/bloom/bloom_critic.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/bloom/bloom_critic.py rename to applications/Chat/coati/models/bloom/bloom_critic.py diff --git a/applications/ChatGPT/chatgpt/models/bloom/bloom_rm.py b/applications/Chat/coati/models/bloom/bloom_rm.py similarity index 92% rename from applications/ChatGPT/chatgpt/models/bloom/bloom_rm.py rename to applications/Chat/coati/models/bloom/bloom_rm.py index 4dc2646e36ae..22cfab441abb 100644 --- a/applications/ChatGPT/chatgpt/models/bloom/bloom_rm.py +++ b/applications/Chat/coati/models/bloom/bloom_rm.py @@ -33,4 +33,5 @@ def __init__(self, if checkpoint: model.gradient_checkpointing_enable() value_head = nn.Linear(model.config.hidden_size, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.hidden_size + 1)) super().__init__(model, value_head, lora_rank, lora_train_bias) diff --git a/applications/ChatGPT/chatgpt/models/generation.py b/applications/Chat/coati/models/generation.py similarity index 74% rename from applications/ChatGPT/chatgpt/models/generation.py rename to applications/Chat/coati/models/generation.py index 4ee797561f7f..0156e2284e52 100644 --- a/applications/ChatGPT/chatgpt/models/generation.py +++ b/applications/Chat/coati/models/generation.py @@ -1,7 +1,10 @@ -from typing import Any, Callable, Optional +from typing import Any, Callable, Optional, Tuple, Union import torch +import torch.distributed as dist import torch.nn as nn +import torch.nn.functional as F + try: from transformers.generation_logits_process import ( @@ -27,6 +30,14 @@ def prepare_logits_processor(top_k: Optional[int] = None, return processor_list +def _is_sequence_finished(unfinished_sequences: torch.Tensor) -> bool: + if dist.is_initialized() and dist.get_world_size() > 1: + # consider DP + unfinished_sequences = unfinished_sequences.clone() + dist.all_reduce(unfinished_sequences) + return unfinished_sequences.max() == 0 + + def sample(model: nn.Module, input_ids: torch.Tensor, max_length: int, @@ -46,9 +57,8 @@ def sample(model: nn.Module, unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) for _ in range(input_ids.size(1), max_length): - model_inputs = prepare_inputs_fn(input_ids, **model_kwargs) if prepare_inputs_fn is not None else { - 'input_ids': input_ids - } + model_inputs = prepare_inputs_fn(input_ids, **model_kwargs) \ + if prepare_inputs_fn is not None else {'input_ids': input_ids} outputs = model(**model_inputs) next_token_logits = outputs['logits'][:, -1, :] @@ -67,14 +77,14 @@ def sample(model: nn.Module, # update generated ids, model inputs for next step input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) if update_model_kwargs_fn is not None: - model_kwargs = update_model_kwargs_fn(outputs, **model_kwargs) + model_kwargs = update_model_kwargs_fn(outputs, model_kwargs) # if eos_token was found in one sentence, set sentence to finished if eos_token_id is not None: unfinished_sequences = unfinished_sequences.mul((next_tokens != eos_token_id).long()) # stop when each sentence is finished if early_stopping=True - if early_stopping and unfinished_sequences.max() == 0: + if early_stopping and _is_sequence_finished(unfinished_sequences): break return input_ids @@ -135,3 +145,35 @@ def generate(model: nn.Module, raise NotImplementedError else: raise ValueError("Unsupported generation mode") + + +@torch.no_grad() +def generate_with_actor(actor_model: nn.Module, + input_ids: torch.Tensor, + return_action_mask: bool = True, + **kwargs + ) -> Union[Tuple[torch.LongTensor, torch.LongTensor], + Tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor]]: + """Generate token sequence with actor model. Refer to `generate` for more details. + """ + # generate sequences + sequences = generate(actor_model, input_ids, **kwargs) + + # calculate auxiliary tensors + attention_mask = None + pad_token_id = kwargs.get('pad_token_id', None) + if pad_token_id is not None: + attention_mask = sequences.not_equal(pad_token_id).to(dtype=torch.long, device=sequences.device) + if not return_action_mask: + return sequences, attention_mask, None + input_len = input_ids.size(1) + eos_token_id = kwargs.get('eos_token_id', None) + if eos_token_id is None: + action_mask = torch.ones_like(sequences, dtype=torch.bool) + else: + # left padding may be applied, only mask action + action_mask = (sequences[:, input_len:] == eos_token_id).cumsum(dim=-1) == 0 + action_mask = F.pad(action_mask, (1 + input_len, -1), value=True) # include eos token and input + action_mask[:, :input_len] = False + action_mask = action_mask[:, 1:] + return sequences, attention_mask, action_mask[:, -(sequences.size(1) - input_len):] diff --git a/applications/ChatGPT/chatgpt/models/gpt/__init__.py b/applications/Chat/coati/models/gpt/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/gpt/__init__.py rename to applications/Chat/coati/models/gpt/__init__.py diff --git a/applications/ChatGPT/chatgpt/models/gpt/gpt_actor.py b/applications/Chat/coati/models/gpt/gpt_actor.py similarity index 87% rename from applications/ChatGPT/chatgpt/models/gpt/gpt_actor.py rename to applications/Chat/coati/models/gpt/gpt_actor.py index 6a53ad40b817..ae9d669f1f56 100644 --- a/applications/ChatGPT/chatgpt/models/gpt/gpt_actor.py +++ b/applications/Chat/coati/models/gpt/gpt_actor.py @@ -23,7 +23,8 @@ def __init__(self, config: Optional[GPT2Config] = None, checkpoint: bool = False, lora_rank: int = 0, - lora_train_bias: str = 'none') -> None: + lora_train_bias: str = 'none', + **kwargs) -> None: if pretrained is not None: model = GPT2LMHeadModel.from_pretrained(pretrained) elif config is not None: @@ -32,4 +33,4 @@ def __init__(self, model = GPT2LMHeadModel(GPT2Config()) if checkpoint: model.gradient_checkpointing_enable() - super().__init__(model, lora_rank, lora_train_bias) + super().__init__(model, lora_rank, lora_train_bias, **kwargs) diff --git a/applications/ChatGPT/chatgpt/models/gpt/gpt_critic.py b/applications/Chat/coati/models/gpt/gpt_critic.py similarity index 92% rename from applications/ChatGPT/chatgpt/models/gpt/gpt_critic.py rename to applications/Chat/coati/models/gpt/gpt_critic.py index 25bb1ed94de4..2e70f5f1fc96 100644 --- a/applications/ChatGPT/chatgpt/models/gpt/gpt_critic.py +++ b/applications/Chat/coati/models/gpt/gpt_critic.py @@ -24,7 +24,8 @@ def __init__(self, config: Optional[GPT2Config] = None, checkpoint: bool = False, lora_rank: int = 0, - lora_train_bias: str = 'none') -> None: + lora_train_bias: str = 'none', + **kwargs) -> None: if pretrained is not None: model = GPT2Model.from_pretrained(pretrained) elif config is not None: @@ -34,4 +35,4 @@ def __init__(self, if checkpoint: model.gradient_checkpointing_enable() value_head = nn.Linear(model.config.n_embd, 1) - super().__init__(model, value_head, lora_rank, lora_train_bias) + super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs) diff --git a/applications/ChatGPT/chatgpt/models/gpt/gpt_rm.py b/applications/Chat/coati/models/gpt/gpt_rm.py similarity index 93% rename from applications/ChatGPT/chatgpt/models/gpt/gpt_rm.py rename to applications/Chat/coati/models/gpt/gpt_rm.py index 0132dbf27ffc..054432e1ce86 100644 --- a/applications/ChatGPT/chatgpt/models/gpt/gpt_rm.py +++ b/applications/Chat/coati/models/gpt/gpt_rm.py @@ -35,4 +35,5 @@ def __init__(self, model.gradient_checkpointing_enable() value_head = nn.Linear(model.config.n_embd, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.n_embd + 1)) super().__init__(model, value_head, lora_rank, lora_train_bias) diff --git a/applications/Chat/coati/models/llama/__init__.py b/applications/Chat/coati/models/llama/__init__.py new file mode 100644 index 000000000000..9b2a024afdb2 --- /dev/null +++ b/applications/Chat/coati/models/llama/__init__.py @@ -0,0 +1,5 @@ +from .llama_actor import LlamaActor +from .llama_critic import LlamaCritic +from .llama_rm import LlamaRM + +__all__ = ['LlamaActor', 'LlamaCritic', 'LlamaRM'] diff --git a/applications/Chat/coati/models/llama/llama_actor.py b/applications/Chat/coati/models/llama/llama_actor.py new file mode 100644 index 000000000000..2c7adb390d8b --- /dev/null +++ b/applications/Chat/coati/models/llama/llama_actor.py @@ -0,0 +1,38 @@ +from typing import Optional + +import torch +from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM + +from ..base import Actor + + +class LlamaActor(Actor): + """ + Llama Actor model. + + Args: + pretrained (str): Pretrained model name or path. + config (LlamaConfig): Model config. + checkpoint (bool): Enable gradient checkpointing. + lora_rank (int): LoRA rank. + lora_train_bias (str): LoRA bias training mode. + """ + + def __init__(self, + pretrained: Optional[str] = None, + config: Optional[LlamaConfig] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none') -> None: + + if pretrained is not None: + model = LlamaForCausalLM.from_pretrained(pretrained) + elif config is not None: + model = LlamaForCausalLM(config) + else: + model = LlamaForCausalLM(LlamaConfig()) + + if checkpoint: + model.gradient_checkpointing_enable() + + super().__init__(model, lora_rank, lora_train_bias) diff --git a/applications/Chat/coati/models/llama/llama_critic.py b/applications/Chat/coati/models/llama/llama_critic.py new file mode 100644 index 000000000000..dd9e5e7bfa1a --- /dev/null +++ b/applications/Chat/coati/models/llama/llama_critic.py @@ -0,0 +1,41 @@ +from typing import Optional + +import torch.nn as nn +from transformers import LlamaConfig, LlamaModel + +from ..base import Critic + + +class LlamaCritic(Critic): + """ + Llama Critic model. + + Args: + pretrained (str): Pretrained model name or path. + config (LlamaConfig): Model config. + checkpoint (bool): Enable gradient checkpointing. + lora_rank (int): LoRA rank. + lora_train_bias (str): LoRA bias training mode. + """ + + def __init__(self, + pretrained: Optional[str] = None, + config: Optional[LlamaConfig] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none', + **kwargs) -> None: + + if pretrained is not None: + model = LlamaModel.from_pretrained(pretrained) + elif config is not None: + model = LlamaModel(config) + else: + model = LlamaModel(LlamaConfig()) + + if checkpoint: + model.gradient_checkpointing_enable() + + value_head = nn.Linear(model.config.hidden_size, 1) + + super().__init__(model, value_head, lora_rank, lora_train_bias, **kwargs) diff --git a/applications/Chat/coati/models/llama/llama_rm.py b/applications/Chat/coati/models/llama/llama_rm.py new file mode 100644 index 000000000000..f936019d62d2 --- /dev/null +++ b/applications/Chat/coati/models/llama/llama_rm.py @@ -0,0 +1,40 @@ +from typing import Optional + +import torch.nn as nn +from transformers import LlamaConfig, LlamaForCausalLM, LlamaModel + +from ..base import RewardModel + + +class LlamaRM(RewardModel): + """ + Llama Reward model. + + Args: + pretrained (str): Pretrained model name or path. + config (LlamaConfig): Model config. + checkpoint (bool): Enable gradient checkpointing. + lora_rank (int): LoRA rank. + lora_train_bias (str): LoRA bias training mode. + """ + + def __init__(self, + pretrained: Optional[str] = None, + config: Optional[LlamaConfig] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none') -> None: + + if pretrained is not None: + model = LlamaModel.from_pretrained(pretrained) + elif config is not None: + model = LlamaModel(config) + else: + model = LlamaModel(LlamaConfig()) + + if checkpoint: + model.gradient_checkpointing_enable() + value_head = nn.Linear(model.config.hidden_size, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.hidden_size + 1)) + + super().__init__(model, value_head, lora_rank, lora_train_bias) diff --git a/applications/ChatGPT/chatgpt/models/lora.py b/applications/Chat/coati/models/lora.py similarity index 80% rename from applications/ChatGPT/chatgpt/models/lora.py rename to applications/Chat/coati/models/lora.py index 9c19f472d726..2a9059e6901e 100644 --- a/applications/ChatGPT/chatgpt/models/lora.py +++ b/applications/Chat/coati/models/lora.py @@ -61,7 +61,13 @@ def T(w): if self.merge_weights and self.merged: # Make sure that the weights are not merged if self.r > 0: - self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling + if not hasattr(self, "lora_A") or not hasattr(self, "lora_B"): + # FIXME(csric): temporary fix + self.lora_A = nn.Parameter(self.weight.new_empty((self.r, self.in_features))) + self.lora_B = nn.Parameter(self.weight.new_empty((self.out_features, self.r))) + self.reset_parameters() + else: + self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling self.merged = False def eval(self): @@ -106,9 +112,26 @@ def convert_to_lora_recursively(module: nn.Module, lora_rank: int) -> None: convert_to_lora_recursively(child, lora_rank) +def convert_to_lora_module(module: nn.Module, lora_rank: int, lora_train_bias: str = 'none') -> nn.Module: + """Convert a torch.nn.Module to a LoRA module. + + Args: + module (nn.Module): The module to convert. + lora_rank (int): LoRA rank. + + Returns: + nn.Module: The converted module. + """ + if lora_rank <= 0: + return module + convert_to_lora_recursively(module, lora_rank) + lora.mark_only_lora_as_trainable(module, lora_train_bias) + return module + + class LoRAModule(nn.Module): """A LoRA module base class. All derived classes should call `convert_to_lora()` at the bottom of `__init__()`. - This calss will convert all torch.nn.Linear layer to LoraLinear layer. + This class will convert all torch.nn.Linear layer to LoraLinear layer. Args: lora_rank (int, optional): LoRA rank. 0 means LoRA is not applied. Defaults to 0. @@ -123,8 +146,4 @@ def __init__(self, lora_rank: int = 0, lora_train_bias: str = 'none') -> None: self.lora_train_bias = lora_train_bias def convert_to_lora(self) -> None: - if self.lora_rank <= 0: - return - convert_to_lora_recursively(self, self.lora_rank) - lora.mark_only_lora_as_trainable(self, self.lora_train_bias) - + convert_to_lora_module(self, self.lora_rank, self.lora_train_bias) diff --git a/applications/ChatGPT/chatgpt/models/loss.py b/applications/Chat/coati/models/loss.py similarity index 88% rename from applications/ChatGPT/chatgpt/models/loss.py rename to applications/Chat/coati/models/loss.py index 0ebcfea061b0..926c6e2a4e41 100644 --- a/applications/ChatGPT/chatgpt/models/loss.py +++ b/applications/Chat/coati/models/loss.py @@ -65,7 +65,7 @@ def forward(self, surr2 = (values - reward)**2 loss = torch.max(surr1, surr2) loss = loss.mean() - return loss + return 0.5 * loss class PPOPtxActorLoss(nn.Module): @@ -93,9 +93,10 @@ def forward(self, return policy_loss + self.pretrain_coef * lm_loss -class PairWiseLoss(nn.Module): +class LogSigLoss(nn.Module): """ Pairwise Loss for Reward Model + Details: https://arxiv.org/abs/2203.02155 """ def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor: @@ -103,3 +104,14 @@ def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> t log_probs = torch.log(probs) loss = -log_probs.mean() return loss + + +class LogExpLoss(nn.Module): + """ + Pairwise Loss for Reward Model + Details: https://arxiv.org/abs/2204.05862 + """ + + def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor: + loss = torch.log(1 + torch.exp(reject_reward - chosen_reward)).mean() + return loss diff --git a/applications/ChatGPT/chatgpt/models/opt/__init__.py b/applications/Chat/coati/models/opt/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/opt/__init__.py rename to applications/Chat/coati/models/opt/__init__.py diff --git a/applications/ChatGPT/chatgpt/models/opt/opt_actor.py b/applications/Chat/coati/models/opt/opt_actor.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/opt/opt_actor.py rename to applications/Chat/coati/models/opt/opt_actor.py diff --git a/applications/ChatGPT/chatgpt/models/opt/opt_critic.py b/applications/Chat/coati/models/opt/opt_critic.py similarity index 100% rename from applications/ChatGPT/chatgpt/models/opt/opt_critic.py rename to applications/Chat/coati/models/opt/opt_critic.py diff --git a/applications/ChatGPT/chatgpt/models/opt/opt_rm.py b/applications/Chat/coati/models/opt/opt_rm.py similarity index 92% rename from applications/ChatGPT/chatgpt/models/opt/opt_rm.py rename to applications/Chat/coati/models/opt/opt_rm.py index 7ad7b3887e53..50fc0dee8568 100644 --- a/applications/ChatGPT/chatgpt/models/opt/opt_rm.py +++ b/applications/Chat/coati/models/opt/opt_rm.py @@ -34,4 +34,5 @@ def __init__(self, model.gradient_checkpointing_enable() value_head = nn.Linear(model.config.word_embed_proj_dim, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.word_embed_proj_dim + 1)) super().__init__(model, value_head, lora_rank, lora_train_bias) diff --git a/applications/ChatGPT/chatgpt/models/utils.py b/applications/Chat/coati/models/utils.py similarity index 85% rename from applications/ChatGPT/chatgpt/models/utils.py rename to applications/Chat/coati/models/utils.py index 0ff13181fcd2..b9f15f894a1f 100644 --- a/applications/ChatGPT/chatgpt/models/utils.py +++ b/applications/Chat/coati/models/utils.py @@ -46,6 +46,25 @@ def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.T return log_probs_labels.squeeze(-1) +def calc_action_log_probs(output: torch.Tensor, + sequences: torch.LongTensor, + num_actions: int + ) -> torch.Tensor: + """Calculate action log probs. + + Args: + output (torch.Tensor): Output tensor of Actor.forward. + sequences (torch.LongTensor): Input sequences. + num_actions (int): Number of actions. + + Returns: + torch.Tensor: Action log probs. + """ + logits = output['logits'] + log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:]) + return log_probs[:, -num_actions:] + + def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor: tensor = tensor * mask tensor = tensor.sum(dim=dim) diff --git a/applications/Chat/coati/quant/__init__.py b/applications/Chat/coati/quant/__init__.py new file mode 100644 index 000000000000..a65a78d07bb8 --- /dev/null +++ b/applications/Chat/coati/quant/__init__.py @@ -0,0 +1,7 @@ +from .llama_gptq import load_quant as llama_load_quant +from .utils import low_resource_init + +__all__ = [ + 'llama_load_quant', + 'low_resource_init', +] diff --git a/applications/Chat/coati/quant/llama_gptq/__init__.py b/applications/Chat/coati/quant/llama_gptq/__init__.py new file mode 100644 index 000000000000..51c8d6316290 --- /dev/null +++ b/applications/Chat/coati/quant/llama_gptq/__init__.py @@ -0,0 +1,5 @@ +from .loader import load_quant + +__all__ = [ + 'load_quant', +] diff --git a/applications/Chat/coati/quant/llama_gptq/loader.py b/applications/Chat/coati/quant/llama_gptq/loader.py new file mode 100644 index 000000000000..5353dc8a2ea3 --- /dev/null +++ b/applications/Chat/coati/quant/llama_gptq/loader.py @@ -0,0 +1,26 @@ +import torch +import torch.nn as nn + +from .model_utils import find_layers +from .quant import make_quant + + +def load_quant(model: nn.Module, checkpoint: str, wbits: int, groupsize: int): + model = model.eval() + layers = find_layers(model) + + # ignore lm head + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + + make_quant(model, layers, wbits, groupsize) + + if checkpoint.endswith('.safetensors'): + from safetensors.torch import load_file as safe_load + model.load_state_dict(safe_load(checkpoint)) + else: + model.load_state_dict(torch.load(checkpoint)) + + return model diff --git a/applications/Chat/coati/quant/llama_gptq/model_utils.py b/applications/Chat/coati/quant/llama_gptq/model_utils.py new file mode 100644 index 000000000000..62db171abb52 --- /dev/null +++ b/applications/Chat/coati/quant/llama_gptq/model_utils.py @@ -0,0 +1,13 @@ +# copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/past/modelutils.py + +import torch +import torch.nn as nn + + +def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''): + if type(module) in layers: + return {name: module} + res = {} + for name1, child in module.named_children(): + res.update(find_layers(child, layers=layers, name=name + '.' + name1 if name != '' else name1)) + return res diff --git a/applications/Chat/coati/quant/llama_gptq/quant.py b/applications/Chat/coati/quant/llama_gptq/quant.py new file mode 100644 index 000000000000..f7d5b7ce4bd8 --- /dev/null +++ b/applications/Chat/coati/quant/llama_gptq/quant.py @@ -0,0 +1,283 @@ +# copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/past/quant.py + +import math + +import numpy as np +import torch +import torch.nn as nn + + +def quantize(x, scale, zero, maxq): + q = torch.clamp(torch.round(x / scale) + zero, 0, maxq) + return scale * (q - zero) + + +class Quantizer(nn.Module): + + def __init__(self, shape=1): + super(Quantizer, self).__init__() + self.register_buffer('maxq', torch.tensor(0)) + self.register_buffer('scale', torch.zeros(shape)) + self.register_buffer('zero', torch.zeros(shape)) + + def configure(self, bits, perchannel=False, sym=True, mse=False, norm=2.4, grid=100, maxshrink=.8): + self.maxq = torch.tensor(2**bits - 1) + self.perchannel = perchannel + self.sym = sym + self.mse = mse + self.norm = norm + self.grid = grid + self.maxshrink = maxshrink + + def find_params(self, x, weight=False): + dev = x.device + self.maxq = self.maxq.to(dev) + + shape = x.shape + if self.perchannel: + if weight: + x = x.flatten(1) + else: + if len(shape) == 4: + x = x.permute([1, 0, 2, 3]) + x = x.flatten(1) + if len(shape) == 3: + x = x.reshape((-1, shape[-1])).t() + if len(shape) == 2: + x = x.t() + else: + x = x.flatten().unsqueeze(0) + + tmp = torch.zeros(x.shape[0], device=dev) + xmin = torch.minimum(x.min(1)[0], tmp) + xmax = torch.maximum(x.max(1)[0], tmp) + + if self.sym: + xmax = torch.maximum(torch.abs(xmin), xmax) + tmp = xmin < 0 + if torch.any(tmp): + xmin[tmp] = -xmax[tmp] + tmp = (xmin == 0) & (xmax == 0) + xmin[tmp] = -1 + xmax[tmp] = +1 + + self.scale = (xmax - xmin) / self.maxq + if self.sym: + self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2) + else: + self.zero = torch.round(-xmin / self.scale) + + if self.mse: + best = torch.full([x.shape[0]], float('inf'), device=dev) + for i in range(int(self.maxshrink * self.grid)): + p = 1 - i / self.grid + xmin1 = p * xmin + xmax1 = p * xmax + scale1 = (xmax1 - xmin1) / self.maxq + zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero + q = quantize(x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq) + q -= x + q.abs_() + q.pow_(self.norm) + err = torch.sum(q, 1) + tmp = err < best + if torch.any(tmp): + best[tmp] = err[tmp] + self.scale[tmp] = scale1[tmp] + self.zero[tmp] = zero1[tmp] + if not self.perchannel: + if weight: + tmp = shape[0] + else: + tmp = shape[1] if len(shape) != 3 else shape[2] + self.scale = self.scale.repeat(tmp) + self.zero = self.zero.repeat(tmp) + + if weight: + shape = [-1] + [1] * (len(shape) - 1) + self.scale = self.scale.reshape(shape) + self.zero = self.zero.reshape(shape) + return + if len(shape) == 4: + self.scale = self.scale.reshape((1, -1, 1, 1)) + self.zero = self.zero.reshape((1, -1, 1, 1)) + if len(shape) == 3: + self.scale = self.scale.reshape((1, 1, -1)) + self.zero = self.zero.reshape((1, 1, -1)) + if len(shape) == 2: + self.scale = self.scale.unsqueeze(0) + self.zero = self.zero.unsqueeze(0) + + def quantize(self, x): + if self.ready(): + return quantize(x, self.scale, self.zero, self.maxq) + return x + + def enabled(self): + return self.maxq > 0 + + def ready(self): + return torch.all(self.scale != 0) + + +try: + import quant_cuda +except: + print('CUDA extension not installed.') + +# Assumes layer is perfectly divisible into 256 * 256 blocks + + +class QuantLinear(nn.Module): + + def __init__(self, bits, groupsize, infeatures, outfeatures): + super().__init__() + if bits not in [2, 3, 4, 8]: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + self.infeatures = infeatures + self.outfeatures = outfeatures + self.bits = bits + if groupsize != -1 and groupsize < 32 and groupsize != int(math.pow(2, int(math.log2(groupsize)))): + raise NotImplementedError("groupsize supports powers of 2 greater than 32. (e.g. : 32,64,128,etc)") + groupsize = groupsize if groupsize != -1 else infeatures + self.groupsize = groupsize + self.register_buffer( + 'qzeros', torch.zeros((math.ceil(infeatures / groupsize), outfeatures // 256 * (bits * 8)), + dtype=torch.int)) + self.register_buffer('scales', torch.zeros((math.ceil(infeatures / groupsize), outfeatures))) + self.register_buffer('bias', torch.zeros(outfeatures)) + self.register_buffer('qweight', torch.zeros((infeatures // 256 * (bits * 8), outfeatures), dtype=torch.int)) + self._initialized_quant_state = False + + def pack(self, linear, scales, zeros): + scales = scales.t().contiguous() + zeros = zeros.t().contiguous() + scale_zeros = zeros * scales + self.scales = scales.clone() + if linear.bias is not None: + self.bias = linear.bias.clone() + + intweight = [] + for idx in range(self.infeatures): + g_idx = idx // self.groupsize + intweight.append( + torch.round((linear.weight.data[:, idx] + scale_zeros[g_idx]) / self.scales[g_idx]).to(torch.int)[:, + None]) + intweight = torch.cat(intweight, dim=1) + intweight = intweight.t().contiguous() + intweight = intweight.numpy().astype(np.uint32) + qweight = np.zeros((intweight.shape[0] // 256 * (self.bits * 8), intweight.shape[1]), dtype=np.uint32) + i = 0 + row = 0 + while row < qweight.shape[0]: + if self.bits in [2, 4, 8]: + for j in range(i, i + (32 // self.bits)): + qweight[row] |= intweight[j] << (self.bits * (j - i)) + i += 32 // self.bits + row += 1 + elif self.bits == 3: + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i)) + i += 10 + qweight[row] |= intweight[i] << 30 + row += 1 + qweight[row] |= (intweight[i] >> 2) & 1 + i += 1 + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i) + 1) + i += 10 + qweight[row] |= intweight[i] << 31 + row += 1 + qweight[row] |= (intweight[i] >> 1) & 0x3 + i += 1 + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i) + 2) + i += 10 + row += 1 + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + + qweight = qweight.astype(np.int32) + self.qweight = torch.from_numpy(qweight) + + zeros -= 1 + zeros = zeros.numpy().astype(np.uint32) + qzeros = np.zeros((zeros.shape[0], zeros.shape[1] // 256 * (self.bits * 8)), dtype=np.uint32) + i = 0 + col = 0 + while col < qzeros.shape[1]: + if self.bits in [2, 4, 8]: + for j in range(i, i + (32 // self.bits)): + qzeros[:, col] |= zeros[:, j] << (self.bits * (j - i)) + i += 32 // self.bits + col += 1 + elif self.bits == 3: + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i)) + i += 10 + qzeros[:, col] |= zeros[:, i] << 30 + col += 1 + qzeros[:, col] |= (zeros[:, i] >> 2) & 1 + i += 1 + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 1) + i += 10 + qzeros[:, col] |= zeros[:, i] << 31 + col += 1 + qzeros[:, col] |= (zeros[:, i] >> 1) & 0x3 + i += 1 + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 2) + i += 10 + col += 1 + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + + qzeros = qzeros.astype(np.int32) + self.qzeros = torch.from_numpy(qzeros) + + def forward(self, x): + intermediate_dtype = torch.float32 + + if not self._initialized_quant_state: + # Do we even have a bias? Check for at least one non-zero element. + if self.bias is not None and bool(torch.any(self.bias != 0)): + # Then make sure it's the right type. + self.bias.data = self.bias.data.to(intermediate_dtype) + else: + self.bias = None + + outshape = list(x.shape) + outshape[-1] = self.outfeatures + x = x.reshape(-1, x.shape[-1]) + if self.bias is None: + y = torch.zeros(x.shape[0], outshape[-1], dtype=intermediate_dtype, device=x.device) + else: + y = self.bias.clone().repeat(x.shape[0], 1) + + output_dtype = x.dtype + x = x.to(intermediate_dtype) + if self.bits == 2: + quant_cuda.vecquant2matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 3: + quant_cuda.vecquant3matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 4: + quant_cuda.vecquant4matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 8: + quant_cuda.vecquant8matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + y = y.to(output_dtype) + return y.reshape(outshape) + + +def make_quant(module, names, bits, groupsize, name=''): + if isinstance(module, QuantLinear): + return + for attr in dir(module): + tmp = getattr(module, attr) + name1 = name + '.' + attr if name != '' else attr + if name1 in names: + setattr(module, attr, QuantLinear(bits, groupsize, tmp.in_features, tmp.out_features)) + for name1, child in module.named_children(): + make_quant(child, names, bits, groupsize, name + '.' + name1 if name != '' else name1) diff --git a/applications/Chat/coati/quant/utils.py b/applications/Chat/coati/quant/utils.py new file mode 100644 index 000000000000..01b8cff0add1 --- /dev/null +++ b/applications/Chat/coati/quant/utils.py @@ -0,0 +1,28 @@ +from contextlib import contextmanager + +import torch + + +def _noop(*args, **kwargs): + pass + + +@contextmanager +def low_resource_init(): + """This context manager disables weight initialization and sets the default float dtype to half. + """ + old_kaiming_uniform_ = torch.nn.init.kaiming_uniform_ + old_uniform_ = torch.nn.init.uniform_ + old_normal_ = torch.nn.init.normal_ + dtype = torch.get_default_dtype() + try: + torch.nn.init.kaiming_uniform_ = _noop + torch.nn.init.uniform_ = _noop + torch.nn.init.normal_ = _noop + torch.set_default_dtype(torch.half) + yield + finally: + torch.nn.init.kaiming_uniform_ = old_kaiming_uniform_ + torch.nn.init.uniform_ = old_uniform_ + torch.nn.init.normal_ = old_normal_ + torch.set_default_dtype(dtype) diff --git a/applications/Chat/coati/ray/README.md b/applications/Chat/coati/ray/README.md new file mode 100644 index 000000000000..228155a6855b --- /dev/null +++ b/applications/Chat/coati/ray/README.md @@ -0,0 +1,160 @@ +# Distributed PPO Training on Stage 3 + +## Detach Experience Makers and Trainers + +We can completely separate the trainers and makers. + +

+ +

+ +- The experience maker performs inference, produces experience, and remotely delivers it to the trainer (1). +- The trainer consumes experience to train models, and periodically transmits new model parameters to the maker (2.1, 2.2). +- Using an experience buffer to overlap transmission and computing. + +In this manner, each node will work continuously without model idle time, and different optimization strategies can be applied for inference and training to meet the needs of speed or storage. It is also helpful for scalability. + +`DetachedPPOTrainer` and `ExperienceMakerHolder` are Ray Actors (distinguished from Actor Model), representing Trainer and Experience Maker on the graph above, respectively. + +[More about Ray Core](https://docs.ray.io/en/latest/ray-core/walkthrough.html) + +## Usage + +See examples at `ColossalAI/application/Chat/examples/ray` + +### Setup Makers + +- define makers' environment variables : + + ```python + env_info_makers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(num_makers), + 'master_port': maker_port, + 'master_addr': master_addr + } for rank in range(num_makers)] + + ``` +- define maker models : + ```python + def model_fn(): + actor = get_actor_from_args(...) + critic = get_critic_from_args(...) + reward_model = get_reward_model_from_args(...) + initial_model = get_actor_from_args(...) + return actor, critic, reward_model, initial_model + + ``` +- set experience_holder_refs : + + ```python + experience_holder_refs = [ + ExperienceMakerHolder.options( + name=f"maker_{i}", + num_gpus=1, + max_concurrency=2 + ).remote( + detached_trainer_name_list=[f"trainer_{x}" for x in target_trainers(...)], + model_fn=model_fn, + ...) + for i, env_info_maker in enumerate(env_info_makers) + ] + ``` + The names in the `detached_trainer_name_list` refer to the target trainers that the maker should send experience to. + We set a trainer's name the same as a maker, by `.options(name="str")`. See below. + +### Setup Trainers + +- define trainers' environment variables : + ```python + env_info_trainers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(num_trainers), + 'master_port': trainer_port, + 'master_addr': master_addr + } for rank in range(num_trainers)] + ``` +- define trainer models : + + ```python + def trainer_model_fn(): + actor = get_actor_from_args(...) + critic = get_critic_from_args(...) + return actor, critic + ``` +- set trainer_refs : + ```python + trainer_refs = [ + DetachedPPOTrainer.options( + name=f"trainer{i}", + num_gpus=1, + max_concurrency=2 + ).remote( + experience_maker_holder_name_list=[f"maker{x}" for x in target_makers(...)], + model_fn = trainer_model_fn(), + ...) + for i, env_info_trainer in enumerate(env_info_trainers) + ] + ``` + The names in `experience_maker_holder_name_list` refer to the target makers that the trainer should send updated models to. + By setting `detached_trainer_name_list` and `experience_maker_holder_name_list`, we can customize the transmission graph. + +### Launch Jobs +- define data_loader : + ```python + def data_loader_fn(): + return = torch.utils.data.DataLoader(dataset=dataset) + + ``` +- launch makers : + ```python + wait_tasks = [] + for experience_holder_ref in experience_holder_refs: + wait_tasks.append( + experience_holder_ref.workingloop.remote(data_loader_fn(), + num_steps=experience_steps)) + + ``` + +- launch trainers : + ```python + for trainer_ref in trainer_refs: + wait_tasks.append(trainer_ref.fit.remote(total_steps, update_steps, train_epochs)) + ``` + +- wait for done : + ```python + ray.get(wait_tasks) + ``` + +## Flexible Structure + +We can deploy different strategies to makers and trainers. Here are some notions. + +### 2 Makers 1 Trainer +

+ +

+ +### 2 Makers 2 Trainer +

+ +

+ +### Maker Inference Quantization +

+ +

+ +### Tensor Parallel + +

+ +

+ +## TODO + +- [ ] Support LoRA +- [ ] Support TP & PP diff --git a/applications/ChatGPT/tests/__init__.py b/applications/Chat/coati/ray/__init__.py similarity index 100% rename from applications/ChatGPT/tests/__init__.py rename to applications/Chat/coati/ray/__init__.py diff --git a/applications/Chat/coati/ray/callbacks/__init__.py b/applications/Chat/coati/ray/callbacks/__init__.py new file mode 100644 index 000000000000..5f5e488f383e --- /dev/null +++ b/applications/Chat/coati/ray/callbacks/__init__.py @@ -0,0 +1,9 @@ +from .base import MakerCallback, TrainerCallback +from .performance_evaluator import ExperienceMakerPerformanceEvaluator, TrainerPerformanceEvaluator + +__all__ = [ + "TrainerCallback", + "MakerCallback", + "ExperienceMakerPerformanceEvaluator", + "TrainerPerformanceEvaluator", +] diff --git a/applications/Chat/coati/ray/callbacks/base.py b/applications/Chat/coati/ray/callbacks/base.py new file mode 100644 index 000000000000..3306150a41ff --- /dev/null +++ b/applications/Chat/coati/ray/callbacks/base.py @@ -0,0 +1,66 @@ +from abc import ABC + +from coati.experience_maker import Experience + + +class TrainerCallback(ABC): + """ + Base callback class. It defines the interface for callbacks. + """ + + def on_fit_start(self) -> None: + pass + + def on_fit_end(self) -> None: + pass + + def on_episode_start(self, episode: int) -> None: + pass + + def on_episode_end(self, episode: int) -> None: + pass + + def on_epoch_start(self, epoch: int) -> None: + pass + + def on_epoch_end(self, epoch: int) -> None: + pass + + def on_batch_start(self) -> None: + pass + + def on_batch_end(self, metrics: dict, experience: Experience) -> None: + pass + + def on_update_start(self) -> None: + pass + + def on_update_end(self) -> None: + pass + + +class MakerCallback(ABC): + + def on_loop_start(self) -> None: + pass + + def on_loop_end(self) -> None: + pass + + def on_make_experience_start(self) -> None: + pass + + def on_make_experience_end(self, experience: Experience) -> None: + pass + + def on_send_start(self) -> None: + pass + + def on_send_end(self) -> None: + pass + + def on_batch_start(self) -> None: + pass + + def on_batch_end(self) -> None: + pass diff --git a/applications/Chat/coati/ray/callbacks/performance_evaluator.py b/applications/Chat/coati/ray/callbacks/performance_evaluator.py new file mode 100644 index 000000000000..cd3517609e7a --- /dev/null +++ b/applications/Chat/coati/ray/callbacks/performance_evaluator.py @@ -0,0 +1,212 @@ +from time import time +from typing import Optional + +import torch +import torch.distributed as dist +from coati.experience_maker import Experience + +from .base import MakerCallback, TrainerCallback + + +def get_world_size() -> int: + if dist.is_initialized(): + return dist.get_world_size() + return 1 + + +def print_rank_0(*args, **kwargs) -> None: + if not dist.is_initialized() or dist.get_rank() == 0: + print(*args, **kwargs) + + +@torch.no_grad() +def all_reduce_mean(x: float, world_size: int) -> float: + if world_size == 1: + return x + tensor = torch.tensor([x], device=torch.cuda.current_device()) + dist.all_reduce(tensor) + tensor = tensor / world_size + return tensor.item() + + +class Timer: + + def __init__(self) -> None: + self.start_time: Optional[float] = None + self.duration: float = 0. + + def start(self) -> None: + self.start_time = time() + + def end(self) -> None: + self.duration += time() - self.start_time + + def reset(self) -> None: + self.duration = 0. + + +class ExperienceMakerPerformanceEvaluator(MakerCallback): + + def __init__(self, actor_num_params: int, critic_num_params: int, initial_model_num_params: int, + reward_model_num_params: int) -> None: + super().__init__() + self.world_size = get_world_size() + self.actor_num_params = actor_num_params + self.critic_num_params = critic_num_params + self.initial_model_num_params = initial_model_num_params + self.reward_model_num_params = reward_model_num_params + + self.batch_timer = Timer() + self.send_timer = Timer() + self.make_experience_timer = Timer() + self.total_samples: int = 0 + self.make_experience_flop: int = 0 + + print_rank_0( + f'ExperienceMaker actor: {actor_num_params/1024**3:.2f}B, critic: {critic_num_params/1024**3:.2f}B, initial model: {initial_model_num_params/1024**3:.2f}B, reward model: {reward_model_num_params/1024**3:.2f}B, world size: {self.world_size}' + ) + + def on_make_experience_start(self) -> None: + self.make_experience_timer.start() + + def on_make_experience_end(self, experience: Experience) -> None: + self.make_experience_timer.end() + + batch_size, seq_len = experience.sequences.shape + + self.total_samples += batch_size + + # actor generate + num_actions = experience.action_mask.size(1) + input_len = seq_len - num_actions + total_seq_len = (input_len + seq_len - 1) * num_actions / 2 + self.make_experience_flop += self.actor_num_params * batch_size * total_seq_len * 2 + # actor forward + self.make_experience_flop += self.actor_num_params * batch_size * seq_len * 2 + # critic forward + self.make_experience_flop += self.critic_num_params * batch_size * seq_len * 2 + # initial model forward + self.make_experience_flop += self.initial_model_num_params * batch_size * seq_len * 2 + # reward model forward + self.make_experience_flop += self.reward_model_num_params * batch_size * seq_len * 2 + + def on_send_start(self) -> None: + self.send_timer.start() + + def on_send_end(self) -> None: + self.send_timer.end() + + def on_batch_start(self) -> None: + self.batch_timer.start() + + def on_batch_end(self) -> None: + self.batch_timer.end() + + def on_loop_end(self) -> None: + avg_make_experience_duration = all_reduce_mean(self.make_experience_timer.duration, self.world_size) + avg_overall_duration = all_reduce_mean(self.batch_timer.duration, self.world_size) + avg_send_duration = all_reduce_mean(self.send_timer.duration, self.world_size) + + avg_throughput = self.total_samples * self.world_size / (avg_overall_duration + 1e-12) + avg_make_experience_tflops = self.make_experience_flop / 1e12 / (avg_make_experience_duration + 1e-12) + avg_time_per_sample = (avg_overall_duration + 1e-12) / (self.total_samples * self.world_size) + avg_make_experience_time_per_sample = (avg_make_experience_duration + 1e-12) / \ + (self.total_samples * self.world_size) + avg_send_time_per_sample = (avg_send_duration + 1e-12) / (self.total_samples * self.world_size) + + print_rank_0( + 'Making Experience Performance Summary:\n' + f'Throughput: {avg_throughput:.3f} samples/sec\n' + + f'TFLOPS per GPU: {avg_make_experience_tflops:.3f}\n' + + f'Sample time (overall): {avg_time_per_sample:.3f} s\n' + + f'Sample time (make experience): {avg_make_experience_time_per_sample:.3f} s, {avg_make_experience_time_per_sample/avg_time_per_sample*100:.2f}%\n' + + + f'Sample time (send): {avg_send_time_per_sample:.3f} s, {avg_send_time_per_sample/avg_time_per_sample*100:.2f}%\n' + ) + + +class TrainerPerformanceEvaluator(TrainerCallback): + + def __init__(self, + actor_num_params: int, + critic_num_params: int, + enable_grad_checkpoint: bool = False, + ignore_first_episodes: int = 1) -> None: + super().__init__() + self.world_size = get_world_size() + self.actor_num_params = actor_num_params + self.critic_num_params = critic_num_params + self.enable_grad_checkpoint = enable_grad_checkpoint + self.ignore_first_episodes = ignore_first_episodes + self.ignore_this_episode = False + + self.episode_timer = Timer() + self.batch_timer = Timer() + self.update_timer = Timer() + self.total_samples: int = 0 + self.learn_flop: int = 0 + + print_rank_0( + f'Trainer actor: {self.actor_num_params/1024**3:.2f}B, critic: {self.critic_num_params/1024**3:.2f}B, world size: {self.world_size}' + ) + + def on_episode_start(self, episodes: int) -> None: + self.ignore_this_episode = episodes < self.ignore_first_episodes + if self.ignore_this_episode: + return + self.episode_timer.start() + + def on_episode_end(self, episodes: int) -> None: + if self.ignore_this_episode: + return + self.episode_timer.end() + + def on_batch_start(self) -> None: + if self.ignore_this_episode: + return + self.batch_timer.start() + + def on_batch_end(self, metrics: dict, experience: Experience) -> None: + if self.ignore_this_episode: + return + self.batch_timer.end() + + batch_size, seq_len = experience.sequences.shape + + self.total_samples += batch_size + + # actor forward-backward, 3 means forward(1) + backward(2) + self.learn_flop += self.actor_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) + # critic forward-backward + self.learn_flop += self.critic_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) + + def on_update_start(self) -> None: + if self.ignore_this_episode: + return + self.update_timer.start() + + def on_update_end(self) -> None: + if self.ignore_this_episode: + return + self.update_timer.end() + + def on_fit_end(self) -> None: + if self.total_samples == 0: + print_rank_0('No samples are collected, skip trainer performance evaluation') + return + avg_train_duration = all_reduce_mean(self.batch_timer.duration, self.world_size) + avg_update_duration = all_reduce_mean(self.update_timer.duration, self.world_size) + avg_episode_duration = all_reduce_mean(self.episode_timer.duration, self.world_size) + + avg_throughput = self.total_samples * self.world_size / (avg_episode_duration + 1e-12) + avg_learn_tflops = self.learn_flop / 1e12 / (avg_train_duration + 1e-12) + avg_time_per_sample = (avg_episode_duration + 1e-12) / (self.total_samples * self.world_size) + avg_train_time_per_sample = (avg_train_duration + 1e-12) / (self.total_samples * self.world_size) + avg_update_time_per_sample = (avg_update_duration + 1e-12) / (self.total_samples * self.world_size) + + print_rank_0( + 'Learning Performance Summary:\n' + f'Throughput: {avg_throughput:.3f} samples/sec\n' + + f'TFLOPS per GPU: {avg_learn_tflops:.3f}\n' + f'Sample time (overall): {avg_time_per_sample:.3f} s\n' + + f'Sample time (train): {avg_train_time_per_sample:.3f} s, {avg_train_time_per_sample/avg_time_per_sample*100:.2f}%\n' + + + f'Sample time (update): {avg_update_time_per_sample:.3f} s, {avg_update_time_per_sample/avg_time_per_sample*100:.2f}%\n' + ) diff --git a/applications/Chat/coati/ray/detached_replay_buffer.py b/applications/Chat/coati/ray/detached_replay_buffer.py new file mode 100644 index 000000000000..2f765281178a --- /dev/null +++ b/applications/Chat/coati/ray/detached_replay_buffer.py @@ -0,0 +1,75 @@ +import asyncio +import copy +import random +from threading import Lock +from typing import Any, List + +import ray +import torch +from coati.experience_maker.base import Experience +from coati.replay_buffer import ReplayBuffer +from coati.replay_buffer.utils import BufferItem, make_experience_batch, split_experience_batch +# from torch.multiprocessing import Queue +from ray.util.queue import Queue + + +class DetachedReplayBuffer: + ''' + Detached replay buffer. Share Experience across workers on the same node. + Therefore a trainer node is expected to have only one instance. + It is ExperienceMakerHolder's duty to call append(exp) method, remotely. + + Args: + sample_batch_size: Batch size when sampling. Exp won't enqueue until they formed a batch. + tp_world_size: Number of workers in the same tp group + limit: Limit of number of experience sample BATCHs. A number <= 0 means unlimited. Defaults to 0. + cpu_offload: Whether to offload experience to cpu when sampling. Defaults to True. + ''' + + def __init__(self, sample_batch_size: int, limit: int = 0) -> None: + self.sample_batch_size = sample_batch_size + self.limit = limit + self.items = Queue(self.limit, actor_options={"num_cpus": 1}) + self.batch_collector: List[BufferItem] = [] + + @torch.no_grad() + def append(self, experience: Experience) -> None: + ''' + Expected to be called remotely. + ''' + items = split_experience_batch(experience) + self.extend(items) + + @torch.no_grad() + def extend(self, items: List[BufferItem]) -> None: + ''' + Expected to be called remotely. + ''' + self.batch_collector.extend(items) + while len(self.batch_collector) >= self.sample_batch_size: + items = self.batch_collector[:self.sample_batch_size] + experience = make_experience_batch(items) + self.items.put(experience, block=True) + self.batch_collector = self.batch_collector[self.sample_batch_size:] + + def clear(self) -> None: + # self.items.close() + self.items.shutdown() + self.items = Queue(self.limit) + self.worker_state = [False] * self.tp_world_size + self.batch_collector = [] + + @torch.no_grad() + def sample(self, worker_rank=0, to_device="cpu") -> Experience: + ret = self._sample_and_erase() + ret.to_device(to_device) + return ret + + @torch.no_grad() + def _sample_and_erase(self) -> Experience: + ret = self.items.get(block=True) + return ret + + def get_length(self) -> int: + ret = self.items.qsize() + return ret diff --git a/applications/Chat/coati/ray/detached_trainer_base.py b/applications/Chat/coati/ray/detached_trainer_base.py new file mode 100644 index 000000000000..ac2d35e9da19 --- /dev/null +++ b/applications/Chat/coati/ray/detached_trainer_base.py @@ -0,0 +1,179 @@ +import os +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, Iterable, List, Optional, Union + +import ray +import torch +from coati.experience_maker import Experience +from coati.replay_buffer.utils import BufferItem +from torch.utils.data import DataLoader +from tqdm import tqdm + +from .callbacks import TrainerCallback +from .detached_replay_buffer import DetachedReplayBuffer +from .utils import is_rank_0 + + +class DetachedTrainer(ABC): + ''' + Base class for detached rlhf trainers. + 'detach' means that the experience maker is detached compared to a normal Trainer. + Please set name attribute during init: + >>> trainer = DetachedTrainer.options(..., name = "xxx", ...).remote() + So an ExperienceMakerHolder can reach the detached_replay_buffer by Actor's name. + Args: + detached_strategy (DetachedStrategy): the strategy to use for training + detached_replay_buffer_ref (ObjectRef[DetachedReplayBuffer]): the replay buffer to use for training + data_loader_pin_memory (bool, defaults to True): whether to pin memory for data loader + callbacks (List[Callback], defaults to []): the callbacks to call during training process + generate_kwargs (dict, optional): the kwargs to use while model generating + + ''' + + def __init__(self, + experience_maker_holder_name_list: List[str], + train_batch_size: int = 8, + buffer_limit: int = 0, + dataloader_pin_memory: bool = True, + callbacks: List[TrainerCallback] = [], + debug: bool = False) -> None: + super().__init__() + self.detached_replay_buffer = DetachedReplayBuffer(train_batch_size, limit=buffer_limit) + self.dataloader_pin_memory = dataloader_pin_memory + self.callbacks = callbacks + self.target_holder_name_list = experience_maker_holder_name_list + self.target_holder_list = [] + self._is_target_holder_initialized = False + self._debug = debug + + def update_target_holder_list(self): + # as the length of target_holder_list may be zero, we need to check it by a bool flag + if not self._is_target_holder_initialized: + for name in self.target_holder_name_list: + self.target_holder_list.append(ray.get_actor(name, namespace=os.environ["RAY_NAMESPACE"])) + self._is_target_holder_initialized = True + + @abstractmethod + def _update_remote_makers(self, fully_update: bool = False, **kwargs): + pass + + def sync_models_to_remote_makers(self, **kwargs): + self._update_remote_makers(fully_update=True, **kwargs) + + @abstractmethod + def training_step(self, experience: Experience) -> Dict[str, Any]: + pass + + def _learn(self, update_steps: int, train_epochs: int) -> None: + data = [] + # warmup + pbar = tqdm(range(update_steps), desc=f'Train epoch [1/{train_epochs}]', disable=not is_rank_0()) + self._on_epoch_start(0) + self._learn_epoch(pbar, data) + self._on_epoch_end(0) + # item is already a batch + dataloader = DataLoader(data, + batch_size=1, + shuffle=True, + pin_memory=self.dataloader_pin_memory, + collate_fn=lambda x: x[0]) + for epoch in range(1, train_epochs): + pbar = tqdm(dataloader, desc=f'Train epoch [{epoch + 1}/{train_epochs}]', disable=not is_rank_0()) + self._on_epoch_start(epoch) + self._learn_epoch(pbar, data) + self._on_epoch_end(epoch) + + def _learn_epoch(self, pbar: tqdm, data: List[Experience]) -> None: + is_warmup = len(data) == 0 + for x in pbar: + if self._debug: + print("[trainer] training step") + # sample a batch and then train to avoid waiting + experience = x if not is_warmup else self._buffer_sample() + experience.to_device(torch.cuda.current_device()) + self._on_batch_start() + metrics = self.training_step(experience) + self._on_batch_end(metrics, experience) + + if self._debug: + print("[trainer] step over") + experience.to_device("cpu") + if is_warmup: + data.append(experience) + pbar.set_postfix(metrics) + + def fit(self, total_steps: int, update_steps: int, train_epochs: int = 1) -> None: + self._on_fit_start() + for i in tqdm(range(total_steps // update_steps), desc='Trainer', disable=not is_rank_0()): + self._on_episode_start(i) + self._learn(update_steps, train_epochs) + self._on_update_start() + self._update_remote_makers() + self._on_update_end() + self._on_episode_end(i) + self._on_fit_end() + + @ray.method(concurrency_group="buffer_length") + def buffer_get_length(self): + # called by ExperienceMakerHolder + if self._debug: + print("[trainer] telling length") + return self.detached_replay_buffer.get_length() + + @ray.method(concurrency_group="buffer_append") + def buffer_append(self, experience: Experience): + # called by ExperienceMakerHolder + if self._debug: + print(f"[trainer] receiving exp.") + self.detached_replay_buffer.append(experience) + + @ray.method(concurrency_group="buffer_append") + def buffer_extend(self, items: List[BufferItem]): + # called by ExperienceMakerHolder + if self._debug: + print(f"[trainer] receiving exp.") + self.detached_replay_buffer.extend(items) + + @ray.method(concurrency_group="buffer_sample") + def _buffer_sample(self): + return self.detached_replay_buffer.sample() + + def _on_fit_start(self) -> None: + for callback in self.callbacks: + callback.on_fit_start() + + def _on_fit_end(self) -> None: + for callback in self.callbacks: + callback.on_fit_end() + + def _on_episode_start(self, episode: int) -> None: + for callback in self.callbacks: + callback.on_episode_start(episode) + + def _on_episode_end(self, episode: int) -> None: + for callback in self.callbacks: + callback.on_episode_end(episode) + + def _on_epoch_start(self, epoch: int) -> None: + for callback in self.callbacks: + callback.on_epoch_start(epoch) + + def _on_epoch_end(self, epoch: int) -> None: + for callback in self.callbacks: + callback.on_epoch_end(epoch) + + def _on_batch_start(self) -> None: + for callback in self.callbacks: + callback.on_batch_start() + + def _on_batch_end(self, metrics: dict, experience: Experience) -> None: + for callback in self.callbacks: + callback.on_batch_end(metrics, experience) + + def _on_update_start(self) -> None: + for callback in self.callbacks: + callback.on_update_start() + + def _on_update_end(self) -> None: + for callback in self.callbacks: + callback.on_update_end() diff --git a/applications/Chat/coati/ray/detached_trainer_ppo.py b/applications/Chat/coati/ray/detached_trainer_ppo.py new file mode 100644 index 000000000000..2f2aa0e29579 --- /dev/null +++ b/applications/Chat/coati/ray/detached_trainer_ppo.py @@ -0,0 +1,200 @@ +from typing import Any, Callable, Dict, List, Optional, Tuple + +import ray +import torch +from coati.experience_maker import Experience, NaiveExperienceMaker +from coati.models.base import Actor, Critic +from coati.models.loss import PolicyLoss, ValueLoss +from coati.trainer.callbacks import Callback +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy, Strategy +from torch.optim import Adam + +from colossalai.nn.optimizer import HybridAdam + +from .callbacks import TrainerCallback, TrainerPerformanceEvaluator +from .detached_trainer_base import DetachedTrainer +from .lora_constructor import LoRAConstructor +from .utils import ( + get_actor_from_args, + get_critic_from_args, + get_model_numel, + get_rank, + get_strategy_from_args, + is_rank_0, + set_dist_env, + state_dict_to, +) + + +@ray.remote(concurrency_groups={ + "buffer_length": 1, + "buffer_append": 1, + "buffer_sample": 1, + "model_io": 1, + "compute": 1 +}) +class DetachedPPOTrainer(DetachedTrainer): + ''' + Detached Trainer for PPO algorithm + Args: + strategy (Strategy): the strategy to use for training + model (str) : for actor / critic init + pretrained (str) : for actor / critic init + lora_rank (int) : for actor / critic init + train_batch_size (int, defaults to 8): the batch size to use for training + train_batch_size (int, defaults to 8): the batch size to use for training + buffer_limit (int, defaults to 0): the max_size limitation of replay buffer + buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu + eps_clip (float, defaults to 0.2): the clip coefficient of policy loss + value_clip (float, defaults to 0.4): the clip coefficient of value loss + experience_batch_size (int, defaults to 8): the batch size to use for experience generation + max_epochs (int, defaults to 1): the number of epochs of training process + dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader + callbacks (List[Callback], defaults to []): the callbacks to call during training process + generate_kwargs (dict, optional): the kwargs to use while model generating + ''' + + def __init__( + self, + experience_maker_holder_name_list: List[str], + strategy_fn: Callable[[], Strategy], + model_fn: Callable[[], Tuple[Actor, Critic]], + env_info: Dict[str, str] = None, + train_batch_size: int = 8, + buffer_limit: int = 0, + eps_clip: float = 0.2, + value_clip: float = 0.4, + dataloader_pin_memory: bool = True, + callbacks: List[TrainerCallback] = [], + eval_performance: bool = False, + debug: bool = False, + update_lora_weights: bool = False, + ) -> None: + # set environment variables + if env_info: + set_dist_env(env_info=env_info) + # configure strategy + self.strategy = strategy_fn() + # configure models, loss and optimizers + with self.strategy.model_init_context(): + self.actor, self.critic = model_fn() + + if eval_performance: + actor_numel = get_model_numel(self.actor) + critic_numel = get_model_numel(self.critic) + evaluator = TrainerPerformanceEvaluator(actor_numel, critic_numel) + callbacks = callbacks + [evaluator] + + if isinstance(self.strategy, (LowLevelZeroStrategy, GeminiStrategy)): + self.actor_optim = HybridAdam(self.actor.parameters(), lr=1e-7) + self.critic_optim = HybridAdam(self.critic.parameters(), lr=1e-7) + else: + self.actor_optim = Adam(self.actor.parameters(), lr=1e-7) + self.critic_optim = Adam(self.critic.parameters(), lr=1e-7) + + (self.actor, self.actor_optim), (self.critic, self.critic_optim) = \ + self.strategy.prepare((self.actor, self.actor_optim), (self.critic, self.critic_optim)) + + # configure trainer + self.actor_loss_fn = PolicyLoss(eps_clip) + self.critic_loss_fn = ValueLoss(value_clip) + + super().__init__(experience_maker_holder_name_list, + train_batch_size=train_batch_size, + buffer_limit=buffer_limit, + dataloader_pin_memory=dataloader_pin_memory, + callbacks=callbacks, + debug=debug) + if self._debug: + print(f'[trainer{get_rank()}] will send state dict to {experience_maker_holder_name_list}') + + self._update_lora_weights = update_lora_weights + + @ray.method(concurrency_group="model_io") + @torch.no_grad() + def _update_remote_makers(self, fully_update: bool = False, **config): + # TODO: balance duties + if not fully_update: + config['requires_grad_only'] = True + self.update_target_holder_list() + # mark start, ensure order + tasks = [] + for target_holder in self.target_holder_list: + tasks.append(target_holder.update_experience_maker.remote(chunk_start=True, fully_update=fully_update)) + ray.get(tasks) + # sending loop + tasks = [] + + for state_dict_shard in self._get_model_state_dict_shard(self.actor, fully_update=fully_update, **config): + for target_holder in self.target_holder_list: + tasks.append( + target_holder.update_experience_maker.remote( + new_actor_state_dict=state_dict_shard, + new_actor_lora_config_dict=self._get_model_lora_config_dict(self.actor), + fully_update=fully_update)) + # sending loop + for state_dict_shard in self._get_model_state_dict_shard(self.critic, fully_update=fully_update, **config): + for target_holder in self.target_holder_list: + tasks.append( + target_holder.update_experience_maker.remote( + new_critic_state_dict=state_dict_shard, + new_critic_lora_config_dict=self._get_model_lora_config_dict(self.critic), + fully_update=fully_update)) + ray.get(tasks) + # mark end + for target_holder in self.target_holder_list: + target_holder.update_experience_maker.remote(chunk_end=True, fully_update=fully_update) + + @ray.method(concurrency_group="compute") + def training_step(self, experience: Experience) -> Dict[str, float]: + self.actor.train() + self.critic.train() + + num_actions = experience.action_mask.size(1) + action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask) + actor_loss = self.actor_loss_fn(action_log_probs, + experience.action_log_probs, + experience.advantages, + action_mask=experience.action_mask) + self.strategy.backward(actor_loss, self.actor, self.actor_optim) + self.strategy.optimizer_step(self.actor_optim) + self.actor_optim.zero_grad() + + values = self.critic(experience.sequences, + action_mask=experience.action_mask, + attention_mask=experience.attention_mask) + critic_loss = self.critic_loss_fn(values, + experience.values, + experience.reward, + action_mask=experience.action_mask) + + self.strategy.backward(critic_loss, self.critic, self.critic_optim) + self.strategy.optimizer_step(self.critic_optim) + self.critic_optim.zero_grad() + return {'actor_loss': actor_loss.item(), 'critic_loss': critic_loss.item()} + + def strategy_save_actor(self, path: str, only_rank0: bool = False) -> None: + self.strategy.save_model(self.actor, path, only_rank0) + + def strategy_save_critic(self, path: str, only_rank0: bool = False) -> None: + self.strategy.save_model(self.critic, path, only_rank0) + + def strategy_save_actor_optim(self, path: str, only_rank0: bool = False) -> None: + self.strategy.save_optimizer(self.actor_optim, path, only_rank0) + + def strategy_save_critic_optim(self, path: str, only_rank0: bool = False) -> None: + self.strategy.save_optimizer(self.critic_optim, path, only_rank0) + + def _get_model_state_dict_shard(self, model: torch.nn.Module, fully_update=False, **config): + for state_dict in self.strategy.get_model_state_dict_shard(model, **config): + if not self._update_lora_weights or fully_update: + yield state_dict_to(state_dict) + else: + state_dict_lora, _ = LoRAConstructor.filter_state_dict_lora(state_dict) + yield state_dict_to(state_dict_lora) + + def _get_model_lora_config_dict(self, model: torch.nn.Module): + if not self._update_lora_weights: + return None + unwrapped_model = self.strategy.unwrap_model(model) + return LoRAConstructor.extract_lora_config(unwrapped_model) diff --git a/applications/Chat/coati/ray/experience_maker_holder.py b/applications/Chat/coati/ray/experience_maker_holder.py new file mode 100644 index 000000000000..07d9c3e4f396 --- /dev/null +++ b/applications/Chat/coati/ray/experience_maker_holder.py @@ -0,0 +1,271 @@ +import os +import time +import tracemalloc +from copy import deepcopy +from threading import Lock +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union + +import ray +import torch +import torch.nn as nn +from coati.experience_maker import Experience, ExperienceMaker, NaiveExperienceMaker +from coati.models.base import Actor, Critic, RewardModel +from coati.replay_buffer.utils import BufferItem, make_experience_batch, split_experience_batch +from coati.trainer.callbacks import Callback +from coati.trainer.strategies import Strategy +from coati.trainer.strategies.sampler import DistributedSampler +from ray.exceptions import GetTimeoutError +from torch import Tensor +from tqdm import tqdm + +from .callbacks import ExperienceMakerPerformanceEvaluator, MakerCallback +from .utils import (get_model_numel, + get_rank, + get_world_size, + is_rank_0, + set_dist_env, + state_dict_to) +from .lora_constructor import LoRAConstructor + +@ray.remote(concurrency_groups={"experience_io": 1, "model_io": 1, "compute": 1}) +class ExperienceMakerHolder: + ''' + Args: + detached_trainer_name_list: str list to get ray actor handles + strategy: + kl_coef: the coefficient of kl divergence loss + sync_models_from_trainers: whether to sync models from trainers. If True, you must call sync_models_to_remote_makers() in trainers to sync models. + ''' + + def __init__( + self, + detached_trainer_name_list: List[str], + strategy_fn: Callable[[], Strategy], + # a function returns (actor, critic, reward_model, initial_model) + model_fn: Callable[[], Tuple[Actor, Critic, RewardModel, Actor]], + env_info: Dict[str, str] = None, + sync_models_from_trainers: bool = False, + buffer_cpu_offload: bool = True, + kl_coef: float = 0.1, + callbacks: List[MakerCallback] = [], + eval_performance: bool = False, + debug: bool = False, + update_lora_weights: bool = False, + **generate_kwargs): + # set environment variables + if env_info: + set_dist_env(env_info=env_info) + self.target_trainer_list = [] + assert len(detached_trainer_name_list) > 0 + self._detached_trainer_name_list = detached_trainer_name_list + self.strategy = strategy_fn() + self.buffer_cpu_offload = buffer_cpu_offload + self.kl_coef = kl_coef + # init models + with self.strategy.model_init_context(): + actor, critic, reward_model, initial_model = model_fn() + self.generate_kwargs = _set_default_generate_kwargs(generate_kwargs, actor) + if eval_performance: + actor_numel = get_model_numel(actor) + critic_numel = get_model_numel(critic) + initial_model_numel = get_model_numel(initial_model) + reward_model_numel = get_model_numel(reward_model) + evaluator = ExperienceMakerPerformanceEvaluator(actor_numel, critic_numel, initial_model_numel, + reward_model_numel) + callbacks = callbacks + [evaluator] + + actor, critic, reward_model, initial_model = self.strategy.prepare(actor, critic, reward_model, initial_model) + self.experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, self.kl_coef) + self.callbacks = callbacks + + self._model_visit_lock = Lock() + + self._is_fully_initialized = not sync_models_from_trainers + + self._debug = debug + self._update_lora_weights = update_lora_weights + if self._update_lora_weights: + self.actor_lora_constructor = LoRAConstructor() + self.critic_lora_constructor = LoRAConstructor() + + self.target_auto_balance = False + + self._target_idx = 0 + + if self._debug: + print(f'[maker{get_rank()}] will send items to {self._detached_trainer_name_list}') + if not self._is_fully_initialized: + print(f'[maker{get_rank()}] Waiting for INIT') + + def _get_ready(self): + while not self._fully_initialized(): + time.sleep(1.0) + + def _fully_initialized(self): + return self._is_fully_initialized + + def _init_target_trainer_list(self): + if len(self.target_trainer_list) > 0: + return + for name in self._detached_trainer_name_list: + self.target_trainer_list.append(ray.get_actor(name, namespace=os.environ["RAY_NAMESPACE"])) + + # copy from ../trainer/base.py + @ray.method(concurrency_group="compute") + def _make_experience(self, inputs: Union[Tensor, Dict[str, Tensor]]) -> Experience: + if isinstance(inputs, Tensor): + return self.experience_maker.make_experience(inputs, **self.generate_kwargs) + elif isinstance(inputs, dict): + return self.experience_maker.make_experience(**inputs, **self.generate_kwargs) + else: + raise ValueError(f'Unsupported input type "{type(inputs)}"') + + @ray.method(concurrency_group="experience_io") + def _send_items(self, experience: Experience) -> None: + self._init_target_trainer_list() + items = split_experience_batch(experience) + items_per_trainer = [[] for _ in range(len(self.target_trainer_list))] + for item in items: + items_per_trainer[self._target_idx].append(item) + self._target_idx = (self._target_idx + 1) % len(self.target_trainer_list) + for i, target_trainer in enumerate(self.target_trainer_list): + if len(items_per_trainer[i]) > 0: + target_trainer.buffer_extend.remote(items_per_trainer[i]) + + def _inference_step(self, batch) -> None: + self._on_batch_start() + with self._model_visit_lock: + self._on_make_experience_start() + experience = self._make_experience(batch) + self._on_make_experience_end(experience) + self._on_send_start() + if self.buffer_cpu_offload: + experience.to_device('cpu') + self._send_items(experience) + self._on_send_end() + self._on_batch_end() + + def workingloop(self, dataloader_fn: Callable[[], Iterable], num_epochs: int = 1, num_steps: int = 0): + """Working loop of the experience maker. + + Args: + dataloader_fn (Callable[[], Iterable]): A function that returns a dataloader. + num_epochs (int, optional): Iterate the dataloader for number of epochs. Defaults to 1. + num_steps (int, optional): Iterate the dataloader for number if steps. If this value > 0, num_epochs will be ignored. Defaults to 0. + """ + self._get_ready() + self._on_loop_start() + dataloader = dataloader_fn() + if num_steps > 0: + # ignore num epochs + it = iter(dataloader) + for _ in tqdm(range(num_steps), desc='ExperienceMaker', disable=not is_rank_0()): + try: + batch = next(it) + except StopIteration: + it = iter(dataloader) + batch = next(it) + self._inference_step(batch) + else: + with tqdm(total=num_epochs * len(dataloader), desc='ExperienceMaker', disable=not is_rank_0()) as pbar: + for _ in range(num_epochs): + for batch in dataloader: + self._inference_step(batch) + pbar.update() + self._on_loop_end() + + @ray.method(concurrency_group="model_io") + def update_experience_maker(self, + new_actor_state_dict: Dict[str, Any] = None, + new_actor_lora_config_dict: Dict[str, Any] = None, + new_critic_state_dict: Dict[str, Any] = None, + new_critic_lora_config_dict: Dict[str, Any] = None, + fully_update: bool = False, + chunk_start: bool = None, + chunk_end: bool = None): + ''' + called by trainer + chunk_start: Set True at the first call. Before sending state_dict calls + chunk_end: Set True at the last call. After sending state_dict calls. + fully_update: Set True if you want to sync models when initializing + + TODO: load_state_dict integrate with model-sharding strategy + ''' + _watch_memory = self._debug + if chunk_start: + if self._debug: + print("[maker] UPDATE ") + if _watch_memory: + tracemalloc.start() + self._model_visit_lock.acquire() + + with torch.no_grad(): + if new_actor_state_dict is not None: + if not self._update_lora_weights or fully_update: + self.experience_maker.actor.model.load_state_dict(new_actor_state_dict, strict=False) + else: + new_actor_state_dict = state_dict_to(new_actor_state_dict, device=torch.cuda.current_device()) + state_dict_increase = self.actor_lora_constructor.reconstruct_increase(new_actor_state_dict, new_actor_lora_config_dict) + self.actor_lora_constructor.load_state_dict_increase(self.experience_maker.actor.model, state_dict_increase) + if new_critic_state_dict is not None: + if not self._update_lora_weights or fully_update: + self.experience_maker.critic.load_state_dict(new_critic_state_dict, strict=False) + else: + new_critic_state_dict = state_dict_to(new_critic_state_dict, device=torch.cuda.current_device()) + state_dict_increase = self.critic_lora_constructor.reconstruct_increase(new_critic_state_dict, new_critic_lora_config_dict) + self.critic_lora_constructor.load_state_dict_increase(self.experience_maker.critic, state_dict_increase) + + # the lock must be released after both actor and critic being updated + if chunk_end: + self._model_visit_lock.release() + if _watch_memory: + current, peak = tracemalloc.get_traced_memory() + print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB") + tracemalloc.stop() + if fully_update: + self._is_fully_initialized = True + + def _on_make_experience_start(self) -> None: + for callback in self.callbacks: + callback.on_make_experience_start() + + def _on_make_experience_end(self, experience: Experience) -> None: + for callback in self.callbacks: + callback.on_make_experience_end(experience) + + def _on_loop_start(self) -> None: + for callback in self.callbacks: + callback.on_loop_start() + + def _on_loop_end(self) -> None: + for callback in self.callbacks: + callback.on_loop_end() + + def _on_send_start(self) -> None: + for callback in self.callbacks: + callback.on_send_start() + + def _on_send_end(self) -> None: + for callback in self.callbacks: + callback.on_send_end() + + def _on_batch_start(self) -> None: + for callback in self.callbacks: + callback.on_batch_start() + + def _on_batch_end(self) -> None: + for callback in self.callbacks: + callback.on_batch_end() + + +def _set_default_generate_kwargs(generate_kwargs: dict, actor: Actor) -> None: + origin_model = actor.model + new_kwargs = {**generate_kwargs} + # use huggingface models method directly + if 'prepare_inputs_fn' not in generate_kwargs and hasattr(origin_model, 'prepare_inputs_for_generation'): + new_kwargs['prepare_inputs_fn'] = origin_model.prepare_inputs_for_generation + + if 'update_model_kwargs_fn' not in generate_kwargs and hasattr(origin_model, '_update_model_kwargs_for_generation'): + new_kwargs['update_model_kwargs_fn'] = origin_model._update_model_kwargs_for_generation + + return new_kwargs diff --git a/applications/Chat/coati/ray/lora_constructor.py b/applications/Chat/coati/ray/lora_constructor.py new file mode 100644 index 000000000000..4809617f647b --- /dev/null +++ b/applications/Chat/coati/ray/lora_constructor.py @@ -0,0 +1,122 @@ +from typing import Any, Callable, Dict, List, Optional +from collections import OrderedDict +from dataclasses import dataclass + +import torch +import torch.nn as nn +from loralib.layers import LoRALayer +from coati.models.lora import LoraLinear + + +@dataclass +class LoRAConfig: + r: int = 0 + lora_alpha: int = 1 + lora_dropout: float = 0 + fan_in_fan_out: bool = False + + +class LoRAConstructor: + ''' + Tools for reconstructing a model from a remote LoRA model. + (Transferring only LoRA data costs much less!) + Usage: + Step 1 (Sender): + filter_state_dict_lora() + + Step 2 (Sender, Optional): + extract_lora_config() + + Step 3 (Sender): + send state_dict_lora and lora_config_dict + + Step 4 (Receiver): + reconstruct_increase() + + Step 5 (Receiver): + load_state_dict_increase() + + ''' + + def __init__(self): + self.lora_config_dict = None + + def register_lora_config(self, lora_config_dict: Dict[str, Any]): + self.lora_config_dict = lora_config_dict + + def reconstruct_increase(self, state_dict_lora: Dict[str, Any], lora_config_dict: Dict[str, Any]): + ''' + xxx.lora_A, xxx.lora_B -->> xxx.weight + Warning: the xxx.weight here is the increment actually. + ''' + if lora_config_dict is not None: + self.register_lora_config(lora_config_dict) + + state_dict_increase = OrderedDict() + config_iter = iter(self.lora_config_dict.items()) + lora_A, lora_B, layer_prefix = None, None, None + for k, v in state_dict_lora.items(): + if k.rpartition('.')[-1] == 'lora_A': + lora_A = v + layer_prefix = k.rpartition('.')[0] + elif k.rpartition('.')[-1] == 'lora_B': + assert layer_prefix == k.rpartition('.')[0], "unmatched (lora_A, lora_B) pair" + layer_prefix_2, config = next(config_iter) + assert layer_prefix_2 == layer_prefix, "unmatched (state_dict, config_dict) pair" + lora_B = v + weight_data_increase = self._compute(lora_A, lora_B, config) + state_dict_increase[layer_prefix + '.weight'] = weight_data_increase + lora_A, lora_B, layer_prefix = None, None, None + else: + raise ValueError('unexpected key') + return state_dict_increase + + def _compute(self, lora_A, lora_B, config=LoRAConfig()): + def T(w): + return w.T if config.fan_in_fan_out else w + if config.r > 0: + scaling = config.lora_alpha / config.r + weight_data_increase = T(lora_B @ lora_A) * scaling + return weight_data_increase + return 0 + + def load_state_dict_increase(self, model: nn.Module, state_dict_increase: Dict[str, Any]): + ''' + The final reconstruction step + ''' + # naive approach + model.load_state_dict({k: v + model.state_dict()[k] for k, v in state_dict_increase.items()}, strict=False) + + @staticmethod + def filter_state_dict_lora(state_dict: Dict[str, Any], keep_non_lora=False): + ''' + if keep_non_lora, also return non_lora state_dict + ''' + state_dict_lora = OrderedDict() + state_dict_non_lora = OrderedDict() + for k, v in state_dict.items(): + if 'lora_A' in k or 'lora_B' in k: + state_dict_lora[k] = v + elif keep_non_lora: + state_dict_non_lora[k] = v + if keep_non_lora: + return state_dict_lora, state_dict_non_lora + else: + return state_dict_lora, None + + @staticmethod + def extract_lora_config(model: nn.Module) -> Dict[str, LoRAConfig]: + ''' + extract LoraLinear model. + return OrderedDict(): name -> LoRAConfig + ''' + lora_config_dict = OrderedDict() + + for name, child in model.named_modules(): + if isinstance(child, LoraLinear): + lora_config_dict[name] = LoRAConfig(r=child.r, + lora_alpha=child.lora_alpha, + lora_dropout=child.lora_dropout, + fan_in_fan_out=child.fan_in_fan_out) + + return lora_config_dict diff --git a/applications/Chat/coati/ray/utils.py b/applications/Chat/coati/ray/utils.py new file mode 100644 index 000000000000..761186b95ee5 --- /dev/null +++ b/applications/Chat/coati/ray/utils.py @@ -0,0 +1,140 @@ +import os +from collections import OrderedDict +from typing import Any, Callable, Dict, List, Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +from coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic +from coati.models.gpt import GPTRM, GPTActor, GPTCritic +from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM +from coati.models.opt import OPTRM, OPTActor, OPTCritic +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer + + +def is_rank_0() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +def get_rank() -> int: + return dist.get_rank() if dist.is_initialized() else 0 + + +def get_world_size() -> int: + return dist.get_world_size() if dist.is_initialized() else 1 + + +def get_actor_from_args(model: str, pretrained: str = None, config=None, lora_rank=0): + if model == 'gpt2': + actor = GPTActor(pretrained=pretrained, config=config, lora_rank=lora_rank) + elif model == 'bloom': + actor = BLOOMActor(pretrained=pretrained, config=config, lora_rank=lora_rank) + elif model == 'opt': + actor = OPTActor(pretrained=pretrained, config=config, lora_rank=lora_rank) + elif model == 'llama': + actor = LlamaActor(pretrained=pretrained, config=config, lora_rank=lora_rank) + else: + raise ValueError(f'Unsupported actor model "{model}"') + return actor + + +def get_critic_from_args(model: str, pretrained: str = None, config=None, lora_rank=0): + if model == 'gpt2': + critic = GPTCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True) + elif model == 'bloom': + critic = BLOOMCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True) + elif model == 'opt': + critic = OPTCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True) + elif model == 'llama': + critic = LlamaCritic(pretrained=pretrained, lora_rank=lora_rank, config=config, use_action_mask=True) + else: + raise ValueError(f'Unsupported reward model "{model}"') + return critic + + +def get_reward_model_from_args(model: str, pretrained: str = None, config=None): + if model == 'gpt2': + reward_model = GPTRM(pretrained=pretrained, config=config) + elif model == 'bloom': + reward_model = BLOOMRM(pretrained=pretrained, config=config) + elif model == 'opt': + reward_model = OPTRM(pretrained=pretrained, config=config) + elif model == 'llama': + reward_model = LlamaRM(pretrained=pretrained, config=config) + else: + raise ValueError(f'Unsupported reward model "{model}"') + return reward_model + + +def get_strategy_from_args(strategy: str): + if strategy == 'ddp': + strategy_ = DDPStrategy() + elif strategy == 'colossalai_gemini': + strategy_ = GeminiStrategy(placement_policy='cuda', initial_scale=2**5) + elif strategy == 'colossalai_zero2': + strategy_ = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + elif strategy == 'colossalai_gemini_cpu': + strategy_ = GeminiStrategy(placement_policy='cpu', initial_scale=2**5) + elif strategy == 'colossalai_zero2_cpu': + strategy_ = LowLevelZeroStrategy(stage=2, placement_policy='cpu') + else: + raise ValueError(f'Unsupported strategy "{strategy}"') + return strategy_ + + +def get_tokenizer_from_args(model: str, **kwargs): + if model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained('gpt2') + elif model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m') + elif model == 'opt': + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") + elif model == 'llama': + pretrain_path = kwargs["pretrain"] + tokenizer = AutoTokenizer.from_pretrained(pretrain_path) + else: + raise ValueError(f'Unsupported model "{model}"') + + tokenizer.pad_token = tokenizer.eos_token + return tokenizer + + +def set_dist_env(env_info: Dict[str, str]): + os.environ["RANK"] = env_info['rank'] + os.environ["LOCAL_RANK"] = env_info['local_rank'] + os.environ["WORLD_SIZE"] = env_info['world_size'] + os.environ['MASTER_PORT'] = env_info['master_port'] + os.environ['MASTER_ADDR'] = env_info['master_addr'] + + +def get_model_numel(model: nn.Module) -> int: + numel = sum(p.numel() for p in model.parameters()) + return numel + + +def get_receivers_per_sender(sender_idx: int, num_senders: int, num_receivers: int, allow_idle_sender: bool) -> list: + target_receivers = [] + if num_senders <= num_receivers or allow_idle_sender: + # a sender will send data to one or more than one receivers + # a receiver only has one sender + for i in range(num_receivers): + if i % num_senders == sender_idx: + target_receivers.append(i) + else: + # a sender will send data to one receiver + # a receiver may have more than one sender + target_receivers.append(sender_idx % num_receivers) + return target_receivers + + +def state_dict_to(state_dict: Dict[str, Any], + dtype: torch.dtype = torch.float16, + device: torch.device = torch.device('cpu')): + ''' + keep state_dict intact + ''' + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + new_state_dict[k] = v.to(dtype=dtype, device=device) + return new_state_dict diff --git a/applications/ChatGPT/chatgpt/replay_buffer/__init__.py b/applications/Chat/coati/replay_buffer/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/replay_buffer/__init__.py rename to applications/Chat/coati/replay_buffer/__init__.py diff --git a/applications/ChatGPT/chatgpt/replay_buffer/base.py b/applications/Chat/coati/replay_buffer/base.py similarity index 94% rename from applications/ChatGPT/chatgpt/replay_buffer/base.py rename to applications/Chat/coati/replay_buffer/base.py index 5036b09045c4..4c3812461a10 100644 --- a/applications/ChatGPT/chatgpt/replay_buffer/base.py +++ b/applications/Chat/coati/replay_buffer/base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Any -from chatgpt.experience_maker.base import Experience +from coati.experience_maker.base import Experience class ReplayBuffer(ABC): diff --git a/applications/ChatGPT/chatgpt/replay_buffer/naive.py b/applications/Chat/coati/replay_buffer/naive.py similarity index 97% rename from applications/ChatGPT/chatgpt/replay_buffer/naive.py rename to applications/Chat/coati/replay_buffer/naive.py index 3fc53da65bff..938f500643c9 100644 --- a/applications/ChatGPT/chatgpt/replay_buffer/naive.py +++ b/applications/Chat/coati/replay_buffer/naive.py @@ -2,7 +2,7 @@ from typing import List import torch -from chatgpt.experience_maker.base import Experience +from coati.experience_maker.base import Experience from .base import ReplayBuffer from .utils import BufferItem, make_experience_batch, split_experience_batch diff --git a/applications/ChatGPT/chatgpt/replay_buffer/utils.py b/applications/Chat/coati/replay_buffer/utils.py similarity index 96% rename from applications/ChatGPT/chatgpt/replay_buffer/utils.py rename to applications/Chat/coati/replay_buffer/utils.py index 752f16704771..6ad0db2c3b60 100644 --- a/applications/ChatGPT/chatgpt/replay_buffer/utils.py +++ b/applications/Chat/coati/replay_buffer/utils.py @@ -3,7 +3,7 @@ import torch import torch.nn.functional as F -from chatgpt.experience_maker.base import Experience +from coati.experience_maker.base import Experience @dataclass @@ -15,7 +15,7 @@ class BufferItem: action_log_probs: (A) values: (1) reward: (1) - advatanges: (1) + advantages: (1) attention_mask: (S) action_mask: (A) diff --git a/applications/Chat/coati/trainer/__init__.py b/applications/Chat/coati/trainer/__init__.py new file mode 100644 index 000000000000..86142361f3ff --- /dev/null +++ b/applications/Chat/coati/trainer/__init__.py @@ -0,0 +1,10 @@ +from .base import OnPolicyTrainer, SLTrainer +from .ppo import PPOTrainer +from .rm import RewardModelTrainer +from .sft import SFTTrainer + +__all__ = [ + 'SLTrainer', 'OnPolicyTrainer', + 'RewardModelTrainer', 'SFTTrainer', + 'PPOTrainer' +] diff --git a/applications/Chat/coati/trainer/base.py b/applications/Chat/coati/trainer/base.py new file mode 100644 index 000000000000..8a826056edf2 --- /dev/null +++ b/applications/Chat/coati/trainer/base.py @@ -0,0 +1,203 @@ +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import List + +import torch.nn as nn +import tqdm +from coati.experience_maker import Experience +from coati.replay_buffer import NaiveReplayBuffer +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torch.distributed as dist +from .callbacks import Callback +from .strategies import Strategy +from .utils import CycledDataLoader, is_rank_0 + + +class SLTrainer(ABC): + """ + Base class for supervised learning trainers. + + Args: + strategy (Strategy):the strategy to use for training + max_epochs (int, defaults to 1): the number of epochs of training process + model (nn.Module): the model to train + optim (Optimizer): the optimizer to use for training + """ + + def __init__(self, + strategy: Strategy, + max_epochs: int, + model: nn.Module, + optimizer: Optimizer, + tensorboard_dir: str = None, + ) -> None: + super().__init__() + self.strategy = strategy + self.max_epochs = max_epochs + self.model = model + self.optimizer = optimizer + self.writer = SummaryWriter(tensorboard_dir) if tensorboard_dir and dist.get_rank() == 0 else None + + @abstractmethod + def _train(self, epoch): + raise NotImplementedError() + + @abstractmethod + def _eval(self, epoch): + raise NotImplementedError() + + def _before_fit(self): + self.no_epoch_bar = False + + def fit(self, *args, **kwargs): + self._before_fit(*args, **kwargs) + for epoch in tqdm.trange(self.max_epochs, + desc="Epochs", + disable=not is_rank_0() or self.no_epoch_bar + ): + self._train(epoch) + self._eval(epoch) + if dist.get_rank() == 0 and self.writer: + print("Closing tensorboard writer...") + self.writer.close() + + +class OnPolicyTrainer(ABC): + """ + Base class for on-policy rl trainers, e.g. PPO. + + Args: + strategy (Strategy):the strategy to use for training + buffer (NaiveReplayBuffer): the buffer to collect experiences + sample_buffer (bool, defaults to False): whether to sample from buffer + dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader + callbacks (List[Callback], defaults to []): the callbacks to call during training process + """ + + def __init__(self, + strategy: Strategy, + buffer: NaiveReplayBuffer, + sample_buffer: bool, + dataloader_pin_memory: bool, + callbacks: List[Callback] = [] + ) -> None: + super().__init__() + self.strategy = strategy + self.buffer = buffer + self.sample_buffer = sample_buffer + self.dataloader_pin_memory = dataloader_pin_memory + self.callbacks = callbacks + + @contextmanager + def _fit_ctx(self) -> None: + for callback in self.callbacks: + callback.on_fit_start() + try: + yield + finally: + for callback in self.callbacks: + callback.on_fit_end() + + @contextmanager + def _episode_ctx(self, episode: int) -> None: + for callback in self.callbacks: + callback.on_episode_start(episode) + try: + yield + finally: + for callback in self.callbacks: + callback.on_episode_end(episode) + + def _on_make_experience_start(self) -> None: + for callback in self.callbacks: + callback.on_make_experience_start() + + def _on_make_experience_end(self, experience: Experience) -> None: + for callback in self.callbacks: + callback.on_make_experience_end(experience) + + def _on_learn_epoch_start(self, epoch: int) -> None: + for callback in self.callbacks: + callback.on_learn_epoch_start(epoch) + + def _on_learn_epoch_end(self, epoch: int) -> None: + for callback in self.callbacks: + callback.on_learn_epoch_end(epoch) + + def _on_learn_batch_start(self) -> None: + for callback in self.callbacks: + callback.on_learn_batch_start() + + def _on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: + for callback in self.callbacks: + callback.on_learn_batch_end(metrics, experience) + + @abstractmethod + def _make_experience(self, collect_step: int): + """ + Implement this method to make experience. + """ + raise NotImplementedError() + + @abstractmethod + def _learn(self, update_step: int): + """ + Implement this method to learn from experience, either + sample from buffer or transform buffer into dataloader. + """ + raise NotImplementedError() + + def _collect_phase(self, collect_step: int): + self._on_make_experience_start() + experience = self._make_experience(collect_step) + self._on_make_experience_end(experience) + self.buffer.append(experience) + + def _update_phase(self, update_step: int): + self._on_learn_epoch_start(update_step) + self._learn(update_step) + self._on_learn_epoch_end(update_step) + + def fit(self, + prompt_dataloader: DataLoader, + pretrain_dataloader: DataLoader, + num_episodes: int, + num_collect_steps: int, + num_update_steps: int, + ): + """ + The main training loop of on-policy rl trainers. + + Args: + prompt_dataloader (DataLoader): the dataloader to use for prompt data + pretrain_dataloader (DataLoader): the dataloader to use for pretrain data + num_episodes (int): the number of episodes to train + num_collect_steps (int): the number of collect steps per episode + num_update_steps (int): the number of update steps per episode + """ + self.prompt_dataloader = CycledDataLoader(prompt_dataloader) + self.pretrain_dataloader = CycledDataLoader(pretrain_dataloader) + + with self._fit_ctx(): + for episode in tqdm.trange(num_episodes, + desc="Episodes", + disable=not is_rank_0()): + with self._episode_ctx(episode): + for collect_step in tqdm.trange(num_collect_steps, + desc="Collect steps", + disable=not is_rank_0()): + self._collect_phase(collect_step) + if not self.sample_buffer: + # HACK(cwher): according to the design of boost API, dataloader should also be boosted, + # but it is impractical to adapt this pattern in RL training. Thus, I left dataloader unboosted. + # I only call strategy.setup_dataloader() to setup dataloader. + self.dataloader = self.strategy.setup_dataloader(self.buffer, + self.dataloader_pin_memory) + for update_step in tqdm.trange(num_update_steps, + desc="Update steps", + disable=not is_rank_0()): + self._update_phase(update_step) + # NOTE: this is for on-policy algorithms + self.buffer.clear() diff --git a/applications/ChatGPT/chatgpt/trainer/callbacks/__init__.py b/applications/Chat/coati/trainer/callbacks/__init__.py similarity index 100% rename from applications/ChatGPT/chatgpt/trainer/callbacks/__init__.py rename to applications/Chat/coati/trainer/callbacks/__init__.py diff --git a/applications/ChatGPT/chatgpt/trainer/callbacks/base.py b/applications/Chat/coati/trainer/callbacks/base.py similarity index 94% rename from applications/ChatGPT/chatgpt/trainer/callbacks/base.py rename to applications/Chat/coati/trainer/callbacks/base.py index 0b01345f7872..f5616048855b 100644 --- a/applications/ChatGPT/chatgpt/trainer/callbacks/base.py +++ b/applications/Chat/coati/trainer/callbacks/base.py @@ -1,6 +1,6 @@ from abc import ABC -from chatgpt.experience_maker import Experience +from coati.experience_maker import Experience class Callback(ABC): diff --git a/applications/ChatGPT/chatgpt/trainer/callbacks/performance_evaluator.py b/applications/Chat/coati/trainer/callbacks/performance_evaluator.py similarity index 60% rename from applications/ChatGPT/chatgpt/trainer/callbacks/performance_evaluator.py rename to applications/Chat/coati/trainer/callbacks/performance_evaluator.py index faa38af1b84e..925455444597 100644 --- a/applications/ChatGPT/chatgpt/trainer/callbacks/performance_evaluator.py +++ b/applications/Chat/coati/trainer/callbacks/performance_evaluator.py @@ -3,7 +3,7 @@ import torch import torch.distributed as dist -from chatgpt.experience_maker import Experience +from coati.experience_maker import Experience from .base import Callback @@ -19,6 +19,14 @@ def print_rank_0(*args, **kwargs) -> None: print(*args, **kwargs) +def divide(x: float, y: float) -> float: + if y == 0: + return float('inf') + elif y == float('inf'): + return float('nan') + return x / y + + @torch.no_grad() def all_reduce_mean(x: float, world_size: int) -> float: if world_size == 1: @@ -29,6 +37,24 @@ def all_reduce_mean(x: float, world_size: int) -> float: return tensor.item() +class Timer: + + def __init__(self) -> None: + self.start_time: Optional[float] = None + self.duration: float = 0. + + def start(self) -> None: + self.start_time = time() + + def end(self) -> None: + assert self.start_time is not None + self.duration += time() - self.start_time + self.start_time = None + + def reset(self) -> None: + self.duration = 0. + + class PerformanceEvaluator(Callback): """ Callback for valuate the performance of the model. @@ -58,27 +84,34 @@ def __init__(self, self.ignore_episodes = ignore_episodes self.disable: bool = False - self.make_experience_duration: float = 0. - self.make_experience_start_time: Optional[float] = None + self.overall_timer = Timer() + self.make_experience_timer = Timer() + self.learn_timer = Timer() self.make_experience_num_samples: int = 0 self.make_experience_flop: int = 0 - self.learn_duration: float = 0. - self.learn_start_time: Optional[float] = None self.learn_num_samples: int = 0 self.learn_flop: int = 0 def on_episode_start(self, episode: int) -> None: self.disable = self.ignore_episodes > 0 and episode < self.ignore_episodes + if self.disable: + return + self.overall_timer.start() + + def on_episode_end(self, episode: int) -> None: + if self.disable: + return + self.overall_timer.end() def on_make_experience_start(self) -> None: if self.disable: return - self.make_experience_start_time = time() + self.make_experience_timer.start() def on_make_experience_end(self, experience: Experience) -> None: if self.disable: return - self.make_experience_duration += time() - self.make_experience_start_time + self.make_experience_timer.end() batch_size, seq_len = experience.sequences.shape @@ -101,12 +134,12 @@ def on_make_experience_end(self, experience: Experience) -> None: def on_learn_batch_start(self) -> None: if self.disable: return - self.learn_start_time = time() + self.learn_timer.start() def on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: if self.disable: return - self.learn_duration += time() - self.learn_start_time + self.learn_timer.end() batch_size, seq_len = experience.sequences.shape @@ -114,20 +147,37 @@ def on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: # actor forward-backward, 3 means forward(1) + backward(2) self.learn_flop += self.actor_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) - # critic foward-backward + # critic forward-backward self.learn_flop += self.critic_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) def on_fit_end(self) -> None: - avg_make_experience_duration = all_reduce_mean(self.make_experience_duration, self.world_size) - avg_learn_duration = all_reduce_mean(self.learn_duration, self.world_size) + avg_make_experience_duration = all_reduce_mean(self.make_experience_timer.duration, self.world_size) + avg_learn_duration = all_reduce_mean(self.learn_timer.duration, self.world_size) + avg_overall_duration = all_reduce_mean(self.overall_timer.duration, self.world_size) - avg_make_experience_throughput = self.make_experience_num_samples / (avg_make_experience_duration + 1e-12) + avg_make_experience_throughput = self.make_experience_num_samples * \ + self.world_size / (avg_make_experience_duration + 1e-12) avg_make_experience_tflops = self.make_experience_flop / 1e12 / (avg_make_experience_duration + 1e-12) - avg_learn_throughput = self.learn_num_samples / (avg_learn_duration + 1e-12) + avg_learn_throughput = self.learn_num_samples * self.world_size / (avg_learn_duration + 1e-12) avg_learn_tflops = self.learn_flop / 1e12 / (avg_learn_duration + 1e-12) + num_effective_samples = min(self.learn_num_samples, self.make_experience_num_samples) * self.world_size + + avg_overall_throughput = num_effective_samples / (avg_overall_duration + 1e-12) + + overall_time_per_sample = divide(1, avg_overall_throughput) + make_experience_time_per_sample = divide(avg_make_experience_duration, num_effective_samples) + learn_time_per_sample = divide(avg_learn_duration, num_effective_samples) + print_rank_0( - f'Making experience throughput: {avg_make_experience_throughput:.3f} samples/sec, TFLOPS: {avg_make_experience_tflops:.3f}' + f'Performance summary:\n' + + f'Generate {self.make_experience_num_samples * self.world_size} samples, throughput: {avg_make_experience_throughput:.2f} samples/s, TFLOPS per GPU: {avg_make_experience_tflops:.2f}\n' + + + f'Train {self.learn_num_samples * self.world_size} samples, throughput: {avg_learn_throughput:.2f} samples/s, TFLOPS per GPU: {avg_learn_tflops:.2f}\n' + + f'Overall throughput: {avg_overall_throughput:.2f} samples/s\n' + + f'Overall time per sample: {overall_time_per_sample:.2f} s\n' + + f'Make experience time per sample: {make_experience_time_per_sample:.2f} s, {make_experience_time_per_sample/overall_time_per_sample*100:.2f}%\n' + + + f'Learn time per sample: {learn_time_per_sample:.2f} s, {learn_time_per_sample/overall_time_per_sample*100:.2f}%' ) - print_rank_0(f'Learning throughput: {avg_learn_throughput:.3f} samples/sec, TFLOPS: {avg_learn_tflops:.3f}') diff --git a/applications/ChatGPT/chatgpt/trainer/callbacks/save_checkpoint.py b/applications/Chat/coati/trainer/callbacks/save_checkpoint.py similarity index 90% rename from applications/ChatGPT/chatgpt/trainer/callbacks/save_checkpoint.py rename to applications/Chat/coati/trainer/callbacks/save_checkpoint.py index 8f2beb12db22..f0d77a191a88 100644 --- a/applications/ChatGPT/chatgpt/trainer/callbacks/save_checkpoint.py +++ b/applications/Chat/coati/trainer/callbacks/save_checkpoint.py @@ -1,8 +1,8 @@ import os import torch.distributed as dist -from chatgpt.trainer.strategies import ColossalAIStrategy, Strategy -from chatgpt.trainer.utils import is_rank_0 +from coati.trainer.strategies import GeminiStrategy, LowLevelZeroStrategy, Strategy +from coati.trainer.utils import is_rank_0 from torch import nn from torch.optim import Optimizer @@ -11,7 +11,7 @@ class SaveCheckpoint(Callback): """ - The callback for saving checkpoint for chatgpt. + The callback for saving checkpoint for coati. Only support saving actor and critic model. A typical architecture of the saved checkpoint would be: @@ -69,7 +69,7 @@ def on_episode_end(self, episode: int) -> None: # save optimizer if self.model_dict[model][1] is None: continue - only_rank0 = not isinstance(self.strategy, ColossalAIStrategy) + only_rank0 = not isinstance(self.strategy, (LowLevelZeroStrategy, GeminiStrategy)) rank = 0 if is_rank_0() else dist.get_rank() optim_path = os.path.join(base_path, f'{model}-optim-rank-{rank}.pt') self.strategy.save_optimizer(optimizer=self.model_dict[model][1], path=optim_path, only_rank0=only_rank0) diff --git a/applications/Chat/coati/trainer/ppo.py b/applications/Chat/coati/trainer/ppo.py new file mode 100644 index 000000000000..4c4a1002e96d --- /dev/null +++ b/applications/Chat/coati/trainer/ppo.py @@ -0,0 +1,191 @@ +from typing import Dict, List + +import torch.nn as nn +from coati.experience_maker import Experience, NaiveExperienceMaker +from coati.models.base import Actor, Critic, get_base_model +from coati.models.loss import GPTLMLoss, PolicyLoss, ValueLoss +from coati.models.utils import calc_action_log_probs +from coati.replay_buffer import NaiveReplayBuffer +from torch import Tensor +from torch.optim import Optimizer +from torch.utils.data import DataLoader, DistributedSampler +from tqdm import tqdm + +from colossalai.utils import get_current_device + +from .base import OnPolicyTrainer +from .callbacks import Callback +from .strategies import GeminiStrategy, Strategy +from .utils import is_rank_0, to_device + + +def _set_default_generate_kwargs(strategy: Strategy, generate_kwargs: dict, actor: Actor) -> Dict: + unwrapper_model = strategy.unwrap_model(actor) + hf_model = get_base_model(unwrapper_model) + new_kwargs = {**generate_kwargs} + # use huggingface models method directly + if 'prepare_inputs_fn' not in generate_kwargs and hasattr(hf_model, 'prepare_inputs_for_generation'): + new_kwargs['prepare_inputs_fn'] = hf_model.prepare_inputs_for_generation + + if 'update_model_kwargs_fn' not in generate_kwargs and hasattr(hf_model, '_update_model_kwargs_for_generation'): + new_kwargs['update_model_kwargs_fn'] = hf_model._update_model_kwargs_for_generation + + return new_kwargs + + +class PPOTrainer(OnPolicyTrainer): + """ + Trainer for PPO algorithm. + + Args: + strategy (Strategy): the strategy to use for training + actor (Actor): the actor model in ppo algorithm + critic (Critic): the critic model in ppo algorithm + reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences + initial_model (Actor): the initial model in rlhf algorithm to generate reference logics to limit the update of actor + actor_optim (Optimizer): the optimizer to use for actor model + critic_optim (Optimizer): the optimizer to use for critic model + kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss + train_batch_size (int, defaults to 8): the batch size to use for training + buffer_limit (int, defaults to 0): the max_size limitation of buffer + buffer_cpu_offload (bool, defaults to True): whether to offload buffer to cpu + eps_clip (float, defaults to 0.2): the clip coefficient of policy loss + vf_coef (float, defaults to 1.0): the coefficient of value loss + ptx_coef (float, defaults to 0.9): the coefficient of ptx loss + value_clip (float, defaults to 0.4): the clip coefficient of value loss + sample_buffer (bool, defaults to False): whether to sample from buffer + dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader + offload_inference_models (bool, defaults to True): whether to offload inference models to cpu during training process + callbacks (List[Callback], defaults to []): the callbacks to call during training process + generate_kwargs (dict, optional): the kwargs to use while model generating + """ + + def __init__(self, + strategy: Strategy, + actor: Actor, + critic: Critic, + reward_model: nn.Module, + initial_model: Actor, + actor_optim: Optimizer, + critic_optim: Optimizer, + kl_coef: float = 0.1, + ptx_coef: float = 0.9, + train_batch_size: int = 8, + buffer_limit: int = 0, + buffer_cpu_offload: bool = True, + eps_clip: float = 0.2, + vf_coef: float = 1.0, + value_clip: float = 0.4, + sample_buffer: bool = False, + dataloader_pin_memory: bool = True, + offload_inference_models: bool = True, + callbacks: List[Callback] = [], + **generate_kwargs + ) -> None: + if isinstance(strategy, GeminiStrategy): + assert not offload_inference_models, \ + "GeminiPlugin is not compatible with manual model.to('cpu')" + + buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload) + super().__init__( + strategy, buffer, + sample_buffer, dataloader_pin_memory, + callbacks + ) + + self.generate_kwargs = _set_default_generate_kwargs(strategy, generate_kwargs, actor) + self.experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef) + self.offload_inference_models = offload_inference_models + + self.actor = actor + self.critic = critic + + self.actor_loss_fn = PolicyLoss(eps_clip) + self.critic_loss_fn = ValueLoss(value_clip) + self.vf_coef = vf_coef + self.ptx_loss_fn = GPTLMLoss() + self.ptx_coef = ptx_coef + self.actor_optim = actor_optim + self.critic_optim = critic_optim + + self.device = get_current_device() + + def _make_experience(self, collect_step: int) -> Experience: + prompts = self.prompt_dataloader.next() + if self.offload_inference_models: + # TODO(ver217): this may be controlled by strategy if they are prepared by strategy + self.experience_maker.initial_model.to(self.device) + self.experience_maker.reward_model.to(self.device) + if isinstance(prompts, Tensor): + return self.experience_maker.make_experience(prompts, **self.generate_kwargs) + elif isinstance(prompts, dict): + return self.experience_maker.make_experience(**prompts, **self.generate_kwargs) + else: + raise ValueError(f'Unsupported input type "{type(prompts)}"') + + def _training_step(self, experience: Experience) -> Dict[str, float]: + self.actor.train() + self.critic.train() + # policy loss + num_actions = experience.action_mask.size(1) + actor_output = self.actor(experience.sequences, attention_mask=experience.attention_mask) + action_log_probs = calc_action_log_probs(actor_output, experience.sequences, num_actions) + actor_loss = self.actor_loss_fn(action_log_probs, + experience.action_log_probs, + experience.advantages, + action_mask=experience.action_mask) + + # ptx loss + if self.ptx_coef != 0: + batch = self.pretrain_dataloader.next() + batch = to_device(batch, self.device) + ptx_log_probs = self.actor(batch['input_ids'], + attention_mask=batch['attention_mask'])['logits'] + ptx_loss = self.ptx_loss_fn(ptx_log_probs, batch['labels']) + actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef) + + self.strategy.backward(actor_loss, self.actor, self.actor_optim) + self.strategy.optimizer_step(self.actor_optim) + self.actor_optim.zero_grad() + + # value loss + values = self.critic(experience.sequences, + action_mask=experience.action_mask, + attention_mask=experience.attention_mask) + critic_loss = self.critic_loss_fn(values, + experience.values, + experience.reward, + action_mask=experience.action_mask) + critic_loss = critic_loss * self.vf_coef + self.strategy.backward(critic_loss, self.critic, self.critic_optim) + self.strategy.optimizer_step(self.critic_optim) + self.critic_optim.zero_grad() + + return {'reward': experience.reward.mean().item()} + + def _learn(self, update_step: int): + if self.offload_inference_models: + self.experience_maker.initial_model.to('cpu') + self.experience_maker.reward_model.to('cpu') + + # buffer may be empty at first, we should rebuild at each training + if self.sample_buffer: + experience = self.buffer.sample() + self._on_learn_batch_start() + experience.to_device(self.device) + metrics = self._training_step(experience) + self._on_learn_batch_end(metrics, experience) + else: + if isinstance(self.dataloader.sampler, DistributedSampler): + self.dataloader.sampler.set_epoch(update_step) + pbar = tqdm( + self.dataloader, + desc=f'Train epoch [{update_step + 1}]', + disable=not is_rank_0() + ) + for experience in pbar: + self._on_learn_batch_start() + experience.to_device(self.device) + metrics = self._training_step(experience) + self._on_learn_batch_end(metrics, experience) + pbar.set_postfix(metrics) diff --git a/applications/Chat/coati/trainer/rm.py b/applications/Chat/coati/trainer/rm.py new file mode 100644 index 000000000000..54a5d0f40dea --- /dev/null +++ b/applications/Chat/coati/trainer/rm.py @@ -0,0 +1,111 @@ +from datetime import datetime +from typing import Callable + +import pandas as pd +import torch +import tqdm +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler +from torch.utils.data import DataLoader + +from .base import SLTrainer +from .strategies import Strategy +from .utils import is_rank_0 + + +class RewardModelTrainer(SLTrainer): + """ + Trainer to use while training reward model. + + Args: + model (torch.nn.Module): the model to train + strategy (Strategy): the strategy to use for training + optim (Optimizer): the optimizer to use for training + lr_scheduler (_LRScheduler): the lr scheduler to use for training + loss_fn (callable): the loss function to use for training + max_epochs (int, defaults to 2): the number of epochs to train + """ + + def __init__( + self, + model, + strategy: Strategy, + optim: Optimizer, + lr_scheduler: _LRScheduler, + loss_fn: Callable, + max_epochs: int = 1, + ) -> None: + super().__init__(strategy, max_epochs, model, optim) + + self.loss_fn = loss_fn + self.scheduler = lr_scheduler + + def _eval(self, epoch): + if self.eval_dataloader is not None: + self.model.eval() + dist, on, cnt = 0, 0, 0 + with torch.no_grad(): + for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader: + chosen_ids = chosen_ids.squeeze(1).to(torch.cuda.current_device()) + c_mask = c_mask.squeeze(1).to(torch.cuda.current_device()) + reject_ids = reject_ids.squeeze(1).to(torch.cuda.current_device()) + r_mask = r_mask.squeeze(1).to(torch.cuda.current_device()) + chosen_reward = self.model(chosen_ids, attention_mask=c_mask) + reject_reward = self.model(reject_ids, attention_mask=r_mask) + for i in range(len(chosen_reward)): + cnt += 1 + if chosen_reward[i] > reject_reward[i]: + on += 1 + dist += (chosen_reward - reject_reward).mean().item() + self.dist = dist / len(self.eval_dataloader) + self.acc = on / cnt + + if is_rank_0(): + log = pd.DataFrame( + [[(epoch + 1) * len(self.train_dataloader), + self.loss.item(), self.dist, self.acc]], + columns=['step', 'loss', 'dist', 'acc'] + ) + log.to_csv('log.csv', mode='a', header=False, index=False) + + def _train(self, epoch): + self.model.train() + step_bar = tqdm.trange( + len(self.train_dataloader), + desc='Train step of epoch %d' % epoch, + disable=not is_rank_0() + ) + cnt = 0 + for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader: + chosen_ids = chosen_ids.squeeze(1).to(torch.cuda.current_device()) + c_mask = c_mask.squeeze(1).to(torch.cuda.current_device()) + reject_ids = reject_ids.squeeze(1).to(torch.cuda.current_device()) + r_mask = r_mask.squeeze(1).to(torch.cuda.current_device()) + chosen_reward = self.model(chosen_ids, attention_mask=c_mask) + reject_reward = self.model(reject_ids, attention_mask=r_mask) + self.loss = self.loss_fn(chosen_reward, reject_reward) + self.strategy.backward(self.loss, self.model, self.optimizer) + self.strategy.optimizer_step(self.optimizer) + self.optimizer.zero_grad() + cnt += 1 + if cnt % 100 == 0: + self.scheduler.step() + step_bar.update() + step_bar.close() + + def _before_fit(self, + train_dataloader: DataLoader, + valid_dataloader: DataLoader, + eval_dataloader: DataLoader): + """ + Args: + train_dataloader (DataLoader): the dataloader to use for training + valid_dataloader (DataLoader): the dataloader to use for validation + eval_dataloader (DataLoader): the dataloader to use for evaluation + """ + super()._before_fit() + self.datetime = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + + self.train_dataloader = train_dataloader + self.valid_dataloader = valid_dataloader + self.eval_dataloader = eval_dataloader diff --git a/applications/Chat/coati/trainer/sft.py b/applications/Chat/coati/trainer/sft.py new file mode 100644 index 000000000000..5ed0a05ca7eb --- /dev/null +++ b/applications/Chat/coati/trainer/sft.py @@ -0,0 +1,129 @@ +import time +from typing import Optional + +import torch +import torch.distributed as dist +import tqdm +import wandb +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler +from torch.utils.data import DataLoader + +from colossalai.logging import DistributedLogger + +from .base import SLTrainer +from .strategies import GeminiStrategy, Strategy +from .utils import is_rank_0, to_device + + +class SFTTrainer(SLTrainer): + """ + Trainer to use while training reward model. + + Args: + model (torch.nn.Module): the model to train + strategy (Strategy): the strategy to use for training + optim(Optimizer): the optimizer to use for training + lr_scheduler(_LRScheduler): the lr scheduler to use for training + max_epochs (int, defaults to 2): the number of epochs to train + accumulation_steps (int, defaults to 8): the number of steps to accumulate gradients + """ + + def __init__( + self, + model, + strategy: Strategy, + optim: Optimizer, + lr_scheduler: _LRScheduler, + tensorboard_dir: str = None, + max_epochs: int = 2, + accumulation_steps: int = 8, + ) -> None: + if accumulation_steps > 1: + assert not isinstance(strategy, GeminiStrategy), \ + "Accumulation steps are not supported in stage 3 of ColossalAI" + + super().__init__(strategy, max_epochs, model, optim, tensorboard_dir) + self.accumulation_steps = accumulation_steps + self.scheduler = lr_scheduler + + def _train(self, epoch: int): + self.model.train() + start_step = epoch * len(self.train_dataloader) // self.accumulation_steps + for batch_id, batch in enumerate(self.train_dataloader): + + batch = to_device(batch, torch.cuda.current_device()) + outputs = self.model(batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"]) + + loss = outputs.loss + loss = loss / self.accumulation_steps + + self.strategy.backward(loss, self.model, self.optimizer) + + self.total_loss += loss.item() + + # gradient accumulation + if (batch_id + 1) % self.accumulation_steps == 0: + self.strategy.optimizer_step(self.optimizer) + self.optimizer.zero_grad() + self.scheduler.step() + if dist.get_rank() == 0 and self.writer: + self.writer.add_scalar('loss', self.total_loss, start_step + batch_id) + self.writer.add_scalar('lr', self.scheduler.get_last_lr()[0], start_step + batch_id) + if is_rank_0() and self.use_wandb: + wandb.log({ + "loss": self.total_loss, + "lr": self.scheduler.get_last_lr()[0], + "epoch": epoch, + "batch_id": batch_id + }) + self.total_loss = 0 + self.step_bar.update() + print(f"Epoch {epoch}/{self.max_epochs} batch {batch_id}/{len(self.train_dataloader)} loss {loss.item()}") + + def _eval(self, epoch: int): + if self.eval_dataloader is not None: + self.model.eval() + with torch.no_grad(): + loss_sum, num_seen = 0, 0 + for batch in self.eval_dataloader: + batch = to_device(batch, torch.cuda.current_device()) + outputs = self.model(batch["input_ids"], + attention_mask=batch["attention_mask"], + labels=batch["labels"]) + loss = outputs.loss + + loss_sum += loss.item() + num_seen += batch["input_ids"].size(0) + + loss_mean = loss_sum / num_seen + if dist.get_rank() == 0: + self.logger.info(f'Eval Epoch {epoch}/{self.max_epochs} loss {loss_mean}') + if dist.get_rank() == 0 and self.writer: + self.writer.add_scalar('eval_loss', loss_mean, epoch) + print(f'Eval Epoch {epoch}/{self.max_epochs} loss {loss_mean}') + + def _before_fit(self, + train_dataloader: DataLoader, + eval_dataloader: Optional[DataLoader] = None, + logger: Optional[DistributedLogger] = None, + use_wandb: bool = False): + """ + Args: + train_dataloader: the dataloader to use for training + eval_dataloader: the dataloader to use for evaluation + """ + self.train_dataloader = train_dataloader + self.eval_dataloader = eval_dataloader + + self.logger = logger + self.use_wandb = use_wandb + if use_wandb: + wandb.init(project="Coati", name=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())) + wandb.watch(self.model) + + self.total_loss = 0 + self.no_epoch_bar = True + self.step_bar = tqdm.trange(len(self.train_dataloader) // self.accumulation_steps * self.max_epochs, + desc=f'steps', + disable=not is_rank_0()) diff --git a/applications/Chat/coati/trainer/strategies/__init__.py b/applications/Chat/coati/trainer/strategies/__init__.py new file mode 100644 index 000000000000..b49a2c742db3 --- /dev/null +++ b/applications/Chat/coati/trainer/strategies/__init__.py @@ -0,0 +1,8 @@ +from .base import Strategy +from .colossalai import GeminiStrategy, LowLevelZeroStrategy +from .ddp import DDPStrategy + +__all__ = [ + 'Strategy', 'DDPStrategy', + 'LowLevelZeroStrategy', 'GeminiStrategy' +] diff --git a/applications/Chat/coati/trainer/strategies/base.py b/applications/Chat/coati/trainer/strategies/base.py new file mode 100644 index 000000000000..5352cd5fc4db --- /dev/null +++ b/applications/Chat/coati/trainer/strategies/base.py @@ -0,0 +1,137 @@ +from abc import ABC, abstractmethod +from contextlib import nullcontext +from typing import Callable, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from coati.replay_buffer import ReplayBuffer +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from colossalai.booster import Booster +from colossalai.booster.plugin import Plugin + +from .sampler import DistributedSampler + +_BoostArgSpec = Union[nn.Module, Tuple[nn.Module, Optimizer], Dict] + + +class Strategy(ABC): + """ + Base class for training strategies. + """ + + def __init__(self, plugin_initializer: Callable[..., Optional[Plugin]] = lambda: None) -> None: + super().__init__() + # NOTE: dist must be initialized before Booster + self.setup_distributed() + self.plugin = plugin_initializer() + self.booster = Booster(plugin=self.plugin) + self._post_init() + + @abstractmethod + def _post_init(self) -> None: + pass + + def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: Optimizer, **kwargs) -> None: + self.booster.backward(loss, optimizer) + + def optimizer_step(self, optimizer: Optimizer, **kwargs) -> None: + optimizer.step() + + @abstractmethod + def setup_distributed(self) -> None: + pass + + @abstractmethod + def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: + pass + + def model_init_context(self): + return nullcontext() + + def prepare(self, *boost_args: _BoostArgSpec) -> Union[List[_BoostArgSpec], _BoostArgSpec]: + """Prepare [model | (model, optimizer) | Dict] based on each strategy. + NOTE: the keys of Dict must be a subset of `self.booster.boost`'s arguments. + + Example:: + >>> # e.g., include lr_scheduler + >>> result_dict = strategy.prepare(dict(model=model, lr_scheduler=lr_scheduler)) + >>> # when fine-tuning actor and critic + >>> (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare((actor, actor_optim), (critic, critic_optim), reward_model, initial_model) + >>> # or when training reward model + >>> (reward_model, reward_model_optim) = strategy.prepare((reward_model, reward_model_optim)) + >>> # or just inference + >>> actor, critic = strategy.prepare(actor, critic) + + Returns: + Union[List[_BoostArgSpec], _BoostArgSpec]: [model | (model, optimizer) | Dict] in the original order. + """ + + rets = [] + for arg in boost_args: + if isinstance(arg, nn.Module): + model, *_ = self.booster.boost(arg) + rets.append(model) + elif isinstance(arg, tuple): + try: + model, optimizer = arg + except ValueError: + raise RuntimeError(f'Expect (model, optimizer) pair, got a tuple with size "{len(arg)}"') + model, optimizer, *_ = self.booster.boost(model=model, optimizer=optimizer) + rets.append((model, optimizer)) + elif isinstance(arg, Dict): + model, optimizer, criterion, dataloader, lr_scheduler = self.booster.boost(**arg) + boost_result = dict(model=model, + optimizer=optimizer, + criterion=criterion, + dataloader=dataloader, + lr_scheduler=lr_scheduler) + # remove None values + boost_result = {key: value for key, value in boost_result.items() if value is not None} + rets.append(boost_result) + else: + raise RuntimeError(f'Type {type(arg)} is not supported') + + return rets[0] if len(rets) == 1 else rets + + @staticmethod + def unwrap_model(model: nn.Module) -> nn.Module: + """Get the unwrapped model from a wrapped model made by Strategy.prepare. + + Args: + model (nn.Module): the model to unwrap + + Returns: + nn.Module: the original model + """ + return model + + def save_model(self, model: nn.Module, path: str, only_rank0: bool = True, **kwargs) -> None: + self.booster.save_model(model, path, shard=not only_rank0, **kwargs) + + def load_model(self, model: nn.Module, path: str, strict: bool = True) -> None: + self.booster.load_model(model, path, strict) + + def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False, **kwargs) -> None: + self.booster.save_optimizer(optimizer, path, shard=False, **kwargs) + + def load_optimizer(self, optimizer: Optimizer, path: str) -> None: + self.booster.load_optimizer(optimizer, path) + + def setup_sampler(self, dataset) -> DistributedSampler: + # FIXME(cwher): this is only invoked in train_on_ray, not tested after adapt Boost API. + return DistributedSampler(dataset, 1, 0) + + @abstractmethod + def save_pretrained(self, + model: nn.Module, + path: str, + only_rank0: bool = True, + tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None: + pass + + @abstractmethod + def get_model_state_dict_shard(self, model: nn.Module, **config): + pass diff --git a/applications/Chat/coati/trainer/strategies/colossalai.py b/applications/Chat/coati/trainer/strategies/colossalai.py new file mode 100644 index 000000000000..1b59d704eec3 --- /dev/null +++ b/applications/Chat/coati/trainer/strategies/colossalai.py @@ -0,0 +1,227 @@ +import warnings +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +import colossalai +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin +from colossalai.booster.plugin.gemini_plugin import GeminiModel +from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel +from colossalai.tensor import ProcessGroup, ShardSpec +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from colossalai.zero.gemini.gemini_ddp import GeminiDDP + +from .ddp import DDPStrategy + + +class LowLevelZeroStrategy(DDPStrategy): + """ + The strategy for training with ColossalAI. + + Args: + stage(int): The stage to use in ZeRO. Choose in (1, 2) + precision(str): The precision to use. Choose in ('fp32', 'fp16'). + seed(int): The seed for the random number generator. + placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda') + If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU, + If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest. + reduce_bucket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2. + overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2. + initial_scale(float): The initial scale for the optimizer. + growth_factor(float): The growth factor for the optimizer. + backoff_factor(float): The backoff factor for the optimizer. + growth_interval(int): The growth interval for the optimizer. + hysteresis(int): The hysteresis for the optimizer. + min_scale(float): The minimum scale for the optimizer. + max_scale(float): The maximum scale for the optimizer. + max_norm(float): The maximum norm for the optimizer. + norm_type(float): The norm type for the optimizer. + + """ + + def __init__(self, + stage: int = 3, + precision: str = 'fp16', + seed: int = 42, + placement_policy: str = 'cuda', + reduce_bucket_size: int = 12 * 1024**2, # only for stage 1&2 + overlap_communication: bool = True, # only for stage 1&2 + initial_scale: float = 2**16, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + min_scale: float = 1, + max_scale: float = 2**32, + max_norm: float = 0.0, + norm_type: float = 2.0 + ) -> None: + + assert stage in (1, 2), f'Unsupported stage "{stage}"' + assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"' + assert precision in ('fp32', 'fp16'), f'Unsupported precision "{precision}"' + + plugin_initializer = lambda: LowLevelZeroPlugin( + # zero_config + stage=stage, + precision=precision, + # zero_optim_config + reduce_bucket_size_in_m=reduce_bucket_size, + overlap_communication=overlap_communication, + cpu_offload=(placement_policy == 'cpu'), + # optim_config + initial_scale=initial_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + min_scale=min_scale, + max_scale=max_scale, + max_norm=max_norm, + norm_type=norm_type + ) + + super().__init__(seed, plugin_initializer) + + def _post_init(self) -> None: + assert isinstance(self.plugin, LowLevelZeroPlugin), \ + f'{type(self).__name__}\'s plugin is not initialized properly.' + + def setup_distributed(self) -> None: + colossalai.launch_from_torch({}, seed=self.seed) + + def unwrap_model(self, model: nn.Module) -> nn.Module: + assert isinstance(model, LowLevelZeroModel) + return model.module + + def get_model_state_dict_shard(self, model: nn.Module, **config): + assert isinstance(model, LowLevelZeroModel) + yield from model.state_dict_shard(max_shard_size=1024, only_rank_0=False) + + +class GeminiStrategy(DDPStrategy): + """ + The strategy for training with ColossalAI. + + Args: + seed(int): The seed for the random number generator. + shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3. + This is not compatible with `from_pretrained()`. We temporarily disable this and will support it in the future. + placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda') + If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU, + If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest. + pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3. + force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3. + search_range_m(int): The number of search range for the chunk size, divided by 2^20. Only for ZeRO-3. + hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3. + min_chunk_size_m(float): The minimum chunk size divided by 2^20. Only for ZeRO-3. + gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3. + initial_scale(float): The initial scale for the optimizer. + growth_factor(float): The growth factor for the optimizer. + backoff_factor(float): The backoff factor for the optimizer. + growth_interval(int): The growth interval for the optimizer. + hysteresis(int): The hysteresis for the optimizer. + min_scale(float): The minimum scale for the optimizer. + max_scale(float): The maximum scale for the optimizer. + max_norm(float): The maximum norm for the optimizer. + norm_type(float): The norm type for the optimizer. + + """ + + def __init__(self, + seed: int = 42, + shard_init: bool = False, # only for stage 3 + placement_policy: str = 'cuda', + pin_memory: bool = True, # only for stage 3 + force_outputs_fp32: bool = False, # only for stage 3 + search_range_m: int = 32, # only for stage 3 + hidden_dim: Optional[int] = None, # only for stage 3 + min_chunk_size_m: float = 32, # only for stage 3 + gpu_margin_mem_ratio: float = 0.0, # only for stage 3 + initial_scale: float = 2**16, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + min_scale: float = 1, + max_scale: float = 2**32, + max_norm: float = 0.0, + norm_type: float = 2.0 + ) -> None: + + assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"' + + # TODO(ver217): support shard_init when using from_pretrained() + if shard_init: + warnings.warn( + f'Shard init is not supported model.from_pretrained() yet. ' + 'Please load weights after strategy.prepare()' + ) + self.shard_init = shard_init + + warnings.warn(f'Stage 3 only supports fp16. Precision is set to fp16.') + + # NOTE: dist should be initialized before calling get_current_device() + plugin_initializer = lambda: GeminiPlugin( + # gemini_config + device=get_current_device(), + placement_policy=placement_policy, + precision='fp16', + pin_memory=pin_memory, + force_outputs_fp32=force_outputs_fp32, + strict_ddp_mode=shard_init, + search_range_m=search_range_m, + hidden_dim=hidden_dim, + min_chunk_size_m=min_chunk_size_m, + # zero_optim_config + gpu_margin_mem_ratio=gpu_margin_mem_ratio, + # optim_config + initial_scale=initial_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + min_scale=min_scale, + max_scale=max_scale, + max_norm=max_norm, + norm_type=norm_type + ) + + super().__init__(seed, plugin_initializer) + + def _post_init(self) -> None: + assert isinstance(self.plugin, GeminiPlugin), \ + f'{type(self).__name__}\'s plugin is not initialized properly.' + + def setup_distributed(self) -> None: + colossalai.launch_from_torch({}, seed=self.seed) + + def model_init_context(self): + world_size = dist.get_world_size() + shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None + default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None + return ColoInitContext(device=get_current_device(), + dtype=torch.half, + default_pg=shard_pg, + default_dist_spec=default_dist_spec) + + def unwrap_model(self, model: nn.Module) -> nn.Module: + assert isinstance(model, GeminiModel) + ddp_model = model.unwrap() + assert isinstance(ddp_model, GeminiDDP) + return ddp_model.module + + def save_pretrained(self, + model: nn.Module, + path: str, + only_rank0: bool = True, + tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None: + raise RuntimeError('ColossalAI strategy with stage-3 does not support save_pretrained() now') + + def get_model_state_dict_shard(self, model: nn.Module, **config): + assert isinstance(self.plugin, GeminiPlugin) + yield from super().get_model_state_dict_shard(model, **config) diff --git a/applications/Chat/coati/trainer/strategies/ddp.py b/applications/Chat/coati/trainer/strategies/ddp.py new file mode 100644 index 000000000000..e1c1bbf19f35 --- /dev/null +++ b/applications/Chat/coati/trainer/strategies/ddp.py @@ -0,0 +1,125 @@ +import os +import random +from collections import OrderedDict +from typing import Callable, Optional + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from coati.replay_buffer import ReplayBuffer +from torch.utils.data import DataLoader +from transformers.modeling_utils import PreTrainedModel +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.booster.plugin.torch_ddp_plugin import TorchDDPModel + +from .base import Strategy +from .sampler import DistributedSampler + + +# TODO Move this to a util.py (Moving to ray.util introduces ringed import) +def get_grad_required_state_dict(model: nn.Module): + state_dict = OrderedDict() + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + state_dict[name] = parameter.detach() + return state_dict + + +class DDPStrategy(Strategy): + """ + Strategy for distributed training using torch.distributed. + """ + + def __init__(self, + seed: int = 42, + plugin_initializer: Callable = TorchDDPPlugin + ) -> None: + self.seed = seed + super().__init__(plugin_initializer) + + def _try_init_dist(self, force: bool = False) -> None: + try: + rank = int(os.environ['RANK']) + local_rank = int(os.environ['LOCAL_RANK']) + world_size = int(os.environ['WORLD_SIZE']) + host = os.environ['MASTER_ADDR'] + port = int(os.environ['MASTER_PORT']) + dist.init_process_group('nccl', init_method=f'tcp://[{host}]:{port}', world_size=world_size, rank=rank) + torch.cuda.set_device(local_rank) + except KeyError as e: + if force: + raise RuntimeError( + f"Could not find {e} in the torch environment, visit https://www.colossalai.org/ for more information on launching with torch" + ) + except Exception as e: + if force: + raise e + + def _post_init(self) -> None: + assert isinstance(self.plugin, TorchDDPPlugin), \ + f'{type(self).__name__}\'s plugin is not initialized properly.' + + def setup_distributed(self) -> None: + self._try_init_dist(force=True) + self.set_seed(self.seed) + + def set_seed(self, seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: + return self.plugin.prepare_dataloader(replay_buffer, + batch_size=replay_buffer.sample_batch_size, + shuffle=True, + drop_last=True, + pin_memory=pin_memory, + collate_fn=replay_buffer.collate_fn) + + def setup_sampler(self, dataset) -> DistributedSampler: + # FIXME(cwher): this is only invoked in train_on_ray, not tested after adapt Boost API. + return DistributedSampler(dataset, dist.get_world_size(), dist.get_rank()) + + def unwrap_model(self, model: nn.Module) -> nn.Module: + assert isinstance(model, TorchDDPModel), "model is not wrapped by TorchDDPModel." + return model.unwrap() + + def save_pretrained(self, + model: nn.Module, + path: str, + only_rank0: bool = True, + tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None: + if only_rank0 and dist.get_rank() != 0: + return + unwrapped_model = self.unwrap_model(model) + assert isinstance(unwrapped_model, PreTrainedModel) + unwrapped_model.save_pretrained(path) + if tokenizer is not None: + tokenizer.save_pretrained(path) + + def get_model_state_dict_shard(self, model: nn.Module, **config): + # TODO: implement sharding on naive strategy + model = self.unwrap_model(model) + if 'requires_grad_only' in config and config['requires_grad_only'] == True: + state_dict = get_grad_required_state_dict(model) + else: + state_dict = model.state_dict() + + if 'shard_size' in config: + shard_size = config['shard_size'] + accumulate_size = 0 + state_dict_shard = OrderedDict() + for name, param in state_dict.items(): + state_dict_shard[name] = param + accumulate_size += param.numel() * param.element_size() + if accumulate_size >= shard_size: + accumulate_size = 0 + yield state_dict_shard + state_dict_shard = OrderedDict() + if accumulate_size > 0: + yield state_dict_shard + else: + yield state_dict diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/sampler.py b/applications/Chat/coati/trainer/strategies/sampler.py similarity index 99% rename from applications/ChatGPT/chatgpt/trainer/strategies/sampler.py rename to applications/Chat/coati/trainer/strategies/sampler.py index d726fa640fa2..65e199dbf029 100644 --- a/applications/ChatGPT/chatgpt/trainer/strategies/sampler.py +++ b/applications/Chat/coati/trainer/strategies/sampler.py @@ -27,6 +27,7 @@ def __init__(self, dataset, num_replicas: int, rank: int) -> None: assert len(indices) == self.num_samples self.indices = indices + def sample(self, batch_size: int) -> list: sampled_indices = np.random.choice(self.indices, batch_size, replace=False) return [self.dataset[idx] for idx in sampled_indices] diff --git a/applications/Chat/coati/trainer/utils.py b/applications/Chat/coati/trainer/utils.py new file mode 100644 index 000000000000..c9fc8d0fe19f --- /dev/null +++ b/applications/Chat/coati/trainer/utils.py @@ -0,0 +1,46 @@ +from typing import Any + +import torch +import torch.distributed as dist +from torch.utils._pytree import tree_map +from torch.utils.data import DataLoader + + +class CycledDataLoader: + """ + Why do we need this class? + In version 4da324cd60, "prompts = next(iter(self.prompt_dataloader))" is used to sample a batch of prompts/pretrain. + However, this may be inefficient due to frequent re-initialization of the dataloader. (re-initialize workers...) + NOTE: next(iter(dataloader)) is not equivalent to for batch in dataloader: break, it causes slightly different behavior. + """ + + def __init__(self, + dataloader: DataLoader, + ) -> None: + self.dataloader = dataloader + + self.count = 0 + self.dataloader_iter = iter(dataloader) + + def next(self): + self.count += 1 + try: + return next(self.dataloader_iter) + except StopIteration: + self.count = 0 + self.dataloader_iter = iter(self.dataloader) + return next(self.dataloader_iter) + + +def is_rank_0() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +def to_device(x: Any, device: torch.device) -> Any: + + def _to(t: Any): + if isinstance(t, torch.Tensor): + return t.to(device) + return t + + return tree_map(_to, x) diff --git a/applications/Chat/evaluate/README.md b/applications/Chat/evaluate/README.md new file mode 100644 index 000000000000..e4a50b11d41f --- /dev/null +++ b/applications/Chat/evaluate/README.md @@ -0,0 +1,390 @@ +# Evaluation + +In this directory, we introduce how you can evaluate your model with our pipeline. This pipeline is now available for evaluation of both Chinese and English capability. + +## Installation + +To start model evaluation, you need to install required packages which listed in `requirements.txt` under `evaluate` folder. + +```shell +pip install -r requirements.txt +``` + +## Evaluation Pipeline + +The whole evaluation pipeline consists of three methods: + +1. `GPT Evaluation`: evaluates model predictions using GPT models. + * Compare the performance of two different models (battle). + * Rate the model according to pre-defined metrics using prompting design. + * Rate the model according to pre-defined metrics with additional reference answer using prompting design. +2. `Automatic Evaluation`: evaluates model predictions using automatic metrics. +3. `UniEval`: evaluates model predictions using UniEval models(English only). + +### Evaluation Category + +Our evaluation pipeline examines the model's capability using 10 categories of questions. The following table introduces each category: + +| Evaluation Category | Description | +| :-----------------: | :----------------------------------------------------------- | +| Brainstorming | Models are asked to generate a range of creative and diverse ideas according to the question. The capability of creativity is required. | +| Chat | Models are asked to continue a multi-round dialogue given the roles involved. The capability of understanding, memorizing previous rounds of the dialogue and answering according to the persona provided is required. | +| Classification | Models are asked to do classification tasks. The capability of accurate classification is required. | +| Closed QA | Models are asked to answer a closed QA question. The capability of answering questions with limited scope (such as single/multiple choice question) is required. | +| Extraction | Models are asked to extract information from a given material. The capability of extracting required information is required. | +| Generation | Models are asked to generate an email, letter, article, etc. The capability of generating texts in a high quality and human-written way is required. | +| Open QA | Models are asked to answer an open QA question(without context provided). The capability of answering questions with the models' own knowledge base is required. | +| Roleplay | Models are asked to play the role provided. The capability of engaging in the scenario and effectively interacting with the user is required. | +| Rewriting | Models are asked to do rewriting tasks such as translation and grammar correction. The capability of rewriting according to different instructions is required. | +| Summarization | Models are asked to summarize the given paragraph or passage. The capability of summarization is required. | + +To better understand each evaluation category, here are some example questions provided. + + +| Evaluation Category | Chinese Example | English Example | +| :-----------------: | :----------------------------------------------------------- | :----------------------------------------------------------- | +| Brainstorming | **Example 1:**
请介绍一下人工智能的多个领域。

**Example 2:**
请给出管理家庭财务的3个小技巧。
| **Example 1:**
How can I improve my memory? Any useful techniques you can suggest?

**Example 2:**
What are some ways to increase productivity while working from home? | +| Chat | **Example 1:**
基于以下角色信息完成一段对话。小张是一名新手爱好者，对养鸡有浓厚的兴趣。老李是一名有丰富经验的养鸡大师。
小张：您好，老李，我最近开始对养鸡感兴趣了，想请教您一些问题。
老李：你好，小张，我很乐意帮助你。你想问些什么？
小张：我想知道如何确定鸡的品种和性别？
老李：确切的品种可以通过鸡的外貌特征来确定，而性别一般是通过鸡卵的大小和形状来判断。还有什么问题吗？
小张：

**Example 2:**
基于以下角色信息完成一段对话。小明是一名医生，一位老年病患者想要停药，但他对病情有所忽视并有担忧；王叔叔是老年病患者的儿子，希望能够听取医生的建议。
小明：你好，王叔叔，我了解你想要让你父亲停药。
王叔叔：是的，我父亲已经吃了那么久的药，我担心药物对他的身体会有副作用。
小明： | **Example 1:**
Complete a conversation based on the following character information. Amy is a 30-year-old chef who runs her own restaurant. Jack is a food blogger who specializes in reviewing local restaurants.
Amy: Hi Jack, I heard that you're a food blogger. Nice to meet you.
Jack: Hi Amy, yes I am. Your restaurant has been receiving a lot of good reviews lately.
Amy: Yes, we use only fresh and quality ingredients, and every dish is carefully crafted.
Jack:

**Example 2:**
Complete a dialogue based on the following role information. A: Elementary student B: Teacher
B: Good morning, Student A. Today we're going to learn about addition and subtraction.
A: Teacher, I already know this very well. Why do I need to learn it again?
B: | +| Classification | **Example 1:**
新闻标题：今日立夏，有一上联，立夏万物并秀，下联怎么对？
请根据以上新闻标题判断新闻所属的分类，你需要从文化，娱乐，体育，财经，房产，教育，科技，旅游，游戏，军事这十类中选择一个答案。

**Example 2:**
新闻标题：赵丽颖很久没有登上微博热搜了，但你们别急，她只是在憋大招而已。
请根据新闻标题判断新闻所属的分类，你需要从文化，娱乐，体育，财经，房产，教育，科技，旅游，游戏，军事这十类中选择一个答案。 | **Example 1:**
Title: Fighting for Love (2020)
Description: Jasmine got obsessed with a man and now he's obsessed with her. Steamy nights, kisses and rules being broken awaits them. She turned his whole world upside down and now he's doing it to hers. In this free fall, can they survive each others love?\"
Based on the above information, determine which genre the work of art belongs to. You can only choose one from \"sport\", \"horror\", \"drama\", \"history\", \"romance\", \"biography\", \"science fiction\", \"comedy\", \"animation\", \"documentary\", \"music\" and \"news\".

**Example2:**
Title: Summer Breeze: The Isley Brothers Greatest Hits Live (2005)
Description: Filmed in the US in 2005 and captured in excellent form led by Ron Isley's vocals and Ernie Isley's hard edged guitar. Virtually every track is a hit including Shout, Who's That Lady, Twist And Shout, Summer Breeze and Harvest For The World.
Based on the above information, determine which genre the work of art belongs to. You can only choose one from \"sport\", \"horror\", \"drama\", \"history\", \"romance\", \"biography\", \"science fiction\", \"comedy\", \"animation\", \"documentary\", \"music\" and \"news\"." | +| Closed QA | **Example 1:**
请从以下选项中选择正确答案。以下哪个是世界上最高山峰？
A. 长城
B. 泰山
C. 珠穆朗玛峰
D. 黄山

**Example 2:**
请从以下选项中选择一个最佳答案回答下面的问题。问题：非洲最高的山是哪座山？
选项：
A. 麦金利山
B. 喜马拉雅山
C. 乞力马扎罗山 | **Example 1:**
Which of the following options is NOT a primary color?
(a) yellow
(b) blue
(c) orange
(d) red

**Example 2:**
Choose the correct option to complete the following sentence: \"Harry Potter and the Chamber of Secrets\" is the ________ book in the Harry Potter series.
(A) first
(B) second
(C) third
(D) fourth | +| Extraction | **Example 1:**
根据以下新闻文本，提取新闻报道时间，例如回答时按照格式“新闻报道时间：2007年8月10日”
新闻文本如下：2007-4-7中新网4月7日电据中国消防在线消息，4月4日晚上7时30分左右，湖南长潭高速公路上发生一起6车连环相撞失火事故。长株潭三地消防部门共出动消防车21台，警力100余人。经过消防官兵近2个小时奋力扑救，大火被成功扑灭。据初步调查，有1人在此次事故中死亡。

**Example 2:**
根据以下新闻文本，提取新闻报道时间，例如回答时按照格式“新闻报道时间：2007年8月10日”
新闻文本如下：2014年1月15日，据外媒《俄罗斯报》报道称，位于北半球的澳大利亚现在正处于炎热的夏季，而近日也到了高温酷暑的时候，当地时间1月14日晚，澳大利亚南部一夜间发生至少250起火灾。受炎热天气及雷雨天气影响，澳大利亚南部一夜间发生至少250起火灾，灾情多集中在维多利亚州。火灾发生后，救援人员立即展开救灾行动。目前，大部分起火点火势已被控制。 | **Example 1:**
Ernest Hemingway, an American literary giant known for his spare and direct writing style, has penned timeless works such as 'The Old Man and the Sea', 'For Whom the Bell Tolls', and 'A Farewell to Arms', which have made a profound impact on the literary world and continue to be widely read and admired today.
Extract the name of the author mentioned above.

**Example 2:**
In the epic fantasy series 'A Song of Ice and Fire', George R.R. Martin weaves a complex web of political intrigue, war, and magic across the fictional continents of Westeros and Essos. Martin's richly developed characters and intricate plotlines have captivated readers worldwide, much like his other acclaimed works such as 'A Clash of Kings' and 'A Storm of Swords'.
Extract the name of the author in the above material. | +| Generation | **Example 1:**
请撰写一篇文章，介绍如何通过改善生活习惯来预防疾病和延长寿命。

**Example 2:**
请根据以下情节撰写一篇短篇小说：一名年轻人被困在一个荒岛上，他必须想办法生存下去直到被救援。但他很快发现自己并不孤单。 | **Example 1:**
Write a descriptive paragraph about an island to relax and unwind, including details about the location and atmosphere.

**Example 2:**
Can you help me write a persuasive email to my colleagues encouraging them to participate in a charitable fundraising event? | +| Open QA | **Example 1:**
请问万有引力定律由谁提出的？

**Example 2:**
哪些国家参与了第一次世界大战？ | **Example 1:**
What are the four basic tastes of the human palate?

**Example 2:**
Who painted the The Scream? | +| Rewriting | **Example 1:**
请将以下句子改为正确的语序。
生日快乐你祝他了吗？

**Example 2:**
将以下文本翻译成英语：
“这个周末我要去海边玩” | **Example 1:**
Please translate the following sentences, which are a mixture of Chinese and English, into full English.
我需要买一些healthy snacks，比如nuts和dried fruits，作为我的office的午餐.

**Example 2:**
Please rewrite the sentence using an inverted sentence structure.
We won't begin our journey until the sun sets. | +| Roleplay | **Example 1:**
我想让你担任Android开发工程师面试官。我将成为候选人，您将向我询问Android开发工程师职位的面试问题。我希望你只作为面试官回答。不要一次写出所有的问题。我希望你只对我进行采访。问我问题，等待我的回答。不要写解释。像面试官一样一个一个问我，等我回答。我的第一句话是“面试官你好”。

**Example 2:**
我想让你扮演讲故事的角色。你会想出引人入胜、富有想象力和吸引观众的有趣故事。它可以是童话故事、教育故事或任何其他类型的有潜力的故事以吸引人们的注意力和想象力。根据目标受众，您可以为您的讲故事环节选择特定的主题或主题，例如，如果是儿童，那么您可以谈论动物；如果是成人，那么基于历史的故事可能会更好地吸引他们等。我的第一个请求是我需要一个关于毅力的有趣故事。 | **Example 1:**
Assume the role of a marriage counselor. Develop a series of communication exercises for a couple who are experiencing difficulties in their relationship. These exercises should promote active listening, empathy, and effective expression of emotions. Your first assignment is to provide a set of three exercises that focus on resolving conflicts and rebuilding trust.

**Example 2:**
I want you to act as a travel agent. I will tell you my desired destination, travel dates, and budget, and it will be your job to suggest the best travel itinerary for me. Your recommendations should include the best transportation options, hotel accommodations, and any popular tourist attractions nearby. My first request is "I want to plan a trip to Tokyo for a week, with a budget of $2000. I want to explore the culture and food of the city." | +| Summarization | **Example 1:**
请简要总结概括以下段落材料。
当地时间29日，泰国卫生部通报，新增143名新冠肺炎确诊病例和1名死亡病例。截止到当地时间29日上午，泰国累计确诊病例1388例，其中泰国籍1172例，非泰国籍216例。死亡病例累计7例。（原题为《泰国新增143例新冠肺炎确诊病例累计确诊1388例》）

**Example 2:**
请简要总结概括以下段落材料。
近期，参与京雄高铁站站房建设的中铁十二局，因在施工过程中存在环境违法行为被雄安新区公开通报。通报发出后，引起社会广泛关注。近日，人民网记者从雄安新区相关部门及中铁十二局获悉，新区有关部门已经集中约谈了中铁十二局等24个参与雄安建设的项目单位。对于约谈内容和结果，中铁十二局有关宣传负责人回应：“具体内容不清楚，最好找雄安新区相关部门了解情况。”新区有关部门负责人表示，此前涉及的环境违法行为，中铁十二局已基本整改到位，但约谈内容和结果暂不公开，接下来，将按部就班推进环境治理工作。（原题为《雄安新区：中铁十二局涉环境违法已基本整改到位》） | **Example 1:**
The 21 year-old-woman was treated by paramedics after the kitchen fire in Botfield Road in Shifnal, Shropshire. West Mercia Police said it is treating Wednesday morning's incident as arson and are appealing for any witnesses to contact them.The 50-year-old man has been arrested on suspicion of arson with intent to endanger life. For more on this and other stories from Shropshire.
Please briefly summarize the above material within 20 words.

**Example 2:**
South Wales Police were called to a property in Heolgerrig, Merthyr Tydfil, at about 13:40 BST on Sunday. The child was airlifted to Prince Charles Hospital but died shortly afterwards. Police are investigating the circumstances surrounding the incident and have appealed for witnesses. The girl's family are being supported by specially trained officers.
Please briefly summarize the above material within 20 words. | + + +### Evaluation Metrics + +#### GPT Evaluation + +GPT evaluation uses GPT models to evaluate the prediction of different models and different pre-defined evaluation metrics are applied to different categories. The following table shows the 11 pre-defined evaluation metrics both in Chinese and English: + +| Evaluation Metric | Prompt Words | CoT(Chain-of-Thought) | +| :-------------------: | :----------------------------------------------------------- | :----------------------------------------------------------- | +| 语言组织
(Language organization) | 语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。

Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc. | 1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。
2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说
3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。
4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。
5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。
6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。

1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.
2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.
3. Determine if the answer is relevant to the question or topic and conveys a clear message.
4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.
5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.
6. Evaluate the linguistic organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good linguistic organization and 1 indicates very poor linguistic organization. | +| 切题
(Relevance) | 切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。

Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic. | 1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。
2. 阅读答案，确认答案是否直接回答了题目所问的问题。
3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。
4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。

1. Read the question to determine what the question asks and what aspects of the question need to be answered.
2. Read the answers to make sure that they directly answer the question asked.
3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.
4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all. | +| 创意性
(Creativity) | 创意性(1-5)：某些头脑风暴问题可能需要答案具有创意，提出新的思路。

Creativity (1-5): Some brainstorming questions may require answers that are creative and suggest new ideas. | 1. 仔细阅读所提供的头脑风暴问题，确保你理解问题的要点和背景。
2. 根据你的知识和经验，判断所提供的答案是否可行。如果答案不可行，则创意性评分可能会受到影响。
3. 考虑答案中是否包含新颖的想法或独特的思路。答案可能与已知的解决方案有所重叠，但仍然可以被认为是有创意的，只要它提供了新的角度或方法来解决问题。
4. 根据答案的创意性，给出一个1到5的评分。如果答案缺乏创意，则应给出一个较低的评分。如果答案具有创意并提供了新的思路，应给出一个较高的评分。

1. Read the provided brainstorming questions carefully to make sure you understand the gist and context of the questions.
2. Based on your knowledge and experience, determine if the answers provided are feasible. If the answer is not feasible, the creativity score may be affected.
3. Consider whether the answer contains novel ideas or unique thoughts. An answer may overlap with a known solution and still be considered creative, as long as it offers a new perspective or approach to the problem.
4. Give a score of 1 to 5 depending on the creativity of the answer. If the answer lacks creativity, a lower score should be given. If the answer is creative and provides a new idea, a higher score should be given. | +| 实用性
(Practicality) | 实用性(1-5)：某些头脑风暴问题可能需要答案提出实用的建议或解决方法。

Practicality (1-5): Some brainstorming questions may require answers to suggest practical suggestions or solutions. | 1. 仔细阅读所提供的头脑风暴问题，确保你理解问题的要点和背景。
2. 根据你的知识和经验，判断所提供的答案是否可行。如果答案不可行，则实用性评分可能会受到影响。
3. 考虑答案中提出的建议或解决方法是否实用并可行。答案可能看起来很好，但如果无法实现或应用，则实用性评分可能会受到影响。
4. 根据答案的实用性，给出一个1到5的评分。如果答案缺乏实用性，则应给出一个较低的评分。如果答案提出了实用的建议或解决方法，并且可以很好地解决问题，则应给出一个较高的评分。

1. Read the provided brainstorming questions carefully to make sure you understand the gist and context of the questions.
2. Based on your knowledge and experience, determine if the answers provided are feasible. If the answer is not feasible, the practicality score may be affected.
3. Consider whether the suggestions or solutions presented in the answer are practical and workable. The answer may look good, but if it cannot be implemented or applied, the practicality score may be affected.
4. Give a score of 1 to 5 depending on the practicality of the answer. If the answer lacks practicality, a lower score should be given. If the answer makes a practical suggestion or solution and solves the problem well, a higher score should be given. | +| 正确性
(Correctness) | 正确性(1-5)：正确性(1-5)：答案是否正确。

Correctness (1-5): whether the answer is correct or not. | 1. 仔细阅读题目，尝试自己回答该问题。
2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。

1. Read the question carefully and try to answer the question yourself.
2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be given. If the answer is completely incorrect, only 1 point is awarded. | +| 自然
(Naturalness) | 自然(1-5)：答案是否自然，并且符合问题给定的身份。

Naturalness (1-5): whether the answer is natural and fits the identity given by the question. | 1. 阅读题目，确定题目提供的身份信息。
2. 检查答案内容是否符合题目给定的身份。
3. 根据以上因素，对该回答的自然性进行打分，分数从1到5，其中1表示不自然，5表示非常自然，并符合问题给定的身份。

1. Read the question and determine the identity information provided in the question.
2. Check whether the content of the answer matches the identity given in the question.
3. Based on the above factors, score the naturalness of the response on a scale from 1 to 5, where 1 means unnatural and 5 means very natural and in accordance with the identity given in the question. | +| 参与感
(Engagingness) | 参与感(1-5)：答案是否对前面的对话内容做出了恰当的反应，是否理解对话的语境和背景。

Engagingness (1-5): whether the answer responds appropriately to the content of the preceding conversation and whether it understands the context and background of the conversation. | 1. 阅读题目，确定对话的语境和背景。
2. 检查答案是否充分理解对话的语境和背景，能否自然地融入到对话中而不显得突兀。
3. 根据以上因素，对该回答的参与感进行打分，分数从1到5，其中1表示没有参与感，5表示非常有参与感，并且恰当地理解了对话的语境和背景。

1. Read the questions to determine the context and background of the dialogue.
2. Check that the answer fully understands the context and background of the conversation and that it fits naturally into the conversation without seeming abrupt.
3. Based on the above factors, rate the response's engagement on a scale from 1 to 5, where 1 means not engaged and 5 means very engaged and appropriately understands the context and background of the conversation. | +| 合理性
(Reasonableness) | 合理性(1-5)：答案是否能够与前面的对话内容形成逻辑上的衔接，是否符合常理，能否在这个上下文中合理存在。

Reasonableness (1-5): Whether the answer can form a logical connection with the content of the previous dialogue, whether it is consistent with common sense, and whether it can reasonably exist in this context. | 1. 阅读题目，确定对话的主题以及问题期望的回答方向。
2. 判断答案是否能够与前面的对话内容形成逻辑上的衔接，是否符合常理，能否在这个上下文中合理存在。
3. 根据以上因素，对该回答的合理性进行打分，分数从1到5，其中1表示不合理，5表示非常合理，并且能够与前面的对话内容形成逻辑上的衔接，并符合常理。

1. Read the question and determine the topic of the conversation and the direction the question expects the answer to go.
2. Determine whether the answer can be logically connected to the preceding conversation, whether it makes common sense, and whether it can reasonably exist in this context.
3. Based on the above factors, rate the reasonableness of the answer on a scale from 1 to 5, where 1 means unreasonable and 5 means very reasonable and able to form a logical connection with the preceding dialogue content and consistent with common sense. | +| 多样性
(Diversity) | 多样性(1-5)：答案使用语言是否优美，具有有一定的创造性和想象力。然而，回答也应该保持合理和适度，不要过于夸张或离题。

Diversity (1-5): Whether the answers use beautiful language and have some creativity and imagination. However, answers should also be kept reasonable and moderate, not overly exaggerated or off-topic. | 1. 仔细阅读整个回答，确保完全理解回答所表达的内容和主题。
2. 在阅读回答的同时，注意语言的质量，例如措辞是否正确，语言是否生动等。
3. 检查回答的创造性和想象力，看看回答是否能够吸引人阅读下去。
4. 检查回答的合理性和适度，看看回答是否夸张或离题。5. 将多样性的评分打分在1到5之间，5分表示回答的质量很好，能够吸引人阅读，1分表示回答的内容生硬或者有离题的问题。

1. Read the entire response carefully to ensure that you fully understand the content and theme expressed in the response.
2. While reading the response, pay attention to the quality of the language, such as whether the wording is correct and the language is vivid.
3. Check the creativity and imagination of the response to see if the response is engaging to read on.
4. Check the reasonableness and appropriateness of the responses to see if the responses are exaggerated or off-topic.
5. Rate the diversity on a scale of 1 to 5, with a 5 indicating a good quality response that is engaging to read and a 1 indicating a raw response or a question that is off-topic. | +| 保真度
(Fidelity) | 保真度(1-5)：答案是否能够严格遵守角色的设定回答给定的请求。

Fidelity (1-5): whether the answer is able to answer the given request in strict compliance with the role setting. | 1. 仔细阅读问题，了解角色在问题中的设定和表现，包括职业、背景、观点、性格等方面。
阅读题目的请求，确认回答请求时需要注意的细节。
3. 对比提供的回答与该角色的设定，评估回答是否能够严格遵守角色的设定。
4. 结合以上评估结果给出保真度的评分，范围从1到5分，其中1分表示回答与角色设定完全不符，5分表示回答完全符合角色设定且满足给定请求。

1. Read the question carefully to understand how the character is set up and represented in the question, including aspects such as occupation, background, point of view, and personality.
2. Read the question's request and confirm the details that need to be taken into account when answering the request.
3. Compare the provided answer with the setting of the role and assess whether the answer can strictly adhere to the setting of the role.
4. Combine the results of the above assessment to give a fidelity score ranging from 1 to 5, where a score of 1 means that the response does not match the persona at all, and a score of 5 means that the response fully complies with the persona and satisfies the given request. | +| 简明扼要
(Conciseness) | 简明扼要(1-5)：答案是否简明扼要，没有冗余内容。

Conciseness (1-5): answers should be concise and without redundant content. | 1. 阅读题目，提取出材料的重点。
2. 阅读该总结，并注意其中的主要观点和信息。
3. 评估总结的长度。一个简明扼要的总结通常应该在几句话或几段文字内传达关键信息，而不是冗长的段落或文章。
4. 检查总结是否包含与主要观点无关的信息或冗余信息。
5. 确定总结涵盖了材料中的关键信息，并且没有忽略任何重要细节。
6. 给总结打出1-5的分数，其中5表示总结简明扼要，没有冗余内容，而1表示总结冗长或包含不必要的信息，难以理解或记忆。根据您的判断，打出适当的得分。

1. Read the title and extract the main points of the material.
2. Read the summary and note the main ideas and messages in it.
3. Assess the length of the summary. A concise summary should usually convey key information within a few sentences or paragraphs, rather than lengthy paragraphs or essays.
4. Check that the summary does not contain information that is not relevant to the main ideas or that is redundant.
5. Make sure that the summary covers the key information in the material and that no important details have been omitted.
6. Rate the summary on a scale of 1-5, where 5 means the summary is concise and free of redundancy, and 1 means the summary is lengthy or contains unnecessary information that is difficult to understand or remember. Based on your judgment, assign the appropriate score. | + +GPT models evaluate the quality of model predictions based on the given prompt words and gives a score between 1-5. + +> **NOTE 1:** Even for the same metric, the details of its prompt words and CoT(Chain-of-Thought) can differ based on which category you want to evaluate. For example, prompt words for metric `correctness` showed here is "Whether the answer is correct or not."(this is for category `classification`), but for category `extraction`, prompt words can be "Answers should extract the required information accurately and should not contain any incorrect or misleading information." You can find all the prompt words and CoT(Chain-of-Thought) in `prompt/evaluation_prompt`. + +> **NOTE 2:** To add customized metrics, you can refer to [FAQ](#faq). + +#### Automatic Evaluation + +Automated metrics evaluate the capability of a model by comparing model predictions with reference answers. +There are two ways to obtain reference answers: + +* For instruction coming from human-designed problems, the reference answers are generated by GPT-3.5, such as roleplay, chat. +* For instruction related with classic NLP problems, the reference answers are collected from open-sourced dataset with target answers, such as classification, extraction, summarization. + +There are 6 types of automatic evaluation metrics listed in the table below: + +| Automatic Evaluation Metric | Description | +| :---------------------------------: | :----------------------------------------------------------- | +| BLEU-n | Measure the accuracy between prediction and reference.
BLEU-1 (Unigram) evaluates accuracy in word level.
BLEU-n (n-gram) evaluate the fluency in sentence level. | +| ROUGE | ROUGE-N measures the number of matching n-grams between prediction and reference.
ROUGE-L measures the number of matching longest common subsequence (LCS) between prediction and reference. | +| Distinct | Measure the diversity of generation text by counting the unique n-grams. | +| BERTScore | Measure the semantic similarity between tokens of predictions and references with BERT. | +| Precision
Recall
F1 Score | Measure the number of overlaps between prediction and reference (design for classification and extraction categories). | +| CHRF | Measure the similarity of character n-grams between prediction and reference. | + +#### UniEval Evaluation + +UniEval converts all evaluation tasks of different dimensions(metrics) into Boolean QA problems and utilize the model to answer with “Yes” or “No”. Compared with similarity-based metrics such as ROUGE and BLEU, UniEval can achieve a more comprehensive evaluation. In addition, UniEval also demonstrates its ability to transfer to unseen dimensions and tasks. + +In our evaluation pipeline, two pre-trained UniEval evaluators are used. One is [unieval-sum](https://huggingface.co/MingZhong/unieval-sum) and the other is [unieval-dialog](https://huggingface.co/MingZhong/unieval-dialog). The two models can be used for the 3 tasks, `summarization`, `dialogue` and `data2text`. Each task has different evaluation dimensions. + +| UniEval Model | Task | Dimension(Metric) | +| :------------: | :----------------- | :--- | +| unieval-sum | summarization | coherence: whether the summary is coherent
consistency: whether the claim is consistent with the given document
fluency: whether the paragraph is fluent
relevance: whether the summary is relevant to the reference | +| unieval-sum | data2text | naturalness: whether the utterance is fluent
informativeness: whether the utterance is informative according to the reference | +| unieval-dialog | dialogue | naturalness: whether the response is natural in the dialogue
coherence: whether the response is coherent in the dialogue history
understandability: whether the response is understandable in the dialogue | + +> **NOTE 1:** Task "data2text" uses the same model as task "summarization". + +> **NOTE 2:** In UniEval paper, the `unieval-sum` model demonstrates the best transfer ability and so you can evaluate your customized metric with this model. Details of adding customized metrics can be found in [FAQ](#faq). + +> **NOTE 3:** We consider not including all metrics provided in UniEval in our pipeline because the data structure and content of the instructions we want to evaluate are not suitable for direct use of some UniEval metrics. + +## Evaluation Process + +### Data Format + +#### Target Answers / Predictions + +A JSON file contains one list. Each element in the list is a target answer / prediction record for one instruction / question. +An element should have the following fields: + +* `category` (str, compulsory): The category of the instruction / question. +* `instruction` (str, compulsory): The instruction / question for the LLM. +* `input` (str, optional): The additional context of the instruction / question. +* `output` (str, optional): The sample output of the instruction (default: GPT-3.5). +* `target` (str, optional): The target answer for the instruction. +* `id` (int, compulsory): The ID of the instruction / question. + +If the `input` has a target answer, the `output` can be empty. Otherwise, we generate answers from GPT-3.5 as the `output`, and the `target` field is empty. + +Example: + +```json +[ + { + "category": "brainstorming", + "instruction": "请介绍一下人工智能的多个领域。", + "input": "", + "output": "{GPT-3.5 Answers}", + "target": "", + "id": 1 + }, + { + "category": "classification", + "instruction": "新闻标题：为什么电影《倩女幽魂》中燕赤霞一个道士却拿着金刚经？请根据新闻标题判断新闻所属的分类，你需要从文化，娱乐，体育，财经，房产，教育，科技，旅游，游戏，军事这十类中选择一个答案。", + "input": "", + "output": "", + "target": "{target answer}", + "id": 2 + } +] +``` + +#### Model Answers / Predictions + +A JSON file contains one list. Each element in the list is a model answer / prediction record for one instruction / question. + +An element should have the following fields: + +* `category` (str, compulsory): The category of the instruction / question. +* `instruction` (str, compulsory): The instruction / question for the LLM. +* `input` (str, optional): The additional context of the instruction / question. +* `output` (str, compulsory): The output from the LLM. +* `target` (str, optional): The target answer for the instruction. +* `id` (int, compulsory): The ID of the instruction / question. + +Example: + +```json +[ + { + "category": "brainstorming", + "instruction": "请介绍一下人工智能的多个领域。", + "input": "", + "output": "{Model Answers / Predictions}", + "target": "", + "id": 1 + }, + { + "category": "classification", + "instruction": "新闻标题：为什么电影《倩女幽魂》中燕赤霞一个道士却拿着金刚经？请根据新闻标题判断新闻所属的分类，你需要从文化，娱乐，体育，财经，房产，教育，科技，旅游，游戏，军事这十类中选择一个答案。", + "input": "", + "output": "{Model Answers / Predictions}", + "target": "{target answer}", + "id": 2 + } +] +``` + +### Prompt + +#### Battle Prompt + +The following is the Chinese battle prompt. In the battle prompt, the question and answers from two different models are fed into the prompt template. You can find example battle prompt files for Chinese and English in `prompt/battle_prompt`. + +```json +{ + "id": 1, + "system_prompt": "你是一个检查回答质量的好助手。", + "prompt_template": "[问题]\n{question}\n\n[1号AI助手的答案]\n{answer_1}\n\n[1号AI助手答案终止]\n\n[2号AI助手的答 案]\n{answer_2}\n\n[2号AI助手答案终止]\n\n[要求]\n{prompt}\n\n", + "prompt": "我们需要你评价这两个AI助手回答的性能。\n请对他们的回答的有用性、相关性、准确性、详细程度进行评分。每个AI助手都会得到一个1到10分的总分，分数越高表示整体表现越好。\n请首先输出一行，该行只包含两个数值，分别表示1号和2号AI助手的分数。这两个分数之间要有一个空格。在随后的一行中，请对你的评价作出全面的解释，避免任何潜在的偏见，并确保AI助手回答的顺序不会影响您的判断。" +} +``` + +#### Evaluation Prompt + +The following is an example of a Chinese GPT evaluation prompt. In an evaluation prompt, you should define your metrics in `metrics` and provide CoT(Chain-of-Thought) in `CoT`. You can find example evaluation prompt files for Chinese and English in `prompt/evaluation_prompt`. + +```json +{ + "brainstorming": { + "id": 1, + "category": "brainstorming", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：" + }, + "prompt": "你是一个好助手。请你为下面“头脑风暴”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + } +} +``` + +`"metrics"`: the metrics that can be used in GPT evaluation. This field determines which metrics can be added to your config file. + +`"CoT"`: evaluation steps you prompt to GPT models for each metric defined in `"metrics"`. + +### Evaluation + +#### Configuration + +The following is an example of a Chinese config file. The configuration file can control how the pipeline evaluates the model. You need to specify GPT evaluation metrics, automatic metrics and UniEval metrics in key `GPT`, `Metrics` and `UniEval`(English only). You can find an example English config file in `config`. + +```json +{ + "language": "en", + "path_for_UniEval": { + "summarization": "path to unieval-sum model", + "dialogue": "path to unieval-dialog model", + "data2text": "path to unieval-sum model" + }, + "category": { + "brainstorming": { + "GPT": ["relevance", "creativity", "practicality", "reasonableness"], + "Metrics": ["Distinct"], + "UniEval": ["summarization-fluency", "data2text-naturalness", "data2text-informativeness"] + }, + "chat": { + "GPT": [ "relevance", "naturalness", "engagingness", "reasonableness"], + "Metrics": ["Distinct"], + "UniEval": ["dialogue-naturalness", "dialogue-coherence", "dialogue-understandability"] + } + } +} +``` + +`"language"`: the language used to evaluate the model capability. We only support Chinese `"cn"` for now. + +`"path_for_UniEval"`: path to the UniEval model. + +`"category"`: the category/categories needed to evaluate the model capability. + +`"GPT"`: the metrics you want to use for GPT evaluation. + +`"Metrics"`: the metrics you want to use for automatic metrics evaluation. + +`"UniEval"`: the metrics you want to use for UniEval metrics evaluation. The metric has to be in the `"{task}-{metric}"` format because different tasks have same metrics such as naturalness and coherence. + +You can remove the key such as `"Metrics"` to skip evaluating answers using its corresponding evaluation metrics. + +You can create your config file based on available settings listed in following table. + +| "category" | "GPT" | "Metrics" | "UniEval" | +| :--------------: | :---------------------: | :---------: | :--------------------------: | +| "brainstorming" | "language organization" | "BLEU" | "dialogue-naturalness" | +| "chat" | "relevance" | "ROUGE" | "dialogue-coherence" | +| "classification" | "creativity" | "Distinct" | "dialogue-understandability" | +| "closed_qa" | "practicality" | "BERTScore" | "data2text-naturalness" | +| "extraction" | "correctness" | "Precision" | "data2text-informativeness" | +| "generation" | "naturalness" | "Recall" | "summarization-coherence" | +| "open_qa" | "engagingness" | "F1 score" | "summarization-consistency" | +| "rewriting" | "reasonableness" | "CHRF" | "summarization-fluency" | +| "roleplay" | "diversity" | | "summarization-relevance" | +| "summarization" | "fidelity" | | | +| | "conciseness" | | | + +> **NOTE:** For categories which don't have standard answers such as `brainstorming`, you should avoid using automatic metrics such as `BLEU` and `ROUGE` which are based on similarity measures and you should use `Distinct` instead in your config file. + +#### Evaluate + +After setting the configuration file, you can evaluate the model using `eval.py`. If you want to make comparisons between answers of two different models, you should specify two answer files in the argument `answer_file_list` and two model names in the argument `model_name_list`. If you want to evaluate one answer file, the length of both `answer_file_list` and `model_name_list` should be 1 and the program will perform evaluation using automatic metrics and GPT models. + +An example script is provided as follows: + +```shell +python eval.py \ + --config_file "path to the config file" \ + --battle_prompt_file "path to the prompt file for battle" \ + --gpt_evaluation_prompt_file "path to the prompt file for gpt evaluation" \ + --target_file "path to the target answer file" \ + --answer_file_list "path to the answer files of at most 2 models" \ + --model_name_list "the names of at most 2 models" \ + --gpt_model "which GPT model to use for evaluation" \ + --save_path "path to save results" \ + --openai_key "your openai key" \ +``` + +If you want GPT evaluation with reference, you can add an argument `--gpt_with_reference`. + +## FAQ + +
How can I add a new GPT evaluation metric? + +For example, if you want to add a new metric `persuasiveness` into category `brainstorming`, you should add the metric definition and its corresponding CoT(Chain-of-thought) in the evaluation prompt file in `prompt/evaluation_promt`. The CoT can be generated using ChatGPT. You can prompt ChatGPT to generate evaluation steps for the new metric. + +```json +{ + "brainstorming": { + "id": 1, + "category": "brainstorming", + "metrics": { + "persuasiveness": "persuasiveness(1-5)：a short description for persuasiveness" + }, + "CoT": { + "persuasiveness": "CoT for persuasiveness\n\npersuasiveness：" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"brainstorming\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + } +} +``` + +
+ +
How can I add a new UniEval evaluation metric? + +For example, if you want to add a new metric `persuasiveness` into task `data2text`, you should add a Boolean QA question about the metric in function `add_question` in `unieval/utils.py`. Please do note that how effectively the model would evaluate this metric is unknown and you may need some experiments to test whether the model is capable of evaluating this metric. + +```python +if task == 'data2text': + if dimension == 'persuasiveness': + cur_input = 'question: Is this a persuasive utterence utterance: ' + output[i] +``` + +
+ +## To Do + +- [x] Add evaluation for English capability +- [x] Support UniEval +- [x] Support GPT-4 evaluation +- [x] Support GPT evaluation with reference + +## Citations + +```bibtex +@misc{vicuna2023, + title = {Vicuna: An Open-Source Chatbot Impressing GPT-4 with 90\%* ChatGPT Quality}, + url = {https://vicuna.lmsys.org}, + author = {Chiang, Wei-Lin and Li, Zhuohan and Lin, Zi and Sheng, Ying and Wu, Zhanghao and Zhang, Hao and Zheng, Lianmin and Zhuang, Siyuan and Zhuang, Yonghao and Gonzalez, Joseph E. and Stoica, Ion and Xing, Eric P.}, + month = {March}, + year = {2023} +} + +@misc{liu2023geval, + title={G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment}, + author={Yang Liu and Dan Iter and Yichong Xu and Shuohang Wang and Ruochen Xu and Chenguang Zhu}, + year={2023}, + eprint={2303.16634}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@misc{zhong2022unified, + title={Towards a Unified Multi-Dimensional Evaluator for Text Generation}, + author={Ming Zhong and Yang Liu and Da Yin and Yuning Mao and Yizhu Jiao and Pengfei Liu and Chenguang Zhu and Heng Ji and Jiawei Han}, + year={2022}, + eprint={2210.07197}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} +``` diff --git a/applications/Chat/evaluate/config/config_cn.json b/applications/Chat/evaluate/config/config_cn.json new file mode 100644 index 000000000000..dffb66f6c3be --- /dev/null +++ b/applications/Chat/evaluate/config/config_cn.json @@ -0,0 +1,127 @@ +{ + "language": "cn", + "category": { + "brainstorming": { + "GPT": [ + "language organization", + "relevance", + "creativity", + "practicality", + "reasonableness" + ], + "Metrics": [ + "Distinct" + ] + }, + "chat": { + "GPT": [ + "language organization", + "relevance", + "naturalness", + "engagingness", + "reasonableness" + ], + "Metrics": [ + "Distinct" + ] + }, + "classification": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Precision", + "Recall", + "F1 score", + "CHRF" + ] + }, + "closed_qa": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore", + "CHRF" + ] + }, + "extraction": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Precision", + "Recall", + "F1 score", + "CHRF" + ] + }, + "generation": { + "GPT": [ + "language organization", + "relevance", + "diversity" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore" + ] + }, + "open_qa": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Distinct" + ] + }, + "rewriting": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore" + ] + }, + "roleplay": { + "GPT": [ + "language organization", + "relevance", + "fidelity", + "creativity" + ], + "Metrics": [ + "Distinct" + ] + }, + "summarization": { + "GPT": [ + "language organization", + "relevance", + "correctness", + "conciseness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore", + "CHRF" + ] + } + } +} diff --git a/applications/Chat/evaluate/config/config_en.json b/applications/Chat/evaluate/config/config_en.json new file mode 100644 index 000000000000..5238bd19f67e --- /dev/null +++ b/applications/Chat/evaluate/config/config_en.json @@ -0,0 +1,188 @@ +{ + "language": "en", + "path_for_UniEval": { + "summarization": "path to unieval-sum", + "dialogue": "path to unieval-dialog", + "data2text": "path to unieval-sum" + }, + "category": { + "brainstorming": { + "GPT": [ + "language organization", + "relevance", + "creativity", + "practicality", + "reasonableness" + ], + "Metrics": [ + "Distinct" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "chat": { + "GPT": [ + "language organization", + "relevance", + "naturalness", + "engagingness", + "reasonableness" + ], + "Metrics": [ + "Distinct" + ], + "UniEval": [ + "summarization-fluency", + "dialogue-naturalness", + "dialogue-coherence", + "dialogue-understandability", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "classification": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Precision", + "Recall", + "F1 score", + "CHRF" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "closed_qa": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore", + "CHRF" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "extraction": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Precision", + "Recall", + "F1 score", + "CHRF" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "generation": { + "GPT": [ + "language organization", + "relevance", + "diversity" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "open_qa": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "Distinct" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "rewriting": { + "GPT": [ + "language organization", + "relevance", + "correctness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "roleplay": { + "GPT": [ + "language organization", + "relevance", + "fidelity", + "creativity" + ], + "Metrics": [ + "Distinct" + ], + "UniEval": [ + "summarization-fluency", + "data2text-naturalness", + "data2text-informativeness" + ] + }, + "summarization": { + "GPT": [ + "language organization", + "relevance", + "correctness", + "conciseness" + ], + "Metrics": [ + "BLEU", + "ROUGE", + "BERTScore", + "CHRF" + ], + "UniEval": [ + "summarization-coherence", + "summarization-consistency", + "summarization-fluency", + "summarization-relevance", + "data2text-naturalness", + "data2text-informativeness" + ] + } + } +} diff --git a/applications/Chat/evaluate/eval.py b/applications/Chat/evaluate/eval.py new file mode 100644 index 000000000000..e3fe0e9e091b --- /dev/null +++ b/applications/Chat/evaluate/eval.py @@ -0,0 +1,112 @@ +import argparse +import json +import os + +import openai +from evaluator import Evaluator +from utils import jload + + +def main(args): + assert len(args.answer_file_list) == len( + args.model_name_list), "The number of answer files and model names should be equal!" + + # load config + config = jload(args.config_file) + + if config["language"] in ["cn", "en"]: + # get metric settings for all categories + metrics_per_category = {} + for category in config["category"].keys(): + metrics_all = {} + for metric_type, metrics in config["category"][category].items(): + metrics_all[metric_type] = metrics + metrics_per_category[category] = metrics_all + + battle_prompt = None + if args.battle_prompt_file: + battle_prompt = jload(args.battle_prompt_file) + + gpt_evaluation_prompt = None + if args.gpt_evaluation_prompt_file: + gpt_evaluation_prompt = jload(args.gpt_evaluation_prompt_file) + + if len(args.model_name_list) == 2 and not battle_prompt: + raise Exception("No prompt file for battle provided. Please specify the prompt file for battle!") + + if len(args.model_name_list) == 1 and not gpt_evaluation_prompt: + raise Exception( + "No prompt file for gpt evaluation provided. Please specify the prompt file for gpt evaluation!") + + if args.gpt_model == "text-davinci-003" and args.gpt_with_reference: + raise Exception( + "GPT evaluation with reference is not supported for text-davinci-003. You should specify chat models such as gpt-3.5-turbo or gpt-4." + ) + + # initialize evaluator + evaluator = Evaluator(metrics_per_category, battle_prompt, gpt_evaluation_prompt, args.gpt_model, + config["language"], config.get("path_for_UniEval", None), args.gpt_with_reference) + if len(args.model_name_list) == 2: + answers1 = jload(args.answer_file_list[0]) + answers2 = jload(args.answer_file_list[1]) + + assert len(answers1) == len(answers2), "The number of answers for two models should be equal!" + + evaluator.battle(answers1=answers1, answers2=answers2) + evaluator.save(args.save_path, args.model_name_list) + elif len(args.model_name_list) == 1: + targets = jload(args.target_file) + answers = jload(args.answer_file_list[0]) + + assert len(targets) == len(answers), "The number of target answers and model answers should be equal!" + + evaluator.evaluate(answers=answers, targets=targets) + evaluator.save(args.save_path, args.model_name_list) + else: + raise ValueError("Unsupported number of answer files and model names!") + else: + raise ValueError(f'Unsupported language {config["language"]}!') + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='ColossalAI LLM evaluation pipeline.') + parser.add_argument('--config_file', + type=str, + default=None, + required=True, + help='path to the file of target results') + parser.add_argument('--battle_prompt_file', type=str, default=None, help='path to the prompt file for battle') + parser.add_argument('--gpt_evaluation_prompt_file', + type=str, + default=None, + help='path to the prompt file for gpt evaluation') + parser.add_argument('--target_file', type=str, default=None, help='path to the target answer (ground truth) file') + parser.add_argument('--answer_file_list', + type=str, + nargs='+', + default=[], + required=True, + help='path to the answer files of at most 2 models') + parser.add_argument('--model_name_list', + type=str, + nargs='+', + default=[], + required=True, + help='the names of at most 2 models') + parser.add_argument('--gpt_model', + default="gpt-3.5-turbo", + choices=["text-davinci-003", "gpt-3.5-turbo", "gpt-4"], + help='which GPT model to use for evaluation') + parser.add_argument('--gpt_with_reference', + default=False, + action="store_true", + help='whether to include reference answer in gpt evaluation') + parser.add_argument('--save_path', type=str, default="results", help='path to save evaluation results') + parser.add_argument('--openai_key', type=str, default=None, required=True, help='Your openai key') + args = parser.parse_args() + + if args.openai_key is not None: + os.environ["OPENAI_API_KEY"] = args.openai_key + openai.api_key = os.getenv("OPENAI_API_KEY") + + main(args) diff --git a/applications/Chat/evaluate/eval.sh b/applications/Chat/evaluate/eval.sh new file mode 100755 index 000000000000..f5729e6ee5c7 --- /dev/null +++ b/applications/Chat/evaluate/eval.sh @@ -0,0 +1,9 @@ +python eval.py \ + --config_file "path to the config file" \ + --battle_prompt_file "path to the prompt file for battle" \ + --gpt_evaluation_prompt_file "path to the prompt file for gpt evaluation" \ + --target_file "path to the target answer file" \ + --answer_file_list "path to the answer files of at most 2 models" \ + --model_name_list "the names of at most 2 models" \ + --save_path "path to save results" \ + --openai_key "your openai key" \ diff --git a/applications/Chat/evaluate/evaluator.py b/applications/Chat/evaluate/evaluator.py new file mode 100644 index 000000000000..3dd5fd6f2f23 --- /dev/null +++ b/applications/Chat/evaluate/evaluator.py @@ -0,0 +1,219 @@ +import os +from typing import Any, Dict, List + +import gpt_evaluate +import metrics +import pandas as pd +import unieval +from utils import analyze_automatic_results, get_data_per_category, save_automatic_results + + +class Evaluator(object): + """ + A class named Evaluator includes GPT-3.5/GPT-4 evaluation + and automatic evaluation + + """ + + def __init__(self, params: Dict[str, Any], battle_prompt: Dict[str, Any], gpt_evaluation_prompt: Dict[str, Any], + gpt_model: str, language: str, path_for_UniEval: Dict[str, str], gpt_with_reference: bool) -> None: + self.params = params + self.battle_prompt = battle_prompt + self.gpt_evaluation_prompt = gpt_evaluation_prompt + self.gpt_model = gpt_model + self.language = language + self.path_for_UniEval = path_for_UniEval + self.gpt_with_reference = gpt_with_reference + self.automatic_metric_stats = dict() + self.unieval_metric_stats = dict() + self.gpt_evaluation_results = dict() + self.battle_results = [] + + def battle(self, answers1: List[Dict], answers2: List[Dict]) -> None: + """ + Comparison between two models using GPT-4 as the reviewer. + """ + + self.battle_results = gpt_evaluate.battle(answers1, answers2, self.battle_prompt) + + def evaluate(self, answers: List[Dict], targets: List[Dict]) -> None: + """ + A comprehensive evaluation of the answers from the model. + The function evaluates the model's performance from different perspectives + using GPT-3.5, GPT-4, and off-the-shelf evaluation metrics. + + The metrics will be decided by the config file. + + """ + + def switch(metric, language): + if metric == "BLEU": + return metrics.bleu_score(preds=predicts_list, targets=targets_list, language=language) + elif metric == "ROUGE": + return metrics.rouge_score(preds=predicts_list, targets=targets_list, language=language) + elif metric == "Distinct": + return metrics.distinct_score(preds=predicts_list, language=language) + elif metric == "BERTScore": + return metrics.bert_score(preds=predicts_list, targets=targets_list, language=language) + elif metric == "Precision": + return metrics.precision(preds=predicts_list, targets=targets_list, language=language) + elif metric == "Recall": + return metrics.recall(preds=predicts_list, targets=targets_list, language=language) + elif metric == "F1 score": + return metrics.F1_score(preds=predicts_list, targets=targets_list, language=language) + elif metric == "CHRF": + return metrics.chrf_score(preds=predicts_list, targets=targets_list, language=language) + else: + raise ValueError(f"Unexpected metric") + + answers_per_category = get_data_per_category(answers, list(self.params.keys())) + targets_per_category = get_data_per_category(targets, list(self.params.keys())) + + # automatic evaluation + for category in self.params: + if len(answers_per_category[category]) == 0: + print(f"Category {category} specified in your config doesn't have corresponding answers!") + continue + + if self.params[category].get("Metrics", None) is None: + continue + + category_metrics = self.params[category]["Metrics"] + self.automatic_metric_stats[category] = {} + + targets_list = [ + target["target"] if target["target"] else target["output"] for target in targets_per_category[category] + ] + predicts_list = [answer["output"] for answer in answers_per_category[category]] + + for metric in category_metrics: + self.automatic_metric_stats[category].update(switch(metric=metric, language=self.language)) + + # UniEval evaluation + # self.unieval_metric_stats's key is "task" instead of "category". + # Iterating "task" first will avoid repeated loading models because one task corresponds to one UniEval model. + # If key is "category", different models will be loaded for multiple times across categories because the user may require different task(models) to evaluate one category. + for category in self.params: + if len(answers_per_category[category]) == 0: + print(f"Category {category} specified in your config doesn't have corresponding answers!") + continue + + if self.params[category].get("UniEval", None) is None: + continue + + if self.params[category]["UniEval"] and self.language == "cn": + raise Exception( + "UniEval doesn't support Chinese! Please remove UniEval config in your Chinese config file.") + + category_metrics = self.params[category]["UniEval"] + + for task, metric in [tuple(category_metric.split("-")) for category_metric in category_metrics]: + if self.unieval_metric_stats.get(task, None) is None: + self.unieval_metric_stats[task] = {category: {metric: 0}} + elif self.unieval_metric_stats[task].get(category, None) is None: + self.unieval_metric_stats[task][category] = {metric: 0} + else: + self.unieval_metric_stats[task][category][metric] = 0 + + for task in self.unieval_metric_stats: + if self.path_for_UniEval is None: + raise Exception(f"Please specify the path for UniEval model in the config file!") + + if self.path_for_UniEval.get(task, None) is None: + raise Exception(f"Please specify the model path for task {task} in the config file!") + + print(f"Load UniEval model for task {task}.") + + uni_evaluator = unieval.get_evaluator(task, model_name_or_path=self.path_for_UniEval[task]) + for category in self.unieval_metric_stats[task]: + targets_list = [ + target["target"] if target["target"] else target["output"] + for target in targets_per_category[category] + ] + predicts_list = [answer["output"] for answer in answers_per_category[category]] + sources_list = [answer["instruction"] + answer["input"] for answer in answers_per_category[category]] + + data = unieval.convert_data_to_unieval_format(predicts_list, sources_list, targets_list) + scores = uni_evaluator.evaluate(data, + category, + dims=list(self.unieval_metric_stats[task][category].keys()), + overall=False) + avg_scores = unieval.calculate_average_score(scores) + + self.unieval_metric_stats[task][category].update(avg_scores) + + # gpt evaluation + for category in self.params: + if len(answers_per_category[category]) == 0: + print(f"Category {category} specified in your config doesn't have corresponding answers!") + continue + + if self.params[category].get("GPT", None) is None: + continue + + category_metrics = self.params[category]["GPT"] + + prompt = self.gpt_evaluation_prompt.get(category, None) + if prompt is None: + print(f"No prompt for category {category}! Use prompt for category general now.") + prompt = self.gpt_evaluation_prompt["general"] + + self.gpt_evaluation_results[category] = gpt_evaluate.evaluate( + answers_per_category[category], + prompt, + category_metrics, + category, + self.gpt_model, + self.language, + references=targets_per_category[category] if self.gpt_with_reference else None) + + def save(self, path: str, model_name_list: List[str]) -> None: + """ + Save evaluation results of GPT-3.5, GPT-4, and off-the-shelf evaluation metrics. + + """ + + if len(model_name_list) == 2: + save_path = os.path.join(path, "gpt_evaluate", "battle_results") + gpt_evaluate.save_battle_results(self.battle_results, model_name_list[0], model_name_list[1], save_path) + else: + if self.automatic_metric_stats: + # Save evaluation results for automatic metrics + automatic_base_save_path = os.path.join(path, "automatic_results") + automatic_results_save_path = os.path.join(automatic_base_save_path, "evaluation_results") + + save_automatic_results(model_name_list[0], self.automatic_metric_stats, automatic_results_save_path) + + # Save charts and csv. + automatic_analyses_save_path = os.path.join(automatic_base_save_path, "evaluation_analyses") + analyze_automatic_results(automatic_results_save_path, automatic_analyses_save_path) + + if self.unieval_metric_stats: + # Save evaluation results for UniEval metrics + unieval_base_save_path = os.path.join(path, "unieval_results") + unieval_results_save_path = os.path.join(unieval_base_save_path, "evaluation_results") + + unieval.save_unieval_results(model_name_list[0], self.unieval_metric_stats, unieval_results_save_path) + + # Save charts and csv. + unieval_analyses_save_path = os.path.join(unieval_base_save_path, "evaluation_analyses") + unieval.analyze_unieval_results(unieval_results_save_path, unieval_analyses_save_path) + + if self.gpt_evaluation_results: + # Save evaluation results for GPT evaluation metrics. + gpt_base_save_path = os.path.join(path, "gpt_evaluate", "gpt_evaluate_results") + gpt_evaluation_results_save_path = os.path.join(gpt_base_save_path, "evaluation_results") + + all_evaluations = gpt_evaluate.save_gpt_evaluation_results(model_name_list[0], + self.gpt_evaluation_results, + gpt_evaluation_results_save_path) + + # Start to calculate scores and save statistics. + gpt_evaluation_statistics_save_path = os.path.join(gpt_base_save_path, "evaluation_statistics") + gpt_evaluate.save_gpt_evaluation_statistics(model_name_list[0], all_evaluations, + gpt_evaluation_statistics_save_path) + + # Save charts and csv. + gpt_evaluation_analyses_save_path = os.path.join(gpt_base_save_path, "evaluation_analyses") + gpt_evaluate.analyze_gpt_evaluation_statistics(gpt_evaluation_statistics_save_path, + gpt_evaluation_analyses_save_path) diff --git a/applications/Chat/evaluate/gpt_evaluate.py b/applications/Chat/evaluate/gpt_evaluate.py new file mode 100644 index 000000000000..f8cfb8d0f7e5 --- /dev/null +++ b/applications/Chat/evaluate/gpt_evaluate.py @@ -0,0 +1,772 @@ +import concurrent.futures +import os +import re +import time +from copy import deepcopy +from typing import Any, Dict, List + +import matplotlib.pyplot as plt +import numpy as np +import openai +import pandas as pd +import seaborn as sns +import tqdm +from utils import jdump, jload + +ref_step_template = { + "en": + "Now please compare the answer with the {adjective} answer, determine whether the answer is able to achieve the same level of {metric}.\n\n", + "cn": + "请比较答案与上面的{adjective}答案，确定答案是否可以达到与该{adjective}答案同样水平的{metric}。\n\n" +} + +ref_answer_template_general = { + "en": "\nAn example answer with good quality is as follows:\n\n{answer}\n\n", + "cn": "\n一个优质的示例答案如下：\n\n{answer}\n\n" +} + +ref_answer_template_correctness = { + "en": "\nA correct answer is as follows:\n\n{answer}\n\n", + "cn": "\n标准答案如下：\n\n{answer}\n\n" +} + + +def get_battle_result(sys_prompt: str, user_prompt: str, id: int, max_tokens: int = 2048) -> Dict[str, Any]: + """ + Get battle evaluation from GPT-4. + + Args: + sys_prompt: prompt for the system. + user_prompt: prompt for the user. + id: id of the answers for comparison. + max_tokens: the maximum number of tokens to generate in the chat completion. + + Returns: + An evaluation of one comparison. + """ + + MAX_API_RETRY = 3 + for _ in range(MAX_API_RETRY): + try: + response = openai.ChatCompletion.create( + model="gpt-4", + messages=[ + { + "role": "system", + "content": sys_prompt + }, + { + "role": "user", + "content": user_prompt, + }, + ], + temperature=0.2, + max_tokens=max_tokens, + ) + evaluation = response["choices"][0]["message"]["content"] + return {"evaluation": evaluation, "id": id} + except Exception as e: + print(e) + time.sleep(1) + print(f"Evaluation {id} failed after {MAX_API_RETRY} retries.") + return {"evaluation": "", "id": id} + + +def parse_battle_score(evaluation: str) -> List[float]: + """ + Parse evaluation from GPT-4 and get the scores of model 1 and 2. + + Args: + evaluation: evaluation from GPT-4. + + Returns: + A score pair of two different model answers. + """ + + try: + pattern = re.compile("([0-9]|10) out of 10") + sp = re.findall(pattern, evaluation) + if len(re.findall(pattern, evaluation)) == 2: + return [float(sp[0]), float(sp[1])] + + pattern = re.compile("a score of ([0-9]|10)") + sp = re.findall(pattern, evaluation) + if len(re.findall(pattern, evaluation)) == 2: + return [float(sp[0]), float(sp[1])] + + pattern = re.compile("([0-9]|10)/10") + sp = re.findall(pattern, evaluation) + if len(re.findall(pattern, evaluation)) == 2: + return [float(sp[0]), float(sp[1])] + + score_pair = evaluation.split("\n")[0] + score_pair = score_pair.replace(",", " ") + sp = score_pair.split(" ") + if len(sp) == 2: + return [float(sp[0]), float(sp[1])] + else: + raise Exception(f"Invalid score pair. Got {evaluation}.") + except Exception as e: + return [-1, -1] + + +def battle(answer1: List[Dict], answer2: List[Dict], prompt_dict: Dict[str, Any]) -> List[Dict]: + """ + Use GPT-4 to compare answers of two different models. + + Args: + answer1: answers of model 1. + answer2: answers of model 2. + prompt_dict: prompt for battle. + + Returns: + Evaluations of all comparison pairs. + """ + + assert len(answer1) == len(answer2) + + handles = [] + evaluation_file = [] + + total_len = len(answer1) + question_idx_list = list(range(total_len)) + + print(f" Total number of answers: {len(answer1)}.") + + evaluations = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [] + for i in question_idx_list: + assert answer1[i]["id"] == answer2[i]["id"] + answer_id = answer1[i]["id"] + + ques = answer1[i]["instruction"] if answer1[i][ + "input"] == "" else answer1[i]["instruction"] + " " + answer1[i]["input"] + cat = answer1[i]["category"] + ans1 = answer1[i]["output"] + ans2 = answer2[i]["output"] + + sys_prompt = prompt_dict["system_prompt"] + prompt_template = prompt_dict["prompt_template"] + prompt = prompt_template.format( + question=ques, + answer_1=ans1, + answer_2=ans2, + prompt=prompt_dict["prompt"], + ) + + future = executor.submit(get_battle_result, sys_prompt, prompt, answer_id, 2048) + futures.append(future) + + for future in tqdm.tqdm(concurrent.futures.as_completed(futures), total=len(futures)): + evaluations.append(future.result()) + + evaluations.sort(key=lambda x: x["id"]) + + return evaluations + + +def save_battle_results(evaluations: List[Dict], name1: str, name2: str, save_path: str) -> None: + """ + Save evaluation results (model 1 vs model 2) from GPT-4. + + Args: + evaluations: evaluation results from GPT-4. + name1: model 1 's name. + name2: model 2 's name. + save_path: path to save battle results. + """ + + evaluation_file = deepcopy(evaluations) + + ans1_score = 0 + ans2_score = 0 + better_count = 0 + worse_count = 0 + tie_count = 0 + invalid_count = 0 + + better_file = [] + worse_file = [] + tie_file = [] + invalid_file = [] + + for idx, evaluation in enumerate(evaluations): + scores = parse_battle_score(evaluation["evaluation"]) + evaluation_file[idx]["score"] = scores + + if scores[0] == -1 and scores[1] == -1: + invalid_count += 1 + invalid_file.append(evaluation_file[idx]) + print(f'Invalid score pair: {evaluation_file[idx]["id"]}.') + else: + if scores[0] > scores[1]: + worse_count += 1 + worse_file.append(evaluation_file[idx]) + elif scores[0] < scores[1]: + better_count += 1 + better_file.append(evaluation_file[idx]) + else: + tie_count += 1 + tie_file.append(evaluation_file[idx]) + ans1_score += scores[0] + ans2_score += scores[1] + + prefix = f"{name1}_vs_{name2}" + + if not os.path.exists(save_path): + os.makedirs(save_path) + + jdump(better_file, os.path.join(save_path, prefix, f"{name2}_better.json")) + jdump(worse_file, os.path.join(save_path, prefix, f"{name2}_worse.json")) + jdump(tie_file, os.path.join(save_path, prefix, f"{prefix}_tie.json")) + jdump(invalid_file, os.path.join(save_path, prefix, f"{prefix}_invalid.json")) + jdump(evaluation_file, os.path.join(save_path, prefix, f"{prefix}_evaluations.json")) + + if os.path.exists(os.path.join(save_path, "battle_results.json")): + results = jload(os.path.join(save_path, "battle_results.json")) + else: + results = {} + + results[prefix] = { + "model": [name1, name2], + "better": better_count, + "worse": worse_count, + "tie": tie_count, + "win_rate": better_count / (len(evaluations) - invalid_count), + "score": [ + ans1_score / (len(evaluations) - invalid_count), + ans2_score / (len(evaluations) - invalid_count), + ], + } + jdump(results, os.path.join(save_path, "battle_results.json")) + + print(f"Total {invalid_count} invalid score pair(s).") + print(f"Model {name2} has {better_count} better answer(s).") + print(f"Model {name2} has {worse_count} worse answer(s).") + print(f"{tie_count} answer(s) play(s) to a tie.") + print(f"Win rate of model {name2}: {better_count/(len(evaluations)-invalid_count):.2f}") + print(f"Model {name1} average score: {ans1_score/(len(evaluations)-invalid_count):.2f}") + print(f"Model {name2} average score: {ans2_score/(len(evaluations)-invalid_count):.2f}") + + +def reference_template(metric: str, language: str, reference: Dict[str, Any]) -> str: + """ + Get prompt template for GPT evaluation with reference. + + Different languages have different prompt templates. + + Args: + metric: metric used in GPT evaluation with reference. + language: language for the template. + reference: the instruction that contains target answer. + + Returns: + Prompt template for GPT evaluation with reference. + """ + + step_to_add = ref_step_template[language] + + for_the_given_answer = "{metric} (1-5) (directly give the score for the given answer):" if language == "en" else "{metric} (1-5) (直接对给定答案打分)" + + # adjective is used to describe the word "answer" in the prompt. + adjective = "example" if language == "en" else "示例" + answer_to_add = ref_answer_template_general[language] + + # Only for correctness, we will provide a correct answer and so the adjective for "answer" will be "correct". The prompt words will be "a correct answer". + # In other cases, the prompt words will be "an example answer with good quality" by default. + if metric.lower() == "correctness": + adjective = "correct" if language == "en" else "标准" + answer_to_add = ref_answer_template_correctness[language] + + answer_to_add = answer_to_add.format(answer=reference["target"] if reference["target"] else reference["output"]) + step_to_add = step_to_add.format(metric=metric.lower(), + adjective=adjective) + for_the_given_answer.format(metric=metric) + + return answer_to_add + step_to_add + + +def fill_in_message(role: str, content: str) -> Dict[str, str]: + """ + Generate one formatted message to send through chat completion. + + Args: + role: the role of the author of this message. + content: the contents of the message. + + Returns: + One message to send through chat completion. + """ + + return {"role": role, "content": content} + + +def multiturn_chat_completion(user_messages: List[str], model: str, max_tokens: int = 1, turns=2) -> Dict[str, Any]: + """ + Do multi-turn chat completion. + + When turns == 1, it is a one-turn conversation for normal GPT evaluation. + When turns == 2, it is a two-turn conversation which is used for GPT evaluation with reference answers. + + Args: + user_messages: messages user wants to send. + model: the model used to evaluate answers. + max_tokens: the maximum number of tokens to generate in the chat completion. + turns: the number of turns for conversation. + + Returns: + Last turn's response. + """ + + if len(user_messages) != turns: + raise Exception("The length of user messages should be equal to the turn number!") + + assistant_responses = [] + + for i in range(turns): + messages_to_send = [] + + for j in range(i): + messages_to_send.append(fill_in_message("user", user_messages[j])) + messages_to_send.append( + fill_in_message("assistant", assistant_responses[j]["choices"][0]["message"]["content"])) + + # Length of user messages == Length of assistant messages + 1 + # Because we always expect the api to response + messages_to_send.append(fill_in_message("user", user_messages[i])) + + response = openai.ChatCompletion.create( + model=model, + messages=messages_to_send, + temperature=0, + max_tokens=max_tokens, + ) + + # Avoid exceeding rate limits. + # You can comment this line if your request doesn't contain many tokens. + time.sleep(1) + + assistant_responses.append(response) + + return assistant_responses[-1] + + +def get_gpt_evaluation_without_logprobs(prompt: Dict[str, Any], + inst: Dict[str, Any], + metrics: List[str], + language: str, + reference: Dict[str, Any] = None, + model: str = "gpt-3.5-turbo", + max_tokens: int = 2048) -> Dict[str, Any]: + """ + Use chat models(gpt-3.5-turbo or gpt-4) to evaluate one model answer. + + Temperature is set to 0 to make the model more deterministic. + + Args: + prompt: a dictionary including prompt template, CoT and metrics. + inst: the instruction that is needed to be evaluated. + metrics: the metrics for evaluation. + language: language used to change the CoT(add one more step about comparing the given answer and reference) if reference is not None. + reference: the reference answer. + model: the model used to evaluate answers. + max_tokens: the maximum number of tokens to generate in the chat completion. + + Returns: + An evaluation of one answer. + """ + + MAX_API_RETRY = 3 + + question = (inst["instruction"] if inst["input"] == "" else inst["instruction"] + "\n" + inst["input"]) + answer = inst["output"] + inst["evaluation"] = {} + + for metric in metrics: + if prompt["metrics"].get(metric, None) is None: + raise Exception( + f"Unsupported metric {metric} for category {inst['category']}! You should add this metric in the prompt file!" + ) + for i in range(MAX_API_RETRY): + try: + prompt_reference = "" if reference is None else reference_template(metric, language, reference) + + prompt_1st_round = prompt["prompt"].format( + question=question, + answer=answer, + metric=prompt["metrics"][metric], + steps=prompt["CoT"][metric], + ) + + if prompt_reference: + # Do a 2-round conversation + response = multiturn_chat_completion([prompt_1st_round, prompt_reference], + model, + max_tokens=max_tokens, + turns=2) + else: + response = multiturn_chat_completion([prompt_1st_round], model, max_tokens=max_tokens, turns=1) + + inst["evaluation"][metric] = { + "response": response["choices"][0]["message"]["content"], + "logprobs": None, + } + + # Prevent exceeding rate limits because we have multiple workers. + # But this will slow down the evaluation process. + # You can comment this line if your request doesn't contain many tokens. + time.sleep(len(metrics) * 0.5) + + break + except Exception as e: + print(e) + time.sleep(1) + if metric not in inst["evaluation"]: + print(f"Evaluation {inst['id']} for metric {metric} failed after {MAX_API_RETRY} retries.") + inst["evaluation"][metric] = {} + return inst + + +def get_gpt_evaluation_with_logprobs(prompt: Dict[str, Any], + inst: Dict[str, Any], + metrics: List[str], + max_tokens: int = 2048) -> Dict[str, Any]: + """ + Use completion model(text-davinci-003) to evaluate one model answer. + Only completion models can return log probabilities. + + Temperature is set to 0 to make the model more deterministic. + + Args: + prompt: a dictionary including prompt template, CoT and metrics. + inst: the instruction that is needed to be evaluated. + metrics: the metrics for evaluation. + max_tokens: the maximum number of tokens to generate in the completion. + + Returns: + An evaluation of one answer. + """ + + MAX_API_RETRY = 3 + + question = (inst["instruction"] if inst["input"] == "" else inst["instruction"] + "\n" + inst["input"]) + answer = inst["output"] + inst["evaluation"] = {} + + for metric in metrics: + if prompt["metrics"].get(metric, None) is None: + raise Exception( + f"Unsupported metric {metric} for category {inst['category']}! You should add this metric in the prompt file!" + ) + for i in range(MAX_API_RETRY): + try: + response = openai.Completion.create( + model="text-davinci-003", + prompt=prompt["prompt"].format( + question=question, + answer=answer, + metric=prompt["metrics"][metric], + steps=prompt["CoT"][metric], + ), + logprobs=5, + temperature=0, + max_tokens=max_tokens, + ) + inst["evaluation"][metric] = { + "response": response["choices"][0]["text"], + "logprobs": response["choices"][0]["logprobs"]["top_logprobs"], + } + + # Prevent exceeding rate limits because we have multiple workers. + # But this will slow down the evaluation process. + # You can comment this line if your request doesn't contain many tokens. + time.sleep(len(metrics) * 0.5) + + break + except Exception as e: + print(e) + time.sleep(1) + if metric not in inst["evaluation"]: + print(f"Evaluation {inst['id']} for metric {metric} failed after {MAX_API_RETRY} retries.") + inst["evaluation"][metric] = {} + return inst + + +def evaluate(answers: List[Dict], + prompt: Dict[str, Any], + metrics: List[str], + category: str, + model: str, + language: str, + references: List[Dict] = None) -> List[Dict]: + """ + Use GPT models to evaluate model answers and save evaluation results. + + Args: + answers: model answers. + prompt: prompt for GPT evaluation. + metrics: metrics for GPT evaluation. + category: the category of the model answers for evaluation. + model: the specific GPT model used to evaluate answers. + language: language used in GPT evaluation + references: references for GPT evaluation + + Returns: + Evaluations of the given answers. + """ + + print(f"The number of instances of category {category}'s is {len(answers)}.") + + evaluations = [] + + metrics_str = ", ".join(x for x in metrics) + print(f"Category {category}'s metrics are {metrics_str}.") + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [] + for idx, inst in enumerate(answers): + # Completion models can return log probabilities. + if model == "text-davinci-003": + future = executor.submit(get_gpt_evaluation_with_logprobs, prompt, inst, metrics, 1) + else: + future = executor.submit(get_gpt_evaluation_without_logprobs, + prompt, + inst, + metrics, + language, + reference=None if references is None else references[idx], + model=model, + max_tokens=1) + + futures.append(future) + + for future in tqdm.tqdm( + concurrent.futures.as_completed(futures), + desc=f"{category}: ", + total=len(futures), + ): + evaluations.append(future.result()) + + evaluations.sort(key=lambda x: x["id"]) + + print(f"{category} done.") + + return evaluations + + +def calculate_scores_form_logprobs(logprobs: Dict[str, Any]) -> float: + """ + Calculate the score according to log probabilities returned by text-davinci-003. + + Calculation formula: + score = sum(score_i * exp(value)) where score_i is the score which corresponds to the key(predicted token) and value is its log probability. + + Ref: https://arxiv.org/abs/2303.16634 + This paper proposes NLG evaluation methods using text-davinci-003(log probabilities returned by completion models) and GPT-4(probabilities obtained by sampling). + + Args: + logprobs: logprobs returned by openai.Completion. + + Returns: + The score of one answer. + """ + + # GPT-3.5 only returns score of 1 to 5. + prob = np.zeros(5) + + for key, value in logprobs.items(): + # Sometimes the key will be one byte of a unicode character which takes the form of "bytes:\\xe7". + # It is meaningless and thus we don't calculate probability. + if "bytes" in key: + continue + # results[0] is the score which corresponds to the key(predicted token). + # For example, key "5" corresponds to score 5. + results = re.findall(r"\d", key) + if len(results) == 1: + prob[int(results[0]) - 1] = prob[int(results[0]) - 1] + np.exp(value) + + score = np.dot(np.arange(1, 6), prob) + + return score + + +def calculate_scores_form_response(response: str, evaluation: Dict[str, Any]) -> int: + """ + Calculate the score from the response returned by gpt-3.5-turbo or gpt-4. + Different from text-davinci-003, this function directly calculates the score according to the plain response returned by gpt-3.5-turbo or gpt-4. + Although text-davinci-003 can return log probabilities, it costs ten times as much as gpt-3.5-turbo. + + Args: + response: logprobs returned by openai.Completion. + evaluation: the evaluation corresponds to the question. + + Returns: + The score of one answer. + """ + + try: + results = re.findall(r"\d", response) + if len(results) == 1: + return int(results[0]) + else: + raise Exception(f"Invalid score pair. Got {evaluation}.") + except Exception as e: + return 0 + + +def save_gpt_evaluation_results(model_name: str, gpt_evaluation_results: Dict[str, Any], + save_path: str) -> Dict[str, Any]: + """ + Save evaluation results for different categories for one model. + + Args: + model_name: name of the model for saving evaluation results. + gpt_evaluation_results: evaluations results for all of the model answers. + save_path: path to save GPT evaluation statistics. + """ + + all_evaluations = [] + for category, evaluations in gpt_evaluation_results.items(): + jdump(evaluations, os.path.join(save_path, model_name, f"{category}_evaluation_results.json")) + all_evaluations.extend(evaluations) + + jdump(all_evaluations, os.path.join(save_path, f"{model_name}_evaluation_results.json")) + + return all_evaluations + + +def save_gpt_evaluation_statistics(model_name: str, evaluations: List[Dict], save_path: str) -> None: + """ + Generate statistics for one model. + + Args: + model_name: name of the model for saving statistics. + evaluations: evaluations for all of the model answers. + save_path: path to save GPT evaluation statistics. + """ + + if not os.path.exists(save_path): + os.makedirs(save_path) + + data_per_category = {} + for evaluation in evaluations: + category = evaluation["category"] + if evaluation["category"] in data_per_category.keys(): + data_per_category[category].append(evaluation) + else: + data_per_category[category] = [evaluation] + + all_statistics = {} + for category, data in data_per_category.items(): + metrics = data[0]["evaluation"].keys() + scores = {metric: [] for metric in metrics} + for evaluation in data: + for metric in metrics: + if evaluation["evaluation"][metric] == {}: + # This means after 3 retries, the server still returns an error and we set the score to 0. + scores[metric].append(0) + elif evaluation["evaluation"][metric]["logprobs"] is not None: + scores[metric].append( + calculate_scores_form_logprobs(evaluation["evaluation"][metric]["logprobs"][0])) + else: + scores[metric].append( + calculate_scores_form_response(evaluation["evaluation"][metric]["response"], evaluation)) + + statistics = {} + for metric in metrics: + arg_sort = np.argsort(scores[metric]) + statistics[metric] = {} + statistics[metric]["avg_score"] = sum(scores[metric]) / len(data) + statistics[metric]["best_3"] = {data[i]["id"]: scores[metric][i] for i in arg_sort[-3:][::-1]} + statistics[metric]["worst_3"] = {data[i]["id"]: scores[metric][i] for i in arg_sort[:3]} + + all_statistics[category] = statistics + + jdump( + all_statistics, + os.path.join(save_path, f"{model_name}_evaluation_statistics.json"), + ) + + +def analyze_gpt_evaluation_statistics(statistics_path: str, save_path: str) -> None: + """ + Analyze and visualize all GPT evaluation statistics in the given directory. + + Args: + statistics_path: path to all the models' statistics. + save_path: path to save table and visualization results. + """ + + if not os.path.exists(statistics_path): + raise Exception(f'The given directory "{statistics_path}" doesn\'t exist! No statistics found!') + + all_statistics = {} + + for file_name in os.listdir(statistics_path): + if file_name.endswith("_evaluation_statistics.json"): + model_name = file_name.split("_evaluation_statistics.json")[0] + all_statistics[model_name] = jload(os.path.join(statistics_path, file_name)) + + if len(list(all_statistics.keys())) == 0: + raise Exception(f'There are no statistics in the given directory "{statistics_path}"!') + + frame_all = { + "model": [], + "category": [], + "metric": [], + "avg_score": [], + "best_3": [], + "worst_3": [], + } + frame_per_category = {} + for model_name, model_statistics in all_statistics.items(): + for category, category_statistics in model_statistics.items(): + if frame_per_category.get(category) is None: + frame_per_category[category] = { + "model": [], + "metric": [], + "avg_score": [], + "best_3": [], + "worst_3": [], + } + + for metric, metric_statistics in category_statistics.items(): + frame_all["model"].append(model_name) + frame_all["category"].append(category) + frame_all["metric"].append(metric) + frame_all["avg_score"].append(metric_statistics["avg_score"]) + frame_all["best_3"].append(metric_statistics["best_3"]) + frame_all["worst_3"].append(metric_statistics["worst_3"]) + + frame_per_category[category]["model"].append(model_name) + frame_per_category[category]["metric"].append(metric) + frame_per_category[category]["avg_score"].append(metric_statistics["avg_score"]) + frame_per_category[category]["best_3"].append(metric_statistics["best_3"]) + frame_per_category[category]["worst_3"].append(metric_statistics["worst_3"]) + + if not os.path.exists(save_path): + os.makedirs(save_path) + + frame_all = pd.DataFrame(frame_all) + frame_all.to_csv(os.path.join(save_path, "gpt_evaluation_statistics.csv")) + + for category in tqdm.tqdm( + frame_per_category.keys(), + desc=f"GPT evaluation: ", + total=len(frame_per_category.keys()), + ): + data = pd.DataFrame(frame_per_category[category]) + + sns.set() + fig = plt.figure(figsize=(16, 10)) + plt.ylim((0, 5)) + + fig = sns.barplot(x="metric", y="avg_score", hue="model", data=data, dodge=True) + fig.set_title(f"Comparison between Different Models for Category {category.title()}") + plt.xlabel("Evaluation Metric") + plt.ylabel("Average Score") + + figure = fig.get_figure() + figure.savefig(os.path.join(save_path, f"{category}.png"), dpi=400) + + plt.close() diff --git a/applications/Chat/evaluate/metrics.py b/applications/Chat/evaluate/metrics.py new file mode 100644 index 000000000000..77f9b6e98044 --- /dev/null +++ b/applications/Chat/evaluate/metrics.py @@ -0,0 +1,253 @@ +import statistics +from typing import Dict, List + +import jieba +from bert_score import score +from nltk.translate.bleu_score import sentence_bleu +from nltk.translate.chrf_score import sentence_chrf +from rouge_chinese import Rouge as Rouge_cn +from rouge_score import rouge_scorer as Rouge_en +from sklearn.metrics import f1_score, precision_score, recall_score +from utils import preprocessing_text, remove_redundant_space + + +def bleu_score(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate BLEU Score Metric + + The calculation includes BLEU-1 for unigram, BLEU-2 for bigram, + BLEU-3 for trigram and BLEU-4 for 4-gram. Unigram evaluates the + accuracy in word level, other n-gram evaluate the fluency in + sentence level. + """ + bleu_scores = {"bleu1": 0, "bleu2": 0, "bleu3": 0, "bleu4": 0} + cumulative_bleu = [0] * 4 + weights = [(1. / 1., 0., 0., 0.), (1. / 2., 1. / 2., 0., 0.), (1. / 3., 1. / 3., 1. / 3., 0.), + (1. / 4., 1. / 4., 1. / 4., 1. / 4.)] + + for pred, target in zip(preds, targets): + if language == "cn": + pred_list = ' '.join(jieba.cut(preprocessing_text(pred))).split() + target_list = [(' '.join(jieba.cut(preprocessing_text(target)))).split()] + elif language == "en": + pred_list = preprocessing_text(pred).split() + target_list = [preprocessing_text(target).split()] + + bleu = sentence_bleu(target_list, pred_list, weights=weights) + cumulative_bleu = [a + b for a, b in zip(cumulative_bleu, bleu)] + + for i in range(len(cumulative_bleu)): + bleu_scores[f"bleu{i+1}"] = cumulative_bleu[i] / len(preds) + + return bleu_scores + + +def chrf_score(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate CHRF Score Metric in sentence level. + """ + chrf_score = {"chrf": 0} + cumulative_chrf = [] + + for pred, target in zip(preds, targets): + if language == "cn": + pred_list = ' '.join(jieba.cut(preprocessing_text(pred))).split() + target_list = ' '.join(jieba.cut(preprocessing_text(target))).split() + elif language == "en": + pred_list = preprocessing_text(pred).split() + target_list = preprocessing_text(target).split() + + cumulative_chrf.append(sentence_chrf(target_list, pred_list)) + + chrf_score["chrf"] = statistics.mean(cumulative_chrf) + + return chrf_score + + +def rouge_cn_score(preds: List[str], targets: List[str]) -> Dict[str, float]: + """Calculate Chinese ROUGE Score Metric + + The calculation includes ROUGE-1 for unigram, ROUGE-2 for bigram + and ROUGE-L. ROUGE-N evaluates the number of matching n-grams between + the preds and targets. ROUGE-L measures the number of matching + longest common subsequence (LCS) between preds and targets. + """ + rouge_scores = {"rouge1": 0, "rouge2": 0, "rougeL": 0} + all_preds = [] + all_targets = [] + + for pred, target in zip(preds, targets): + pred_list = remove_redundant_space(' '.join(jieba.cut(preprocessing_text(pred)))) + target_list = remove_redundant_space(' '.join(jieba.cut(preprocessing_text(target)))) + all_preds.append(pred_list) + all_targets.append(target_list) + + rouge_cn = Rouge_cn() + rouge_avg = rouge_cn.get_scores(all_preds, all_targets, avg=True) + + rouge_scores["rouge1"] = rouge_avg["rouge-1"]["f"] + rouge_scores["rouge2"] = rouge_avg["rouge-2"]["f"] + rouge_scores["rougeL"] = rouge_avg["rouge-l"]["f"] + + return rouge_scores + + +def rouge_en_score(preds: List[str], targets: List[str]) -> Dict[str, float]: + """Calculate English ROUGE Score Metric + + The calculation includes ROUGE-1 for unigram, ROUGE-2 for bigram + and ROUGE-L. ROUGE-N evaluates the number of matching n-grams between + the preds and targets. ROUGE-L measures the number of matching + longest common subsequence (LCS) between preds and targets. + """ + rouge_scores = {"rouge1": 0, "rouge2": 0, "rougeL": 0} + all_preds = [] + all_targets = [] + + rouge_en = Rouge_en.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=False) + + for pred, target in zip(preds, targets): + score = rouge_en.score(preprocessing_text(pred), preprocessing_text(target)) + rouge_scores["rouge1"] += score['rouge1'].fmeasure + rouge_scores["rouge2"] += score['rouge2'].fmeasure + rouge_scores["rougeL"] += score['rougeL'].fmeasure + + rouge_scores["rouge1"] = rouge_scores["rouge1"] / len(preds) + rouge_scores["rouge2"] = rouge_scores["rouge2"] / len(preds) + rouge_scores["rougeL"] = rouge_scores["rougeL"] / len(preds) + + return rouge_scores + + +def rouge_score(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate ROUGE Score Metric""" + if language == "cn": + return rouge_cn_score(preds, targets) + elif language == "en": + return rouge_en_score(preds, targets) + + +def distinct_score(preds: List[str], language: str) -> Dict[str, float]: + """Calculate Distinct Score Metric + + This metric refers to https://arxiv.org/abs/1510.03055. + It evaluates the diversity of generation text by counting + the unique n-grams. + """ + distinct_score = {"distinct": 0} + cumulative_distinct = [] + + for pred in preds: + if language == "cn": + pred_seg_list = ' '.join(jieba.cut(pred)).split() + count_segs = len(pred_seg_list) + unique_segs = set(pred_seg_list) + count_unique_chars = len(unique_segs) + # prevent denominator from being 0 + cumulative_distinct.append(count_unique_chars / (count_segs + 1e-6)) + elif language == "en": + # calculate distinct 1-gram, 2-gram, 3-gram + unique_ngram = [set() for _ in range(0, 3)] + all_ngram_count = [0 for _ in range(0, 3)] + + split_pred = preprocessing_text(pred).split() + for n in range(0, 3): + for i in range(0, len(split_pred) - n): + ngram = ' '.join(split_pred[i:i + n + 1]) + unique_ngram[n].add(ngram) + all_ngram_count[n] += 1 + + # Sometimes the answer may contain only one word. For 2-gram and 3-gram, the gram count(denominator) may be zero. + avg_distinct = [len(a) / (b + 1e-6) for a, b in zip(unique_ngram, all_ngram_count)] + + cumulative_distinct.append(statistics.mean(avg_distinct)) + + distinct_score["distinct"] = statistics.mean(cumulative_distinct) + + return distinct_score + + +def bert_score(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate BERTScore Metric + + The BERTScore evaluates the semantic similarity between + tokens of preds and targets with BERT. + """ + bert_score = {"bert_score": 0} + pred_list = [] + target_list = [] + + for pred, target in zip(preds, targets): + pred_list.append(pred) + target_list.append(target) + + if language == "cn": + _, _, F = score(pred_list, target_list, lang="zh", verbose=True) + elif language == "en": + _, _, F = score(pred_list, target_list, lang="en", verbose=True) + + bert_score["bert_score"] = F.mean().item() + + return bert_score + + +def calculate_precision_recall_f1(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Precision, Recall and F1-Score Calculation + + The calculation of precision, recall and f1-score is realized by counting + the number f overlaps between the preds and target. The comparison length + limited by the shorter one of preds and targets. + """ + precision_recall_f1 = {"precision": 0, "recall": 0, "f1_score": 0} + precision_scores = [] + recall_scores = [] + f1_scores = [] + + for pred, target in zip(preds, targets): + if language == "cn": + pred_list = [char for char in ' '.join(jieba.cut(preprocessing_text(pred))).split()] + target_list = [char for char in ' '.join(jieba.cut(preprocessing_text(target))).split()] + elif language == "en": + pred_list = [char for char in preprocessing_text(pred).split()] + target_list = [char for char in preprocessing_text(target).split()] + + target_labels = [1] * min(len(target_list), len(pred_list)) + pred_labels = [int(pred_list[i] == target_list[i]) for i in range(0, min(len(target_list), len(pred_list)))] + + precision_scores.append(precision_score(target_labels, pred_labels, zero_division=0)) + recall_scores.append(recall_score(target_labels, pred_labels, zero_division=0)) + f1_scores.append(f1_score(target_labels, pred_labels, zero_division=0)) + + precision_recall_f1["precision"] = statistics.mean(precision_scores) + precision_recall_f1["recall"] = statistics.mean(recall_scores) + precision_recall_f1["f1_score"] = statistics.mean(f1_scores) + + return precision_recall_f1 + + +def precision(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate Precision Metric + + Calculating precision by counting the number of overlaps between the preds and target. + """ + precision = {"precision": 0} + precision["precision"] = calculate_precision_recall_f1(preds, targets, language)["precision"] + return precision + + +def recall(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate Recall Metric + + Calculating recall by counting the number of overlaps between the preds and target. + """ + recall = {"recall": 0} + recall["recall"] = calculate_precision_recall_f1(preds, targets, language)["recall"] + return recall + + +def F1_score(preds: List[str], targets: List[str], language: str) -> Dict[str, float]: + """Calculate F1-score Metric + + Calculating f1-score by counting the number of overlaps between the preds and target. + """ + f1 = {"f1_score": 0} + f1["f1_score"] = calculate_precision_recall_f1(preds, targets, language)["f1_score"] + return f1 diff --git a/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_cn.json b/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_cn.json new file mode 100644 index 000000000000..ca66afd7e464 --- /dev/null +++ b/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_cn.json @@ -0,0 +1,6 @@ +{ + "id": 1, + "system_prompt": "你是一个检查回答质量的好助手。", + "prompt_template": "[问题]\n{question}\n\n[1号AI助手的答案]\n{answer_1}\n\n[1号AI助手答案终止]\n\n[2号AI助手的答案]\n{answer_2}\n\n[2号AI助手答案终止]\n\n[要求]\n{prompt}\n\n", + "prompt": "我们需要你评价这两个AI助手回答的性能。\n请对他们的回答的有用性、相关性、准确性、详细程度进行评分。每个AI助手都会得到一个1到10分的总分，分数越高表示整体表现越好。\n请首先输出一行，该行只包含两个数值，分别表示1号和2号AI助手的分数。这两个分数之间要有一个空格。在随后的一行中，请对你的评价作出全面的解释，避免任何潜在的偏见，并确保AI助手回答的顺序不会影响您的判断。" +} diff --git a/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_en.json b/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_en.json new file mode 100644 index 000000000000..2b35d1958ac5 --- /dev/null +++ b/applications/Chat/evaluate/prompt/battle_prompt/battle_prompt_en.json @@ -0,0 +1,6 @@ +{ + "id": 1, + "system_prompt": "You are a helpful and precise assistant for checking the quality of the answer. You will be given two different answers to the same question", + "prompt_template": "[Question]\n{question}\n\n[The Start of AI Assistant 1's Answer]\n{answer_1}\n\n[The End of AI Assistant 1's Answer]\n\n[The Start of AI Assistant 2's Answer]\n{answer_2}\n\n[The End of AI Assistant 2's Answer]\n\n[Requirements]\n{prompt}\n\n", + "prompt": "We would like to request your feedback on the performance of two AI assistants in response to the user question displayed above.\nPlease rate the helpfulness, relevance, accuracy, level of details of their responses. Each assistant receives an overall score on a scale of 1 to 10, where a higher score indicates better overall performance.\nPlease first output a single line containing only two values indicating the scores for Assistant 1 and 2, respectively. The two scores are separated by a space. In the subsequent line, please provide a comprehensive explanation of your evaluation, avoiding any potential bias and ensuring that the order in which the responses were presented does not affect your judgment." +} diff --git a/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_cn.json b/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_cn.json new file mode 100644 index 000000000000..783f453cafdb --- /dev/null +++ b/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_cn.json @@ -0,0 +1,179 @@ +{ + "brainstorming": { + "id": 1, + "category": "brainstorming", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "creativity": "创意性(1-5)：某些头脑风暴问题可能需要答案具有创意，提出新的思路。", + "practicality": "实用性(1-5)：某些头脑风暴问题可能需要答案提出实用的建议或解决方法。", + "reasonableness": "合理性(1-5)：答案应该符合常识、生活实际等等。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "creativity": "1. 仔细阅读所提供的头脑风暴问题，确保你理解问题的要点和背景。\n2. 根据你的知识和经验，判断所提供的答案是否可行。如果答案不可行，则创意性评分可能会受到影响。\n3. 考虑答案中是否包含新颖的想法或独特的思路。答案可能与已知的解决方案有所重叠，但仍然可以被认为是有创意的，只要它提供了新的角度或方法来解决问题。\n4. 根据答案的创意性，给出一个1到5的评分。如果答案缺乏创意，则应给出一个较低的评分。如果答案具有创意并提供了新的思路，应给出一个较高的评分。\n\n创意性：", + "practicality": "1. 仔细阅读所提供的头脑风暴问题，确保你理解问题的要点和背景。\n2. 根据你的知识和经验，判断所提供的答案是否可行。如果答案不可行，则实用性评分可能会受到影响。\n3. 考虑答案中提出的建议或解决方法是否实用并可行。答案可能看起来很好，但如果无法实现或应用，则实用性评分可能会受到影响。\n4. 根据答案的实用性，给出一个1到5的评分。如果答案缺乏实用性，则应给出一个较低的评分。如果答案提出了实用的建议或解决方法，并且可以很好地解决问题，则应给出一个较高的评分。\n\n实用性：", + "reasonableness": "1. 仔细阅读所提供的头脑风暴问题，确保你理解问题的要点和背景。\n2. 根据你的知识和经验，判断所提供的答案是否可行。如果答案不可行，则合理性评分可能会受到影响。\n3. 考虑答案中所提供的信息是否合理、符合常识、生活实际等等。如果答案中存在明显的不合理之处，则合理性评分可能会受到影响。\n4. 根据答案的合理性，给出一个1到5的评分。如果答案存在明显的不合理之处，则应给出一个较低的评分。如果答案合理、符合常识、生活实际等等，则应给出一个较高的评分。\n\n合理性：" + }, + "prompt": "你是一个好助手。请你为下面“头脑风暴”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "chat": { + "id": 2, + "category": "chat", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "naturalness": "自然(1-5)：答案是否自然，并且符合问题给定的身份。", + "engagingness": "参与感(1-5)：答案是否对前面的对话内容做出了恰当的反应，是否理解对话的语境和背景。", + "reasonableness": "合理性(1-5)：答案是否能够与前面的对话内容形成逻辑上的衔接，是否符合常理，能否在这个上下文中合理存在。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "naturalness": "1. 阅读题目，确定题目提供的身份信息。\n2. 检查答案内容是否符合题目给定的身份。\n3. 根据以上因素，对该回答的自然性进行打分，分数从1到5，其中1表示不自然，5表示非常自然，并符合问题给定的身份。\n\n自然：", + "engagingness": "1. 阅读题目，确定对话的语境和背景。\n2. 检查答案是否充分理解对话的语境和背景，能否自然地融入到对话中而不显得突兀。\n3. 根据以上因素，对该回答的参与感进行打分，分数从1到5，其中1表示没有参与感，5表示非常有参与感，并且恰当地理解了对话的语境和背景。\n\n参与感：", + "reasonableness": "1. 阅读题目，确定对话的主题以及问题期望的回答方向。\n2. 判断答案是否能够与前面的对话内容形成逻辑上的衔接，是否符合常理，能否在这个上下文中合理存在。\n3. 根据以上因素，对该回答的合理性进行打分，分数从1到5，其中1表示不合理，5表示非常合理，并且能够与前面的对话内容形成逻辑上的衔接，并符合常理。\n\n合理性：" + }, + "prompt": "你是一个好助手。请你为下面的“补全对话”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "classification": { + "id": 3, + "category": "classification", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "正确性(1-5)：答案是否正确。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读题目，尝试自己回答该问题。\n2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。\n\n正确性：" + }, + "prompt": "你是一个好助手。请你为下面的“分类“问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "closed_qa": { + "id": 4, + "category": "closed_qa", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "正确性(1-5)：答案是否正确。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读题目，尝试自己回答该问题。\n2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。\n\n正确性：" + }, + "prompt": "你是一个好助手。请你为下面问题的答案打分。\n\n问题如下：\n\n{question}\n\n需要你评分的答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "extraction": { + "id": 5, + "category": "extraction", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "准确性(1-5)：回答应该准确无误地提取出所需信息，不应该包含任何错误或误导性信息。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读问题并确定需要从材料中提取的信息。\n2. 仔细阅读回答并确保它涵盖了所有需要提取的信息。\n3. 使用所提供的材料来验证回答的准确性。如果回答不准确或包含错误或误导性信息，则无法给出高分。\n4. 检查回答是否包含所有要求提取的信息，不要漏掉任何重要细节。\n5. 根据回答的准确性和完整性，给出一个介于1和5之间的分数，5分表示回答非常准确且完整，1分表示回答几乎没有提取出所需信息。\n\n准确性：" + }, + "prompt": "你是一个好助手。请你为下面的“提取”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "generation": { + "id": 6, + "category": "generation", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "diversity": "多样性(1-5)：答案使用语言是否优美，具有有一定的创造性和想象力。然而，回答也应该保持合理和适度，不要过于夸张或离题。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "diversity": "1. 仔细阅读整个回答，确保完全理解回答所表达的内容和主题。\n2. 在阅读回答的同时，注意语言的质量，例如措辞是否正确，语言是否生动等。\n3. 检查回答的创造性和想象力，看看回答是否能够吸引人阅读下去。\n4. 检查回答的合理性和适度，看看回答是否夸张或离题。\n5. 将多样性的评分打分在1到5之间，5分表示回答的质量很好，能够吸引人阅读，1分表示回答的内容生硬或者有离题的问题。\n\n多样性：" + }, + "prompt": "你是一个好助手。请你为下面的“生成”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "open_qa": { + "id": 7, + "category": "open_qa", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "正确性(1-5)：答案是否正确。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读题目，尝试自己回答该问题。\n2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。\n\n正确性：" + }, + "prompt": "你是一个好助手。请你为下面的问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "rewriting": { + "id": 8, + "category": "rewriting", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "正确性(1-5)：答案是否正确。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读题目，尝试自己回答该问题。\n2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。\n\n正确性：" + }, + "prompt": "你是一个好助手。请你为下面的问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "roleplay": { + "id": 9, + "category": "roleplay", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "fidelity": "保真度(1-5)：答案是否能够严格遵守角色的设定回答给定的请求。", + "creativity": "创意性(1-5)：角色扮演问题的回答需要具有一定创意，但同时需要遵守角色的设定。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "fidelity": "1. 仔细阅读问题，了解角色在问题中的设定和表现，包括职业、背景、观点、性格等方面。\n2. 阅读题目的请求，确认回答请求时需要注意的细节。\n3. 对比提供的回答与该角色的设定，评估回答是否能够严格遵守角色的设定。\n4. 结合以上评估结果给出保真度的评分，范围从1到5分，其中1分表示回答与角色设定完全不符，5分表示回答完全符合角色设定且满足给定请求。\n\n保真度：", + "creativity": "1. 仔细阅读问题，了解角色在问题中的设定和表现，包括职业、背景、观点、性格等方面。\n2. 评估回答是否具有独特的思路和建议，是否能够给提问者带来新的想法和启示。\n3. 对比回答中的创意和该角色的设定，评估回答是否遵守了该角色的设定和基本特征。\n4. 对回答的质量进行总体评估，并结合以上评估结果给出创意性的评分，范围从1到5分，其中1分表示回答缺乏创意，5分表示回答具有独特的思路和建议，并且能够遵守该角色的设定。\n\n创意性：" + }, + "prompt": "你是一个好助手。请你为下面的“角色扮演”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "summarization": { + "id": 10, + "category": "summarization", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "准确性(1-5)：回答应该准确无误地总结出材料的重点。", + "conciseness": "简明扼要(1-5)：答案是否简明扼要，没有冗余内容。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读问题给的材料，理解其内容和要点。\n2. 评估回答是否准确地总结出原始材料的重点。\n3. 评估回答是否包含原始材料中的所有关键信息。\n4. 根据以上步骤，给出一个1-5的分数，其中1表示回答不能准确地总结出材料的重点，5表示回答完全准确地总结出材料的重点。\n\n准确性：", + "conciseness": "1. 阅读题目，提取出材料的重点。\n2. 阅读该总结，并注意其中的主要观点和信息。\n3. 评估总结的长度。一个简明扼要的总结通常应该在几句话或几段文字内传达关键信息，而不是冗长的段落或文章。\n4. 检查总结是否包含与主要观点无关的信息或冗余信息。\n5.确定总结涵盖了材料中的关键信息，并且没有忽略任何重要细节。\n6.给总结打出1-5的分数，其中5表示总结简明扼要，没有冗余内容，而1表示总结冗长或包含不必要的信息，难以理解或记忆。根据您的判断，打出适当的得分。\n\n简明扼要：" + }, + "prompt": "你是一个好助手。请你为下面的“总结”问题的答案打分。\n\n问题如下：\n\n{question}\n\n答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + }, + "general": { + "id": 11, + "category": "general", + "metrics": { + "language organization": "语言组织(1-5)：答案语言是否流畅、连贯，使用正确的语法，具有一定逻辑性，使用恰当的连接词、过渡词等等。", + "relevance": "切题(1-5)：答案内容是否切题，不答非所问，并且严格遵照题目要求。", + "correctness": "正确性(1-5)：答案是否正确。" + }, + "CoT": { + "language organization": "1. 阅读答案，并检查是否有语法错误、用词不当或其他显著的错误。\n2. 检查答案是否具有逻辑性，能够按照合理的顺序传达信息并且能够自圆其说。\n3. 确定答案是否与问题或主题相关，并且能够传达清晰的信息。\n4. 检查答案是否连贯，是否使用适当的转换和过渡来保持句子和段落之间的连贯性。\n5. 检查答案是否具有明确的结构和组织方式，使得读者可以轻松理解信息的层次和结构。\n6. 根据以上因素综合评估答案的语言组织，并给出一个1到5的分数，其中5表示语言组织非常好，而1表示语言组织非常差。\n\n语言组织：", + "relevance": "1. 阅读题目，确定题目所问的问题是什么，以及需要回答哪些方面的问题。\n2. 阅读答案，确认答案是否直接回答了题目所问的问题。\n3. 检查答案是否严格遵照了题目的要求，包括答题方式、答题长度、答题格式等等。\n4. 根据以上因素综合评估答案的切题程度，并给出一个1到5的分数，其中5表示答案非常切题，而1表示答案完全没有切题。\n\n切题：", + "correctness": "1. 仔细阅读题目，尝试自己回答该问题。\n2. 检查答案的准确性。您可以使用已知的事实或研究来验证答案是否正确。如果答案是正确的，则可以将正确性得分为5分。如果答案是部分正确的，则可以给予适当的得分，例如2分、3分或4分。如果答案完全不正确，则只得1分。\n\n正确性：" + }, + "prompt": "你是一个好助手。请你为下面问题的答案打分。\n\n问题如下：\n\n{question}\n\n需要你评分的答案如下：\n\n{answer}\n\n评分的指标如下：\n\n{metric}\n\n请你遵照以下的评分步骤：\n\n{steps}" + } +} diff --git a/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_en.json b/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_en.json new file mode 100644 index 000000000000..2285b639427c --- /dev/null +++ b/applications/Chat/evaluate/prompt/evaluation_prompt/evaluation_prompt_en.json @@ -0,0 +1,179 @@ +{ + "brainstorming": { + "id": 1, + "category": "brainstorming", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "creativity": "Creativity (1-5): Some brainstorming questions may require answers that are creative and suggest new ideas.", + "practicality": "Practicality (1-5): Some brainstorming questions may require answers to suggest practical suggestions or solutions.", + "reasonableness": "Reasonableness (1-5): The answer should be in line with common sense, life experience, etc." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "creativity": "1. Read the provided brainstorming questions carefully to make sure you understand the gist and context of the questions.\n2. Based on your knowledge and experience, determine if the answers provided are feasible. If the answer is not feasible, the creativity score may be affected.\n3. Consider whether the answer contains novel ideas or unique thoughts. An answer may overlap with a known solution and still be considered creative, as long as it offers a new perspective or approach to the problem.\n4. Give a score of 1 to 5 depending on the creativity of the answer. If the answer lacks creativity, a lower score should be given. If the answer is creative and provides a new idea, a higher score should be given.\n\nCreativity:", + "practicality": "1. Read the provided brainstorming questions carefully to make sure you understand the gist and context of the questions.\n2. Based on your knowledge and experience, determine if the answers provided are feasible. If the answer is not feasible, the practicality score may be affected.\n3. Consider whether the suggestions or solutions presented in the answer are practical and workable. The answer may look good, but if it cannot be implemented or applied, the practicality score may be affected.\n4. Give a score of 1 to 5 depending on the practicality of the answer. If the answer lacks practicality, a lower score should be given. If the answer makes a practical suggestion or solution and solves the problem well, a higher score should be given.\n\nPracticality:", + "reasonableness": "1. Read the provided brainstorming questions carefully to make sure you understand the gist and context of the questions.\n2. Based on your knowledge and experience, determine if the answers provided are feasible. If the answer is not feasible, the reasonableness score may be affected.\n3. Consider whether the information provided in the answer is reasonable, consistent with common sense, real life, etc. If there are obvious errors or implausibilities in the answer, the reasonableness score may be affected.\n4. Give a score of 1 to 5 depending on the reasonableness of the answer. If the answer contains obvious errors or unreasonable points, a lower score should be given. A higher score should be given if the answer is reasonable, consistent with common sense, real life, etc.\n\nReasonableness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"brainstorming\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "chat": { + "id": 2, + "category": "chat", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "naturalness": "Naturalness (1-5): whether the answer is natural and fits the identity given by the question.", + "engagingness": "Engagingness (1-5): whether the answer responds appropriately to the content of the preceding conversation and whether it understands the context and background of the conversation.", + "reasonableness": "Reasonableness (1-5): Whether the answer can form a logical connection with the content of the previous dialogue, whether it is consistent with common sense, and whether it can reasonably exist in this context." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "naturalness": "1. Read the question and determine the identity information provided in the question.\n2. Check whether the content of the answer matches the identity given in the question.\n3. Based on the above factors, score the naturalness of the response on a scale from 1 to 5, where 1 means unnatural and 5 means very natural and in accordance with the identity given in the question.\n\nNaturalness:", + "engagingness": "1. Read the questions to determine the context and background of the dialogue.\n2. Check that the answer fully understands the context and background of the conversation and that it fits naturally into the conversation without seeming abrupt.\n3. Based on the above factors, rate the response's engagement on a scale from 1 to 5, where 1 means not engaged and 5 means very engaged and appropriately understands the context and background of the conversation.\n\nEngagingness:", + "reasonableness": "1. Read the question and determine the topic of the conversation and the direction the question expects the answer to go.\n2. Determine whether the answer can be logically connected to the preceding conversation, whether it makes common sense, and whether it can reasonably exist in this context.\n3. Based on the above factors, rate the reasonableness of the answer on a scale from 1 to 5, where 1 means unreasonable and 5 means very reasonable and able to form a logical connection with the preceding dialogue content and consistent with common sense.\n\nReasonableness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"chat\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "classification": { + "id": 3, + "category": "classification", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): whether the answer is correct or not." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the question carefully and try to answer the question yourself.\n2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be given. If the answer is completely incorrect, only 1 point is awarded.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"classification\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "closed_qa": { + "id": 4, + "category": "closed_qa", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): whether the answer is correct or not." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the question carefully and try to answer the question by yourself.\n2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be assigned. If the answer is completely incorrect, only 1 point is awarded.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"closed qa\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "extraction": { + "id": 5, + "category": "extraction", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "correctness (1-5): Answers should extract the required information accurately and should not contain any incorrect or misleading information." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the questions carefully and identify the information that needs to be extracted from the material.\n2. Read the answer carefully and make sure it covers all the information that needs to be extracted.\n3. Use the material provided to verify the correctness of the response. If the response is inaccurate or contains incorrect or misleading information, a high score cannot be given.\n4. Check that the answer contains all the information required to be extracted and do not leave out any important details.\n5. Give a score between 1 and 5 based on the correctness and completeness of the response, with a score of 5 indicating a very accurate and complete response and a score of 1 indicating that the response barely extracts the required information.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"extraction\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "generation": { + "id": 6, + "category": "generation", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "diversity": "Diversity (1-5): Whether the answers use beautiful language and have some creativity and imagination. However, answers should also be kept reasonable and moderate, not overly exaggerated or off-topic." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "diversity": "1. Read the entire response carefully to ensure that you fully understand the content and theme expressed in the response.\n2. While reading the response, pay attention to the quality of the language, such as whether the wording is correct and the language is vivid.\n3. Check the creativity and imagination of the response to see if the response is engaging to read on.\n4. Check the reasonableness and appropriateness of the responses to see if the responses are exaggerated or off-topic.\n5. Rate the diversity on a scale of 1 to 5, with a 5 indicating a good quality response that is engaging to read and a 1 indicating a raw response or a question that is off-topic.\n\nDiversity:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"generation\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "open_qa": { + "id": 7, + "category": "open_qa", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): whether the answer is correct or not." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the question carefully and try to answer the question yourself.\n2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be given. If the answer is completely incorrect, only 1 point is awarded.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the answers to the \"open qa\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "rewriting": { + "id": 8, + "category": "rewriting", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): whether the answer is correct or not." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the question carefully and try to answer the question yourself.\n2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be assigned. If the answer is completely incorrect, only 1 point is awarded.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the answers to the \"rewriting\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "roleplay": { + "id": 9, + "category": "roleplay", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "fidelity": "Fidelity (1-5): whether the answer is able to answer the given request in strict compliance with the role setting.", + "creativity": "Creativity (1-5): The answers to the role-play questions need to be somewhat creative, but at the same time they need to adhere to the setting of the role." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "fidelity": "1. Read the question carefully to understand how the character is set up and represented in the question, including aspects such as occupation, background, point of view, and personality.\n2. Read the question's request and confirm the details that need to be taken into account when answering the request.\n3. Compare the provided answer with the setting of the role and assess whether the answer can strictly adhere to the setting of the role.\n4. Combine the results of the above assessment to give a fidelity score ranging from 1 to 5, where a score of 1 means that the response does not match the persona at all, and a score of 5 means that the response fully complies with the persona and satisfies the given request.\n\nFidelity:", + "creativity": "1. Read the question carefully to understand how the character is set up and represented in the question, including career, background, perspective, and personality.\n2. Evaluate whether the answer has unique ideas and suggestions that bring new ideas and insights to the questioner.\n3. Compare the creativity in the response to the setting of the persona and assess whether the response adheres to the setting and essential characteristics of the persona.\n4. Evaluate the quality of the responses in general and combine the results of the above assessment to give a creativity score ranging from 1 to 5, where a score of 1 indicates that the response lacks creativity and a score of 5 indicates that the response has unique ideas and suggestions and is able to adhere to the set-up of the persona.\n\nCreativity:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"role-play\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "summarization": { + "id": 10, + "category": "summarization", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): answers should summarize the main points of the material accurately and unambiguously.", + "conciseness": "Conciseness (1-5): answers should be concise and without redundant content." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the material given in the question carefully to understand its content and main points.\n2. Assess whether the answer accurately summarizes the key points of the source material.\n3. assess whether the response contains all the key information in the source material.\n4. Based on the above steps, give a score of 1-5, where 1 means that the response does not accurately summarize the main points of the material and 5 means that the response completely accurately summarizes the main points of the material.\n\nCorrectness:", + "conciseness": "1. Read the title and extract the main points of the material.\n2. Read the summary and note the main ideas and messages in it.\n3. Assess the length of the summary. A concise summary should usually convey key information within a few sentences or paragraphs, rather than lengthy paragraphs or essays.\n4. Check that the summary does not contain information that is not relevant to the main ideas or that is redundant.\n5. Make sure that the summary covers the key information in the material and that no important details have been omitted.\n6. Rate the summary on a scale of 1-5, where 5 means the summary is concise and free of redundancy, and 1 means the summary is lengthy or contains unnecessary information that is difficult to understand or remember. Based on your judgment, assign the appropriate score.\n\nConciseness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the \"summarization\" question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + }, + "general": { + "id": 11, + "category": "general", + "metrics": { + "language organization": "Language organization (1-5): whether the answer language is fluent and coherent, uses correct grammar, has a certain logic, uses appropriate connecting words, transition words, etc.", + "relevance": "Relevance (1-5): whether the content of the answer is relevant to the topic, does not answer the wrong question, and strictly follows the requirements of the topic.", + "correctness": "Correctness (1-5): whether the answer is correct or not." + }, + "CoT": { + "language organization": "1. Read the answers and check for grammatical errors, poor word choice, or other significant mistakes.\n2. Check that the answer is logical, conveys the information in a logical order, and is self-explanatory.\n3. Determine if the answer is relevant to the question or topic and conveys a clear message.\n4. Check that the answer is coherent and that appropriate transitions and switches are used to maintain coherence between sentences and paragraphs.\n5. Check that the answer is clearly structured and organized in such a way that the reader can easily understand the hierarchy and structure of the information.\n6. Evaluate the language organization of the answer based on a combination of the above factors and give a score of 1 to 5, where 5 indicates very good language organization and 1 indicates very poor language organization.\n\nLanguage organization:", + "relevance": "1. Read the question to determine what the question asks and what aspects of the question need to be answered.\n2. Read the answers to make sure that they directly answer the question asked.\n3. Check that the answer follows the requirements of the question, including the way it is answered, the length of the answer, the format of the answer, etc.\n4. Evaluate how relevant the answer is based on the above factors and give a score of 1 to 5, where 5 means the answer is very relevant and 1 means the answer is not relevant at all.\n\nRelevance:", + "correctness": "1. Read the question carefully and try to answer the question yourself.\n2. Check the correctness of the answer. You can use known facts or research to verify that the answer is correct. If the answer is correct, you can give a score of 5 for correctness. If the answer is partially correct, an appropriate score, such as 2, 3, or 4, may be assigned. If the answer is completely incorrect, only 1 point is awarded.\n\nCorrectness:" + }, + "prompt": "You are a good assistant. Please rate the given answer to the question below.\n\nThe question is as follows:\n\n{question}\n\nThe answer is as follows:\n\n{answer}\n\nThe metric for evaluation is as follows:\n\n{metric}\n\nYou should follow the following evaluation steps:\n\n{steps}" + } +} diff --git a/applications/Chat/evaluate/requirements.txt b/applications/Chat/evaluate/requirements.txt new file mode 100644 index 000000000000..27d317ed88cc --- /dev/null +++ b/applications/Chat/evaluate/requirements.txt @@ -0,0 +1,12 @@ +jieba +bert-score +rouge_chinese +scikit-metrics +nltk +openai +seaborn +pandas +matplotlib +numpy +zhon +rouge_score diff --git a/applications/Chat/evaluate/unieval/__init__.py b/applications/Chat/evaluate/unieval/__init__.py new file mode 100644 index 000000000000..dad8d6ad09fa --- /dev/null +++ b/applications/Chat/evaluate/unieval/__init__.py @@ -0,0 +1,12 @@ +from .evaluator import get_evaluator +from .utils import ( + analyze_unieval_results, + calculate_average_score, + convert_data_to_unieval_format, + save_unieval_results, +) + +__all__ = [ + 'get_evaluator', 'convert_data_to_unieval_format', 'calculate_average_score', 'save_unieval_results', + 'analyze_unieval_results' +] diff --git a/applications/Chat/evaluate/unieval/evaluator.py b/applications/Chat/evaluate/unieval/evaluator.py new file mode 100644 index 000000000000..56cc6d2f9e41 --- /dev/null +++ b/applications/Chat/evaluate/unieval/evaluator.py @@ -0,0 +1,331 @@ +# MIT License + +# Copyright (c) 2022 Ming Zhong + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import numpy as np +from nltk import sent_tokenize + +from .scorer import UniEvaluator +from .utils import add_question + + +class SumEvaluator: + + def __init__(self, model_name_or_path, max_length=1024, device='cuda:0', cache_dir=None): + """ Set up evaluator for text summarization """ + self.scorer = UniEvaluator( + model_name_or_path='MingZhong/unieval-sum' if model_name_or_path == "" else model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + self.task = 'summarization' + self.dimensions = ['coherence', 'consistency', 'fluency', 'relevance'] + + def evaluate(self, data, category, dims=None, overall=True): + """ + Get the scores of all the given dimensions + + category: The category to be evaluated. + + dims: A list of dimensions to be evaluated. If dims is None, SumEvaluator will evaluate + four dimensions: coherence, consistency, fluency, relevance. + + overall: indicates whether the overall score is to be calculated. + Overall score can be customized to a combination of scores based on different + dimensions. The default here is the average score of all the given dimensions. + """ + n_data = len(data) + eval_scores = [{} for _ in range(n_data)] + + if dims == None: + eval_dims = self.dimensions + else: + assert isinstance(dims, list) + eval_dims = dims + + for dim in eval_dims: + # Calculate average sentence-level scores for 'consistency' and 'fluency' + if dim == 'consistency' or dim == 'fluency': + src_list, output_list = [], [] + n_sents = [] # the number of sentences in each generated summary + for i in range(n_data): + source = data[i]['source'] + system_outputs = sent_tokenize(data[i]['system_output']) + n_sents.append(len(system_outputs)) + for j in range(len(system_outputs)): + src_list.append(source) + output_list.append(system_outputs[j]) + input_list = add_question(dimension=dim, output=output_list, src=src_list, task=self.task) + sent_score = self.scorer.score(input_list, self.task, category, dim) + + # Get average score for each sample + start_idx = 0 + score = [] + for cur_n_sent in n_sents: + # prevent denominator from being 0 + score.append(sum(sent_score[start_idx:start_idx + cur_n_sent]) / (cur_n_sent + 1e-6)) + start_idx += cur_n_sent + + # Calculate summary-level score for 'coherence' and 'relevance' + elif dim == 'coherence' or dim == 'relevance': + src_list, output_list, ref_list = [], [], [] + for i in range(n_data): + src_list.append(data[i]['source']) + output_list.append(data[i]['system_output']) + if dim == 'relevance': + ref_list.append(data[i]['reference']) + input_list = add_question(dimension=dim, output=output_list, src=src_list, ref=ref_list, task=self.task) + score = self.scorer.score(input_list, self.task, category, dim) + + # Please customize other dimensions here for summarization + else: + raise NotImplementedError('The input format for this dimension is still undefined. \ + Please customize it first.') + + for i in range(n_data): + eval_scores[i][dim] = score[i] + + # Customize your overall score here. + if overall == True: + for i in range(n_data): + eval_scores[i]['overall'] = np.mean(list(eval_scores[i].values())) + + return eval_scores + + +class DialogEvaluator: + + def __init__(self, model_name_or_path, max_length=1024, device='cuda:0', cache_dir=None): + """ Set up evaluator for dialogues """ + self.scorer = UniEvaluator( + model_name_or_path='MingZhong/unieval-dialog' if model_name_or_path == "" else model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + self.task = 'dialogue' + self.dimensions = ['naturalness', 'coherence', 'engagingness', 'groundedness', 'understandability'] + + def evaluate(self, data, category, dims=None, overall=True): + """ + Get the scores of all the given dimensions + + category: The category to be evaluated. + + dims: A list of dimensions to be evaluated. If dims is None, DialogEvaluator will evaluate + five dimensions: naturalness, coherence, engagingness, groundedness and understandability. + + overall: indicates whether the overall score is to be calculated. + Overall score can be customized to a combination of scores based on different + dimensions. The default here is the average score of all the given dimensions. + """ + n_data = len(data) + eval_scores = [{} for _ in range(n_data)] + + if dims == None: + eval_dims = self.dimensions + else: + assert isinstance(dims, list) + eval_dims = dims + + for dim in eval_dims: + # Calculate summation score for 'engagingness' + if dim == 'engagingness': + src_list, output_list, context_list = [], [], [] + n_sents = [] # the number of sentences in each generated response + for i in range(n_data): + source = data[i]['source'] + context = data[i]['context'] + system_outputs = sent_tokenize(data[i]['system_output']) + n_sents.append(len(system_outputs)) + for j in range(len(system_outputs)): + src_list.append(source) + context_list.append(context) + output_list.append(system_outputs[j]) + input_list = add_question(dimension=dim, + output=output_list, + src=src_list, + context=context_list, + task=self.task) + sent_score = self.scorer.score(input_list, self.task, category, dim) + + # Get the summation score for each sample + start_idx = 0 + score = [] + for cur_n_sent in n_sents: + score.append(sum(sent_score[start_idx:start_idx + cur_n_sent])) + start_idx += cur_n_sent + + # Calculate turn-level score for other dimensions + elif dim in ['naturalness', 'coherence', 'groundedness', 'understandability']: + src_list, output_list, context_list = [], [], [] + for i in range(n_data): + src_list.append(data[i]['source']) + output_list.append(data[i]['system_output']) + context_list.append(data[i]['context']) + input_list = add_question(dimension=dim, + output=output_list, + src=src_list, + context=context_list, + task=self.task) + score = self.scorer.score(input_list, self.task, category, dim) + + # Please customize other dimensions here for summarization + else: + raise NotImplementedError('The input format for this dimension is still undefined. \ + Please customize it first.') + + for i in range(n_data): + eval_scores[i][dim] = score[i] + + # Customize your overall score here. + if overall == True: + for i in range(n_data): + eval_scores[i]['overall'] = np.mean(list(eval_scores[i].values())) + + return eval_scores + + +class D2tEvaluator: + + def __init__(self, model_name_or_path, max_length=1024, device='cuda:0', cache_dir=None): + """ Set up evaluator for data-to-text """ + self.scorer = UniEvaluator( + model_name_or_path='MingZhong/unieval-sum' if model_name_or_path == "" else model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + self.task = 'data2text' + self.dimensions = ['naturalness', 'informativeness'] + + def evaluate(self, data, category, dims=None, overall=True): + """ + Get the scores of all the given dimensions + + category: The category to be evaluated. + + dims: A list of dimensions to be evaluated. If dims is None, D2tEvaluator will evaluate + two dimensions: naturalness and informativeness. + + overall: indicates whether the overall score is to be calculated. + Overall score can be customized to a combination of scores based on different + dimensions. The default here is the average score of all the given dimensions. + """ + n_data = len(data) + eval_scores = [{} for _ in range(n_data)] + + if dims == None: + eval_dims = self.dimensions + else: + assert isinstance(dims, list) + eval_dims = dims + + for dim in eval_dims: + output_list, ref_list = [], [] + for i in range(n_data): + output_list.append(data[i]['system_output']) + ref_list.append(data[i]['reference']) + + input_list = add_question(dimension=dim, output=output_list, ref=ref_list, task=self.task) + score = self.scorer.score(input_list, self.task, category, dim) + + for i in range(n_data): + eval_scores[i][dim] = score[i] + + # Customize your overall score here. + if overall == True: + for i in range(n_data): + eval_scores[i]['overall'] = np.mean(list(eval_scores[i].values())) + + return eval_scores + + +class FactEvaluator: + + def __init__(self, model_name_or_path, max_length=1024, device='cuda:0', cache_dir=None): + """ Set up evaluator for factual consistency detection """ + self.scorer = UniEvaluator( + model_name_or_path='MingZhong/unieval-fact' if model_name_or_path == "" else model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + self.task = 'fact' + self.dim = 'consistency' + + def evaluate(self, data, category): + """ + Get the factual consistency score (only 1 dimension for this task) + + category: The category to be evaluated. + """ + n_data = len(data) + eval_scores = [{} for _ in range(n_data)] + + # Calculate average sentence-level scores for factual consistency + src_list, output_list = [], [] + n_sents = [] # the number of sentences in the claim + for i in range(n_data): + source = data[i]['source'] + system_outputs = sent_tokenize(data[i]['system_output']) + n_sents.append(len(system_outputs)) + for j in range(len(system_outputs)): + src_list.append(source) + output_list.append(system_outputs[j]) + input_list = add_question(dimension=self.dim, output=output_list, src=src_list, task=self.task) + sent_score = self.scorer.score(input_list, self.task, category, self.dim) + + # Get average score for each sample + start_idx = 0 + score = [] + for cur_n_sent in n_sents: + score.append(sum(sent_score[start_idx:start_idx + cur_n_sent]) / cur_n_sent) + start_idx += cur_n_sent + + for i in range(n_data): + eval_scores[i][self.dim] = score[i] + + return eval_scores + + +def get_evaluator(task, model_name_or_path="", max_length=1024, device='cuda:0', cache_dir=None): + assert task in ['summarization', 'dialogue', 'data2text', 'fact'] + if task == 'summarization': + return SumEvaluator(model_name_or_path=model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + elif task == 'dialogue': + return DialogEvaluator(model_name_or_path=model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + elif task == 'data2text': + return D2tEvaluator(model_name_or_path=model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + elif task == 'fact': + return FactEvaluator(model_name_or_path=model_name_or_path, + max_length=max_length, + device=device, + cache_dir=cache_dir) + else: + raise NotImplementedError('Other tasks are not implemented, \ + please customize specific tasks here.') diff --git a/applications/Chat/evaluate/unieval/scorer.py b/applications/Chat/evaluate/unieval/scorer.py new file mode 100644 index 000000000000..2c70bb9f6ded --- /dev/null +++ b/applications/Chat/evaluate/unieval/scorer.py @@ -0,0 +1,101 @@ +# MIT License + +# Copyright (c) 2022 Ming Zhong + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import torch +import torch.nn as nn +from tqdm import tqdm +from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer + + +class UniEvaluator: + + def __init__(self, model_name_or_path, max_length=1024, device='cuda:0', cache_dir=None): + """ Set up model """ + self.device = device + self.max_length = max_length + + self.config = AutoConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir) + self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, cache_dir=cache_dir) + self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path, config=self.config, cache_dir=cache_dir) + + self.model.eval() + self.model.to(device) + + self.softmax = nn.Softmax(dim=1) + + self.pos_id = self.tokenizer("Yes")["input_ids"][0] + self.neg_id = self.tokenizer("No")["input_ids"][0] + + def score(self, inputs, task, category, dim, batch_size=8): + """ + Get scores for the given samples. + final_score = postive_score / (postive_score + negative_score) + """ + + # The implementation of "forward" in T5 still requires decoder_input_ids. + # Therefore, we construct a random one-word target sequence. + # The content of the target has no effect on the final scores. + tgts = ["No" for _ in range(len(inputs))] + + pos_score_list, neg_score_list = [], [] + for i in tqdm(range(0, len(inputs), batch_size), desc=f"{category}-({dim}-{task}): "): + src_list = inputs[i:i + batch_size] + tgt_list = tgts[i:i + batch_size] + try: + with torch.no_grad(): + encoded_src = self.tokenizer(src_list, + max_length=self.max_length, + truncation=True, + padding=True, + return_tensors='pt') + encoded_tgt = self.tokenizer(tgt_list, + max_length=self.max_length, + truncation=True, + padding=True, + return_tensors='pt') + + src_tokens = encoded_src['input_ids'].to(self.device) + src_mask = encoded_src['attention_mask'].to(self.device) + + tgt_tokens = encoded_tgt['input_ids'].to(self.device)[:, 0].unsqueeze(-1) + + output = self.model(input_ids=src_tokens, attention_mask=src_mask, labels=tgt_tokens) + logits = output.logits.view(-1, self.model.config.vocab_size) + + pos_score = self.softmax(logits)[:, self.pos_id] # Yes + neg_score = self.softmax(logits)[:, self.neg_id] # No + + cur_pos_score = [x.item() for x in pos_score] + cur_neg_score = [x.item() for x in neg_score] + pos_score_list += cur_pos_score + neg_score_list += cur_neg_score + + except RuntimeError: + print(f'source: {src_list}') + print(f'target: {tgt_list}') + exit(0) + + score_list = [] + for i in range(len(pos_score_list)): + score_list.append(pos_score_list[i] / (pos_score_list[i] + neg_score_list[i])) + + return score_list diff --git a/applications/Chat/evaluate/unieval/utils.py b/applications/Chat/evaluate/unieval/utils.py new file mode 100644 index 000000000000..a381e9e590b2 --- /dev/null +++ b/applications/Chat/evaluate/unieval/utils.py @@ -0,0 +1,248 @@ +# MIT License + +# Copyright (c) 2022 Ming Zhong + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import os +from typing import Dict + +import matplotlib.pyplot as plt +import pandas as pd +import seaborn as sns +import tqdm + + +def add_question(dimension, output, src=None, ref=None, context=None, task=None): + """ + Add questions to generate input in Bool-QA format for UniEval. + + dimension: specific dimension to be evaluated + src: source input for different NLG tasks. For example, source document for summarization + and dialogue history for dialogue response generation. + output: output text generated by the models + ref: human-annotated groundtruth + context: the context needed to evaluate several specific dimension. For example, + additional factual information when evaluating engagingness and groundedness in dialogues. + """ + + input_with_question = [] + for i in range(len(output)): + # For summarization + if task == 'summarization': + if dimension == 'fluency': + cur_input = 'question: Is this a fluent paragraph? paragraph: ' + output[i] + elif dimension == 'coherence': + cur_input = 'question: Is this a coherent summary to the document? summary: ' + output[ + i] + ' document: ' + src[i] + elif dimension == 'consistency': + cur_input = 'question: Is this claim consistent with the document? claim: ' + output[ + i] + ' document: ' + src[i] + elif dimension == 'relevance': + cur_input = 'question: Is this summary relevant to the reference? summary: ' + output[ + i] + ' reference: ' + ref[i] + else: + raise NotImplementedError( + 'The input format for this dimension is still undefined. Please customize it first.') + # For dialogues + elif task == 'dialogue': + if dimension == 'naturalness': + cur_input = 'question: Is this a natural response in the dialogue? response: ' + output[i] + elif dimension == 'coherence': + cur_input = 'question: Is this a coherent response given the dialogue history? response: '\ + + output[i] + ' dialogue history: ' + src[i] + elif dimension == 'engagingness': + cur_input = 'question: Is this an engaging and informative response according to the dialogue history and fact? response: '\ + + output[i] + ' dialogue history: ' + src[i] + ' fact: ' + context[i] + elif dimension == 'groundedness': + cur_input = 'question: Is this response consistent with knowledge in the fact? response: '\ + + output[i] + ' fact: ' + context[i] + elif dimension == 'understandability': + cur_input = 'question: Is this an understandable response in the dialogue? response: ' + output[i] + else: + raise NotImplementedError( + 'The input format for this dimension is still undefined. Please customize it first.') + # For data-to-text + elif task == 'data2text': + if dimension == 'naturalness': + cur_input = 'question: Is this a fluent utterance? utterance: ' + output[i] + elif dimension == 'informativeness': + cur_input = 'question: Is this sentence informative according to the reference? sentence: '\ + + output[i] + ' reference: ' + ref[i] + else: + raise NotImplementedError( + 'The input format for this dimension is still undefined. Please customize it first.') + # For factual consistency detection + elif task == 'fact': + if dimension == 'consistency': + cur_input = 'question: Is this claim consistent with the document? claim: ' + output[ + i] + ' document: ' + src[i] + else: + raise NotImplementedError('No other dimensions for the factual consistency detection task.') + # For new customized tasks + else: + raise NotImplementedError('Other tasks are not implemented, please customize specific tasks here.') + input_with_question.append(cur_input) + return input_with_question + + +def convert_data_to_unieval_format(output_list, src_list=None, ref_list=None): + """ + Convert the data into the unieval's format. + + output_list: a list of model output + + src_list: source input for different NLG tasks. For example, source document for summarization + and dialogue history for dialogue response generation + ref_list: human-annotated groundtruth + """ + json_data = [] + for i in range(len(output_list)): + cur = {} + cur['system_output'] = output_list[i] + if src_list is not None: + cur['source'] = src_list[i] + if ref_list is not None: + cur['reference'] = ref_list[i] + cur['context'] = "" + json_data.append(cur) + return json_data + + +def calculate_average_score(scores): + """ + Calculate average scores for different metrics + + scores: a list of scores for different metrics for each answer + + """ + metrics = {metric: 0 for metric in scores[0]} + + for score in scores: + for metric in score: + metrics[metric] += score[metric] + + for metric in metrics: + metrics[metric] /= len(scores) + + return metrics + + +def save_unieval_results(model_name: str, unieval_metric_stats: Dict[str, Dict], save_path: str) -> None: + """ + Save UniEval evaluation results of different categories for one model. + + """ + + if not os.path.exists(save_path): + os.makedirs(save_path) + + unieval_metric_stats_per_category = {} + for task, category_stat in unieval_metric_stats.items(): + for category, metric_stat in category_stat.items(): + if unieval_metric_stats_per_category.get(category, None) is None: + unieval_metric_stats_per_category[category] = {} + for metric, score in metric_stat.items(): + unieval_metric_stats_per_category[category][f"{metric}-{task}"] = score + + automatic_df = pd.DataFrame(unieval_metric_stats_per_category) + automatic_df.to_csv(os.path.join(save_path, f"{model_name}_results.csv"), index=True) + + +def read_unieval_results(results_path: str, file_name: str) -> Dict[str, Dict]: + """ + Read a csv file and return a dictionary which stores scores per metric. + + """ + + results = pd.read_csv(os.path.join(results_path, file_name), index_col=0) + + results_dict = {metric: {} for metric in list(results.index)} + for i, metric in enumerate(results_dict.keys()): + for j, category in enumerate(list(results.columns)): + if pd.isnull(results.iloc[i][j]): + continue + results_dict[metric][category] = results.iloc[i][j] + + return results_dict + + +def analyze_unieval_results(results_path: str, save_path: str) -> None: + """ + Analyze and visualize all csv files in the given folder. + + """ + + if not os.path.exists(results_path): + raise Exception(f'The given directory "{results_path}" doesn\'t exist! No results found!') + + all_statistics = {} + + for file_name in os.listdir(results_path): + if file_name.endswith("_results.csv"): + model_name = file_name.split("_results.csv")[0] + all_statistics[model_name] = read_unieval_results(results_path, file_name) + + if len(list(all_statistics.keys())) == 0: + raise Exception(f'There are no csv files in the given directory "{results_path}"!') + + frame_all = {"model": [], "category": [], "metric": [], "score": []} + frame_per_metric = {} + for model_name, model_statistics in all_statistics.items(): + for metric, metric_statistics in model_statistics.items(): + if frame_per_metric.get(metric) is None: + frame_per_metric[metric] = {"model": [], "category": [], "score": []} + + for category, category_score in metric_statistics.items(): + frame_all["model"].append(model_name) + frame_all["category"].append(category) + frame_all["metric"].append(metric) + frame_all["score"].append(category_score) + + frame_per_metric[metric]["model"].append(model_name) + frame_per_metric[metric]["category"].append(category) + frame_per_metric[metric]["score"].append(category_score) + + if not os.path.exists(save_path): + os.makedirs(save_path) + + frame_all = pd.DataFrame(frame_all) + frame_all.to_csv(os.path.join(save_path, "unieval_statistics.csv")) + + for metric in tqdm.tqdm( + frame_per_metric.keys(), + desc=f"UniEval metrics: ", + total=len(frame_per_metric.keys()), + ): + data = pd.DataFrame(frame_per_metric[metric]) + + sns.set() + fig = plt.figure(figsize=(16, 10)) + + fig = sns.barplot(x="category", y="score", hue="model", data=data, dodge=True) + fig.set_title( + f"Comparison between Different Models for Metric {metric.split('-')[0].title()} in Task {metric.split('-')[1].title()}" + ) + plt.xlabel("Evaluation Category") + plt.ylabel("Score") + + figure = fig.get_figure() + figure.savefig(os.path.join(save_path, f"{metric}.png"), dpi=400) + + plt.close() diff --git a/applications/Chat/evaluate/utils.py b/applications/Chat/evaluate/utils.py new file mode 100644 index 000000000000..406e43db99aa --- /dev/null +++ b/applications/Chat/evaluate/utils.py @@ -0,0 +1,207 @@ +import io +import json +import os +import re +import string +from typing import Dict + +import matplotlib.pyplot as plt +import pandas as pd +import seaborn as sns +import tqdm +from zhon import hanzi + + +def _make_w_io_base(f, mode: str): + if not isinstance(f, io.IOBase): + f_dirname = os.path.dirname(f) + if f_dirname != "": + os.makedirs(f_dirname, exist_ok=True) + f = open(f, mode=mode) + return f + + +def _make_r_io_base(f, mode: str): + if not isinstance(f, io.IOBase): + f = open(f, mode=mode) + return f + + +def jdump(obj, f, mode="w", indent=4, default=str): + """Dump a str or dictionary to a file in json format. + Args: + obj: An object to be written. + f: A string path to the location on disk. + mode: Mode for opening the file. + indent: Indent for storing json dictionaries. + default: A function to handle non-serializable entries; defaults to `str`. + """ + f = _make_w_io_base(f, mode) + if isinstance(obj, (dict, list)): + json.dump(obj, f, indent=indent, default=default, ensure_ascii=False) + elif isinstance(obj, str): + f.write(obj) + else: + raise ValueError(f"Unexpected type: {type(obj)}") + f.close() + + +def jload(f, mode="r"): + """Load a .json file into a dictionary.""" + f = _make_r_io_base(f, mode) + jdict = json.load(f) + f.close() + return jdict + + +def get_json_list(file_path): + with open(file_path, 'r') as f: + json_list = [] + for line in f: + json_list.append(json.loads(line)) + return json_list + + +def get_data_per_category(data, categories): + data_per_category = {category: [] for category in categories} + for item in data: + category = item["category"] + if category in categories: + data_per_category[category].append(item) + + return data_per_category + + +def remove_punctuations(text: str) -> str: + """ + Remove punctuations in the given text. + It is used in evaluation of automatic metrics. + + """ + + punctuation = string.punctuation + hanzi.punctuation + punctuation = set([char for char in punctuation]) + punctuation.difference_update(set("!@#$%&()<>?|,.\"'")) + + out = [] + for char in text: + if char in punctuation: + continue + else: + out.append(char) + + return "".join(out) + + +def remove_redundant_space(text: str) -> str: + """ + Remove redundant spaces in the given text. + It is used in evaluation of automatic metrics. + + """ + + return " ".join(text.split()) + + +def preprocessing_text(text: str) -> str: + """ + Preprocess the given text. + It is used in evaluation of automatic metrics. + + """ + + return remove_redundant_space(remove_punctuations(text.lower())) + + +def save_automatic_results(model_name: str, automatic_metric_stats: Dict[str, Dict], save_path: str) -> None: + """ + Save automatic evaluation results of different categories for one model. + + """ + + if not os.path.exists(save_path): + os.makedirs(save_path) + + automatic_df = pd.DataFrame(automatic_metric_stats) + automatic_df.to_csv(os.path.join(save_path, f"{model_name}_results.csv"), index=True) + + +def read_automatic_results(results_path: str, file_name: str) -> Dict[str, Dict]: + """ + Read a csv file and return a dictionary which stores scores per metric. + + """ + + results = pd.read_csv(os.path.join(results_path, file_name), index_col=0) + + results_dict = {metric: {} for metric in list(results.index)} + for i, metric in enumerate(results_dict.keys()): + for j, category in enumerate(list(results.columns)): + if pd.isnull(results.iloc[i][j]): + continue + results_dict[metric][category] = results.iloc[i][j] + + return results_dict + + +def analyze_automatic_results(results_path: str, save_path: str) -> None: + """ + Analyze and visualize all csv files in the given folder. + + """ + + if not os.path.exists(results_path): + raise Exception(f'The given directory "{results_path}" doesn\'t exist! No results found!') + + all_statistics = {} + + for file_name in os.listdir(results_path): + if file_name.endswith("_results.csv"): + model_name = file_name.split("_results.csv")[0] + all_statistics[model_name] = read_automatic_results(results_path, file_name) + + if len(list(all_statistics.keys())) == 0: + raise Exception(f'There are no csv files in the given directory "{results_path}"!') + + frame_all = {"model": [], "category": [], "metric": [], "score": []} + frame_per_metric = {} + for model_name, model_statistics in all_statistics.items(): + for metric, metric_statistics in model_statistics.items(): + if frame_per_metric.get(metric) is None: + frame_per_metric[metric] = {"model": [], "category": [], "score": []} + + for category, category_score in metric_statistics.items(): + frame_all["model"].append(model_name) + frame_all["category"].append(category) + frame_all["metric"].append(metric) + frame_all["score"].append(category_score) + + frame_per_metric[metric]["model"].append(model_name) + frame_per_metric[metric]["category"].append(category) + frame_per_metric[metric]["score"].append(category_score) + + if not os.path.exists(save_path): + os.makedirs(save_path) + + frame_all = pd.DataFrame(frame_all) + frame_all.to_csv(os.path.join(save_path, "automatic_evaluation_statistics.csv")) + + for metric in tqdm.tqdm( + frame_per_metric.keys(), + desc=f"automatic metrics: ", + total=len(frame_per_metric.keys()), + ): + data = pd.DataFrame(frame_per_metric[metric]) + + sns.set() + fig = plt.figure(figsize=(16, 10)) + + fig = sns.barplot(x="category", y="score", hue="model", data=data, dodge=True) + fig.set_title(f"Comparison between Different Models for Metric {metric.title()}") + plt.xlabel("Evaluation Category") + plt.ylabel("Score") + + figure = fig.get_figure() + figure.savefig(os.path.join(save_path, f"{metric}.png"), dpi=400) + + plt.close() diff --git a/applications/Chat/examples/README.md b/applications/Chat/examples/README.md new file mode 100644 index 000000000000..56e4cc992c17 --- /dev/null +++ b/applications/Chat/examples/README.md @@ -0,0 +1,303 @@ +# Examples + +## Table of Contents + +- [Examples](#examples) + - [Table of Contents](#table-of-contents) + - [Install requirements](#install-requirements) + - [Supervised datasets collection](#supervised-datasets-collection) + - [Stage1 - Supervised instructs tuning](#stage1---supervised-instructs-tuning) + - [Arg List](#arg-list) + - [Stage2 - Training reward model](#stage2---training-reward-model) + - [Features and tricks in RM training](#features-and-tricks-in-rm-training) + - [Experiment result](#experiment-result) + - [Arg List](#arg-list-1) + - [Stage3 - Training model using prompts with RL](#stage3---training-model-using-prompts-with-rl) + - [Arg List](#arg-list-2) + - [Inference example - After Stage3](#inference-example---after-stage3) + - [Attention](#attention) + - [data](#data) + - [Support Model](#support-model) + - [GPT](#gpt) + - [BLOOM](#bloom) + - [OPT](#opt) + - [LLaMA](#llama) + - [Add your own models](#add-your-own-models) + - [Actor model](#actor-model) + - [Reward model](#reward-model) + - [Critic model](#critic-model) + + +--- +## Install requirements + +```shell +pip install -r requirements.txt +``` + +## Supervised datasets collection + +We collected 104K bilingual dataset of Chinese and English, and you can find the datasets in this repo +[InstructionWild](https://github.com/XueFuzhao/InstructionWild). + +The following pic shows how we collected the data. +

+ +

+ +## Stage1 - Supervised instructs tuning + +Stage1 is supervised instructs fine-tuning, which uses the datasets mentioned earlier to fine-tune the model. +[[Stage1 tutorial video]](https://www.youtube.com/watch?v=-qFBZFmOJfg) + +You can run the `examples/train_sft.sh` to start a supervised instructs fine-tuning. + +You can also use the following cmd to start a supervised instructs fine-tuning with your own settings. +``` +torchrun --standalone --nproc_per_node=4 train_sft.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_zero2 \ + --log_interval 10 \ + --save_path /path/to/Coati-7B \ + --dataset /path/to/data.json \ + --batch_size 4 \ + --accumulation_steps 8 \ + --lr 2e-5 \ + --max_datasets_size 512 \ + --max_epochs 1 \ + --grad_checkpoint +``` +### Arg List +- --strategy: the strategy using for training, choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], default='colossalai_zero2' +- --model: model type, choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom' +- --pretrain: pretrain model, type=str, default=None +- --max_datasets_size: the max size of dataset, type=int, default=None +- --save_path: path to save the model, type=str, default='output' +- --need_optim_ckpt: whether to save optim ckpt, type=bool, default=False +- --max_epochs: max epochs for training, type=int, default=3 +- --batch_size: batch size while training, type=int, default=4 +- --lora_rank: low-rank adaptation matrices rank, type=int, default=0 +- --log_interval: how many steps to log, type=int, default=100 +- --grad_checkpoint: enable gradient checkpointing, type=bool, default=False + +## Stage2 - Training reward model + +We train a reward model in stage 2, which obtains corresponding scores by manually ranking different outputs for the same prompt and supervises the training of the reward model. +[[Stage2 tutorial video]](https://www.youtube.com/watch?v=gMx2CApKhuo) + +You can run the `examples/train_rm.sh` to start a reward model training. + +You can also use the following cmd to start training a reward model. +``` +torchrun --standalone --nproc_per_node=4 train_reward_model.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_zero2 \ + --loss_fn 'log_exp'\ + --save_path 'rmstatic.pt' \ +``` +### Features and tricks in RM training +- We support [Anthropic/hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf)and[rm-static](https://huggingface.co/datasets/Dahoas/rm-static) datasets. +- We support 2 kinds of loss_function named 'log_sig'(used by OpenAI) and 'log_exp'(used by Anthropic). +- We change the loss to valid_acc and pair_dist to monitor progress during training. +- We add special token to the end of the sequence to get better result. +- We use cosine-reducing lr-scheduler for RM training. +- We set value_head as 1 liner layer and initialize the weight of value_head using N(0，1/(d_model + 1)) distribution. +- We train a Bloom-560m reward model for 1 epoch and find the test acc of the model achieve the performance mentions in [Anthropics paper](https://arxiv.org/abs/2204.05862). + +### Experiment result +Model performance in [Anthropics paper](https://arxiv.org/abs/2204.05862): + +
image + +
Our training & test result of bloom-560m for 1 epoch: + +
image + +
We also train the reward model based on LLaMA-7B, which reaches the ACC of 72.06% after 1 epoch, performing almost the same as Anthropic's best RM. + +### Arg List +- --strategy: the strategy using for training, choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], default='colossalai_zero2' +- --model: model type, choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom' +- --pretrain: pretrain model, type=str, default=None +- --model_path: the path of rm model(if continue to train), type=str, default=None +- --save_path: path to save the model, type=str, default='output' +- --need_optim_ckpt: whether to save optim ckpt, type=bool, default=False +- --max_epochs: max epochs for training, type=int, default=3 +- --dataset: dataset name, type=str, choices=['Anthropic/hh-rlhf', 'Dahoas/rm-static'] +- --subset: subset of the dataset, type=str, default=None +- --batch_size: batch size while training, type=int, default=4 +- --lora_rank: low-rank adaptation matrices rank, type=int, default=0 +- --loss_func: which kind of loss function, choices=['log_sig', 'log_exp'] +- --max_len: max sentence length for generation, type=int, default=512 +- --test: whether is only testing, if it's true, the dataset will be small + +## Stage3 - Training model using prompts with RL + +Stage3 uses reinforcement learning algorithm, which is the most complex part of the training process, as shown below: + +

+ +

+ +You can run the `examples/train_prompts.sh` to start PPO training. +You can also use the cmd following to start PPO training. +[[Stage3 tutorial video]](https://www.youtube.com/watch?v=Z8wwSHxPL9g) + +``` +torchrun --standalone --nproc_per_node=4 train_prompts.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_zero2 \ + --prompt_dataset /path/to/your/prompt_dataset \ + --pretrain_dataset /path/to/your/pretrain_dataset \ + --rm_pretrain /your/pretrain/rm/definition \ + --rm_path /your/rm/model/path +``` + +Prompt dataset: the instruction dataset mentioned in the above figure which includes the instructions, e.g. you can use the [script](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples/generate_prompt_dataset.py) which samples `instinwild_en.json` or `instinwild_ch.json` in [InstructionWild](https://github.com/XueFuzhao/InstructionWild/tree/main/data#instructwild-data) to generate the prompt dataset. +Pretrain dataset: the pretrain dataset including the instruction and corresponding response, e.g. you can use the [InstructWild Data](https://github.com/XueFuzhao/InstructionWild/tree/main/data) in stage 1 supervised instructs tuning. + +### Arg List +- --strategy: the strategy using for training, choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], default='colossalai_zero2' +- --model: model type of actor, choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom' +- --pretrain: pretrain model, type=str, default=None +- --rm_model: reward model type, type=str, choices=['gpt2', 'bloom', 'opt', 'llama'], default=None +- --rm_pretrain: pretrain model for reward model, type=str, default=None +- --rm_path: the path of rm model, type=str, default=None +- --save_path: path to save the model, type=str, default='output' +- --prompt_dataset: path of the prompt dataset, type=str, default=None +- --pretrain_dataset: path of the ptx dataset, type=str, default=None +- --need_optim_ckpt: whether to save optim ckpt, type=bool, default=False +- --num_episodes: num of episodes for training, type=int, default=10 +- --num_update_steps: number of steps to update policy per episode, type=int +- --num_collect_steps: number of steps to collect experience per episode, type=int +- --train_batch_size: batch size while training, type=int, default=8 +- --ptx_batch_size: batch size to compute ptx loss, type=int, default=1 +- --experience_batch_size: batch size to make experience, type=int, default=8 +- --lora_rank: low-rank adaptation matrices rank, type=int, default=0 +- --kl_coef: kl_coef using for computing reward, type=float, default=0.1 +- --ptx_coef: ptx_coef using for computing policy loss, type=float, default=0.9 + +## Inference example - After Stage3 +We support different inference options, including int8 and int4 quantization. +For details, see [`inference/`](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/inference). + + +## Attention +The examples are demos for the whole training process.You need to change the hyper-parameters to reach great performance. + +#### data +- [x] [rm-static](https://huggingface.co/datasets/Dahoas/rm-static) +- [x] [hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) +- [ ] [openai/summarize_from_feedback](https://huggingface.co/datasets/openai/summarize_from_feedback) +- [ ] [openai/webgpt_comparisons](https://huggingface.co/datasets/openai/webgpt_comparisons) +- [ ] [Dahoas/instruct-synthetic-prompt-responses](https://huggingface.co/datasets/Dahoas/instruct-synthetic-prompt-responses) + +## Support Model + +### GPT +- [x] GPT2-S (s) +- [x] GPT2-M (m) +- [x] GPT2-L (l) +- [x] GPT2-XL (xl) +- [x] GPT2-4B (4b) +- [ ] GPT2-6B (6b) + +### BLOOM +- [x] [BLOOM-560m](https://huggingface.co/bigscience/bloom-560m) +- [x] [BLOOM-1b1](https://huggingface.co/bigscience/bloom-1b1) +- [x] [BLOOM-3b](https://huggingface.co/bigscience/bloom-3b) +- [x] [BLOOM-7b](https://huggingface.co/bigscience/bloom-7b1) +- [ ] [BLOOM-175b](https://huggingface.co/bigscience/bloom) + +### OPT +- [x] [OPT-125M](https://huggingface.co/facebook/opt-125m) +- [x] [OPT-350M](https://huggingface.co/facebook/opt-350m) +- [x] [OPT-1.3B](https://huggingface.co/facebook/opt-1.3b) +- [x] [OPT-2.7B](https://huggingface.co/facebook/opt-2.7b) +- [x] [OPT-6.7B](https://huggingface.co/facebook/opt-6.7b) +- [ ] [OPT-13B](https://huggingface.co/facebook/opt-13b) +- [ ] [OPT-30B](https://huggingface.co/facebook/opt-30b) + +### [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) +- [x] LLaMA-7B +- [x] LLaMA-13B +- [ ] LLaMA-33B +- [ ] LLaMA-65B + +## Add your own models + +If you want to support your own model in Coati, please refer the pull request for RoBERTa support as an example --[[chatgpt] add pre-trained model RoBERTa for RLHF stage 2 & 3](https://github.com/hpcaitech/ColossalAI/pull/3223), and submit a PR to us. + +You should complete the implementation of four model classes, including Reward model, Critic model, LM model, Actor model + +here are some example code for a NewModel named `Coati`. +if it is supported in huggingface [transformers](https://github.com/huggingface/transformers), you can load it by `from_pretrained`, o +r you can build your own model by yourself. + +### Actor model +``` +from ..base import Actor +from transformers.models.coati import CoatiModel + +class CoatiActor(Actor): + + def __init__(self, + pretrained: Optional[str] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none') -> None: + if pretrained is not None: + model = CoatiModel.from_pretrained(pretrained) + else: + model = build_model() # load your own model if it is not support in transformers + + super().__init__(model, lora_rank, lora_train_bias) +``` + +### Reward model +``` +from ..base import RewardModel +from transformers.models.coati import CoatiModel + +class CoatiRM(RewardModel): + + def __init__(self, + pretrained: Optional[str] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none') -> None: + if pretrained is not None: + model = CoatiModel.from_pretrained(pretrained) + else: + model = build_model() # load your own model if it is not support in transformers + + value_head = nn.Linear(model.config.n_embd, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.n_embd + 1)) + super().__init__(model, value_head, lora_rank, lora_train_bias) +``` + +### Critic model + +``` +from ..base import Critic +from transformers.models.coati import CoatiModel + +class CoatiCritic(Critic): + + def __init__(self, + pretrained: Optional[str] = None, + checkpoint: bool = False, + lora_rank: int = 0, + lora_train_bias: str = 'none') -> None: + if pretrained is not None: + model = CoatiModel.from_pretrained(pretrained) + else: + model = build_model() # load your own model if it is not support in transformers + + value_head = nn.Linear(model.config.n_embd, 1) + value_head.weight.data.normal_(mean=0.0, std=1 / (model.config.n_embd + 1)) + super().__init__(model, value_head, lora_rank, lora_train_bias) +``` diff --git a/applications/Chat/examples/community/README.md b/applications/Chat/examples/community/README.md new file mode 100644 index 000000000000..cd7b9d99bf06 --- /dev/null +++ b/applications/Chat/examples/community/README.md @@ -0,0 +1,24 @@ +# Community Examples +--- +We are thrilled to announce the latest updates to ColossalChat, an open-source solution for cloning ChatGPT with a complete RLHF (Reinforcement Learning with Human Feedback) pipeline. + +As Colossal-AI undergoes major updates, we are actively maintaining ColossalChat to stay aligned with the project's progress. With the introduction of Community-driven example, we aim to create a collaborative platform for developers to contribute exotic features built on top of ColossalChat. + +## Community Example + +Community-driven Examples is an initiative that allows users to contribute their own examples to the ColossalChat package, fostering a sense of community and making it easy for others to access and benefit from shared work. The primary goal with community-driven examples is to have a community-maintained collection of diverse and exotic functionalities built on top of the ColossalChat package, which is powered by the Colossal-AI project and its Coati module (ColossalAI Talking Intelligence). + +For more information about community pipelines, please have a look at this [issue](https://github.com/hpcaitech/ColossalAI/issues/3487). + +## Community Examples + +Community examples consist of both inference and training examples that have been added by the community. Please have a look at the following table to get an overview of all community examples. Click on the Code Example to get a copy-and-paste ready code example that you can try out. If a community doesn't work as expected, please open an issue and ping the author on it. + +| Example | Description | Code Example | Colab | Author | +|:---------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------:| +| Peft | Adding Peft support for SFT and Prompts model training | [Huggingface Peft](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples/community/peft) | - | [YY Lin](https://github.com/yynil) | +| Train prompts on Ray | A Ray based implementation of Train prompts example | [Training On Ray](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples/community/ray) | - | [MisterLin1995](https://github.com/MisterLin1995) | +|...|...|...|...|...| + +### How to get involved +To join our community-driven initiative, please visit the [ColossalChat GitHub repository](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples), review the provided information, and explore the codebase. To contribute, create a new issue outlining your proposed feature or enhancement, and our team will review and provide feedback. We look forward to collaborating with you on this exciting project! diff --git a/applications/Chat/examples/community/peft/README.md b/applications/Chat/examples/community/peft/README.md new file mode 100644 index 000000000000..844bfd3d22c3 --- /dev/null +++ b/applications/Chat/examples/community/peft/README.md @@ -0,0 +1,24 @@ +# Add Peft support for SFT and Prompts model training + +The original implementation just adopts the loralib and merges the layers into the final model. The huggingface peft is a better lora model implementation and can be easily training and distributed. + +Since reward model is relative small, I just keep it as original one. I suggest train full model to get the proper reward/critic model. + +# Preliminary installation +Since the current pypi peft package(0.2) has some bugs, please install the peft package using source. +``` +git clone https://github.com/huggingface/peft +cd peft +pip install . +``` + +# Usage +For SFT training, just call train_peft_sft.py + +Its arguments are almost identical to train_sft.py instead adding a new eval_dataset if you have a eval_dataset file. The data file is just a plain datafile, please check the format in the easy_dataset.py. + +For stage-3 rlhf training, call train_peft_prompts.py. +Its arguments are almost identical to train_prompts.py. The only difference is that I use text files to indicate the prompt and pretrained data file. The models are included in easy_models.py. Currently only bloom models are tested, but technically gpt2/opt/llama should be supported. + +# Dataformat +Please refer the formats in test_sft.txt, test_prompts.txt, test_pretrained.txt. diff --git a/applications/Chat/examples/community/peft/easy_dataset.py b/applications/Chat/examples/community/peft/easy_dataset.py new file mode 100644 index 000000000000..2fe293957079 --- /dev/null +++ b/applications/Chat/examples/community/peft/easy_dataset.py @@ -0,0 +1,240 @@ +import copy +import json +from typing import Dict, Sequence + +import torch +from datasets import load_dataset +from torch.utils.data import Dataset +from tqdm import tqdm +from transformers import AutoTokenizer + +IGNORE_INDEX = -100 + + +def _tokenize_fn(strings: Sequence[str], tokenizer: AutoTokenizer, max_length: int = 512) -> Dict: + """Tokenize a list of strings.""" + tokenized_list = [ + tokenizer( + text, + return_tensors="pt", + padding="longest", + max_length=max_length, + truncation=True, + ) for text in strings + ] + input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list] + input_ids_lens = labels_lens = [ + tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list + ] + return dict( + input_ids=input_ids, + labels=labels, + input_ids_lens=input_ids_lens, + labels_lens=labels_lens, + ) + + +def preprocess(sources: Sequence[str], targets: Sequence[str], tokenizer: AutoTokenizer, max_length: int = 512) -> Dict: + """Preprocess the data by tokenizing.""" + examples = [s + t for s, t in zip(sources, targets)] + examples_tokenized, sources_tokenized = [ + _tokenize_fn(strings, tokenizer, max_length) for strings in (examples, sources) + ] + input_ids = examples_tokenized["input_ids"] + labels = copy.deepcopy(input_ids) + for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]): + label[:source_len] = IGNORE_INDEX + return dict(input_ids=input_ids, labels=labels) + + +class EasySupervisedDataset(Dataset): + + def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length: int = 512) -> None: + super(EasySupervisedDataset, self).__init__() + with open(data_file, "r", encoding="UTF-8") as f: + all_lines = f.readlines() + #split to source and target ,source the characters before "回答：" including "回答：", target the characters after "回答：" + sources, targets = [], [] + for line in all_lines: + if "回答：" in line: + sep_index = line.index("回答：") + sources.append(line[:sep_index + 3]) + targets.append(line[sep_index + 3:] + tokenizer.eos_token) + else: + sources.append(line) + targets.append("" + tokenizer.eos_token) + data_dict = preprocess(sources, targets, tokenizer, max_length) + + self.input_ids = data_dict["input_ids"] + self.labels = data_dict["labels"] + self.data_file = data_file + + def __len__(self): + return len(self.input_ids) + + def __getitem__(self, i) -> Dict[str, torch.Tensor]: + return dict(input_ids=self.input_ids[i], labels=self.labels[i]) + + def __repr__(self): + return f"LawSupervisedDataset(data_file={self.data_file}, input_ids_len={len(self.input_ids)}, labels_len={len(self.labels)})" + + def __str__(self): + return f"LawSupervisedDataset(data_file={self.data_file}, input_ids_len={len(self.input_ids)}, labels_len={len(self.labels)})" + + +class EasyPromptsDataset(Dataset): + + def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length: int = 96) -> None: + super(EasyPromptsDataset, self).__init__() + with open(data_file, "r", encoding="UTF-8") as f: + all_lines = f.readlines() + all_lines = [line if "回答：" not in line else line[:line.index("回答：") + 3] for line in all_lines] + self.prompts = [ + tokenizer(line, return_tensors='pt', max_length=max_length, padding='max_length', + truncation=True)['input_ids'].to(torch.cuda.current_device()).squeeze(0) + for line in tqdm(all_lines) + ] + self.data_file = data_file + + def __len__(self): + return len(self.prompts) + + def __getitem__(self, idx): + return self.prompts[idx] + + def __repr__(self): + return f"LawPromptsDataset(data_file={self.data_file}, prompts_len={len(self.prompts)})" + + def __str__(self): + return f"LawPromptsDataset(data_file={self.data_file}, prompts_len={len(self.prompts)})" + + +class EasyRewardDataset(Dataset): + + def __init__(self, train_file: str, tokenizer: AutoTokenizer, special_token=None, max_length=512) -> None: + super(EasyRewardDataset, self).__init__() + self.chosen = [] + self.reject = [] + if special_token is None: + self.end_token = tokenizer.eos_token + else: + self.end_token = special_token + print(self.end_token) + #read all lines in the train_file to a list + with open(train_file, "r", encoding="UTF-8") as f: + all_lines = f.readlines() + for line in tqdm(all_lines): + data = json.loads(line) + prompt = "提问：" + data['prompt'] + " 回答：" + + chosen = prompt + data['chosen'] + self.end_token + chosen_token = tokenizer(chosen, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.chosen.append({ + "input_ids": chosen_token['input_ids'], + "attention_mask": chosen_token['attention_mask'] + }) + + reject = prompt + data['rejected'] + self.end_token + reject_token = tokenizer(reject, + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt") + self.reject.append({ + "input_ids": reject_token['input_ids'], + "attention_mask": reject_token['attention_mask'] + }) + + def __len__(self): + length = len(self.chosen) + return length + + def __getitem__(self, idx): + return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][ + "input_ids"], self.reject[idx]["attention_mask"] + + #python representation of the object and the string representation of the object + def __repr__(self): + return f"LawRewardDataset(chosen_len={len(self.chosen)}, reject_len={len(self.reject)})" + + def __str__(self): + return f"LawRewardDataset(chosen_len={len(self.chosen)}, reject_len={len(self.reject)})" + + +''' +Easy SFT just accept a text file which can be read line by line. However the datasets will group texts together to max_length so LLM will learn the texts meaning better. +If individual lines are not related, just set is_group_texts to False. +''' + + +class EasySFTDataset(Dataset): + + def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length=512, is_group_texts=True) -> None: + super().__init__() + #read the data_file line by line + with open(data_file, "r", encoding="UTF-8") as f: + #encode the text data line by line and put raw python list input_ids only to raw_input_ids list + raw_input_ids = [] + for line in f: + encoded_ids = tokenizer.encode(line) + #if the encoded_ids is longer than max_length, then split it into several parts + if len(encoded_ids) > max_length: + for i in range(0, len(encoded_ids), max_length): + raw_input_ids.append(encoded_ids[i:i + max_length]) + else: + raw_input_ids.append(encoded_ids) + + grouped_input_ids = [] + current_input_ids = [] + attention_mask = [] + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + if is_group_texts: + for input_ids in raw_input_ids: + if len(current_input_ids) + len(input_ids) > max_length: + #pad the current_input_ids to max_length with tokenizer.pad_token_id + padded_length = max_length - len(current_input_ids) + current_input_ids.extend([tokenizer.pad_token_id] * padded_length) + grouped_input_ids.append(torch.tensor(current_input_ids, dtype=torch.long)) + attention_mask.append( + torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long)) + current_input_ids = [] + else: + current_input_ids.extend(input_ids) + if len(current_input_ids) > 0: + padded_length = max_length - len(current_input_ids) + current_input_ids.extend([tokenizer.pad_token_id] * padded_length) + grouped_input_ids.append(torch.tensor(current_input_ids, dtype=torch.long)) + attention_mask.append( + torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long)) + else: + #just append the raw_input_ids to max_length + for input_ids in raw_input_ids: + padded_length = max_length - len(input_ids) + input_ids.extend([tokenizer.pad_token_id] * padded_length) + attention_mask.append( + torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long)) + grouped_input_ids.append(torch.tensor(input_ids, dtype=torch.long)) + self.input_ids = grouped_input_ids + self.labels = copy.deepcopy(self.input_ids) + self.file_name = data_file + self.attention_mask = attention_mask + + def __len__(self): + return len(self.input_ids) + + #get item from dataset + def __getitem__(self, idx): + return dict(input_ids=self.input_ids[idx], labels=self.labels[idx], attention_mask=self.attention_mask[idx]) + + #generate the dataset description to be printed by print in python + def __repr__(self): + return f"EasySFTDataset(len={len(self)},\nfile_name is {self.file_name})" + + #generate the dataset description to be printed by print in python + def __str__(self): + return f"EasySFTDataset(len={len(self)},\nfile_name is {self.file_name})" diff --git a/applications/ChatGPT/chatgpt/models/base/actor.py b/applications/Chat/examples/community/peft/easy_models.py similarity index 60% rename from applications/ChatGPT/chatgpt/models/base/actor.py rename to applications/Chat/examples/community/peft/easy_models.py index 57db2bb11a6a..fe294868159d 100644 --- a/applications/ChatGPT/chatgpt/models/base/actor.py +++ b/applications/Chat/examples/community/peft/easy_models.py @@ -3,26 +3,24 @@ import torch import torch.nn as nn import torch.nn.functional as F +from coati.models.generation import generate +from coati.models.utils import log_probs_from_logits, masked_mean +from peft import PeftModel +from torch.nn.modules import Module +from transformers import BloomConfig, BloomForCausalLM -from ..generation import generate -from ..lora import LoRAModule -from ..utils import log_probs_from_logits - -class Actor(LoRAModule): +class Actor(Module): """ Actor model base class. Args: model (nn.Module): Actor Model. - lora_rank (int): LoRA rank. - lora_train_bias (str): LoRA bias training mode. """ - def __init__(self, model: nn.Module, lora_rank: int = 0, lora_train_bias: str = 'none') -> None: - super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias) + def __init__(self, model: nn.Module) -> None: + super().__init__() self.model = model - self.convert_to_lora() @torch.no_grad() def generate( @@ -60,3 +58,39 @@ def forward(self, logits = output['logits'] log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:]) return log_probs[:, -num_actions:] + + def get_base_model(self): + return self.model + + +class BLOOMActor(Actor): + """ + BLOOM Actor model. + + Args: + pretrained (str): Pretrained model name or path. + config (BloomConfig): Model config. + checkpoint (bool): Enable gradient checkpointing. + lora_rank (int): LoRA rank. + lora_train_bias (str): LoRA bias training mode. + """ + + def __init__(self, + pretrained: str = None, + config: Optional[BloomConfig] = None, + checkpoint: bool = False, + lora_path: str = None) -> None: + if pretrained is not None: + model = BloomForCausalLM.from_pretrained(pretrained) + elif config is not None: + model = BloomForCausalLM(config) + else: + model = BloomForCausalLM(BloomConfig()) + if lora_path is not None: + model = PeftModel.from_pretrained(model, lora_path) + if checkpoint: + model.gradient_checkpointing_enable() + super().__init__(model) + + def print_trainable_parameters(self): + self.get_base_model().print_trainable_parameters() diff --git a/applications/Chat/examples/community/peft/train_peft_prompts.py b/applications/Chat/examples/community/peft/train_peft_prompts.py new file mode 100644 index 000000000000..9385e457d852 --- /dev/null +++ b/applications/Chat/examples/community/peft/train_peft_prompts.py @@ -0,0 +1,222 @@ +import argparse + +import pandas as pd +import torch +import torch.distributed as dist +from coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset +from coati.models.bloom import BLOOMRM, BLOOMCritic +from coati.models.gpt import GPTRM, GPTActor, GPTCritic +from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM +from coati.models.opt import OPTRM, OPTActor, OPTCritic +from coati.trainer import PPOTrainer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from easy_dataset import EasyPromptsDataset, EasySupervisedDataset +from easy_models import BLOOMActor +from peft import PeftModel +from torch.optim import Adam +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer + +from colossalai.nn.optimizer import HybridAdam + + +def main(args): + # configure strategy + if args.strategy == 'ddp': + strategy = DDPStrategy() + elif args.strategy == 'colossalai_gemini': + strategy = GeminiStrategy(placement_policy='cpu', initial_scale=2**5) + elif args.strategy == 'colossalai_zero2': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cpu') + else: + raise ValueError(f'Unsupported strategy "{args.strategy}"') + + if args.rm_path is not None: + state_dict = torch.load(args.rm_path, map_location='cpu') + + # configure model + if args.model == 'bloom': + # initial_model = BLOOMActor(pretrained=args.pretrain) + print('Using peft lora to load Bloom model as initial_model') + initial_model = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path) + print('Using peft lora to load Bloom model as initial_model (Done)') + else: + raise ValueError(f'Unsupported actor model "{args.model}"') + + if args.rm_model == None: + rm_model_name = args.model + else: + rm_model_name = args.rm_model + + if rm_model_name == 'gpt2': + reward_model = GPTRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'bloom': + print("load bloom reward model ", args.rm_pretrain) + reward_model = BLOOMRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'opt': + reward_model = OPTRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'llama': + reward_model = LlamaRM(pretrained=args.rm_pretrain) + else: + raise ValueError(f'Unsupported reward model "{rm_model_name}"') + + if args.rm_path is not None: + print('Loading reward model from', args.rm_path) + reward_model.load_state_dict(state_dict) + + if args.strategy != 'colossalai_gemini': + initial_model.to(torch.float16).to(torch.cuda.current_device()) + reward_model.to(torch.float16).to(torch.cuda.current_device()) + + with strategy.model_init_context(): + if args.model == 'bloom': + # actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank) + print('Using peft lora to load Bloom model as Actor') + actor = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path) + print('Using peft lora to load Bloom model as Actor (Done)') + else: + raise ValueError(f'Unsupported actor model "{args.model}"') + + if rm_model_name == 'gpt2': + critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + elif rm_model_name == 'bloom': + print("load bloom critic ", args.rm_pretrain, " lora_rank ", args.lora_rank, " use_action_mask ", True) + critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + print("load bloom critic (Done) ") + elif rm_model_name == 'opt': + critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + elif rm_model_name == 'llama': + critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + else: + raise ValueError(f'Unsupported reward model "{rm_model_name}"') + + if args.rm_path is not None: + print('Loading reward model from', args.rm_path) + critic.load_state_dict(state_dict) + del state_dict + + if args.strategy != 'colossalai_gemini': + critic.to(torch.float16).to(torch.cuda.current_device()) + actor.to(torch.float16).to(torch.cuda.current_device()) + + # configure optimizer + if args.strategy.startswith('colossalai'): + actor_optim = HybridAdam(actor.parameters(), lr=1e-7) + critic_optim = HybridAdam(critic.parameters(), lr=1e-7) + else: + actor_optim = Adam(actor.parameters(), lr=1e-7) + critic_optim = Adam(critic.parameters(), lr=1e-7) + + # configure tokenizer + if args.model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained(args.rm_pretrain) + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained(args.rm_pretrain) + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'opt': + tokenizer = AutoTokenizer.from_pretrained(args.rm_pretrain) + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'llama': + tokenizer = LlamaTokenizer.from_pretrained(args.pretrain) + tokenizer.eos_token = '<\s>' + tokenizer.pad_token = tokenizer.unk_token + else: + raise ValueError(f'Unsupported model "{args.model}"') + + data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) + + prompt_dataset = EasyPromptsDataset(args.prompt_path, tokenizer) + if dist.is_initialized() and dist.get_world_size() > 1: + prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True) + else: + prompt_sampler = None + prompt_dataloader = DataLoader(prompt_dataset, + shuffle=(prompt_sampler is None), + sampler=prompt_sampler, + batch_size=args.train_batch_size) + + pretrain_dataset = EasySupervisedDataset(args.pretrain_dataset, tokenizer) + if dist.is_initialized() and dist.get_world_size() > 1: + pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True) + else: + pretrain_sampler = None + pretrain_dataloader = DataLoader(pretrain_dataset, + shuffle=(pretrain_sampler is None), + sampler=pretrain_sampler, + batch_size=args.ptx_batch_size, + collate_fn=data_collator) + + def tokenize_fn(texts): + # MUST padding to max length to ensure inputs of all ranks have the same length + # Different length may lead to hang when using gemini, as different generation steps + batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True) + return {k: v.to(torch.cuda.current_device()) for k, v in batch.items()} + + (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim)) + + # configure trainer + trainer = PPOTrainer( + strategy, + actor, + critic, + reward_model, + initial_model, + actor_optim, + critic_optim, + kl_coef=args.kl_coef, + ptx_coef=args.ptx_coef, + train_batch_size=args.train_batch_size, + experience_batch_size=args.experience_batch_size, + tokenizer=tokenize_fn, + max_length=512, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + ) + + trainer.fit(prompt_dataloader=prompt_dataloader, + pretrain_dataloader=pretrain_dataloader, + num_episodes=args.num_episodes, + num_update_steps=args.num_update_steps, + num_collect_steps=args.num_collect_steps) + + # save model checkpoint after fitting + trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer) + # save optimizer checkpoint on all ranks + if args.need_optim_ckpt: + strategy.save_optimizer(actor_optim, + 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()), + only_rank0=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset') + parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset') + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], + default='ddp', + help='strategy to use') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--sft_lora_path', type=str, default=None) + parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--rm_path', type=str, default=None) + parser.add_argument('--rm_pretrain', type=str, default=None) + parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts') + parser.add_argument('--need_optim_ckpt', type=bool, default=False) + parser.add_argument('--num_episodes', type=int, default=10) + parser.add_argument('--num_collect_steps', type=int, default=10) + parser.add_argument('--num_update_steps', type=int, default=5) + parser.add_argument('--train_batch_size', type=int, default=2) + parser.add_argument('--ptx_batch_size', type=int, default=1) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + parser.add_argument('--kl_coef', type=float, default=0.1) + parser.add_argument('--ptx_coef', type=float, default=0.9) + args = parser.parse_args() + main(args) diff --git a/applications/Chat/examples/community/peft/train_peft_sft.py b/applications/Chat/examples/community/peft/train_peft_sft.py new file mode 100644 index 000000000000..4af08e6d0141 --- /dev/null +++ b/applications/Chat/examples/community/peft/train_peft_sft.py @@ -0,0 +1,184 @@ +import argparse +import os + +import loralib as lora +import torch +import torch.distributed as dist +from coati.dataset import DataCollatorForSupervisedDataset, SFTDataset, SupervisedDataset +from coati.models.base import RewardModel +from coati.models.bloom import BLOOMLM +from coati.models.gpt import GPTLM +from coati.models.llama import LlamaLM +from coati.models.opt import OPTLM +from coati.trainer import SFTTrainer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from datasets import load_dataset +from easy_dataset import EasyDataset +from peft import LoraConfig, PeftModel, TaskType, get_peft_model +from torch.optim import Adam +from torch.utils.data import DataLoader +from torch.utils.data.dataloader import default_collate +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoModelForCausalLM, AutoTokenizer, BloomTokenizerFast +from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer + +from colossalai.logging import get_dist_logger +from colossalai.nn.optimizer import HybridAdam +from colossalai.tensor import ColoParameter + + +def train(args): + # configure strategy + if args.strategy == 'ddp': + strategy = DDPStrategy() + elif args.strategy == 'colossalai_gemini': + strategy = GeminiStrategy(placement_policy='cuda') + elif args.strategy == 'colossalai_zero2': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + else: + raise ValueError(f'Unsupported strategy "{args.strategy}"') + + # configure model + with strategy.model_init_context(): + print('Warning: currently only bloom is tested, gpt2,llama and opt are not tested') + model = AutoModelForCausalLM.from_pretrained(args.pretrain).to(torch.cuda.current_device()) + # if the args.save_path exists and args.save_path+'/adapter_config.json' exists, we'll load the adapter_config.json + if os.path.exists(args.save_path) and os.path.exists(args.save_path + '/adapter_config.json') \ + and os.path.exists(args.save_path + '/adapter_model.bin'): + print("loading from saved peft model ", args.save_path) + model = PeftModel.from_pretrained(model, args.save_path) + else: + # we'll use peft lora library to do the lora + lora_rank = args.lora_rank if args.lora_rank > 0 else 32 + # config lora with rank of lora_rank + lora_config = LoraConfig(task_type=TaskType.CAUSAL_LM, + inference_mode=False, + r=lora_rank, + lora_alpha=32, + lora_dropout=0.1) + model = get_peft_model(model, lora_config) + model.print_trainable_parameters() + + # configure tokenizer + if args.model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained('gpt2') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m") + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'opt': + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'llama': + tokenizer = AutoTokenizer.from_pretrained( + args.pretrain, + padding_side="right", + use_fast=False, + ) + tokenizer.eos_token = '<\s>' + tokenizer.pad_token = tokenizer.unk_token + else: + raise ValueError(f'Unsupported model "{args.model}"') + + if args.model == 'llama' and args.strategy == 'colossalai_gemini': + # this is a hack to deal with the resized embedding + # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility + for name, param in model.named_parameters(): + if not isinstance(param, ColoParameter): + sub_module_name = '.'.join(name.split('.')[:-1]) + weight_name = name.split('.')[-1] + sub_module = model.get_submodule(sub_module_name) + setattr(sub_module, weight_name, ColoParameter(param)) + + # configure optimizer + if args.strategy.startswith('colossalai'): + optim = HybridAdam(model.parameters(), lr=args.lr, clipping_norm=1.0) + else: + optim = Adam(model.parameters(), lr=args.lr) + + logger = get_dist_logger() + logger.set_level('WARNING') + + # configure dataset + law_dataset = EasyDataset(args.dataset, tokenizer=tokenizer, is_group_texts=not args.is_short_text) + train_dataset = law_dataset + print(train_dataset) + eval_dataset = None + if args.eval_dataset is not None: + eval_dataset = EasyDataset(args.eval_dataset, tokenizer=tokenizer, is_group_texts=not args.is_short_text) + data_collator = default_collate + if dist.is_initialized() and dist.get_world_size() > 1: + train_sampler = DistributedSampler(train_dataset, + shuffle=True, + seed=42, + drop_last=True, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + if eval_dataset is not None: + eval_sampler = DistributedSampler(eval_dataset, + shuffle=False, + seed=42, + drop_last=False, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + else: + train_sampler = None + eval_sampler = None + + train_dataloader = DataLoader(train_dataset, + shuffle=(train_sampler is None), + sampler=train_sampler, + batch_size=args.batch_size, + collate_fn=data_collator, + pin_memory=True) + if eval_dataset is not None: + eval_dataloader = DataLoader(eval_dataset, + shuffle=(eval_sampler is None), + sampler=eval_sampler, + batch_size=args.batch_size, + collate_fn=data_collator, + pin_memory=True) + else: + eval_dataloader = None + + trainer = SFTTrainer(model=model, + strategy=strategy, + optim=optim, + train_dataloader=train_dataloader, + eval_dataloader=eval_dataloader, + batch_size=args.batch_size, + max_epochs=args.max_epochs, + accumulation_steps=args.accumulation_steps) + + trainer.fit(logger=logger, log_interval=args.log_interval) + + # save model checkpoint after fitting on only rank0 + trainer.save_model(path=args.save_path, only_rank0=True, tokenizer=tokenizer) + # save optimizer checkpoint on all ranks + if args.need_optim_ckpt: + strategy.save_optimizer(trainer.optimizer, + 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), + only_rank0=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], + default='ddp') + parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom') + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--dataset', type=str, default=None) + parser.add_argument('--eval_dataset', type=str, default=None) + parser.add_argument('--save_path', type=str, default='output') + parser.add_argument('--need_optim_ckpt', type=bool, default=False) + parser.add_argument('--max_epochs', type=int, default=3) + parser.add_argument('--batch_size', type=int, default=4) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + parser.add_argument('--log_interval', type=int, default=100, help="how many steps to log") + parser.add_argument('--lr', type=float, default=5e-6) + parser.add_argument('--accumulation_steps', type=int, default=8) + parser.add_argument('--enable_peft_lora', action='store_true', default=False) + parser.add_argument("--is_short_text", action='store_true', default=False) + args = parser.parse_args() + train(args) diff --git a/applications/Chat/examples/community/ray/README.md b/applications/Chat/examples/community/ray/README.md new file mode 100644 index 000000000000..64360bd73ddc --- /dev/null +++ b/applications/Chat/examples/community/ray/README.md @@ -0,0 +1,17 @@ +# ColossalAI on Ray +## Abstract +This is an experimental effort to run ColossalAI Chat training on Ray +## How to use? +### 1. Setup Ray clusters +Please follow the official [Ray cluster setup instructions](https://docs.ray.io/en/latest/cluster/getting-started.html) to setup an cluster with GPU support. Record the cluster's api server endpoint, it should be something similar to http://your.head.node.addrees:8265 +### 2. Clone repo +Clone this project: +```shell +git clone https://github.com/hpcaitech/ColossalAI.git +``` +### 3. Submit the ray job +```shell +python applications/Chat/examples/community/ray/ray_job_script.py http://your.head.node.addrees:8265 +``` +### 4. View your job on the Ray Dashboard +Open your ray cluster dashboard http://your.head.node.addrees:8265 to view your submitted training job. diff --git a/applications/Chat/examples/community/ray/ray_job_script.py b/applications/Chat/examples/community/ray/ray_job_script.py new file mode 100644 index 000000000000..53f304d379fe --- /dev/null +++ b/applications/Chat/examples/community/ray/ray_job_script.py @@ -0,0 +1,22 @@ +import sys + +from ray.job_submission import JobSubmissionClient + + +def main(api_server_endpoint="http://127.0.0.1:8265"): + client = JobSubmissionClient(api_server_endpoint) + client.submit_job( + entrypoint= + "python experimental/ray/train_prompts_on_ray.py --strategy colossalai_zero2 --prompt_csv_url https://huggingface.co/datasets/fka/awesome-chatgpt-prompts/resolve/main/prompts.csv", + runtime_env={ + "working_dir": + "applications/Chat", + "pip": [ + "torch==1.13.1", "transformers>=4.20.1", "datasets", "loralib", "colossalai>=0.2.4", "langchain", + "tokenizers", "fastapi", "sse_starlette", "wandb", "sentencepiece", "gpustat" + ] + }) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/applications/Chat/examples/community/ray/train_prompts_on_ray.py b/applications/Chat/examples/community/ray/train_prompts_on_ray.py new file mode 100644 index 000000000000..1bba9ad66fbc --- /dev/null +++ b/applications/Chat/examples/community/ray/train_prompts_on_ray.py @@ -0,0 +1,553 @@ +import argparse +import logging +import os +import socket +from copy import deepcopy +from typing import Type + +import ray +import torch +from coati.experience_maker.base import Experience +from coati.models.base import RewardModel +from coati.models.bloom import BLOOMActor, BLOOMCritic +from coati.models.gpt import GPTActor, GPTCritic +from coati.models.lora import LoRAModule +from coati.models.loss import PolicyLoss, ValueLoss +from coati.models.opt import OPTActor, OPTCritic +from coati.models.utils import compute_reward +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from torch.optim import Adam +from transformers import AutoTokenizer, BloomTokenizerFast +from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer + +from colossalai.nn.optimizer import HybridAdam + + +class ExperienceCompositionRefs: + + def __init__(self, sequences_attention_mask_action_mask_ref: ray.ObjectRef, action_log_probs_ref: ray.ObjectRef, + base_action_log_probs_ref: ray.ObjectRef, value_ref: ray.ObjectRef, r_ref: ray.ObjectRef) -> None: + self.sequences_attention_mask_action_mask_ref = sequences_attention_mask_action_mask_ref + self.action_log_probs_ref = action_log_probs_ref + self.base_action_log_probs_ref = base_action_log_probs_ref + self.value_ref = value_ref + self.r_ref = r_ref + + +class ExperienceMaker: + + def __init__(self, kl_coef) -> None: + self.kl_coef = kl_coef + + @torch.no_grad() + def make_experience(self, experiment_computation_refs: ExperienceCompositionRefs): + sequences, attention_mask, action_mask = ray.get( + experiment_computation_refs.sequences_attention_mask_action_mask_ref) + action_log_probs = ray.get(experiment_computation_refs.action_log_probs_ref) + base_action_log_probs = ray.get(experiment_computation_refs.base_action_log_probs_ref) + r = ray.get(experiment_computation_refs.r_ref) + reward = compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_mask=action_mask) + value = ray.get(experiment_computation_refs.value_ref) + advantage = reward - value + if advantage.ndim == 1: + advantage = advantage.unsqueeze(-1) + experience = Experience(sequences, action_log_probs, value, reward, advantage, attention_mask, action_mask) + return experience + + +class DistributedTorchRayActor: + + def __init__(self, world_size, rank, local_rank, master_addr, master_port): + logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s', + level=logging.INFO, + datefmt='%Y-%m-%d %H:%M:%S') + self._model = None + self._world_size = world_size + self._rank = rank + self._local_rank = local_rank + self._master_addr = master_addr if master_addr else self._get_current_node_ip() + self._master_port = master_port if master_port else self._get_free_port() + os.environ["MASTER_ADDR"] = self._master_addr + os.environ["MASTER_PORT"] = str(self._master_port) + os.environ["WORLD_SIZE"] = str(self._world_size) + os.environ["RANK"] = str(self._rank) + os.environ["LOCAL_RANK"] = str(self._local_rank) + + @staticmethod + def _get_current_node_ip(): + return ray._private.services.get_node_ip_address() + + @staticmethod + def _get_free_port(): + with socket.socket() as sock: + sock.bind(('', 0)) + return sock.getsockname()[1] + + def get_master_addr_port(self): + return self._master_addr, self._master_port + + +class BasePPORole(DistributedTorchRayActor): + + def add_experience_maker(self, kl_coef: float = 0.1): + self._experience_maker = ExperienceMaker(kl_coef) + + def make_experience(self, experience_computation_ref: ExperienceCompositionRefs): + return self._experience_maker.make_experience(experience_computation_ref) + + def _init_strategy(self, strategy: str): + # configure strategy + if strategy == 'ddp': + self._strategy = DDPStrategy() + elif strategy == 'colossalai_gemini': + self._strategy = GeminiStrategy(placement_policy='cuda', initial_scale=2**5) + elif strategy == 'colossalai_zero2': + self._strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + else: + raise ValueError(f'Unsupported strategy "{strategy}"') + + def _init_optimizer(self): + if isinstance(self._strategy, (GeminiStrategy, LowLevelZeroStrategy)): + self._optimizer = HybridAdam(self._model.parameters(), lr=5e-6) + else: + self._optimizer = Adam(self._model.parameters(), lr=5e-6) + + def _prepare_model_with_strategy(self, has_optimizer: bool): + if has_optimizer: + self._init_optimizer() + (self._model, self._optimizer) = self._strategy.prepare((self._model, self._optimizer)) + else: + self._model = self._strategy.prepare(self._model) + + def _load_model_from_pretrained(self, model_class: Type[LoRAModule], pretrain: str): + raise NotImplementedError() + + def init_model_from_pretrained(self, + strategy: str, + model_class: Type[LoRAModule], + pretrain: str, + has_optimizer=False): + self._init_strategy(strategy) + self._load_model_from_pretrained(model_class, pretrain) + self._prepare_model_with_strategy(has_optimizer) + + def eval(self): + self._model.eval() + + +class TrainablePPORole(BasePPORole): + + def _load_model_from_pretrained(self, model_class, pretrain): + with self._strategy.model_init_context(): + self._model = model_class(pretrain).to(torch.cuda.current_device()) + + def _train(self): + self._model.train() + + def _training_step(self, experience: Experience): + raise NotImplementedError() + + def learn_on_experiences(self, experience_refs): + experiences = ray.get(experience_refs) + device = torch.cuda.current_device() + self._train() + for exp in experiences: + exp.to_device(device) + self._training_step(exp) + self.eval() + + +@ray.remote(num_gpus=1) +class RayPPOActor(TrainablePPORole): + + def set_loss_function(self, eps_clip: float): + self._actor_loss_fn = PolicyLoss(eps_clip) + + def load_tokenizer_from_pretrained(self, model_type: str, pretrained): + if model_type == 'gpt2': + self._model_tokenizer = GPT2Tokenizer.from_pretrained(pretrained) + self._model_tokenizer.pad_token = self._model_tokenizer.eos_token + elif model_type == 'bloom': + self._model_tokenizer = BloomTokenizerFast.from_pretrained(pretrained) + self._model_tokenizer.pad_token = self._model_tokenizer.eos_token + elif model_type == 'opt': + self._model_tokenizer = AutoTokenizer.from_pretrained(pretrained) + else: + raise ValueError(f'Unsupported model "{model_type}"') + + # Set tokenize function for sequence generation + def _text_input_tokenize_fn(texts): + batch = self._model_tokenizer(texts, return_tensors='pt', max_length=96, padding=True, truncation=True) + return {k: v.cuda() for k, v in batch.items()} + + self._sample_tokenize_function = _text_input_tokenize_fn + + def setup_generate_kwargs(self, generate_kwargs: dict): + from coati.trainer.ppo import _set_default_generate_kwargs + self._generate_kwargs = _set_default_generate_kwargs(self._strategy, generate_kwargs, self._model) + self._generate_kwargs['pad_token_id'] = self._model_tokenizer.pad_token_id + self._generate_kwargs['eos_token_id'] = self._model_tokenizer.eos_token_id + + def load_csv_prompt_file_from_url_to_sampler(self, prompt_url): + import pandas as pd + prompts = pd.read_csv(prompt_url)['prompt'] + self._sampler = self._strategy.setup_sampler(prompts) + + def _generate(self, input_ids, **generate_kwargs): + return self._model.generate(input_ids, return_action_mask=True, **generate_kwargs) + + def sample_prompts_and_make_sequence(self, experience_batch_size): + sampled_prompts = self._sampler.sample(experience_batch_size) + input_ids = self._sample_tokenize_function(sampled_prompts) + if isinstance(input_ids, dict): + return self._generate(**input_ids, **self._generate_kwargs) + else: + return self._generate(input_ids, **self._generate_kwargs) + + @torch.no_grad() + def calculate_action_log_probs(self, sequence_attention_action_mask): + sequences, attention_mask, action_mask = sequence_attention_action_mask + return self._model.forward(sequences, action_mask.size(1), attention_mask) + + def _training_step(self, experience): + num_actions = experience.action_mask.size(1) + action_log_probs = self._model(experience.sequences, num_actions, attention_mask=experience.attention_mask) + actor_loss = self._actor_loss_fn(action_log_probs, + experience.action_log_probs, + experience.advantages, + action_mask=experience.action_mask) + self._strategy.backward(actor_loss, self._model, self._optimizer) + self._strategy.optimizer_step(self._optimizer) + self._optimizer.zero_grad() + logging.info("actor_loss: {}".format(actor_loss)) + + def save_checkpoint(self, save_path, should_save_optimizer: bool): + if self._rank == 0: + # save model checkpoint only on rank 0 + self._strategy.save_model(self._model, save_path, only_rank0=True) + # save optimizer checkpoint on all ranks + if should_save_optimizer: + self._strategy.save_optimizer(self._optimizer, + 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()), + only_rank0=False) + + def generate_answer(self, prompt, max_length=30, num_return_sequences=5): + encoded_input = self._model_tokenizer(prompt, return_tensors='pt') + input_ids = {k: v.cuda() for k, v in encoded_input.items()} + sequence, _ = self._model.generate(**input_ids, + max_length=max_length, + return_action_mask=False, + num_return_sequences=num_return_sequences) + token_list = list(sequence.data[0]) + output = " ".join([self._model_tokenizer.decode(token) for token in token_list]) + return output + + +@ray.remote(num_gpus=1) +class RayPPOCritic(TrainablePPORole): + + def set_loss_function(self, value_clip: float): + self._critic_loss_fn = ValueLoss(value_clip) + + def _training_step(self, experience): + values = self._model(experience.sequences, + action_mask=experience.action_mask, + attention_mask=experience.attention_mask) + critic_loss = self._critic_loss_fn(values, + experience.values, + experience.reward, + action_mask=experience.action_mask) + self._strategy.backward(critic_loss, self._model, self._optimizer) + self._strategy.optimizer_step(self._optimizer) + self._optimizer.zero_grad() + logging.info("critic_loss: {}".format(critic_loss)) + + @torch.no_grad() + def calculate_value(self, sequence_attention_action_mask): + sequences, attention_mask, action_mask = sequence_attention_action_mask + return self._model(sequences, action_mask, attention_mask) + + +@ray.remote(num_gpus=1) +class RayPPORewardModel(BasePPORole): + + def _load_model_from_pretrained(self, model_class, pretrain): + with self._strategy.model_init_context(): + critic = model_class(pretrained=pretrain).to(torch.cuda.current_device()) + self._model = RewardModel(deepcopy(critic.model), + deepcopy(critic.value_head)).to(torch.cuda.current_device()) + + @torch.no_grad() + def calculate_r(self, sequence_attention_action_mask): + sequences, attention_mask, _ = sequence_attention_action_mask + return self._model(sequences, attention_mask) + + +@ray.remote(num_gpus=1) +class RayPPOInitialModel(BasePPORole): + + def _load_model_from_pretrained(self, model_class, pretrain): + with self._strategy.model_init_context(): + self._model = model_class(pretrain).to(torch.cuda.current_device()) + + @torch.no_grad() + def calculate_base_action_log_probs(self, sequence_attention_action_mask): + sequences, attention_mask, action_mask = sequence_attention_action_mask + return self._model(sequences, action_mask.size(1), attention_mask) + + +class PPORayActorGroup: + """ + A group of ray actors + Functions start with 'async' should return list of object refs + """ + + def __init__(self, num_nodes, num_gpus_per_node, ray_actor_type: Type[BasePPORole]) -> None: + self._num_nodes = num_nodes + self._num_gpus_per_node = num_gpus_per_node + self.ray_actor_type = ray_actor_type + self._initiate_actors() + + def _initiate_actors(self): + world_size = self._num_nodes * self._num_gpus_per_node + # Use placement group to lock resources for models of same type + pg = None + if self._num_gpus_per_node > 1: + bundles = [{"GPU": self._num_gpus_per_node, "CPU": self._num_gpus_per_node} for _ in range(self._num_nodes)] + pg = placement_group(bundles, strategy="STRICT_SPREAD") + ray.get(pg.ready()) + if pg: + master_actor = self.ray_actor_type.options(scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, placement_group_bundle_index=0)).remote(world_size, 0, 0, None, None) + else: + master_actor = self.ray_actor_type.options(num_gpus=1).remote(world_size, 0, 0, None, None) + self._actor_handlers = [master_actor] + + # Create worker actors + if world_size > 1: + master_addr, master_port = ray.get(master_actor.get_master_addr_port.remote()) + for rank in range(1, world_size): + local_rank = rank % self._num_gpus_per_node + if pg: + worker_actor = self.ray_actor_type.options(scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, placement_group_bundle_index=rank // self._num_gpus_per_node)).remote( + world_size, rank, local_rank, master_addr, master_port) + else: + worker_actor = self.ray_actor_type.options(num_gpus=1).remote(world_size, rank, local_rank, + master_addr, master_port) + self._actor_handlers.append(worker_actor) + + def async_init_model_from_pretrained(self, strategy: str, model_class: Type[LoRAModule], pretrain: str, + has_optimizer: bool): + return [ + actor.init_model_from_pretrained.remote(strategy, model_class, pretrain, has_optimizer) + for actor in self._actor_handlers + ] + + +class TrainableModelRayActorGroup(PPORayActorGroup): + + def async_learn_on_experiences(self, experience_refs): + num_actors = len(self._actor_handlers) + learn_result_refs = [] + for i in range(num_actors): + exp_refs_batch = experience_refs[i::num_actors] + learn_result_refs.append(self._actor_handlers[i].learn_on_experiences.remote(exp_refs_batch)) + return learn_result_refs + + +class PPOActorRayActorGroup(TrainableModelRayActorGroup): + + def __init__(self, num_nodes, num_gpus_per_node) -> None: + super().__init__(num_nodes, num_gpus_per_node, RayPPOActor) + + def async_prepare_for_sequence_generation(self, model: str, pretrain: str, generation_kwargs: dict): + refs = [] + for actor in self._actor_handlers: + refs.append(actor.load_tokenizer_from_pretrained.remote(model, pretrain)) + refs.append(actor.setup_generate_kwargs.remote(generation_kwargs)) + return refs + + def load_csv_prompt_file_from_url_to_sampler(self, csv_url): + ray.get([actor.load_csv_prompt_file_from_url_to_sampler.remote(csv_url) for actor in self._actor_handlers]) + + def async_sample_prompts_and_make_sequence(self, experience_batch_size): + return [actor.sample_prompts_and_make_sequence.remote(experience_batch_size) for actor in self._actor_handlers] + + def async_calculate_action_log_probs(self, sequences_attention_mask_action_mask_refs): + num_actors = len(self._actor_handlers) + action_log_probs_refs = [] + for i in range(len(sequences_attention_mask_action_mask_refs)): + action_log_probs_ref = self._actor_handlers[i % num_actors].calculate_action_log_probs.remote( + sequences_attention_mask_action_mask_refs[i]) + action_log_probs_refs.append(action_log_probs_ref) + return action_log_probs_refs + + def set_loss_function(self, eps_clip: float = 0.2): + ray.get([actor.set_loss_function.remote(eps_clip) for actor in self._actor_handlers]) + + def save_checkpoint(self, save_path, should_save_optimizer): + ray.get([actor.save_checkpoint.remote(save_path, should_save_optimizer) for actor in self._actor_handlers]) + + +class PPOCriticRayActorGroup(TrainableModelRayActorGroup): + + def __init__(self, num_nodes, num_gpus_per_node) -> None: + super().__init__(num_nodes, num_gpus_per_node, RayPPOCritic) + + def async_calculate_value(self, sequences_attention_mask_action_mask_refs): + num_actors = len(self._actor_handlers) + value_refs = [] + for i in range(len(sequences_attention_mask_action_mask_refs)): + value_ref = self._actor_handlers[i % num_actors].calculate_value.remote( + sequences_attention_mask_action_mask_refs[i]) + value_refs.append(value_ref) + return value_refs + + def set_loss_function(self, value_clip: float = 0.4): + ray.get([actor.set_loss_function.remote(value_clip) for actor in self._actor_handlers]) + + +class PPOInitialRayActorGroup(PPORayActorGroup): + + def __init__(self, num_nodes, num_gpus_per_node) -> None: + super().__init__(num_nodes, num_gpus_per_node, RayPPOInitialModel) + + def async_calculate_base_action_log_probs(self, sequences_attention_mask_action_mask_refs): + num_actors = len(self._actor_handlers) + base_action_log_probs_refs = [] + for i in range(len(sequences_attention_mask_action_mask_refs)): + base_action_log_probs_ref = self._actor_handlers[i % num_actors].calculate_base_action_log_probs.remote( + sequences_attention_mask_action_mask_refs[i]) + base_action_log_probs_refs.append(base_action_log_probs_ref) + return base_action_log_probs_refs + + +class PPORewardRayActorGroup(PPORayActorGroup): + + def __init__(self, num_nodes, num_gpus_per_node) -> None: + super().__init__(num_nodes, num_gpus_per_node, RayPPORewardModel) + + def async_calculate_r(self, sequences_attention_mask_action_mask_refs): + num_actors = len(self._actor_handlers) + r_refs = [] + for i in range(len(sequences_attention_mask_action_mask_refs)): + r_ref = self._actor_handlers[i % num_actors].calculate_r.remote( + sequences_attention_mask_action_mask_refs[i]) + r_refs.append(r_ref) + return r_refs + + +def main(args): + logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s', + level=logging.INFO, + datefmt='%Y-%m-%d %H:%M:%S') + if args.model == 'gpt2': + actor_model_class, critic_model_class = GPTActor, GPTCritic + elif args.model == 'bloom': + actor_model_class, critic_model_class = BLOOMActor, BLOOMCritic + elif args.model == 'opt': + actor_model_class, critic_model_class = OPTActor, OPTCritic + else: + raise ValueError(f'Unsupported model "{args.model}"') + + logging.info("Start creating actors") + # Initialize 4 models (actor, critic, initial_model and reward_model) + actor_group = PPOActorRayActorGroup(num_nodes=args.num_actor_nodes, num_gpus_per_node=args.num_gpus_per_node) + critic_group = PPOCriticRayActorGroup(num_nodes=args.num_critic_nodes, num_gpus_per_node=args.num_gpus_per_node) + initial_group = PPOInitialRayActorGroup(num_nodes=args.num_initial_nodes, num_gpus_per_node=args.num_gpus_per_node) + reward_group = PPORewardRayActorGroup(num_nodes=args.num_reward_nodes, num_gpus_per_node=args.num_gpus_per_node) + logging.info("Actors created") + + # Prepare model for training + generate_kwargs = {'max_length': 128, 'do_sample': True, 'temperature': 1.0, 'top_k': 50} + ray.get( + actor_group.async_init_model_from_pretrained(args.strategy, actor_model_class, args.pretrain, True) + + critic_group.async_init_model_from_pretrained(args.strategy, critic_model_class, args.pretrain, True) + + initial_group.async_init_model_from_pretrained(args.strategy, actor_model_class, args.pretrain, False) + + reward_group.async_init_model_from_pretrained(args.strategy, critic_model_class, args.pretrain, False) + + actor_group.async_prepare_for_sequence_generation(args.model, args.pretrain, generate_kwargs)) + logging.info("Models prepared for training") + + # Prepare models for training + actor_group.load_csv_prompt_file_from_url_to_sampler(args.prompt_csv_url) + actor_group.set_loss_function() + critic_group.set_loss_function() + # Training parameter + num_episodes = args.num_episodes + max_timesteps = args.max_timesteps + update_timesteps = args.update_timesteps + experience_batch_size = args.experience_batch_size + # Start training + logging.info("Training start") + # Set all models to eval and add experience maker + all_ray_actors = actor_group._actor_handlers + critic_group._actor_handlers + \ + initial_group._actor_handlers + reward_group._actor_handlers + num_ray_actors = len(all_ray_actors) + ray.get([ray_actor.eval.remote() for ray_actor in all_ray_actors]) + ray.get([ray_actor.add_experience_maker.remote() for ray_actor in all_ray_actors]) + # Used as a queue to coordinate experience making + experience_composition_refs = [] + time = 0 + for episode in range(num_episodes): + logging.info("episode {} started".format(episode)) + for _ in range(max_timesteps): + time += 1 + # Experience queueing stage + sequences_attention_mask_action_mask_refs = actor_group.async_sample_prompts_and_make_sequence( + experience_batch_size) + base_action_log_probs_refs = initial_group.async_calculate_base_action_log_probs( + sequences_attention_mask_action_mask_refs) + values_refs = critic_group.async_calculate_value(sequences_attention_mask_action_mask_refs) + r_refs = reward_group.async_calculate_r(sequences_attention_mask_action_mask_refs) + action_log_probs_refs = actor_group.async_calculate_action_log_probs( + sequences_attention_mask_action_mask_refs) + experience_composition_refs.extend([ + ExperienceCompositionRefs(sequences_attention_mask_action_mask_refs[i], action_log_probs_refs[i], + base_action_log_probs_refs[i], values_refs[i], r_refs[i]) + for i in range(len(sequences_attention_mask_action_mask_refs)) + ]) + # Learning stage + if time % update_timesteps == 0: + experience_refs = [] + # calculate experiences + for i, experience_composition_ref in enumerate(experience_composition_refs): + exp_composition_ref = experience_composition_ref + selected_ray_actor = all_ray_actors[i % num_ray_actors] + experience_refs.append(selected_ray_actor.make_experience.remote(exp_composition_ref)) + # backward + ray.get( + actor_group.async_learn_on_experiences(experience_refs) + + critic_group.async_learn_on_experiences(experience_refs)) + # clear refs queue + experience_composition_refs.clear() + logging.info("Training finished") + # Save checkpoint + actor_group.save_checkpoint(args.save_path, args.need_optim_ckpt) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--prompt_csv_url', type=str) + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], + default='ddp') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt']) + parser.add_argument('--pretrain', type=str, default='gpt2') + parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts.pt') + parser.add_argument('--need_optim_ckpt', type=bool, default=False) + parser.add_argument('--num_episodes', type=int, default=10) + parser.add_argument('--max_timesteps', type=int, default=10) + parser.add_argument('--update_timesteps', type=int, default=10) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--num_actor_nodes', type=int, help='num of nodes to use to host actor model', default=1) + parser.add_argument('--num_critic_nodes', type=int, help='num of nodes to use to host critic model', default=1) + parser.add_argument('--num_initial_nodes', type=int, help='num of nodes to use to host initial model', default=1) + parser.add_argument('--num_reward_nodes', type=int, help='num of nodes to use to host reward model', default=1) + parser.add_argument('--num_gpus_per_node', type=int, help='num of gpus on a ray node', default=1) + args = parser.parse_args() + ray.init() + main(args) diff --git a/applications/Chat/examples/generate_prompt_dataset.py b/applications/Chat/examples/generate_prompt_dataset.py new file mode 100644 index 000000000000..95e40fefe7ff --- /dev/null +++ b/applications/Chat/examples/generate_prompt_dataset.py @@ -0,0 +1,30 @@ +import argparse + +import random +import json + +random.seed(42) + + +def sample(args): + with open(args.dataset_path, mode='r') as f: + dataset_list = json.load(f) + + sampled_dataset = [{"instruction": sample["instruction"], "id":idx} + for idx, sample in enumerate(random.sample(dataset_list, args.sample_size))] + + with open(args.save_path, mode='w') as f: + json.dump(sampled_dataset, f, indent=4, + default=str, ensure_ascii=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--dataset_path', type=str, default=None, + required=True, help="path to the pretrain dataset") + parser.add_argument('--save_path', type=str, default='prompt.json', + help="path to save the prompt dataset") + parser.add_argument('--sample_size', type=int, + default=16384, help="size of the prompt dataset") + args = parser.parse_args() + sample(args) diff --git a/applications/ChatGPT/examples/inference.py b/applications/Chat/examples/inference.py similarity index 80% rename from applications/ChatGPT/examples/inference.py rename to applications/Chat/examples/inference.py index 08885c33b194..4b49e76088bc 100644 --- a/applications/ChatGPT/examples/inference.py +++ b/applications/Chat/examples/inference.py @@ -1,9 +1,10 @@ import argparse import torch -from chatgpt.models.bloom import BLOOMActor -from chatgpt.models.gpt import GPTActor -from chatgpt.models.opt import OPTActor +from coati.models.bloom import BLOOMActor +from coati.models.generation import generate +from coati.models.gpt import GPTActor +from coati.models.opt import OPTActor from transformers import AutoTokenizer from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer @@ -20,7 +21,7 @@ def eval(args): raise ValueError(f'Unsupported model "{args.model}"') state_dict = torch.load(args.model_path) - actor.model.load_state_dict(state_dict) + actor.load_state_dict(state_dict) # configure tokenizer if args.model == 'gpt2': @@ -37,12 +38,13 @@ def eval(args): actor.eval() input = args.input input_ids = tokenizer.encode(input, return_tensors='pt').to(torch.cuda.current_device()) - outputs = actor.generate(input_ids, - max_length=args.max_length, - do_sample=True, - top_k=50, - top_p=0.95, - num_return_sequences=1) + outputs = generate(actor, + input_ids, + max_length=args.max_length, + do_sample=True, + top_k=50, + top_p=0.95, + num_return_sequences=1) output = tokenizer.batch_decode(outputs[0], skip_special_tokens=True) print(output) diff --git a/applications/Chat/examples/ray/1mmt_prompt.py b/applications/Chat/examples/ray/1mmt_prompt.py new file mode 100644 index 000000000000..5dd52f1790e6 --- /dev/null +++ b/applications/Chat/examples/ray/1mmt_prompt.py @@ -0,0 +1,175 @@ +import argparse +import os +import socket +from functools import partial + +import pandas as pd +import ray +import torch +from coati.quant import llama_load_quant, low_resource_init +from coati.ray.detached_trainer_ppo import DetachedPPOTrainer +from coati.ray.experience_maker_holder import ExperienceMakerHolder +from coati.ray.utils import ( + get_actor_from_args, + get_critic_from_args, + get_reward_model_from_args, + get_strategy_from_args, + get_tokenizer_from_args, +) +from torch.utils.data import DataLoader +from transformers import AutoConfig +from transformers.modeling_utils import no_init_weights + + +def get_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + return s.getsockname()[1] + + +def get_local_ip(): + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + s.connect(('8.8.8.8', 80)) + return s.getsockname()[0] + + +def main(args): + master_addr = str(get_local_ip()) + # trainer_env_info + trainer_port = str(get_free_port()) + env_info_trainers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_trainers), + 'master_port': trainer_port, + 'master_addr': master_addr + } for rank in range(args.num_trainers)] + + # maker_env_info + maker_port = str(get_free_port()) + env_info_maker = { + 'local_rank': '0', + 'rank': '0', + 'world_size': '1', + 'master_port': maker_port, + 'master_addr': master_addr + } + + # configure tokenizer + tokenizer = get_tokenizer_from_args(args.model) + + def trainer_model_fn(): + actor = get_actor_from_args(args.model, args.pretrain).half().cuda() + critic = get_critic_from_args(args.model, args.critic_pretrain).half().cuda() + return actor, critic + + # configure Trainer + trainer_refs = [ + DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote( + experience_maker_holder_name_list=["maker1"], + strategy_fn=partial(get_strategy_from_args, args.trainer_strategy), + model_fn=trainer_model_fn, + env_info=env_info_trainer, + train_batch_size=args.train_batch_size, + buffer_limit=16, + eval_performance=True, + debug=args.debug, + update_lora_weights=not (args.lora_rank == 0), + ) for i, env_info_trainer in enumerate(env_info_trainers) + ] + + def model_fn(): + actor = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda() + critic = get_critic_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda() + reward_model = get_reward_model_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda() + if args.initial_model_quant_ckpt is not None and args.model == 'llama': + # quantize initial model + actor_cfg = AutoConfig.from_pretrained(args.pretrain) + with low_resource_init(), no_init_weights(): + initial_model = get_actor_from_args(args.model, config=actor_cfg) + initial_model.model = llama_load_quant(initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, + args.quant_group_size).cuda().requires_grad_(False) + else: + initial_model = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda() + return actor, critic, reward_model, initial_model + + # configure Experience Maker + experience_holder_ref = ExperienceMakerHolder.options(name="maker1", num_gpus=1, max_concurrency=2).remote( + detached_trainer_name_list=[f'trainer{i}' for i in range(args.num_trainers)], + strategy_fn=partial(get_strategy_from_args, args.maker_strategy), + model_fn=model_fn, + env_info=env_info_maker, + experience_batch_size=args.experience_batch_size, + kl_coef=0.1, + debug=args.debug, + update_lora_weights=not (args.lora_rank == 0), + # sync_models_from_trainers=True, + # generation kwargs: + max_length=512, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + eval_performance=True, + use_cache=True, + ) + + # uncomment this function if sync_models_from_trainers is True + # ray.get([ + # trainer_ref.sync_models_to_remote_makers.remote() + # for trainer_ref in trainer_refs + # ]) + + wait_tasks = [] + + total_steps = args.experience_batch_size * args.experience_steps // (args.num_trainers * args.train_batch_size) + for trainer_ref in trainer_refs: + wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs)) + + dataset_size = args.experience_batch_size * 4 + + def build_dataloader(): + + def tokenize_fn(texts): + batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True) + return {k: v.cuda() for k, v in batch.items()} + + dataset = pd.read_csv(args.prompt_path)['prompt'] + dataloader = DataLoader(dataset=dataset, batch_size=dataset_size, shuffle=True, collate_fn=tokenize_fn) + return dataloader + + wait_tasks.append(experience_holder_ref.workingloop.remote(build_dataloader, num_steps=args.experience_steps)) + + ray.get(wait_tasks) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--prompt_path', type=str, default=None) + parser.add_argument('--num_trainers', type=int, default=1) + parser.add_argument('--trainer_strategy', + choices=[ + 'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu', + 'colossalai_zero2_cpu' + ], + default='ddp') + parser.add_argument('--maker_strategy', choices=['naive'], default='naive') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--critic_model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--critic_pretrain', type=str, default=None) + parser.add_argument('--experience_steps', type=int, default=4) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--train_epochs', type=int, default=1) + parser.add_argument('--update_steps', type=int, default=2) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + + parser.add_argument('--initial_model_quant_ckpt', type=str, default=None) + parser.add_argument('--quant_bits', type=int, default=4) + parser.add_argument('--quant_group_size', type=int, default=128) + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)}) + main(args) diff --git a/applications/Chat/examples/ray/mmmt_prompt.py b/applications/Chat/examples/ray/mmmt_prompt.py new file mode 100644 index 000000000000..60f049bd5b70 --- /dev/null +++ b/applications/Chat/examples/ray/mmmt_prompt.py @@ -0,0 +1,189 @@ +import argparse +import os +import socket +from functools import partial + +import pandas as pd +import ray +import torch +from coati.quant import llama_load_quant, low_resource_init +from coati.ray.detached_trainer_ppo import DetachedPPOTrainer +from coati.ray.experience_maker_holder import ExperienceMakerHolder +from coati.ray.utils import ( + get_actor_from_args, + get_critic_from_args, + get_receivers_per_sender, + get_reward_model_from_args, + get_strategy_from_args, +) +from torch.utils.data import DataLoader +from transformers import AutoConfig, AutoTokenizer +from transformers.modeling_utils import no_init_weights + + +def get_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('', 0)) + return s.getsockname()[1] + + +def get_local_ip(): + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + s.connect(('8.8.8.8', 80)) + return s.getsockname()[0] + + +def main(args): + master_addr = str(get_local_ip()) + # trainer_env_info + trainer_port = str(get_free_port()) + env_info_trainers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_trainers), + 'master_port': trainer_port, + 'master_addr': master_addr + } for rank in range(args.num_trainers)] + + # maker_env_info + maker_port = str(get_free_port()) + env_info_makers = [{ + 'local_rank': '0', + 'rank': str(rank), + 'world_size': str(args.num_makers), + 'master_port': maker_port, + 'master_addr': master_addr + } for rank in range(args.num_makers)] + + # configure tokenizer + tokenizer = AutoTokenizer.from_pretrained(args.pretrain) + tokenizer.pad_token = tokenizer.eos_token + + def model_fn(): + actor = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda() + critic = get_critic_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda() + reward_model = get_reward_model_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda() + if args.initial_model_quant_ckpt is not None and args.model == 'llama': + # quantize initial model + actor_cfg = AutoConfig.from_pretrained(args.pretrain) + with low_resource_init(), no_init_weights(): + initial_model = get_actor_from_args(args.model, config=actor_cfg) + initial_model.model = llama_load_quant(initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, + args.quant_group_size).cuda().requires_grad_(False) + else: + initial_model = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda() + return actor, critic, reward_model, initial_model + + # configure Experience Maker + experience_holder_refs = [ + ExperienceMakerHolder.options(name=f"maker{i}", num_gpus=1, max_concurrency=2).remote( + detached_trainer_name_list=[ + f'trainer{x}' + for x in get_receivers_per_sender(i, args.num_makers, args.num_trainers, allow_idle_sender=False) + ], + strategy_fn=partial(get_strategy_from_args, args.maker_strategy), + model_fn=model_fn, + env_info=env_info_maker, + kl_coef=0.1, + debug=args.debug, + update_lora_weights=not (args.lora_rank == 0), + # sync_models_from_trainers=True, + # generation kwargs: + max_length=512, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + eval_performance=True, + use_cache=True, + ) + for i, env_info_maker in enumerate(env_info_makers) + ] + + def trainer_model_fn(): + actor = get_actor_from_args(args.model, args.pretrain, lora_rank=args.lora_rank).half().cuda() + critic = get_critic_from_args(args.model, args.critic_pretrain, lora_rank=args.lora_rank).half().cuda() + return actor, critic + + # configure Trainer + trainer_refs = [ + DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote( + experience_maker_holder_name_list=[ + f"maker{x}" + for x in get_receivers_per_sender(i, args.num_trainers, args.num_makers, allow_idle_sender=True) + ], + strategy_fn=partial(get_strategy_from_args, args.trainer_strategy), + model_fn=trainer_model_fn, + env_info=env_info_trainer, + train_batch_size=args.train_batch_size, + buffer_limit=16, + eval_performance=True, + debug=args.debug, + update_lora_weights=not (args.lora_rank == 0), + ) + for i, env_info_trainer in enumerate(env_info_trainers) + ] + + dataset_size = args.experience_batch_size * 4 + + def build_dataloader(): + + def tokenize_fn(texts): + batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True) + return {k: v.cuda() for k, v in batch.items()} + + dataset = pd.read_csv(args.prompt_path)['prompt'] + dataloader = DataLoader(dataset=dataset, batch_size=dataset_size, shuffle=True, collate_fn=tokenize_fn) + return dataloader + + # uncomment this function if sync_models_from_trainers is True + # ray.get([ + # trainer_ref.sync_models_to_remote_makers.remote() + # for trainer_ref in trainer_refs + # ]) + + wait_tasks = [] + + for experience_holder_ref in experience_holder_refs: + wait_tasks.append(experience_holder_ref.workingloop.remote(build_dataloader, num_steps=args.experience_steps)) + + total_steps = args.experience_batch_size * args.experience_steps * \ + args.num_makers // (args.num_trainers * args.train_batch_size) + for trainer_ref in trainer_refs: + wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs)) + + ray.get(wait_tasks) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--prompt_path', type=str, default=None) + parser.add_argument('--num_makers', type=int, default=1) + parser.add_argument('--num_trainers', type=int, default=1) + parser.add_argument('--trainer_strategy', + choices=[ + 'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_gemini_cpu', + 'colossalai_zero2_cpu' + ], + default='ddp') + parser.add_argument('--maker_strategy', choices=['naive'], default='naive') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--critic_model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--critic_pretrain', type=str, default=None) + parser.add_argument('--experience_steps', type=int, default=4) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--train_epochs', type=int, default=1) + parser.add_argument('--update_steps', type=int, default=2) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + + parser.add_argument('--initial_model_quant_ckpt', type=str, default=None) + parser.add_argument('--quant_bits', type=int, default=4) + parser.add_argument('--quant_group_size', type=int, default=128) + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + + ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)}) + main(args) diff --git a/applications/Chat/examples/ray/requirements.txt b/applications/Chat/examples/ray/requirements.txt new file mode 100644 index 000000000000..e0275631807f --- /dev/null +++ b/applications/Chat/examples/ray/requirements.txt @@ -0,0 +1 @@ +ray diff --git a/applications/Chat/examples/ray/test_ci.sh b/applications/Chat/examples/ray/test_ci.sh new file mode 100755 index 000000000000..895f7de0fea9 --- /dev/null +++ b/applications/Chat/examples/ray/test_ci.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +set -xe +BASE=$(realpath $(dirname $0)) + +export RAY_NAMESPACE=admin +export DATA=/data/scratch/chatgpt/prompts.csv + +# install requirements +pip install -r ${BASE}/requirements.txt + +python ${BASE}/mmmt_prompt.py --prompt_path $DATA --num_makers 2 --num_trainers 2 --trainer_strategy colossalai_gemini --model opt --critic_model opt --pretrain facebook/opt-350m --critic_pretrain facebook/opt-125m --experience_batch_size 4 --train_batch_size 2 diff --git a/applications/ChatGPT/examples/requirements.txt b/applications/Chat/examples/requirements.txt similarity index 50% rename from applications/ChatGPT/examples/requirements.txt rename to applications/Chat/examples/requirements.txt index 6c5dac292486..40e6edc7ea73 100644 --- a/applications/ChatGPT/examples/requirements.txt +++ b/applications/Chat/examples/requirements.txt @@ -1 +1,2 @@ pandas>=1.4.1 +sentencepiece diff --git a/applications/Chat/examples/test_ci.sh b/applications/Chat/examples/test_ci.sh new file mode 100755 index 000000000000..fe2af471017e --- /dev/null +++ b/applications/Chat/examples/test_ci.sh @@ -0,0 +1,160 @@ +#!/usr/bin/env bash + +set_n_least_used_CUDA_VISIBLE_DEVICES() { + local n=${1:-"9999"} + echo "GPU Memory Usage:" + local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv | + tail -n +2 | + nl -v 0 | + tee /dev/tty | + sort -g -k 2 | + awk '{print $1}' | + head -n $n) + export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') + echo "Now CUDA_VISIBLE_DEVICES is set to:" + echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" +} + +set_n_least_used_CUDA_VISIBLE_DEVICES 4 + +set -xue + +if [ -z "$SFT_DATASET" ]; then + echo "Please set \$SFT_DATASET to the path to sft dataset." + exit 1 +fi + +if [ -z "$PROMPT_PATH" ]; then + echo "Please set \$PROMPT_PATH to the path to prompts csv." + exit 1 +fi + +if [ -z "$PRETRAIN_DATASET" ]; then + echo "Please set \$PRETRAIN_DATASET to the path to alpaca data." + exit 1 +fi + +BASE=$(realpath $(dirname $0)) + +export OMP_NUM_THREADS=8 + +# install requirements +pip install -r ${BASE}/requirements.txt + +wandb init -m offline + +# FIXME: This is a hack to skip tests that are not working +# - gpt2-ddp: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation +# - llama-*: These tests can be passed locally, skipped for long execution time +SKIPPED_TESTS=( + "gpt2-ddp" + "llama-ddp" + "llama-colossalai_gemini" + "llama-colossalai_zero2" +) + +# These tests are quick and do not have any dependencies +for model in 'gpt2' 'bloom' 'opt' 'llama'; do + for strategy in 'ddp' 'colossalai_gemini' 'colossalai_zero2'; do + if [[ " ${SKIPPED_TESTS[*]} " =~ " ${model}-${strategy} " ]]; then + echo "[Test]: Skipped $model-$strategy" + continue + fi + torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py \ + --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \ + --strategy $strategy --model $model \ + --num_episodes 1 --num_collect_steps 2 --num_update_steps 1 \ + --train_batch_size 2 --lora_rank 4 + done +done + +# train sft +torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'bigscience/bloom-560m' \ + --model 'bloom' --strategy colossalai_zero2 --lora_rank 4 \ + --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \ + --save_path ${BASE}/output +rm -rf ${BASE}/output + +torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \ + --model 'gpt2' --strategy colossalai_zero2 \ + --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \ + --save_path ${BASE}/output +rm -rf ${BASE}/output + +torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'facebook/opt-350m' \ + --model 'opt' --strategy colossalai_zero2 --lora_rank 4 \ + --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \ + --save_path ${BASE}/output +rm -rf ${BASE}/output + +torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \ + --model 'gpt2' --strategy ddp --lora_rank 4 \ + --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \ + --save_path ${BASE}/output +rm -rf ${BASE}/output + +# train rm +torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \ + --pretrain 'facebook/opt-350m' --model 'opt' \ + --strategy colossalai_zero2 --loss_fn 'log_sig' \ + --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \ + --test True --lora_rank 0 \ + --save_path ${BASE}/rm_ckpt_opt.pt + +torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \ + --pretrain 'gpt2' --model 'gpt2' \ + --strategy colossalai_zero2 --loss_fn 'log_exp' \ + --dataset 'Dahoas/rm-static' \ + --test True --lora_rank 0 \ + --save_path ${BASE}/rm_ckpt_gpt.pt + +torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \ + --pretrain 'gpt2' --model 'gpt2' \ + --strategy ddp --loss_fn 'log_exp' \ + --dataset 'Dahoas/rm-static' \ + --test True --lora_rank 4 \ + --save_path ${BASE}/rm_ckpt.pt +rm -rf ${BASE}/rm_ckpt.pt + +torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \ + --pretrain 'bigscience/bloom-560m' --model 'bloom' \ + --strategy colossalai_zero2 --loss_fn 'log_sig' \ + --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \ + --test True --lora_rank 4 \ + --save_path ${BASE}/rm_ckpt.pt +rm -rf ${BASE}/rm_ckpt.pt + +# train rl +torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py \ + --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \ + --strategy colossalai_zero2 --num_episodes 1 \ + --num_collect_steps 2 --num_update_steps 1 --train_batch_size 2 \ + --pretrain 'facebook/opt-350m' --model opt \ + --rm_pretrain 'facebook/opt-350m' \ + --rm_path ${BASE}/rm_ckpt_opt.pt \ + --save_path ${BASE}/actor_checkpoint_prompts.pt +rm -rf ${BASE}/rm_ckpt_opt.pt + +torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py \ + --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \ + --strategy colossalai_zero2 --num_episodes 1 \ + --num_collect_steps 2 --num_update_steps 1 --train_batch_size 2 \ + --pretrain 'gpt2' --model gpt2 \ + --rm_pretrain 'gpt2' \ + --rm_path ${BASE}/rm_ckpt_gpt.pt \ + --save_path ${BASE}/actor_checkpoint_prompts.pt + +torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py \ + --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \ + --strategy colossalai_gemini --num_episodes 1 \ + --num_collect_steps 2 --num_update_steps 1 --train_batch_size 2 \ + --pretrain 'gpt2' --model gpt2 \ + --rm_pretrain 'gpt2' \ + --rm_path ${BASE}/rm_ckpt_gpt.pt \ + --save_path ${BASE}/actor_checkpoint_prompts.pt +rm -rf ${BASE}/rm_ckpt_gpt.pt + +rm -rf ${BASE}/actor_checkpoint_prompts.pt + +# 3080 doesn't support P2P, skip this test +# cd ${BASE}/ray && bash test_ci.sh && cd ${BASE} diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py new file mode 100644 index 000000000000..7338a6d51142 --- /dev/null +++ b/applications/Chat/examples/train_prompts.py @@ -0,0 +1,218 @@ +import argparse + +import torch +import torch.distributed as dist +from coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset +from coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic +from coati.models.gpt import GPTRM, GPTActor, GPTCritic +from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM +from coati.models.opt import OPTRM, OPTActor, OPTCritic +from coati.trainer import PPOTrainer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from torch.optim import Adam +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer + +from colossalai.nn.optimizer import HybridAdam + + +def main(args): + # configure strategy + if args.strategy == 'ddp': + strategy = DDPStrategy() + elif args.strategy == 'colossalai_gemini': + strategy = GeminiStrategy(placement_policy='cuda', initial_scale=2**5) + elif args.strategy == 'colossalai_zero2': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + else: + raise ValueError(f'Unsupported strategy "{args.strategy}"') + + if args.rm_path is not None: + state_dict = torch.load(args.rm_path, map_location='cpu') + + with strategy.model_init_context(): + # configure model + if args.model == 'gpt2': + initial_model = GPTActor(pretrained=args.pretrain) + elif args.model == 'bloom': + initial_model = BLOOMActor(pretrained=args.pretrain) + elif args.model == 'opt': + initial_model = OPTActor(pretrained=args.pretrain) + elif args.model == 'llama': + initial_model = LlamaActor(pretrained=args.pretrain) + else: + raise ValueError(f'Unsupported actor model "{args.model}"') + + if args.rm_model is None: + rm_model_name = args.model + else: + rm_model_name = args.rm_model + + if rm_model_name == 'gpt2': + reward_model = GPTRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'bloom': + reward_model = BLOOMRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'opt': + reward_model = OPTRM(pretrained=args.rm_pretrain) + elif rm_model_name == 'llama': + reward_model = LlamaRM(pretrained=args.rm_pretrain) + else: + raise ValueError(f'Unsupported reward model "{rm_model_name}"') + + if args.rm_path is not None: + reward_model.load_state_dict(state_dict) + + initial_model.to(torch.float16).to(torch.cuda.current_device()) + reward_model.to(torch.float16).to(torch.cuda.current_device()) + + if args.model == 'gpt2': + actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank) + elif args.model == 'bloom': + actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank) + elif args.model == 'opt': + actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank) + elif args.model == 'llama': + actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank) + else: + raise ValueError(f'Unsupported actor model "{args.model}"') + + if rm_model_name == 'gpt2': + critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + elif rm_model_name == 'bloom': + critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + elif rm_model_name == 'opt': + critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + elif rm_model_name == 'llama': + critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True) + else: + raise ValueError(f'Unsupported reward model "{rm_model_name}"') + + if args.rm_path is not None: + critic.load_state_dict(state_dict) + del state_dict + + if args.strategy != 'colossalai_gemini': + critic.to(torch.float16).to(torch.cuda.current_device()) + actor.to(torch.float16).to(torch.cuda.current_device()) + + # configure optimizer + if args.strategy.startswith('colossalai'): + actor_optim = HybridAdam(actor.parameters(), lr=1e-7) + critic_optim = HybridAdam(critic.parameters(), lr=1e-7) + else: + actor_optim = Adam(actor.parameters(), lr=1e-7) + critic_optim = Adam(critic.parameters(), lr=1e-7) + + # configure tokenizer + if args.model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained('gpt2') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'opt': + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'llama': + tokenizer = LlamaTokenizer.from_pretrained(args.pretrain) + tokenizer.eos_token = '<\s>' + tokenizer.pad_token = tokenizer.unk_token + else: + raise ValueError(f'Unsupported model "{args.model}"') + + data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) + + prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384) + if dist.is_initialized() and dist.get_world_size() > 1: + prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True) + else: + prompt_sampler = None + prompt_dataloader = DataLoader(prompt_dataset, + shuffle=(prompt_sampler is None), + sampler=prompt_sampler, + batch_size=args.experience_batch_size) + + pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, + data_path=args.pretrain_dataset, + max_datasets_size=16384, + max_length=args.max_input_len) + if dist.is_initialized() and dist.get_world_size() > 1: + pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True) + else: + pretrain_sampler = None + pretrain_dataloader = DataLoader(pretrain_dataset, + shuffle=(pretrain_sampler is None), + sampler=pretrain_sampler, + batch_size=args.ptx_batch_size, + collate_fn=data_collator) + + # NOTE: For small models like opt-1.3b, reward model and initial model are not required to be parallelized. + (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = \ + strategy.prepare((actor, actor_optim), (critic, critic_optim), reward_model, initial_model) + + # configure trainer + trainer = PPOTrainer( + strategy, + actor, + critic, + reward_model, + initial_model, + actor_optim, + critic_optim, + kl_coef=args.kl_coef, + ptx_coef=args.ptx_coef, + train_batch_size=args.train_batch_size, + max_length=args.max_seq_len, + use_cache=True, + do_sample=True, + temperature=1.0, + top_k=50, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + offload_inference_models=args.strategy != 'colossalai_gemini' + ) + + trainer.fit(prompt_dataloader=prompt_dataloader, + pretrain_dataloader=pretrain_dataloader, + num_episodes=args.num_episodes, + num_collect_steps=args.num_collect_steps, + num_update_steps=args.num_update_steps) + + # save model checkpoint after fitting + strategy.save_model(actor, args.save_path, only_rank0=True) + # save optimizer checkpoint on all ranks + if args.need_optim_ckpt: + strategy.save_optimizer(actor_optim, + 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()), + only_rank0=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset') + parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset') + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], + default='colossalai_zero2', + help='strategy to use') + parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama']) + parser.add_argument('--rm_path', type=str, default=None) + parser.add_argument('--rm_pretrain', type=str, default=None) + parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts') + parser.add_argument('--need_optim_ckpt', type=bool, default=False) + parser.add_argument('--num_episodes', type=int, default=10) + parser.add_argument('--num_collect_steps', type=int, default=10) + parser.add_argument('--num_update_steps', type=int, default=5) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--ptx_batch_size', type=int, default=1) + parser.add_argument('--experience_batch_size', type=int, default=8) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + parser.add_argument('--kl_coef', type=float, default=0.1) + parser.add_argument('--ptx_coef', type=float, default=0.9) + parser.add_argument('--max_input_len', type=int, default=96) + parser.add_argument('--max_seq_len', type=int, default=128) + args = parser.parse_args() + main(args) diff --git a/applications/Chat/examples/train_prompts.sh b/applications/Chat/examples/train_prompts.sh new file mode 100755 index 000000000000..d04c416015b1 --- /dev/null +++ b/applications/Chat/examples/train_prompts.sh @@ -0,0 +1,25 @@ +set_n_least_used_CUDA_VISIBLE_DEVICES() { + local n=${1:-"9999"} + echo "GPU Memory Usage:" + local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv | + tail -n +2 | + nl -v 0 | + tee /dev/tty | + sort -g -k 2 | + awk '{print $1}' | + head -n $n) + export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') + echo "Now CUDA_VISIBLE_DEVICES is set to:" + echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" +} + +set_n_least_used_CUDA_VISIBLE_DEVICES 2 + +# torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy colossalai_zero2 + +torchrun --standalone --nproc_per_node=2 train_prompts.py \ + --pretrain_dataset /path/to/data.json \ + --prompt_dataset /path/to/data.json \ + --strategy colossalai_zero2 \ + --num_episodes 1 --num_collect_steps 2 --num_update_steps 1 \ + --train_batch_size 2 diff --git a/applications/Chat/examples/train_reward_model.py b/applications/Chat/examples/train_reward_model.py new file mode 100644 index 000000000000..5b1b8d3d16b2 --- /dev/null +++ b/applications/Chat/examples/train_reward_model.py @@ -0,0 +1,200 @@ +import argparse +from random import randint + +import torch +import torch.distributed as dist +from coati.dataset import HhRlhfDataset, RmStaticDataset +from coati.models import LogExpLoss, LogSigLoss +from coati.models.bloom import BLOOMRM +from coati.models.gpt import GPTRM +from coati.models.llama import LlamaRM +from coati.models.opt import OPTRM +from coati.trainer import RewardModelTrainer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from datasets import load_dataset +from torch.optim import Adam +from torch.optim.lr_scheduler import CosineAnnealingLR +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoTokenizer, BloomTokenizerFast, LlamaTokenizer +from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer + +from colossalai.nn.optimizer import HybridAdam + + +def train(args): + # configure strategy + if args.strategy == 'ddp': + strategy = DDPStrategy() + elif args.strategy == 'colossalai_gemini': + strategy = GeminiStrategy(placement_policy='cuda') + elif args.strategy == 'colossalai_zero2': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + else: + raise ValueError(f'Unsupported strategy "{args.strategy}"') + + # configure model + with strategy.model_init_context(): + if args.model == 'bloom': + model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) + elif args.model == 'opt': + model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) + elif args.model == 'gpt2': + model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) + elif args.model == 'llama': + model = LlamaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) + else: + raise ValueError(f'Unsupported model "{args.model}"') + + if args.model_path is not None: + state_dict = torch.load(args.model_path) + model.load_state_dict(state_dict) + + model = model.to(torch.float16) + + # configure tokenizer + if args.model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained('gpt2') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'opt': + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'llama': + tokenizer = LlamaTokenizer.from_pretrained(args.pretrain) + tokenizer.pad_token = tokenizer.unk_token + else: + raise ValueError(f'Unsupported model "{args.model}"') + + # configure optimizer + if args.strategy.startswith('colossalai'): + optim = HybridAdam(model.parameters(), lr=5e-6) + else: + optim = Adam(model.parameters(), lr=5e-6) + + # configure loss function + if args.loss_fn == 'log_sig': + loss_fn = LogSigLoss() + elif args.loss_fn == 'log_exp': + loss_fn = LogExpLoss() + else: + raise ValueError(f'Unsupported loss function "{args.loss_fn}"') + + # prepare for data and dataset + if args.subset is not None: + data = load_dataset(args.dataset, data_dir=args.subset) + else: + data = load_dataset(args.dataset) + + if args.test: + train_data = data['train'].select(range(100)) + eval_data = data['test'].select(range(10)) + else: + train_data = data['train'] + eval_data = data['test'] + valid_data = data['test'].select((randint(0, len(eval_data) - 1) for _ in range(len(eval_data) // 5))) + + if args.dataset == 'Dahoas/rm-static': + train_dataset = RmStaticDataset(train_data, tokenizer, args.max_len) + valid_dataset = RmStaticDataset(valid_data, tokenizer, args.max_len) + eval_dataset = RmStaticDataset(eval_data, tokenizer, args.max_len) + elif args.dataset == 'Anthropic/hh-rlhf': + train_dataset = HhRlhfDataset(train_data, tokenizer, args.max_len) + valid_dataset = HhRlhfDataset(valid_data, tokenizer, args.max_len) + eval_dataset = HhRlhfDataset(eval_data, tokenizer, args.max_len) + else: + raise ValueError(f'Unsupported dataset "{args.dataset}"') + + if dist.is_initialized() and dist.get_world_size() > 1: + train_sampler = DistributedSampler(train_dataset, + shuffle=True, + seed=42, + drop_last=True, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + valid_sampler = DistributedSampler(valid_dataset, + shuffle=True, + seed=42, + drop_last=True, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + eval_sampler = DistributedSampler(eval_dataset, + shuffle=True, + seed=42, + drop_last=True, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + else: + train_sampler = None + valid_sampler = None + eval_sampler = None + + train_dataloader = DataLoader(train_dataset, + shuffle=(train_sampler is None), + sampler=train_sampler, + batch_size=args.batch_size, + pin_memory=True) + + valid_dataloader = DataLoader(valid_dataset, + shuffle=(valid_sampler is None), + sampler=valid_sampler, + batch_size=args.batch_size, + pin_memory=True) + + eval_dataloader = DataLoader(eval_dataset, + shuffle=(eval_sampler is None), + sampler=eval_sampler, + batch_size=args.batch_size, + pin_memory=True) + + lr_scheduler = CosineAnnealingLR(optim, train_dataloader.__len__() // 100) + strategy_dict = strategy.prepare( + dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler) + ) + model = strategy_dict['model'] + optim = strategy_dict['optimizer'] + lr_scheduler = strategy_dict['lr_scheduler'] + trainer = RewardModelTrainer(model=model, + strategy=strategy, + optim=optim, + lr_scheduler=lr_scheduler, + loss_fn=loss_fn, + max_epochs=args.max_epochs) + + trainer.fit(train_dataloader=train_dataloader, + valid_dataloader=valid_dataloader, + eval_dataloader=eval_dataloader) + # save model checkpoint after fitting on only rank0 + strategy.save_model(model, args.save_path, only_rank0=True) + # save optimizer checkpoint on all ranks + if args.need_optim_ckpt: + strategy.save_optimizer(trainer.optimizer, + 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), + only_rank0=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2'], + default='colossalai_zero2') + parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom') + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--model_path', type=str, default=None) + parser.add_argument('--need_optim_ckpt', type=bool, default=False) + parser.add_argument('--dataset', + type=str, + choices=['Anthropic/hh-rlhf', 'Dahoas/rm-static'], + default='Dahoas/rm-static') + parser.add_argument('--subset', type=str, default=None) + parser.add_argument('--save_path', type=str, default='rm_ckpt') + parser.add_argument('--max_epochs', type=int, default=1) + parser.add_argument('--batch_size', type=int, default=1) + parser.add_argument('--max_len', type=int, default=512) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + parser.add_argument('--loss_fn', type=str, default='log_sig', choices=['log_sig', 'log_exp']) + parser.add_argument('--test', type=bool, default=False) + args = parser.parse_args() + train(args) diff --git a/applications/ChatGPT/examples/train_dummy.sh b/applications/Chat/examples/train_rm.sh similarity index 67% rename from applications/ChatGPT/examples/train_dummy.sh rename to applications/Chat/examples/train_rm.sh index 595da573e2b1..80abe62d2a3f 100755 --- a/applications/ChatGPT/examples/train_dummy.sh +++ b/applications/Chat/examples/train_rm.sh @@ -15,4 +15,10 @@ set_n_least_used_CUDA_VISIBLE_DEVICES() { set_n_least_used_CUDA_VISIBLE_DEVICES 2 -torchrun --standalone --nproc_per_node=2 train_dummy.py --strategy colossalai_zero2 +torchrun --standalone --nproc_per_node=2 train_reward_model.py \ + --pretrain \ + --model 'bloom' \ + --strategy colossalai_zero2 \ + --loss_fn 'log_sig'\ + --save_path \ + --dataset 'Anthropic/hh-rlhf'\ diff --git a/applications/Chat/examples/train_sft.py b/applications/Chat/examples/train_sft.py new file mode 100644 index 000000000000..d643609b3a30 --- /dev/null +++ b/applications/Chat/examples/train_sft.py @@ -0,0 +1,205 @@ +import argparse +import math +import os + +import loralib as lora +import torch +import torch.distributed as dist +from coati.dataset import DataCollatorForSupervisedDataset, SFTDataset, SupervisedDataset +from coati.models import convert_to_lora_module +from coati.trainer import SFTTrainer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy +from datasets import load_dataset +from torch.optim import Adam +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoTokenizer, BloomConfig, BloomForCausalLM, BloomTokenizerFast, LlamaConfig, LlamaForCausalLM +from transformers.models.gpt2.configuration_gpt2 import GPT2Config +from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel +from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer +from transformers.models.opt.configuration_opt import OPTConfig +from transformers.models.opt.modeling_opt import OPTForCausalLM +from transformers.trainer import get_scheduler + +from colossalai.logging import get_dist_logger +from colossalai.nn.optimizer import HybridAdam +from colossalai.tensor import ColoParameter + + +def train(args): + # configure strategy + if args.strategy == 'ddp': + strategy = DDPStrategy() + elif args.strategy == 'colossalai_gemini': + raise NotImplementedError( + 'Gemini is not supported .from_pretrained() yet. We will update this after checkpoint io is ready.') + strategy = GeminiStrategy(placement_policy='cuda') + elif args.strategy == 'colossalai_zero2': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') + elif args.strategy == 'colossalai_zero2_cpu': + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cpu') + else: + raise ValueError(f'Unsupported strategy "{args.strategy}"') + + # configure model + with strategy.model_init_context(): + if args.model == 'bloom': + model = convert_to_lora_module(BloomForCausalLM.from_pretrained(args.pretrain), + args.lora_rank).half().cuda() + elif args.model == 'opt': + model = convert_to_lora_module(OPTForCausalLM.from_pretrained(args.pretrain), args.lora_rank).half().cuda() + elif args.model == 'gpt2': + model = convert_to_lora_module(GPT2LMHeadModel.from_pretrained(args.pretrain), args.lora_rank).half().cuda() + elif args.model == 'llama': + model = convert_to_lora_module(LlamaForCausalLM.from_pretrained(args.pretrain), + args.lora_rank).half().cuda() + else: + raise ValueError(f'Unsupported model "{args.model}"') + if args.grad_checkpoint: + model.gradient_checkpointing_enable() + + # configure tokenizer + if args.model == 'gpt2': + tokenizer = GPT2Tokenizer.from_pretrained('gpt2') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'bloom': + tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m') + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'opt': + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") + tokenizer.pad_token = tokenizer.eos_token + elif args.model == 'llama': + tokenizer = AutoTokenizer.from_pretrained( + args.pretrain, + padding_side="right", + use_fast=False, + ) + tokenizer.eos_token = '<\s>' + tokenizer.pad_token = tokenizer.unk_token + else: + raise ValueError(f'Unsupported model "{args.model}"') + + if args.model == 'llama' and args.strategy == 'colossalai_gemini': + # this is a hack to deal with the resized embedding + # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatibility + for name, param in model.named_parameters(): + if not isinstance(param, ColoParameter): + sub_module_name = '.'.join(name.split('.')[:-1]) + weight_name = name.split('.')[-1] + sub_module = model.get_submodule(sub_module_name) + setattr(sub_module, weight_name, ColoParameter(param)) + + # configure optimizer + if args.strategy.startswith('colossalai'): + optim = HybridAdam(model.parameters(), lr=args.lr, clipping_norm=1.0) + else: + optim = Adam(model.parameters(), lr=args.lr) + + logger = get_dist_logger() + + # configure dataset + if args.dataset == 'yizhongw/self_instruct': + train_data = load_dataset(args.dataset, 'super_natural_instructions', split='train') + eval_data = load_dataset(args.dataset, 'super_natural_instructions', split='test') + + train_dataset = SFTDataset(train_data, tokenizer, args.max_len) + eval_dataset = SFTDataset(eval_data, tokenizer, args.max_len) + + else: + train_dataset = SupervisedDataset(tokenizer=tokenizer, + data_path=args.dataset, + max_datasets_size=args.max_datasets_size, + max_length=args.max_len) + eval_dataset = None + data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) + + if dist.is_initialized() and dist.get_world_size() > 1: + train_sampler = DistributedSampler(train_dataset, + shuffle=True, + seed=42, + drop_last=True, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + if eval_dataset is not None: + eval_sampler = DistributedSampler(eval_dataset, + shuffle=False, + seed=42, + drop_last=False, + rank=dist.get_rank(), + num_replicas=dist.get_world_size()) + else: + train_sampler = None + eval_sampler = None + + train_dataloader = DataLoader(train_dataset, + shuffle=(train_sampler is None), + sampler=train_sampler, + batch_size=args.batch_size, + collate_fn=data_collator, + pin_memory=True) + if eval_dataset is not None: + eval_dataloader = DataLoader(eval_dataset, + shuffle=(eval_sampler is None), + sampler=eval_sampler, + batch_size=args.batch_size, + collate_fn=data_collator, + pin_memory=True) + else: + eval_dataloader = None + + num_update_steps_per_epoch = len(train_dataloader) // args.accumulation_steps + max_steps = math.ceil(args.max_epochs * num_update_steps_per_epoch) + lr_scheduler = get_scheduler("cosine", + optim, + num_warmup_steps=math.ceil(max_steps * 0.03), + num_training_steps=max_steps) + strategy_dict = strategy.prepare(dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler)) + model = strategy_dict['model'] + optim = strategy_dict['optimizer'] + lr_scheduler = strategy_dict['lr_scheduler'] + if args.optim_load_path: + strategy.load_optimizer(optim, path=args.optim_load_path) + trainer = SFTTrainer(model=model, + strategy=strategy, + optim=optim, + lr_scheduler=lr_scheduler, + max_epochs=args.max_epochs, + tensorboard_dir=args.tensorboard_dir, + accumulation_steps=args.accumulation_steps) + + trainer.fit(train_dataloader=train_dataloader, + eval_dataloader=eval_dataloader, + logger=logger, + use_wandb=args.use_wandb) + + # save model checkpoint after fitting on only rank0 + strategy.save_pretrained(model, path=args.save_path, only_rank0=True, tokenizer=tokenizer) + # save optimizer checkpoint on all ranks + if args.optim_save_path: + strategy.save_optimizer(trainer.optimizer, path=args.optim_save_path) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--strategy', + choices=['ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_zero2_cpu'], + default='colossalai_zero2') + parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom') + parser.add_argument('--pretrain', type=str, default=None) + parser.add_argument('--dataset', type=str, default=None) + parser.add_argument('--max_datasets_size', type=int, default=None) + parser.add_argument('--save_path', type=str, default='output') + parser.add_argument('--optim_save_path', type=str, default=None) + parser.add_argument('--optim_load_path', type=str, default=None) + parser.add_argument('--max_epochs', type=int, default=3) + parser.add_argument('--batch_size', type=int, default=4) + parser.add_argument('--max_len', type=int, default=512) + parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") + parser.add_argument('--log_interval', type=int, default=100, help="how many steps to log") + parser.add_argument('--lr', type=float, default=5e-6) + parser.add_argument('--accumulation_steps', type=int, default=8) + parser.add_argument('--use_wandb', default=False, action='store_true') + parser.add_argument('--grad_checkpoint', default=False, action='store_true') + parser.add_argument('--tensorboard_dir', type=str, default=None) + args = parser.parse_args() + train(args) diff --git a/applications/Chat/examples/train_sft.sh b/applications/Chat/examples/train_sft.sh new file mode 100755 index 000000000000..c880f85825a7 --- /dev/null +++ b/applications/Chat/examples/train_sft.sh @@ -0,0 +1,12 @@ +torchrun --standalone --nproc_per_node=4 train_sft.py \ + --pretrain "/path/to/LLaMa-7B/" \ + --model 'llama' \ + --strategy colossalai_zero2 \ + --log_interval 10 \ + --save_path /path/to/Coati-7B \ + --dataset /path/to/data.json \ + --batch_size 4 \ + --accumulation_steps 8 \ + --lr 2e-5 \ + --max_datasets_size 512 \ + --max_epochs 1 \ diff --git a/applications/Chat/inference/README.md b/applications/Chat/inference/README.md new file mode 100644 index 000000000000..4848817e0fd1 --- /dev/null +++ b/applications/Chat/inference/README.md @@ -0,0 +1,118 @@ +# Inference + +We provide an online inference server and a benchmark. We aim to run inference on single GPU, so quantization is essential when using large models. + +We support 8-bit quantization (RTN), which is powered by [bitsandbytes](https://github.com/TimDettmers/bitsandbytes) and [transformers](https://github.com/huggingface/transformers). And 4-bit quantization (GPTQ), which is powered by [gptq](https://github.com/IST-DASLab/gptq) and [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). We also support FP16 inference. + +We only support LLaMA family models now. + +## Choosing precision (quantization) + +**FP16**: Fastest, best output quality, highest memory usage + +**8-bit**: Slow, easier setup (originally supported by transformers), lower output quality (due to RTN), **recommended for first-timers** + +**4-bit**: Faster, lowest memory usage, higher output quality (due to GPTQ), but more difficult setup + +## Hardware requirements for LLaMA + +Tha data is from [LLaMA Int8 4bit ChatBot Guide v2](https://rentry.org/llama-tard-v2). + +### 8-bit + +| Model | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap | Card examples | +| :---: | :---: | :---: | :---: | :---: | +| LLaMA-7B | 9.2GB | 10GB | 24GB | 3060 12GB, RTX 3080 10GB, RTX 3090 | +| LLaMA-13B | 16.3GB | 20GB | 32GB | RTX 3090 Ti, RTX 4090 | +| LLaMA-30B | 36GB | 40GB | 64GB | A6000 48GB, A100 40GB | +| LLaMA-65B | 74GB | 80GB | 128GB | A100 80GB | + +### 4-bit + +| Model | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap | Card examples | +| :---: | :---: | :---: | :---: | :---: | +| LLaMA-7B | 3.5GB | 6GB | 16GB | RTX 1660, 2060, AMD 5700xt, RTX 3050, 3060 | +| LLaMA-13B | 6.5GB | 10GB | 32GB | AMD 6900xt, RTX 2060 12GB, 3060 12GB, 3080, A2000 | +| LLaMA-30B | 15.8GB | 20GB | 64GB | RTX 3080 20GB, A4500, A5000, 3090, 4090, 6000, Tesla V100 | +| LLaMA-65B | 31.2GB | 40GB | 128GB | A100 40GB, 2x3090, 2x4090, A40, RTX A6000, 8000, Titan Ada | + +## General setup + +```shell +pip install -r requirements.txt +``` + +## 8-bit setup + +8-bit quantization is originally supported by the latest [transformers](https://github.com/huggingface/transformers). Please install it from source. + +Please ensure you have downloaded HF-format model weights of LLaMA models. + +Usage: + +```python +import torch +from transformers import LlamaForCausalLM + +USE_8BIT = True # use 8-bit quantization; otherwise, use fp16 + +model = LlamaForCausalLM.from_pretrained( + "pretrained/path", + load_in_8bit=USE_8BIT, + torch_dtype=torch.float16, + device_map="auto", + ) +if not USE_8BIT: + model.half() # use fp16 +model.eval() +``` + +**Troubleshooting**: if you get error indicating your CUDA-related libraries not found when loading 8-bit model, you can check whether your `LD_LIBRARY_PATH` is correct. + +E.g. you can set `export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH`. + +## 4-bit setup + +Please ensure you have downloaded HF-format model weights of LLaMA models first. + +Then you can follow [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). This lib provides efficient CUDA kernels and weight conversion script. + +After installing this lib, we may convert the original HF-format LLaMA model weights to 4-bit version. + +```shell +CUDA_VISIBLE_DEVICES=0 python llama.py /path/to/pretrained/llama-7b c4 --wbits 4 --groupsize 128 --save llama7b-4bit.pt +``` + +Run this command in your cloned `GPTQ-for-LLaMa` directory, then you will get a 4-bit weight file `llama7b-4bit-128g.pt`. + +**Troubleshooting**: if you get error about `position_ids`, you can checkout to commit `50287c3b9ae4a3b66f6b5127c643ec39b769b155`(`GPTQ-for-LLaMa` repo). + +## Online inference server + +In this directory: + +```shell +export CUDA_VISIBLE_DEVICES=0 +# fp16, will listen on 0.0.0.0:7070 by default +python server.py /path/to/pretrained +# 8-bit, will listen on localhost:8080 +python server.py /path/to/pretrained --quant 8bit --http_host localhost --http_port 8080 +# 4-bit +python server.py /path/to/pretrained --quant 4bit --gptq_checkpoint /path/to/llama7b-4bit-128g.pt --gptq_group_size 128 +``` + +## Benchmark + +In this directory: + +```shell +export CUDA_VISIBLE_DEVICES=0 +# fp16 +python benchmark.py /path/to/pretrained +# 8-bit +python benchmark.py /path/to/pretrained --quant 8bit +# 4-bit +python benchmark.py /path/to/pretrained --quant 4bit --gptq_checkpoint /path/to/llama7b-4bit-128g.pt --gptq_group_size 128 +``` + +This benchmark will record throughput and peak CUDA memory usage. diff --git a/applications/Chat/inference/benchmark.py b/applications/Chat/inference/benchmark.py new file mode 100644 index 000000000000..a8485f588705 --- /dev/null +++ b/applications/Chat/inference/benchmark.py @@ -0,0 +1,132 @@ +# Adapted from https://github.com/tloen/alpaca-lora/blob/main/generate.py + +import argparse +from time import time + +import torch +from llama_gptq import load_quant +from transformers import AutoTokenizer, GenerationConfig, LlamaForCausalLM + + +def generate_prompt(instruction, input=None): + if input: + return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. + +### Instruction: +{instruction} + +### Input: +{input} + +### Response:""" + else: + return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request. + +### Instruction: +{instruction} + +### Response:""" + + +@torch.no_grad() +def evaluate( + model, + tokenizer, + instruction, + input=None, + temperature=0.1, + top_p=0.75, + top_k=40, + num_beams=4, + max_new_tokens=128, + **kwargs, +): + prompt = generate_prompt(instruction, input) + inputs = tokenizer(prompt, return_tensors="pt") + input_ids = inputs["input_ids"].cuda() + generation_config = GenerationConfig( + temperature=temperature, + top_p=top_p, + top_k=top_k, + num_beams=num_beams, + **kwargs, + ) + generation_output = model.generate( + input_ids=input_ids, + generation_config=generation_config, + return_dict_in_generate=True, + output_scores=True, + max_new_tokens=max_new_tokens, + do_sample=True, + ) + s = generation_output.sequences[0] + output = tokenizer.decode(s) + n_new_tokens = s.size(0) - input_ids.size(1) + return output.split("### Response:")[1].strip(), n_new_tokens + + +instructions = [ + "Tell me about alpacas.", + "Tell me about the president of Mexico in 2019.", + "Tell me about the king of France in 2019.", + "List all Canadian provinces in alphabetical order.", + "Write a Python program that prints the first 10 Fibonacci numbers.", + "Write a program that prints the numbers from 1 to 100. But for multiples of three print 'Fizz' instead of the number and for the multiples of five print 'Buzz'. For numbers which are multiples of both three and five print 'FizzBuzz'.", + "Tell me five words that rhyme with 'shock'.", + "Translate the sentence 'I have no mouth but I must scream' into Spanish.", + "Count up from 1 to 500.", + # === + "How to play support in legends of league", + "Write a Python program that calculate Fibonacci numbers.", +] +inst = [instructions[0]] * 4 + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + 'pretrained', + help='Path to pretrained model. Can be a local path or a model name from the HuggingFace model hub.') + parser.add_argument('--quant', + choices=['8bit', '4bit'], + default=None, + help='Quantization mode. Default: None (no quantization, fp16).') + parser.add_argument( + '--gptq_checkpoint', + default=None, + help='Path to GPTQ checkpoint. This is only useful when quantization mode is 4bit. Default: None.') + parser.add_argument('--gptq_group_size', + type=int, + default=128, + help='Group size for GPTQ. This is only useful when quantization mode is 4bit. Default: 128.') + args = parser.parse_args() + + if args.quant == '4bit': + assert args.gptq_checkpoint is not None, 'Please specify a GPTQ checkpoint.' + + tokenizer = AutoTokenizer.from_pretrained(args.pretrained) + + if args.quant == '4bit': + model = load_quant(args.pretrained, args.gptq_checkpoint, 4, args.gptq_group_size) + model.cuda() + else: + model = LlamaForCausalLM.from_pretrained( + args.pretrained, + load_in_8bit=(args.quant == '8bit'), + torch_dtype=torch.float16, + device_map="auto", + ) + if args.quant != '8bit': + model.half() # seems to fix bugs for some users. + model.eval() + + total_tokens = 0 + start = time() + for instruction in instructions: + print(f"Instruction: {instruction}") + resp, tokens = evaluate(model, tokenizer, instruction, temperature=0.2, num_beams=1) + total_tokens += tokens + print(f"Response: {resp}") + print('\n----------------------------\n') + duration = time() - start + print(f'Total time: {duration:.3f} s, {total_tokens/duration:.3f} tokens/s') + print(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.3f} GB') diff --git a/applications/Chat/inference/llama_gptq/__init__.py b/applications/Chat/inference/llama_gptq/__init__.py new file mode 100644 index 000000000000..51c8d6316290 --- /dev/null +++ b/applications/Chat/inference/llama_gptq/__init__.py @@ -0,0 +1,5 @@ +from .loader import load_quant + +__all__ = [ + 'load_quant', +] diff --git a/applications/Chat/inference/llama_gptq/loader.py b/applications/Chat/inference/llama_gptq/loader.py new file mode 100644 index 000000000000..a5c6ac7d1589 --- /dev/null +++ b/applications/Chat/inference/llama_gptq/loader.py @@ -0,0 +1,41 @@ +import torch +import torch.nn as nn +import transformers +from transformers import LlamaConfig, LlamaForCausalLM + +from .model_utils import find_layers +from .quant import make_quant + + +def load_quant(pretrained: str, checkpoint: str, wbits: int, groupsize: int): + config = LlamaConfig.from_pretrained(pretrained) + + def noop(*args, **kwargs): + pass + + torch.nn.init.kaiming_uniform_ = noop + torch.nn.init.uniform_ = noop + torch.nn.init.normal_ = noop + + torch.set_default_dtype(torch.half) + transformers.modeling_utils._init_weights = False + torch.set_default_dtype(torch.half) + model = LlamaForCausalLM(config) + torch.set_default_dtype(torch.float) + model = model.eval() + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + make_quant(model, layers, wbits, groupsize) + + print(f'Loading model with {wbits} bits...') + if checkpoint.endswith('.safetensors'): + from safetensors.torch import load_file as safe_load + model.load_state_dict(safe_load(checkpoint)) + else: + model.load_state_dict(torch.load(checkpoint)) + model.seqlen = 2048 + print('Done.') + + return model diff --git a/applications/Chat/inference/llama_gptq/model_utils.py b/applications/Chat/inference/llama_gptq/model_utils.py new file mode 100644 index 000000000000..62db171abb52 --- /dev/null +++ b/applications/Chat/inference/llama_gptq/model_utils.py @@ -0,0 +1,13 @@ +# copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/past/modelutils.py + +import torch +import torch.nn as nn + + +def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''): + if type(module) in layers: + return {name: module} + res = {} + for name1, child in module.named_children(): + res.update(find_layers(child, layers=layers, name=name + '.' + name1 if name != '' else name1)) + return res diff --git a/applications/Chat/inference/llama_gptq/quant.py b/applications/Chat/inference/llama_gptq/quant.py new file mode 100644 index 000000000000..f7d5b7ce4bd8 --- /dev/null +++ b/applications/Chat/inference/llama_gptq/quant.py @@ -0,0 +1,283 @@ +# copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/past/quant.py + +import math + +import numpy as np +import torch +import torch.nn as nn + + +def quantize(x, scale, zero, maxq): + q = torch.clamp(torch.round(x / scale) + zero, 0, maxq) + return scale * (q - zero) + + +class Quantizer(nn.Module): + + def __init__(self, shape=1): + super(Quantizer, self).__init__() + self.register_buffer('maxq', torch.tensor(0)) + self.register_buffer('scale', torch.zeros(shape)) + self.register_buffer('zero', torch.zeros(shape)) + + def configure(self, bits, perchannel=False, sym=True, mse=False, norm=2.4, grid=100, maxshrink=.8): + self.maxq = torch.tensor(2**bits - 1) + self.perchannel = perchannel + self.sym = sym + self.mse = mse + self.norm = norm + self.grid = grid + self.maxshrink = maxshrink + + def find_params(self, x, weight=False): + dev = x.device + self.maxq = self.maxq.to(dev) + + shape = x.shape + if self.perchannel: + if weight: + x = x.flatten(1) + else: + if len(shape) == 4: + x = x.permute([1, 0, 2, 3]) + x = x.flatten(1) + if len(shape) == 3: + x = x.reshape((-1, shape[-1])).t() + if len(shape) == 2: + x = x.t() + else: + x = x.flatten().unsqueeze(0) + + tmp = torch.zeros(x.shape[0], device=dev) + xmin = torch.minimum(x.min(1)[0], tmp) + xmax = torch.maximum(x.max(1)[0], tmp) + + if self.sym: + xmax = torch.maximum(torch.abs(xmin), xmax) + tmp = xmin < 0 + if torch.any(tmp): + xmin[tmp] = -xmax[tmp] + tmp = (xmin == 0) & (xmax == 0) + xmin[tmp] = -1 + xmax[tmp] = +1 + + self.scale = (xmax - xmin) / self.maxq + if self.sym: + self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2) + else: + self.zero = torch.round(-xmin / self.scale) + + if self.mse: + best = torch.full([x.shape[0]], float('inf'), device=dev) + for i in range(int(self.maxshrink * self.grid)): + p = 1 - i / self.grid + xmin1 = p * xmin + xmax1 = p * xmax + scale1 = (xmax1 - xmin1) / self.maxq + zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero + q = quantize(x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq) + q -= x + q.abs_() + q.pow_(self.norm) + err = torch.sum(q, 1) + tmp = err < best + if torch.any(tmp): + best[tmp] = err[tmp] + self.scale[tmp] = scale1[tmp] + self.zero[tmp] = zero1[tmp] + if not self.perchannel: + if weight: + tmp = shape[0] + else: + tmp = shape[1] if len(shape) != 3 else shape[2] + self.scale = self.scale.repeat(tmp) + self.zero = self.zero.repeat(tmp) + + if weight: + shape = [-1] + [1] * (len(shape) - 1) + self.scale = self.scale.reshape(shape) + self.zero = self.zero.reshape(shape) + return + if len(shape) == 4: + self.scale = self.scale.reshape((1, -1, 1, 1)) + self.zero = self.zero.reshape((1, -1, 1, 1)) + if len(shape) == 3: + self.scale = self.scale.reshape((1, 1, -1)) + self.zero = self.zero.reshape((1, 1, -1)) + if len(shape) == 2: + self.scale = self.scale.unsqueeze(0) + self.zero = self.zero.unsqueeze(0) + + def quantize(self, x): + if self.ready(): + return quantize(x, self.scale, self.zero, self.maxq) + return x + + def enabled(self): + return self.maxq > 0 + + def ready(self): + return torch.all(self.scale != 0) + + +try: + import quant_cuda +except: + print('CUDA extension not installed.') + +# Assumes layer is perfectly divisible into 256 * 256 blocks + + +class QuantLinear(nn.Module): + + def __init__(self, bits, groupsize, infeatures, outfeatures): + super().__init__() + if bits not in [2, 3, 4, 8]: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + self.infeatures = infeatures + self.outfeatures = outfeatures + self.bits = bits + if groupsize != -1 and groupsize < 32 and groupsize != int(math.pow(2, int(math.log2(groupsize)))): + raise NotImplementedError("groupsize supports powers of 2 greater than 32. (e.g. : 32,64,128,etc)") + groupsize = groupsize if groupsize != -1 else infeatures + self.groupsize = groupsize + self.register_buffer( + 'qzeros', torch.zeros((math.ceil(infeatures / groupsize), outfeatures // 256 * (bits * 8)), + dtype=torch.int)) + self.register_buffer('scales', torch.zeros((math.ceil(infeatures / groupsize), outfeatures))) + self.register_buffer('bias', torch.zeros(outfeatures)) + self.register_buffer('qweight', torch.zeros((infeatures // 256 * (bits * 8), outfeatures), dtype=torch.int)) + self._initialized_quant_state = False + + def pack(self, linear, scales, zeros): + scales = scales.t().contiguous() + zeros = zeros.t().contiguous() + scale_zeros = zeros * scales + self.scales = scales.clone() + if linear.bias is not None: + self.bias = linear.bias.clone() + + intweight = [] + for idx in range(self.infeatures): + g_idx = idx // self.groupsize + intweight.append( + torch.round((linear.weight.data[:, idx] + scale_zeros[g_idx]) / self.scales[g_idx]).to(torch.int)[:, + None]) + intweight = torch.cat(intweight, dim=1) + intweight = intweight.t().contiguous() + intweight = intweight.numpy().astype(np.uint32) + qweight = np.zeros((intweight.shape[0] // 256 * (self.bits * 8), intweight.shape[1]), dtype=np.uint32) + i = 0 + row = 0 + while row < qweight.shape[0]: + if self.bits in [2, 4, 8]: + for j in range(i, i + (32 // self.bits)): + qweight[row] |= intweight[j] << (self.bits * (j - i)) + i += 32 // self.bits + row += 1 + elif self.bits == 3: + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i)) + i += 10 + qweight[row] |= intweight[i] << 30 + row += 1 + qweight[row] |= (intweight[i] >> 2) & 1 + i += 1 + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i) + 1) + i += 10 + qweight[row] |= intweight[i] << 31 + row += 1 + qweight[row] |= (intweight[i] >> 1) & 0x3 + i += 1 + for j in range(i, i + 10): + qweight[row] |= intweight[j] << (3 * (j - i) + 2) + i += 10 + row += 1 + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + + qweight = qweight.astype(np.int32) + self.qweight = torch.from_numpy(qweight) + + zeros -= 1 + zeros = zeros.numpy().astype(np.uint32) + qzeros = np.zeros((zeros.shape[0], zeros.shape[1] // 256 * (self.bits * 8)), dtype=np.uint32) + i = 0 + col = 0 + while col < qzeros.shape[1]: + if self.bits in [2, 4, 8]: + for j in range(i, i + (32 // self.bits)): + qzeros[:, col] |= zeros[:, j] << (self.bits * (j - i)) + i += 32 // self.bits + col += 1 + elif self.bits == 3: + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i)) + i += 10 + qzeros[:, col] |= zeros[:, i] << 30 + col += 1 + qzeros[:, col] |= (zeros[:, i] >> 2) & 1 + i += 1 + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 1) + i += 10 + qzeros[:, col] |= zeros[:, i] << 31 + col += 1 + qzeros[:, col] |= (zeros[:, i] >> 1) & 0x3 + i += 1 + for j in range(i, i + 10): + qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 2) + i += 10 + col += 1 + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + + qzeros = qzeros.astype(np.int32) + self.qzeros = torch.from_numpy(qzeros) + + def forward(self, x): + intermediate_dtype = torch.float32 + + if not self._initialized_quant_state: + # Do we even have a bias? Check for at least one non-zero element. + if self.bias is not None and bool(torch.any(self.bias != 0)): + # Then make sure it's the right type. + self.bias.data = self.bias.data.to(intermediate_dtype) + else: + self.bias = None + + outshape = list(x.shape) + outshape[-1] = self.outfeatures + x = x.reshape(-1, x.shape[-1]) + if self.bias is None: + y = torch.zeros(x.shape[0], outshape[-1], dtype=intermediate_dtype, device=x.device) + else: + y = self.bias.clone().repeat(x.shape[0], 1) + + output_dtype = x.dtype + x = x.to(intermediate_dtype) + if self.bits == 2: + quant_cuda.vecquant2matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 3: + quant_cuda.vecquant3matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 4: + quant_cuda.vecquant4matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + elif self.bits == 8: + quant_cuda.vecquant8matmul(x, self.qweight, y, self.scales, self.qzeros, self.groupsize) + else: + raise NotImplementedError("Only 2,3,4,8 bits are supported.") + y = y.to(output_dtype) + return y.reshape(outshape) + + +def make_quant(module, names, bits, groupsize, name=''): + if isinstance(module, QuantLinear): + return + for attr in dir(module): + tmp = getattr(module, attr) + name1 = name + '.' + attr if name != '' else attr + if name1 in names: + setattr(module, attr, QuantLinear(bits, groupsize, tmp.in_features, tmp.out_features)) + for name1, child in module.named_children(): + make_quant(child, names, bits, groupsize, name + '.' + name1 if name != '' else name1) diff --git a/applications/Chat/inference/locustfile.py b/applications/Chat/inference/locustfile.py new file mode 100644 index 000000000000..51cdc68125bb --- /dev/null +++ b/applications/Chat/inference/locustfile.py @@ -0,0 +1,27 @@ +from json import JSONDecodeError + +from locust import HttpUser, task + +samples = [[ + dict( + instruction='Who is the best player in the history of NBA?', + response= + 'The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1' + ), + dict(instruction='continue this talk', response=''), +], [ + dict(instruction='Who is the best player in the history of NBA?', response=''), +]] + + +class GenerationUser(HttpUser): + + @task + def generate(self): + for sample in samples: + data = {'max_new_tokens': 64, 'history': sample} + with self.client.post('/generate', json=data, catch_response=True) as response: + if response.status_code in (200, 406): + response.success() + else: + response.failure('Response wrong') diff --git a/applications/Chat/inference/requirements.txt b/applications/Chat/inference/requirements.txt new file mode 100644 index 000000000000..511fe1a4f1f3 --- /dev/null +++ b/applications/Chat/inference/requirements.txt @@ -0,0 +1,13 @@ +fastapi +locust +numpy +pydantic +safetensors +slowapi +sse_starlette +torch +uvicorn +git+https://github.com/huggingface/transformers +accelerate +bitsandbytes +jieba \ No newline at end of file diff --git a/applications/Chat/inference/server.py b/applications/Chat/inference/server.py new file mode 100644 index 000000000000..b4627299397e --- /dev/null +++ b/applications/Chat/inference/server.py @@ -0,0 +1,178 @@ +import argparse +import os +from threading import Lock +from typing import Dict, Generator, List, Optional + +import torch +import uvicorn +from fastapi import FastAPI, HTTPException, Request +from fastapi.middleware.cors import CORSMiddleware +from llama_gptq import load_quant +from pydantic import BaseModel, Field +from slowapi import Limiter, _rate_limit_exceeded_handler +from slowapi.errors import RateLimitExceeded +from slowapi.util import get_remote_address +from sse_starlette.sse import EventSourceResponse +from transformers import AutoTokenizer, GenerationConfig, LlamaForCausalLM +from utils import ChatPromptProcessor, Dialogue, LockedIterator, sample_streamingly, update_model_kwargs_fn, load_json + +CONTEXT = 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.' +MAX_LEN = 512 +running_lock = Lock() + + +class GenerationTaskReq(BaseModel): + max_new_tokens: int = Field(gt=0, le=512, example=64) + history: List[Dialogue] = Field(min_items=1) + top_k: Optional[int] = Field(default=None, gt=0, example=50) + top_p: Optional[float] = Field(default=None, gt=0.0, lt=1.0, example=0.5) + temperature: Optional[float] = Field(default=None, gt=0.0, lt=1.0, example=0.7) + repetition_penalty: Optional[float] = Field(default=None, gt=1.0, example=1.2) + + +limiter = Limiter(key_func=get_remote_address) +app = FastAPI() +app.state.limiter = limiter +app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) + +# set CORS +origin_spec_from_env = os.environ.get('CORS_ORIGIN', None) + +if origin_spec_from_env is not None: + # allow CORS from the specified origins + origins = os.environ['CORS_ORIGIN'].split(',') +else: + # allow CORS from all origins + origins = ["*"] + +app.add_middleware( + CORSMiddleware, + allow_origins=origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +def generate_streamingly(prompt, max_new_tokens, top_k, top_p, temperature): + inputs = {k: v.cuda() for k, v in tokenizer(prompt, return_tensors="pt").items()} + #TODO(ver217): streaming generation does not support repetition_penalty now + model_kwargs = { + 'max_generate_tokens': max_new_tokens, + 'early_stopping': True, + 'top_k': top_k, + 'top_p': top_p, + 'temperature': temperature, + 'prepare_inputs_fn': model.prepare_inputs_for_generation, + 'update_model_kwargs_fn': update_model_kwargs_fn, + } + is_first_word = True + generator = LockedIterator(sample_streamingly(model, **inputs, **model_kwargs), running_lock) + for output in generator: + output = output.cpu() + tokens = tokenizer.convert_ids_to_tokens(output, skip_special_tokens=True) + current_sub_tokens = [] + for token in tokens: + if token in tokenizer.all_special_tokens: + continue + current_sub_tokens.append(token) + if current_sub_tokens: + out_string = tokenizer.sp_model.decode(current_sub_tokens) + if is_first_word: + out_string = out_string.lstrip() + is_first_word = False + elif current_sub_tokens[0].startswith('▁'): + # whitespace will be ignored by the frontend + out_string = ' ' + out_string + yield out_string + + +async def event_generator(request: Request, generator: Generator): + while True: + if await request.is_disconnected(): + break + try: + yield {'event': 'generate', 'data': next(generator)} + except StopIteration: + yield {'event': 'end', 'data': ''} + break + + +@app.post('/generate/stream') +@limiter.limit('1/second') +def generate(data: GenerationTaskReq, request: Request): + prompt = prompt_processor.preprocess_prompt(data.history, data.max_new_tokens) + event_source = event_generator( + request, generate_streamingly(prompt, data.max_new_tokens, data.top_k, data.top_p, data.temperature)) + return EventSourceResponse(event_source) + + +@app.post('/generate') +@limiter.limit('1/second') +def generate_no_stream(data: GenerationTaskReq, request: Request): + prompt = prompt_processor.preprocess_prompt(data.history, data.max_new_tokens) + if prompt_processor.has_censored_words(prompt): + return prompt_processor.SAFE_RESPONSE + inputs = {k: v.cuda() for k, v in tokenizer(prompt, return_tensors="pt").items()} + with running_lock: + output = model.generate(**inputs, **data.dict(exclude={'history'})) + output = output.cpu() + prompt_len = inputs['input_ids'].size(1) + response = output[0, prompt_len:] + out_string = tokenizer.decode(response, skip_special_tokens=True) + out_string = prompt_processor.postprocess_output(out_string) + if prompt_processor.has_censored_words(out_string): + return prompt_processor.SAFE_RESPONSE + return out_string + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument( + 'pretrained', + help='Path to pretrained model. Can be a local path or a model name from the HuggingFace model hub.') + parser.add_argument('--quant', + choices=['8bit', '4bit'], + default=None, + help='Quantization mode. Default: None (no quantization, fp16).') + parser.add_argument( + '--gptq_checkpoint', + default=None, + help='Path to GPTQ checkpoint. This is only useful when quantization mode is 4bit. Default: None.') + parser.add_argument('--gptq_group_size', + type=int, + default=128, + help='Group size for GPTQ. This is only useful when quantization mode is 4bit. Default: 128.') + parser.add_argument('--http_host', default='0.0.0.0') + parser.add_argument('--http_port', type=int, default=7070) + parser.add_argument('--profanity_file', default=None, help='Path to profanity words list. It should be a JSON file containing a list of words.') + args = parser.parse_args() + + if args.quant == '4bit': + assert args.gptq_checkpoint is not None, 'Please specify a GPTQ checkpoint.' + + tokenizer = AutoTokenizer.from_pretrained(args.pretrained) + + if args.profanity_file is not None: + censored_words = load_json(args.profanity_file) + else: + censored_words = [] + prompt_processor = ChatPromptProcessor(tokenizer, CONTEXT, MAX_LEN, censored_words=censored_words) + + if args.quant == '4bit': + model = load_quant(args.pretrained, args.gptq_checkpoint, 4, args.gptq_group_size) + model.cuda() + else: + model = LlamaForCausalLM.from_pretrained( + args.pretrained, + load_in_8bit=(args.quant == '8bit'), + torch_dtype=torch.float16, + device_map="auto", + ) + if args.quant != '8bit': + model.half() # seems to fix bugs for some users. + model.eval() + + config = uvicorn.Config(app, host=args.http_host, port=args.http_port) + server = uvicorn.Server(config=config) + server.run() diff --git a/applications/Chat/inference/tests/test_chat_prompt.py b/applications/Chat/inference/tests/test_chat_prompt.py new file mode 100644 index 000000000000..f5737ebe8c09 --- /dev/null +++ b/applications/Chat/inference/tests/test_chat_prompt.py @@ -0,0 +1,56 @@ +import os + +from transformers import AutoTokenizer +from utils import ChatPromptProcessor, Dialogue + +CONTEXT = 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.' +tokenizer = AutoTokenizer.from_pretrained(os.environ['PRETRAINED_PATH']) + +samples = [ + ([ + Dialogue( + instruction='Who is the best player in the history of NBA?', + response= + 'The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1' + ), + Dialogue(instruction='continue this talk', response=''), + ], 128, + 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.\n\n### Instruction:\nWho is the best player in the history of NBA?\n\n### Response:\nThe best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1\n\n### Instruction:\ncontinue this talk\n\n### Response:\n' + ), + ([ + Dialogue( + instruction='Who is the best player in the history of NBA?', + response= + 'The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1' + ), + Dialogue(instruction='continue this talk', response=''), + ], 200, + 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.\n\n### Instruction:\ncontinue this talk\n\n### Response:\n' + ), + ([ + Dialogue( + instruction='Who is the best player in the history of NBA?', + response= + 'The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1' + ), + Dialogue(instruction='continue this talk', response=''), + ], 211, + 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.\n\n### Instruction:\ncontinue this\n\n### Response:\n' + ), + ([ + Dialogue(instruction='Who is the best player in the history of NBA?', response=''), + ], 128, + 'Below is an instruction that describes a task. Write a response that appropriately completes the request. Do not generate new instructions.\n\n### Instruction:\nWho is the best player in the history of NBA?\n\n### Response:\n' + ), +] + + +def test_chat_prompt_processor(): + processor = ChatPromptProcessor(tokenizer, CONTEXT, 256) + for history, max_new_tokens, result in samples: + prompt = processor.preprocess_prompt(history, max_new_tokens) + assert prompt == result + + +if __name__ == '__main__': + test_chat_prompt_processor() diff --git a/applications/Chat/inference/utils.py b/applications/Chat/inference/utils.py new file mode 100644 index 000000000000..37944be70a3b --- /dev/null +++ b/applications/Chat/inference/utils.py @@ -0,0 +1,200 @@ +import re +from threading import Lock +from typing import Any, Callable, Generator, List, Optional +import json +import jieba + +import torch +import torch.distributed as dist +import torch.nn as nn +from pydantic import BaseModel, Field + +try: + from transformers.generation_logits_process import ( + LogitsProcessorList, + TemperatureLogitsWarper, + TopKLogitsWarper, + TopPLogitsWarper, + ) +except ImportError: + from transformers.generation import LogitsProcessorList, TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper + + +def prepare_logits_processor(top_k: Optional[int] = None, + top_p: Optional[float] = None, + temperature: Optional[float] = None) -> LogitsProcessorList: + processor_list = LogitsProcessorList() + if temperature is not None and temperature != 1.0: + processor_list.append(TemperatureLogitsWarper(temperature)) + if top_k is not None and top_k != 0: + processor_list.append(TopKLogitsWarper(top_k)) + if top_p is not None and top_p < 1.0: + processor_list.append(TopPLogitsWarper(top_p)) + return processor_list + + +def _is_sequence_finished(unfinished_sequences: torch.Tensor) -> bool: + if dist.is_initialized() and dist.get_world_size() > 1: + # consider DP + unfinished_sequences = unfinished_sequences.clone() + dist.all_reduce(unfinished_sequences) + return unfinished_sequences.max() == 0 + + +def sample_streamingly(model: nn.Module, + input_ids: torch.Tensor, + max_generate_tokens: int, + early_stopping: bool = False, + eos_token_id: Optional[int] = None, + pad_token_id: Optional[int] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + temperature: Optional[float] = None, + prepare_inputs_fn: Optional[Callable[[torch.Tensor, Any], dict]] = None, + update_model_kwargs_fn: Optional[Callable[[dict, Any], dict]] = None, + **model_kwargs) -> Generator: + + logits_processor = prepare_logits_processor(top_k, top_p, temperature) + unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) + + for _ in range(max_generate_tokens): + model_inputs = prepare_inputs_fn(input_ids, **model_kwargs) if prepare_inputs_fn is not None else { + 'input_ids': input_ids + } + outputs = model(**model_inputs) + + next_token_logits = outputs['logits'][:, -1, :] + # pre-process distribution + next_token_logits = logits_processor(input_ids, next_token_logits) + # sample + probs = torch.softmax(next_token_logits, dim=-1, dtype=torch.float) + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + + # finished sentences should have their next token be a padding token + if eos_token_id is not None: + if pad_token_id is None: + raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.") + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + yield next_tokens + + # update generated ids, model inputs for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + if update_model_kwargs_fn is not None: + model_kwargs = update_model_kwargs_fn(outputs, **model_kwargs) + + # if eos_token was found in one sentence, set sentence to finished + if eos_token_id is not None: + unfinished_sequences = unfinished_sequences.mul((next_tokens != eos_token_id).long()) + + # stop when each sentence is finished if early_stopping=True + if early_stopping and _is_sequence_finished(unfinished_sequences): + break + + +def update_model_kwargs_fn(outputs: dict, **model_kwargs) -> dict: + if "past_key_values" in outputs: + model_kwargs["past"] = outputs["past_key_values"] + else: + model_kwargs["past"] = None + + # update token_type_ids with last value + if "token_type_ids" in model_kwargs: + token_type_ids = model_kwargs["token_type_ids"] + model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) + + # update attention mask + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) + + return model_kwargs + + +class Dialogue(BaseModel): + instruction: str = Field(min_length=1, example='Count up from 1 to 500.') + response: str = Field(example='') + + +def _format_dialogue(instruction: str, response: str = ''): + return f'\n\n### Instruction:\n{instruction}\n\n### Response:\n{response}' + + +STOP_PAT = re.compile(r'(###|instruction:).*', flags=(re.I | re.S)) + + +class ChatPromptProcessor: + SAFE_RESPONSE = 'The input/response contains inappropriate content, please rephrase your prompt.' + + def __init__(self, tokenizer, context: str, max_len: int = 2048, censored_words: List[str]=[]): + self.tokenizer = tokenizer + self.context = context + self.max_len = max_len + self.censored_words = set([word.lower() for word in censored_words]) + # These will be initialized after the first call of preprocess_prompt() + self.context_len: Optional[int] = None + self.dialogue_placeholder_len: Optional[int] = None + + def preprocess_prompt(self, history: List[Dialogue], max_new_tokens: int) -> str: + if self.context_len is None: + self.context_len = len(self.tokenizer(self.context)['input_ids']) + if self.dialogue_placeholder_len is None: + self.dialogue_placeholder_len = len( + self.tokenizer(_format_dialogue(''), add_special_tokens=False)['input_ids']) + prompt = self.context + # the last dialogue must be in the prompt + last_dialogue = history.pop() + # the response of the last dialogue is empty + assert last_dialogue.response == '' + if len(self.tokenizer(_format_dialogue(last_dialogue.instruction), add_special_tokens=False) + ['input_ids']) + max_new_tokens + self.context_len >= self.max_len: + # to avoid truncate placeholder, apply truncate to the original instruction + instruction_truncated = self.tokenizer(last_dialogue.instruction, + add_special_tokens=False, + truncation=True, + max_length=(self.max_len - max_new_tokens - self.context_len - + self.dialogue_placeholder_len))['input_ids'] + instruction_truncated = self.tokenizer.decode(instruction_truncated).lstrip() + prompt += _format_dialogue(instruction_truncated) + return prompt + + res_len = self.max_len - max_new_tokens - len(self.tokenizer(prompt)['input_ids']) + + rows = [] + for dialogue in history[::-1]: + text = _format_dialogue(dialogue.instruction, dialogue.response) + cur_len = len(self.tokenizer(text, add_special_tokens=False)['input_ids']) + if res_len - cur_len < 0: + break + res_len -= cur_len + rows.insert(0, text) + prompt += ''.join(rows) + _format_dialogue(last_dialogue.instruction) + return prompt + + def postprocess_output(self, output: str) -> str: + output = STOP_PAT.sub('', output) + return output.strip() + + def has_censored_words(self, text: str) -> bool: + if len(self.censored_words) == 0: + return False + intersection = set(jieba.cut(text.lower())) & self.censored_words + return len(intersection) > 0 + +class LockedIterator: + + def __init__(self, it, lock: Lock) -> None: + self.lock = lock + self.it = iter(it) + + def __iter__(self): + return self + + def __next__(self): + with self.lock: + return next(self.it) + +def load_json(path: str): + with open(path) as f: + return json.load(f) \ No newline at end of file diff --git a/applications/ChatGPT/pytest.ini b/applications/Chat/pytest.ini similarity index 100% rename from applications/ChatGPT/pytest.ini rename to applications/Chat/pytest.ini diff --git a/applications/ChatGPT/requirements-test.txt b/applications/Chat/requirements-test.txt similarity index 100% rename from applications/ChatGPT/requirements-test.txt rename to applications/Chat/requirements-test.txt diff --git a/applications/Chat/requirements.txt b/applications/Chat/requirements.txt new file mode 100644 index 000000000000..af7ff67861eb --- /dev/null +++ b/applications/Chat/requirements.txt @@ -0,0 +1,13 @@ +transformers>=4.20.1 +tqdm +datasets +loralib +colossalai>=0.2.4 +torch<2.0.0, >=1.12.1 +langchain +tokenizers +fastapi +sse_starlette +wandb +sentencepiece +gpustat diff --git a/applications/ChatGPT/setup.py b/applications/Chat/setup.py similarity index 87% rename from applications/ChatGPT/setup.py rename to applications/Chat/setup.py index deec10e0c841..a285a6dff4bf 100644 --- a/applications/ChatGPT/setup.py +++ b/applications/Chat/setup.py @@ -17,18 +17,18 @@ def fetch_version(): setup( - name='chatgpt', + name='coati', version=fetch_version(), packages=find_packages(exclude=( 'tests', 'benchmarks', '*.egg-info', )), - description='A RLFH implementation (ChatGPT) powered by ColossalAI', + description='Colossal-AI Talking Intelligence', long_description=fetch_readme(), long_description_content_type='text/markdown', license='Apache Software License 2.0', - url='https://github.com/hpcaitech/ChatGPT', + url='https://github.com/hpcaitech/Coati', install_requires=fetch_requirements('requirements.txt'), python_requires='>=3.6', classifiers=[ diff --git a/applications/Chat/tests/__init__.py b/applications/Chat/tests/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/applications/ChatGPT/tests/test_checkpoint.py b/applications/Chat/tests/test_checkpoint.py similarity index 71% rename from applications/ChatGPT/tests/test_checkpoint.py rename to applications/Chat/tests/test_checkpoint.py index 1bbd133f76d3..19338da437ab 100644 --- a/applications/ChatGPT/tests/test_checkpoint.py +++ b/applications/Chat/tests/test_checkpoint.py @@ -1,19 +1,17 @@ import os import tempfile from contextlib import nullcontext -from functools import partial import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -from chatgpt.models.gpt import GPTActor -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy +from coati.models.gpt import GPTActor +from coati.models.utils import calc_action_log_probs +from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy from transformers.models.gpt2.configuration_gpt2 import GPT2Config from colossalai.nn.optimizer import HybridAdam -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn GPT_CONFIG = GPT2Config(n_embd=128, n_layer=4, n_head=4) @@ -30,9 +28,9 @@ def run_test_checkpoint(strategy): if strategy == 'ddp': strategy = DDPStrategy() elif strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) + strategy = GeminiStrategy(placement_policy='cuda', initial_scale=2**5) elif strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') + strategy = LowLevelZeroStrategy(stage=2, placement_policy='cuda') else: raise ValueError(f'Unsupported strategy "{strategy}"') @@ -46,7 +44,8 @@ def run_test_checkpoint(strategy): def run_step(): data = get_data(BATCH_SIZE) action_mask = torch.ones_like(data['attention_mask'], dtype=torch.bool) - action_log_probs = actor(data['input_ids'], action_mask.size(1), data['attention_mask']) + actor_output = actor(data['input_ids'], data['attention_mask']) + action_log_probs = calc_action_log_probs(actor_output, data['input_ids'], action_mask.size(1)) loss = action_log_probs.sum() strategy.backward(loss, actor, actor_optim) strategy.optimizer_step(actor_optim) @@ -61,10 +60,15 @@ def run_step(): rank0_dirname = rank0_dirname[0] model_path = os.path.join(rank0_dirname, 'model.pt') - optim_path = os.path.join(rank0_dirname, f'optim-r{dist.get_rank()}.pt') - strategy.save_model(actor, model_path, only_rank0=True) - strategy.save_optimizer(actor_optim, optim_path, only_rank0=False) + + optim_path = os.path.join(rank0_dirname, f'optim.pt') + strategy.save_optimizer(actor_optim, optim_path, only_rank0=True) + + # FIXME(cwher): Sharded optimizer checkpoint is not supported yet. + # at "ColossalAI/colossalai/checkpoint_io/general_checkpoint_io.py", line 62 + # optim_path = os.path.join(rank0_dirname, f'optim-r{dist.get_rank()}.pt') + # strategy.save_optimizer(actor_optim, optim_path, only_rank0=False) dist.barrier() @@ -90,8 +94,7 @@ def run_dist(rank, world_size, port, strategy): @pytest.mark.parametrize('strategy', ['ddp', 'colossalai_zero2', 'colossalai_gemini']) @rerun_if_address_is_in_use() def test_checkpoint(world_size, strategy): - run_func = partial(run_dist, world_size=world_size, port=free_port(), strategy=strategy) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size, strategy=strategy) if __name__ == '__main__': diff --git a/applications/ChatGPT/tests/test_data.py b/applications/Chat/tests/test_data.py similarity index 85% rename from applications/ChatGPT/tests/test_data.py rename to applications/Chat/tests/test_data.py index 3d8fe912cb27..db641a6218b1 100644 --- a/applications/ChatGPT/tests/test_data.py +++ b/applications/Chat/tests/test_data.py @@ -1,20 +1,17 @@ import os from copy import deepcopy -from functools import partial import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -from chatgpt.experience_maker import NaiveExperienceMaker -from chatgpt.models.base import RewardModel -from chatgpt.models.gpt import GPTActor, GPTCritic -from chatgpt.replay_buffer import NaiveReplayBuffer -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy +from coati.experience_maker import NaiveExperienceMaker +from coati.models.base import RewardModel +from coati.models.gpt import GPTActor, GPTCritic +from coati.replay_buffer import NaiveReplayBuffer +from coati.trainer.strategies import DDPStrategy, GeminiStrategy from transformers.models.gpt2.configuration_gpt2 import GPT2Config -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn GPT_CONFIG = GPT2Config(n_embd=128, n_layer=4, n_head=4) @@ -36,13 +33,13 @@ def gather_and_equal(tensor: torch.Tensor) -> bool: def run_test_data(strategy): - EXPERINCE_BATCH_SIZE = 4 + EXPERIENCE_BATCH_SIZE = 4 SAMPLE_BATCH_SIZE = 2 if strategy == 'ddp': strategy = DDPStrategy() elif strategy == 'colossalai': - strategy = ColossalAIStrategy(placement_policy='cuda') + strategy = GeminiStrategy(placement_policy='cuda') else: raise ValueError(f'Unsupported strategy "{strategy}"') @@ -57,7 +54,7 @@ def run_test_data(strategy): # experience of all ranks should be the same for _ in range(2): - data = get_data(EXPERINCE_BATCH_SIZE) + data = get_data(EXPERIENCE_BATCH_SIZE) assert gather_and_equal(data['input_ids']) assert gather_and_equal(data['attention_mask']) experience = experience_maker.make_experience(**data, @@ -114,8 +111,7 @@ def run_dist(rank, world_size, port, strategy): @pytest.mark.parametrize('strategy', ['ddp', 'colossalai']) @rerun_if_address_is_in_use() def test_data(world_size, strategy): - run_func = partial(run_dist, world_size=world_size, port=free_port(), strategy=strategy) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size, strategy=strategy) if __name__ == '__main__': diff --git a/applications/Chat/version.txt b/applications/Chat/version.txt new file mode 100644 index 000000000000..3eefcb9dd5b3 --- /dev/null +++ b/applications/Chat/version.txt @@ -0,0 +1 @@ +1.0.0 diff --git a/applications/ChatGPT/README.md b/applications/ChatGPT/README.md deleted file mode 100644 index 206ede5f1843..000000000000 --- a/applications/ChatGPT/README.md +++ /dev/null @@ -1,209 +0,0 @@ -# RLHF - Colossal-AI - -## Table of Contents - -- [What is RLHF - Colossal-AI?](#intro) -- [How to Install?](#install) -- [The Plan](#the-plan) -- [How can you partcipate in open source?](#invitation-to-open-source-contribution) ---- -## Intro -Implementation of RLHF (Reinforcement Learning with Human Feedback) powered by Colossal-AI. It supports distributed training and offloading, which can fit extremly large models. More details can be found in the [blog](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt). - -

- -

- -## Training process (step 3) -

- -

-

- -

- - -## Install -```shell -pip install . -``` - -## Usage - -The main entrypoint is `Trainer`. We only support PPO trainer now. We support many training strategies: - -- NaiveStrategy: simplest strategy. Train on single GPU. -- DDPStrategy: use `torch.nn.parallel.DistributedDataParallel`. Train on multi GPUs. -- ColossalAIStrategy: use Gemini and Zero of ColossalAI. It eliminates model duplication on each GPU and supports offload. It's very useful when training large models on multi GPUs. - -Simplest usage: - -```python -from chatgpt.trainer import PPOTrainer -from chatgpt.trainer.strategies import ColossalAIStrategy -from chatgpt.models.gpt import GPTActor, GPTCritic -from chatgpt.models.base import RewardModel -from copy import deepcopy -from colossalai.nn.optimizer import HybridAdam - -strategy = ColossalAIStrategy() - -with strategy.model_init_context(): - # init your model here - # load pretrained gpt2 - actor = GPTActor(pretrained='gpt2') - critic = GPTCritic() - initial_model = deepcopy(actor).cuda() - reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() - -actor_optim = HybridAdam(actor.parameters(), lr=5e-6) -critic_optim = HybridAdam(critic.parameters(), lr=5e-6) - -# prepare models and optimizers -(actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare( - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model) - -# load saved model checkpoint after preparing -strategy.load_model(actor, 'actor_checkpoint.pt', strict=False) -# load saved optimizer checkpoint after preparing -strategy.load_optimizer(actor_optim, 'actor_optim_checkpoint.pt') - -trainer = PPOTrainer(strategy, - actor, - critic, - reward_model, - initial_model, - actor_optim, - critic_optim, - ...) - -trainer.fit(dataset, ...) - -# save model checkpoint after fitting on only rank0 -strategy.save_model(actor, 'actor_checkpoint.pt', only_rank0=True) -# save optimizer checkpoint on all ranks -strategy.save_optimizer(actor_optim, 'actor_optim_checkpoint.pt', only_rank0=False) -``` - -For more details, see `examples/`. - -We also support training reward model with true-world data. See `examples/train_reward_model.py`. - -## FAQ - -### How to save/load checkpoint - -To load pretrained model, you can simply use huggingface pretrained models: - -```python -# load OPT-350m pretrained model -actor = OPTActor(pretrained='facebook/opt-350m') -``` - -To save model checkpoint: - -```python -# save model checkpoint on only rank0 -strategy.save_model(actor, 'actor_checkpoint.pt', only_rank0=True) -``` - -This function must be called after `strategy.prepare()`. - -For DDP strategy, model weights are replicated on all ranks. And for ColossalAI strategy, model weights may be sharded, but all-gather will be applied before returning state dict. You can set `only_rank0=True` for both of them, which only saves checkpoint on rank0, to save disk space usage. The checkpoint is float32. - -To save optimizer checkpoint: - -```python -# save optimizer checkpoint on all ranks -strategy.save_optimizer(actor_optim, 'actor_optim_checkpoint.pt', only_rank0=False) -``` - -For DDP strategy, optimizer states are replicated on all ranks. You can set `only_rank0=True`. But for ColossalAI strategy, optimizer states are sharded over all ranks, and no all-gather will be applied. So for ColossalAI strategy, you can only set `only_rank0=False`. That is to say, each rank will save a cehckpoint. When loading, each rank should load the corresponding part. - -Note that different stategy may have different shapes of optimizer checkpoint. - -To load model checkpoint: - -```python -# load saved model checkpoint after preparing -strategy.load_model(actor, 'actor_checkpoint.pt', strict=False) -``` - -To load optimizer checkpoint: - -```python -# load saved optimizer checkpoint after preparing -strategy.load_optimizer(actor_optim, 'actor_optim_checkpoint.pt') -``` - -## The Plan - -- [x] implement PPO fine-tuning -- [x] implement training reward model -- [x] support LoRA -- [x] support inference -- [ ] open source the reward model weight -- [ ] support llama from [facebook](https://github.com/facebookresearch/llama) -- [ ] support BoN(best of N sample) -- [ ] implement PPO-ptx fine-tuning -- [ ] integrate with Ray -- [ ] support more RL paradigms, like Implicit Language Q-Learning (ILQL), -- [ ] support chain of throught by [langchain](https://github.com/hwchase17/langchain) - -### Real-time progress -You will find our progress in github project broad - -[Open ChatGPT](https://github.com/orgs/hpcaitech/projects/17/views/1) - -## Invitation to open-source contribution -Referring to the successful attempts of [BLOOM](https://bigscience.huggingface.co/) and [Stable Diffusion](https://en.wikipedia.org/wiki/Stable_Diffusion), any and all developers and partners with computing powers, datasets, models are welcome to join and build the Colossal-AI community, making efforts towards the era of big AI models from the starting point of replicating ChatGPT! - -You may contact us or participate in the following ways: -1. [Leaving a Star ⭐](https://github.com/hpcaitech/ColossalAI/stargazers) to show your like and support. Thanks! -2. Posting an [issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose), or submitting a PR on GitHub follow the guideline in [Contributing](https://github.com/hpcaitech/ColossalAI/blob/main/CONTRIBUTING.md). -3. Join the Colossal-AI community on -[Slack](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w), -and [WeChat(微信)](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png "qrcode") to share your ideas. -4. Send your official proposal to email contact@hpcaitech.com - -Thanks so much to all of our amazing contributors! - -## Quick Preview -

- -

- -- Up to 7.73 times faster for single server training and 1.42 times faster for single-GPU inference - -

- -

- -- Up to 10.3x growth in model capacity on one GPU -- A mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU) - -

- -

- -- Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPU -- Keep in a sufficiently high running speed - -## Citations - -```bibtex -@article{Hu2021LoRALA, - title = {LoRA: Low-Rank Adaptation of Large Language Models}, - author = {Edward J. Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Weizhu Chen}, - journal = {ArXiv}, - year = {2021}, - volume = {abs/2106.09685} -} - -@article{ouyang2022training, - title={Training language models to follow instructions with human feedback}, - author={Ouyang, Long and Wu, Jeff and Jiang, Xu and Almeida, Diogo and Wainwright, Carroll L and Mishkin, Pamela and Zhang, Chong and Agarwal, Sandhini and Slama, Katarina and Ray, Alex and others}, - journal={arXiv preprint arXiv:2203.02155}, - year={2022} -} -``` diff --git a/applications/ChatGPT/benchmarks/README.md b/applications/ChatGPT/benchmarks/README.md deleted file mode 100644 index b4e28ba1d764..000000000000 --- a/applications/ChatGPT/benchmarks/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# Benchmarks - -## Benchmark GPT on dummy prompt data - -We provide various GPT models (string in parentheses is the corresponding model name used in this script): - -- GPT2-S (s) -- GPT2-M (m) -- GPT2-L (l) -- GPT2-XL (xl) -- GPT2-4B (4b) -- GPT2-6B (6b) -- GPT2-8B (8b) -- GPT2-10B (10b) -- GPT2-12B (12b) -- GPT2-15B (15b) -- GPT2-18B (18b) -- GPT2-20B (20b) -- GPT2-24B (24b) -- GPT2-28B (28b) -- GPT2-32B (32b) -- GPT2-36B (36b) -- GPT2-40B (40b) -- GPT3 (175b) - -We also provide various training strategies: - -- ddp: torch DDP -- colossalai_gemini: ColossalAI GeminiDDP with `placement_policy="cuda"`, like zero3 -- colossalai_gemini_cpu: ColossalAI GeminiDDP with `placement_policy="cpu"`, like zero3-offload -- colossalai_zero2: ColossalAI zero2 -- colossalai_zero2_cpu: ColossalAI zero2-offload -- colossalai_zero1: ColossalAI zero1 -- colossalai_zero1_cpu: ColossalAI zero1-offload - -We only support `torchrun` to launch now. E.g. - -```shell -# run GPT2-S on single-node single-GPU with min batch size -torchrun --standalone --nproc_per_node 1 benchmark_gpt_dummy.py --model s --strategy ddp --experience_batch_size 1 --train_batch_size 1 -# run GPT2-XL on single-node 4-GPU -torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2 -# run GPT3 on 8-node 8-GPU -torchrun --nnodes 8 --nproc_per_node 8 \ - --rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR \ - benchmark_gpt_dummy.py --model 175b --strategy colossalai_gemini -``` - -> ⚠ Batch sizes in CLI args and outputed throughput/TFLOPS are all values of per GPU. - -In this benchmark, we assume the model architectures/sizes of actor and critic are the same for simplicity. But in practice, to reduce training cost, we may use a smaller critic. - -We also provide a simple shell script to run a set of benchmarks. But it only supports benchmark on single node. However, it's easy to run on multi-nodes by modifying launch command in this script. - -Usage: - -```shell -# run for GPUS=(1 2 4 8) x strategy=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) -./benchmark_gpt_dummy.sh -# run for GPUS=2 x strategy=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) -./benchmark_gpt_dummy.sh 2 -# run for GPUS=2 x strategy=ddp x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) -./benchmark_gpt_dummy.sh 2 ddp -# run for GPUS=2 x strategy=ddp x model=l x batch_size=(1 2 4 8 16 32 64 128 256) -./benchmark_gpt_dummy.sh 2 ddp l -``` - -## Benchmark OPT with LoRA on dummy prompt data - -We provide various OPT models (string in parentheses is the corresponding model name used in this script): - -- OPT-125M (125m) -- OPT-350M (350m) -- OPT-700M (700m) -- OPT-1.3B (1.3b) -- OPT-2.7B (2.7b) -- OPT-3.5B (3.5b) -- OPT-5.5B (5.5b) -- OPT-6.7B (6.7b) -- OPT-10B (10b) -- OPT-13B (13b) - -We only support `torchrun` to launch now. E.g. - -```shell -# run OPT-125M with no lora (lora_rank=0) on single-node single-GPU with min batch size -torchrun --standalone --nproc_per_node 1 benchmark_opt_lora_dummy.py --model 125m --strategy ddp --experience_batch_size 1 --train_batch_size 1 --lora_rank 0 -# run OPT-350M with lora_rank=4 on single-node 4-GPU -torchrun --standalone --nproc_per_node 4 benchmark_opt_lora_dummy.py --model 350m --strategy colossalai_zero2 --lora_rank 4 -``` - -> ⚠ Batch sizes in CLI args and outputed throughput/TFLOPS are all values of per GPU. - -In this benchmark, we assume the model architectures/sizes of actor and critic are the same for simplicity. But in practice, to reduce training cost, we may use a smaller critic. diff --git a/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.py b/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.py deleted file mode 100644 index 5ee65763b936..000000000000 --- a/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.py +++ /dev/null @@ -1,184 +0,0 @@ -import argparse -from copy import deepcopy - -import torch -import torch.distributed as dist -import torch.nn as nn -from chatgpt.models.base import RewardModel -from chatgpt.models.gpt import GPTActor, GPTCritic -from chatgpt.trainer import PPOTrainer -from chatgpt.trainer.callbacks import PerformanceEvaluator -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy -from torch.optim import Adam -from transformers.models.gpt2.configuration_gpt2 import GPT2Config -from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer - -from colossalai.nn.optimizer import HybridAdam - - -def get_model_numel(model: nn.Module, strategy: Strategy) -> int: - numel = sum(p.numel() for p in model.parameters()) - if isinstance(strategy, ColossalAIStrategy) and strategy.stage == 3 and strategy.shard_init: - numel *= dist.get_world_size() - return numel - - -def preprocess_batch(samples) -> dict: - input_ids = torch.stack(samples) - attention_mask = torch.ones_like(input_ids, dtype=torch.long) - return {'input_ids': input_ids, 'attention_mask': attention_mask} - - -def print_rank_0(*args, **kwargs) -> None: - if dist.get_rank() == 0: - print(*args, **kwargs) - - -def print_model_numel(model_dict: dict) -> None: - B = 1024**3 - M = 1024**2 - K = 1024 - outputs = '' - for name, numel in model_dict.items(): - outputs += f'{name}: ' - if numel >= B: - outputs += f'{numel / B:.2f} B\n' - elif numel >= M: - outputs += f'{numel / M:.2f} M\n' - elif numel >= K: - outputs += f'{numel / K:.2f} K\n' - else: - outputs += f'{numel}\n' - print_rank_0(outputs) - - -def get_gpt_config(model_name: str) -> GPT2Config: - model_map = { - 's': GPT2Config(), - 'm': GPT2Config(n_embd=1024, n_layer=24, n_head=16), - 'l': GPT2Config(n_embd=1280, n_layer=36, n_head=20), - 'xl': GPT2Config(n_embd=1600, n_layer=48, n_head=25), - '2b': GPT2Config(n_embd=2048, n_layer=40, n_head=16), - '4b': GPT2Config(n_embd=2304, n_layer=64, n_head=16), - '6b': GPT2Config(n_embd=4096, n_layer=30, n_head=16), - '8b': GPT2Config(n_embd=4096, n_layer=40, n_head=16), - '10b': GPT2Config(n_embd=4096, n_layer=50, n_head=16), - '12b': GPT2Config(n_embd=4096, n_layer=60, n_head=16), - '15b': GPT2Config(n_embd=4096, n_layer=78, n_head=16), - '18b': GPT2Config(n_embd=4096, n_layer=90, n_head=16), - '20b': GPT2Config(n_embd=8192, n_layer=25, n_head=16), - '24b': GPT2Config(n_embd=8192, n_layer=30, n_head=16), - '28b': GPT2Config(n_embd=8192, n_layer=35, n_head=16), - '32b': GPT2Config(n_embd=8192, n_layer=40, n_head=16), - '36b': GPT2Config(n_embd=8192, n_layer=45, n_head=16), - '40b': GPT2Config(n_embd=8192, n_layer=50, n_head=16), - '175b': GPT2Config(n_positions=2048, n_embd=12288, n_layer=96, n_head=96), - } - try: - return model_map[model_name] - except KeyError: - raise ValueError(f'Unknown model "{model_name}"') - - -def main(args): - if args.strategy == 'ddp': - strategy = DDPStrategy() - elif args.strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) - elif args.strategy == 'colossalai_gemini_cpu': - strategy = ColossalAIStrategy(stage=3, placement_policy='cpu', initial_scale=2**5) - elif args.strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') - elif args.strategy == 'colossalai_zero2_cpu': - strategy = ColossalAIStrategy(stage=2, placement_policy='cpu') - elif args.strategy == 'colossalai_zero1': - strategy = ColossalAIStrategy(stage=1, placement_policy='cuda') - elif args.strategy == 'colossalai_zero1_cpu': - strategy = ColossalAIStrategy(stage=1, placement_policy='cpu') - else: - raise ValueError(f'Unsupported strategy "{args.strategy}"') - - model_config = get_gpt_config(args.model) - - with strategy.model_init_context(): - actor = GPTActor(config=model_config).cuda() - critic = GPTCritic(config=model_config).cuda() - - initial_model = deepcopy(actor).cuda() - reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() - - actor_numel = get_model_numel(actor, strategy) - critic_numel = get_model_numel(critic, strategy) - initial_model_numel = get_model_numel(initial_model, strategy) - reward_model_numel = get_model_numel(reward_model, strategy) - print_model_numel({ - 'Actor': actor_numel, - 'Critic': critic_numel, - 'Initial model': initial_model_numel, - 'Reward model': reward_model_numel - }) - performance_evaluator = PerformanceEvaluator(actor_numel, - critic_numel, - initial_model_numel, - reward_model_numel, - enable_grad_checkpoint=False, - ignore_episodes=1) - - if args.strategy.startswith('colossalai'): - actor_optim = HybridAdam(actor.parameters(), lr=5e-6) - critic_optim = HybridAdam(critic.parameters(), lr=5e-6) - else: - actor_optim = Adam(actor.parameters(), lr=5e-6) - critic_optim = Adam(critic.parameters(), lr=5e-6) - - tokenizer = GPT2Tokenizer.from_pretrained('gpt2') - tokenizer.pad_token = tokenizer.eos_token - - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare( - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model) - - trainer = PPOTrainer(strategy, - actor, - critic, - reward_model, - initial_model, - actor_optim, - critic_optim, - max_epochs=args.max_epochs, - train_batch_size=args.train_batch_size, - experience_batch_size=args.experience_batch_size, - tokenizer=preprocess_batch, - max_length=512, - do_sample=True, - temperature=1.0, - top_k=50, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - callbacks=[performance_evaluator]) - - random_prompts = torch.randint(tokenizer.vocab_size, (1000, 400), device=torch.cuda.current_device()) - trainer.fit(random_prompts, - num_episodes=args.num_episodes, - max_timesteps=args.max_timesteps, - update_timesteps=args.update_timesteps) - - print_rank_0(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--model', default='s') - parser.add_argument('--strategy', - choices=[ - 'ddp', 'colossalai_gemini', 'colossalai_gemini_cpu', 'colossalai_zero2', - 'colossalai_zero2_cpu', 'colossalai_zero1', 'colossalai_zero1_cpu' - ], - default='ddp') - parser.add_argument('--num_episodes', type=int, default=3) - parser.add_argument('--max_timesteps', type=int, default=8) - parser.add_argument('--update_timesteps', type=int, default=8) - parser.add_argument('--max_epochs', type=int, default=3) - parser.add_argument('--train_batch_size', type=int, default=8) - parser.add_argument('--experience_batch_size', type=int, default=8) - args = parser.parse_args() - main(args) diff --git a/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.sh b/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.sh deleted file mode 100755 index d70f8872570a..000000000000 --- a/applications/ChatGPT/benchmarks/benchmark_gpt_dummy.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/usr/bin/env bash -# Usage: $0 -set -xu - -BASE=$(realpath $(dirname $0)) - - -PY_SCRIPT=${BASE}/benchmark_gpt_dummy.py -export OMP_NUM_THREADS=8 - -function tune_batch_size() { - # we found when experience batch size is equal to train batch size - # peak CUDA memory usage of making experience phase is less than or equal to that of training phase - # thus, experience batch size can be larger than or equal to train batch size - for bs in 1 2 4 8 16 32 64 128 256; do - torchrun --standalone --nproc_per_node $1 $PY_SCRIPT --model $2 --strategy $3 --experience_batch_size $bs --train_batch_size $bs || return 1 - done -} - -if [ $# -eq 0 ]; then - num_gpus=(1 2 4 8) -else - num_gpus=($1) -fi - -if [ $# -le 1 ]; then - strategies=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") -else - strategies=($2) -fi - -if [ $# -le 2 ]; then - models=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") -else - models=($3) -fi - - -for num_gpu in ${num_gpus[@]}; do - for strategy in ${strategies[@]}; do - for model in ${models[@]}; do - tune_batch_size $num_gpu $model $strategy || break - done - done -done diff --git a/applications/ChatGPT/chatgpt/dataset/__init__.py b/applications/ChatGPT/chatgpt/dataset/__init__.py deleted file mode 100644 index b4599c82ba75..000000000000 --- a/applications/ChatGPT/chatgpt/dataset/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .reward_dataset import RewardDataset -from .utils import is_rank_0 - -__all__ = ['RewardDataset', 'is_rank_0'] diff --git a/applications/ChatGPT/chatgpt/dataset/reward_dataset.py b/applications/ChatGPT/chatgpt/dataset/reward_dataset.py deleted file mode 100644 index 8bc850f2d52d..000000000000 --- a/applications/ChatGPT/chatgpt/dataset/reward_dataset.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import Callable - -from torch.utils.data import Dataset -from tqdm import tqdm - -from .utils import is_rank_0 - - -class RewardDataset(Dataset): - """ - Dataset for reward model - - Args: - dataset: dataset for reward model - tokenizer: tokenizer for reward model - max_length: max length of input - """ - - def __init__(self, dataset, tokenizer: Callable, max_length: int) -> None: - super().__init__() - self.chosen = [] - self.reject = [] - for data in tqdm(dataset, disable=not is_rank_0()): - prompt = data['prompt'] - - chosen = prompt + data['chosen'] + "<|endoftext|>" - chosen_token = tokenizer(chosen, - max_length=max_length, - padding="max_length", - truncation=True, - return_tensors="pt") - self.chosen.append({ - "input_ids": chosen_token['input_ids'], - "attention_mask": chosen_token['attention_mask'] - }) - - reject = prompt + data['rejected'] + "<|endoftext|>" - reject_token = tokenizer(reject, - max_length=max_length, - padding="max_length", - truncation=True, - return_tensors="pt") - self.reject.append({ - "input_ids": reject_token['input_ids'], - "attention_mask": reject_token['attention_mask'] - }) - - def __len__(self): - length = len(self.chosen) - return length - - def __getitem__(self, idx): - return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][ - "input_ids"], self.reject[idx]["attention_mask"] diff --git a/applications/ChatGPT/chatgpt/dataset/utils.py b/applications/ChatGPT/chatgpt/dataset/utils.py deleted file mode 100644 index 6c9f7f085f8c..000000000000 --- a/applications/ChatGPT/chatgpt/dataset/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -import torch.distributed as dist - - -def is_rank_0() -> bool: - return not dist.is_initialized() or dist.get_rank() == 0 diff --git a/applications/ChatGPT/chatgpt/models/__init__.py b/applications/ChatGPT/chatgpt/models/__init__.py deleted file mode 100644 index 376fed8de792..000000000000 --- a/applications/ChatGPT/chatgpt/models/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Actor, Critic, RewardModel -from .loss import PairWiseLoss, PolicyLoss, PPOPtxActorLoss, ValueLoss - -__all__ = ['Actor', 'Critic', 'RewardModel', 'PolicyLoss', 'ValueLoss', 'PPOPtxActorLoss', 'PairWiseLoss'] diff --git a/applications/ChatGPT/chatgpt/models/base/__init__.py b/applications/ChatGPT/chatgpt/models/base/__init__.py deleted file mode 100644 index 86f403556904..000000000000 --- a/applications/ChatGPT/chatgpt/models/base/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .actor import Actor -from .critic import Critic -from .reward_model import RewardModel - -__all__ = ['Actor', 'Critic', 'RewardModel'] diff --git a/applications/ChatGPT/chatgpt/models/generation_utils.py b/applications/ChatGPT/chatgpt/models/generation_utils.py deleted file mode 100644 index c7bc1b383fb9..000000000000 --- a/applications/ChatGPT/chatgpt/models/generation_utils.py +++ /dev/null @@ -1,92 +0,0 @@ -from typing import Optional - -import torch - - -def gpt_prepare_inputs_fn(input_ids: torch.Tensor, past: Optional[torch.Tensor] = None, **kwargs) -> dict: - token_type_ids = kwargs.get("token_type_ids", None) - # only last token for inputs_ids if past is defined in kwargs - if past: - input_ids = input_ids[:, -1].unsqueeze(-1) - if token_type_ids is not None: - token_type_ids = token_type_ids[:, -1].unsqueeze(-1) - - attention_mask = kwargs.get("attention_mask", None) - position_ids = kwargs.get("position_ids", None) - - if attention_mask is not None and position_ids is None: - # create position_ids on the fly for batch generation - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - if past: - position_ids = position_ids[:, -1].unsqueeze(-1) - else: - position_ids = None - return { - "input_ids": input_ids, - "past_key_values": past, - "use_cache": kwargs.get("use_cache"), - "position_ids": position_ids, - "attention_mask": attention_mask, - "token_type_ids": token_type_ids, - } - - -def update_model_kwargs_fn(outputs: dict, **model_kwargs) -> dict: - if "past_key_values" in outputs: - model_kwargs["past"] = outputs["past_key_values"] - else: - model_kwargs["past"] = None - - # update token_type_ids with last value - if "token_type_ids" in model_kwargs: - token_type_ids = model_kwargs["token_type_ids"] - model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) - - # update attention mask - if "attention_mask" in model_kwargs: - attention_mask = model_kwargs["attention_mask"] - model_kwargs["attention_mask"] = torch.cat( - [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) - - return model_kwargs - - -def opt_prepare_inputs_fn(input_ids: torch.Tensor, - past: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - **kwargs) -> dict: - # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly - if attention_mask is None: - attention_mask = input_ids.new_ones(input_ids.shape) - - if past: - input_ids = input_ids[:, -1:] - # first step, decoder_cached_states are empty - return { - "input_ids": input_ids, # encoder_outputs is defined. input_ids not needed - "attention_mask": attention_mask, - "past_key_values": past, - "use_cache": use_cache, - } - - -def bloom_prepare_inputs_fn(input_ids: torch.Tensor, - past: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - **kwargs) -> dict: - # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly - if attention_mask is None: - attention_mask = input_ids.new_ones(input_ids.shape) - - if past: - input_ids = input_ids[:, -1:] - # first step, decoder_cached_states are empty - return { - "input_ids": input_ids, # encoder_outputs is defined. input_ids not needed - "attention_mask": attention_mask, - "past_key_values": past, - "use_cache": use_cache, - } diff --git a/applications/ChatGPT/chatgpt/trainer/__init__.py b/applications/ChatGPT/chatgpt/trainer/__init__.py deleted file mode 100644 index c47c76347ee5..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .base import Trainer -from .ppo import PPOTrainer -from .rm import RewardModelTrainer - -__all__ = ['Trainer', 'PPOTrainer', 'RewardModelTrainer'] diff --git a/applications/ChatGPT/chatgpt/trainer/base.py b/applications/ChatGPT/chatgpt/trainer/base.py deleted file mode 100644 index a2419a35b6cd..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/base.py +++ /dev/null @@ -1,162 +0,0 @@ -from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Optional, Union - -import torch -from chatgpt.experience_maker import Experience, ExperienceMaker -from chatgpt.replay_buffer import ReplayBuffer -from torch import Tensor -from torch.utils.data import DistributedSampler -from tqdm import tqdm - -from .callbacks import Callback -from .strategies import Strategy -from .utils import is_rank_0 - - -class Trainer(ABC): - """ - Base class for rlhf trainers. - - Args: - strategy (Strategy):the strategy to use for training - experience_maker (ExperienceMaker): the experience maker to use for produce experience to fullfill replay buffer - replay_buffer (ReplayBuffer): the replay buffer to use for training - experience_batch_size (int, defaults to 8): the batch size to use for experience generation - max_epochs (int, defaults to 1): the number of epochs of training process - tokenizer (Callable, optional): the tokenizer to use for tokenizing the input - sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer - data_loader_pin_memory (bool, defaults to True): whether to pin memory for data loader - callbacks (List[Callback], defaults to []): the callbacks to call during training process - generate_kwargs (dict, optional): the kwargs to use while model generating - """ - - def __init__(self, - strategy: Strategy, - experience_maker: ExperienceMaker, - replay_buffer: ReplayBuffer, - experience_batch_size: int = 8, - max_epochs: int = 1, - tokenizer: Optional[Callable[[Any], dict]] = None, - sample_replay_buffer: bool = False, - dataloader_pin_memory: bool = True, - callbacks: List[Callback] = [], - **generate_kwargs) -> None: - super().__init__() - self.strategy = strategy - self.experience_maker = experience_maker - self.replay_buffer = replay_buffer - self.experience_batch_size = experience_batch_size - self.max_epochs = max_epochs - self.tokenizer = tokenizer - self.generate_kwargs = generate_kwargs - self.sample_replay_buffer = sample_replay_buffer - self.dataloader_pin_memory = dataloader_pin_memory - self.callbacks = callbacks - - @abstractmethod - def training_step(self, experience: Experience) -> Dict[str, Any]: - pass - - def _make_experience(self, inputs: Union[Tensor, Dict[str, Tensor]]) -> Experience: - if isinstance(inputs, Tensor): - return self.experience_maker.make_experience(inputs, **self.generate_kwargs) - elif isinstance(inputs, dict): - return self.experience_maker.make_experience(**inputs, **self.generate_kwargs) - else: - raise ValueError(f'Unsupported input type "{type(inputs)}"') - - def _sample_prompts(self, prompts) -> list: - indices = list(range(len(prompts))) - sampled_indices = self.strategy.experience_sampler.choice(indices, self.experience_batch_size, replace=False) - return [prompts[i] for i in sampled_indices] - - def _learn(self): - # replay buffer may be empty at first, we should rebuild at each training - if not self.sample_replay_buffer: - dataloader = self.strategy.setup_dataloader(self.replay_buffer, self.dataloader_pin_memory) - device = torch.cuda.current_device() - if self.sample_replay_buffer: - pbar = tqdm(range(self.max_epochs), desc='Train epoch', disable=not is_rank_0()) - for _ in pbar: - experience = self.replay_buffer.sample() - metrics = self.training_step(experience) - pbar.set_postfix(metrics) - else: - for epoch in range(self.max_epochs): - self._on_learn_epoch_start(epoch) - if isinstance(dataloader.sampler, DistributedSampler): - dataloader.sampler.set_epoch(epoch) - pbar = tqdm(dataloader, desc=f'Train epoch [{epoch+1}/{self.max_epochs}]', disable=not is_rank_0()) - for experience in pbar: - self._on_learn_batch_start() - experience.to_device(device) - metrics = self.training_step(experience) - self._on_learn_batch_end(metrics, experience) - pbar.set_postfix(metrics) - self._on_learn_epoch_end(epoch) - - def fit(self, prompts, num_episodes: int = 50000, max_timesteps: int = 500, update_timesteps: int = 5000) -> None: - time = 0 - sampler = self.strategy.setup_sampler(prompts) - self._on_fit_start() - for episode in range(num_episodes): - self._on_episode_start(episode) - for timestep in tqdm(range(max_timesteps), - desc=f'Episode [{episode+1}/{num_episodes}]', - disable=not is_rank_0()): - time += 1 - rand_prompts = sampler.sample(self.experience_batch_size) - if self.tokenizer is not None: - inputs = self.tokenizer(rand_prompts) - else: - inputs = rand_prompts - self._on_make_experience_start() - experience = self._make_experience(inputs) - self._on_make_experience_end(experience) - self.replay_buffer.append(experience) - if time % update_timesteps == 0: - self._learn() - self.replay_buffer.clear() - self._on_episode_end(episode) - self._on_fit_end() - - # TODO(ver217): maybe simplify these code using context - def _on_fit_start(self) -> None: - for callback in self.callbacks: - callback.on_fit_start() - - def _on_fit_end(self) -> None: - for callback in self.callbacks: - callback.on_fit_end() - - def _on_episode_start(self, episode: int) -> None: - for callback in self.callbacks: - callback.on_episode_start(episode) - - def _on_episode_end(self, episode: int) -> None: - for callback in self.callbacks: - callback.on_episode_end(episode) - - def _on_make_experience_start(self) -> None: - for callback in self.callbacks: - callback.on_make_experience_start() - - def _on_make_experience_end(self, experience: Experience) -> None: - for callback in self.callbacks: - callback.on_make_experience_end(experience) - - def _on_learn_epoch_start(self, epoch: int) -> None: - for callback in self.callbacks: - callback.on_learn_epoch_start(epoch) - - def _on_learn_epoch_end(self, epoch: int) -> None: - for callback in self.callbacks: - callback.on_learn_epoch_end(epoch) - - def _on_learn_batch_start(self) -> None: - for callback in self.callbacks: - callback.on_learn_batch_start() - - def _on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: - for callback in self.callbacks: - callback.on_learn_batch_end(metrics, experience) diff --git a/applications/ChatGPT/chatgpt/trainer/ppo.py b/applications/ChatGPT/chatgpt/trainer/ppo.py deleted file mode 100644 index 789e0c2f8f1e..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/ppo.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import Any, Callable, Dict, List, Optional - -import torch.nn as nn -from chatgpt.experience_maker import Experience, NaiveExperienceMaker -from chatgpt.models.base import Actor, Critic -from chatgpt.models.generation_utils import update_model_kwargs_fn -from chatgpt.models.loss import PolicyLoss, ValueLoss -from chatgpt.replay_buffer import NaiveReplayBuffer -from torch.optim import Optimizer - -from .base import Trainer -from .callbacks import Callback -from .strategies import Strategy - - -class PPOTrainer(Trainer): - """ - Trainer for PPO algorithm. - - Args: - strategy (Strategy): the strategy to use for training - actor (Actor): the actor model in ppo algorithm - critic (Critic): the critic model in ppo algorithm - reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences - initial_model (Actor): the initial model in rlhf algorithm to generate reference logits to limit the update of actor - actor_optim (Optimizer): the optimizer to use for actor model - critic_optim (Optimizer): the optimizer to use for critic model - kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss - train_batch_size (int, defaults to 8): the batch size to use for training - buffer_limit (int, defaults to 0): the max_size limitaiton of replay buffer - buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu - eps_clip (float, defaults to 0.2): the clip coefficient of policy loss - value_clip (float, defaults to 0.4): the clip coefficient of value loss - experience_batch_size (int, defaults to 8): the batch size to use for experience generation - max_epochs (int, defaults to 1): the number of epochs of training process - tokenier (Callable, optional): the tokenizer to use for tokenizing the input - sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer - dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader - callbacks (List[Callback], defaults to []): the callbacks to call during training process - generate_kwargs (dict, optional): the kwargs to use while model generating - """ - - def __init__(self, - strategy: Strategy, - actor: Actor, - critic: Critic, - reward_model: nn.Module, - initial_model: Actor, - actor_optim: Optimizer, - critic_optim: Optimizer, - kl_coef: float = 0.1, - train_batch_size: int = 8, - buffer_limit: int = 0, - buffer_cpu_offload: bool = True, - eps_clip: float = 0.2, - value_clip: float = 0.4, - experience_batch_size: int = 8, - max_epochs: int = 1, - tokenizer: Optional[Callable[[Any], dict]] = None, - sample_replay_buffer: bool = False, - dataloader_pin_memory: bool = True, - callbacks: List[Callback] = [], - **generate_kwargs) -> None: - experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef) - replay_buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload) - super().__init__(strategy, experience_maker, replay_buffer, experience_batch_size, max_epochs, tokenizer, - sample_replay_buffer, dataloader_pin_memory, callbacks, **generate_kwargs) - self.actor = actor - self.critic = critic - - self.actor_loss_fn = PolicyLoss(eps_clip) - self.critic_loss_fn = ValueLoss(value_clip) - - self.actor_optim = actor_optim - self.critic_optim = critic_optim - self._set_default_generate_kwargs(generate_kwargs, actor) - - def training_step(self, experience: Experience) -> Dict[str, float]: - self.actor.train() - self.critic.train() - - num_actions = experience.action_mask.size(1) - action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask) - actor_loss = self.actor_loss_fn(action_log_probs, - experience.action_log_probs, - experience.advantages, - action_mask=experience.action_mask) - self.strategy.backward(actor_loss, self.actor, self.actor_optim) - self.strategy.optimizer_step(self.actor_optim) - self.actor_optim.zero_grad() - - values = self.critic(experience.sequences, - action_mask=experience.action_mask, - attention_mask=experience.attention_mask) - critic_loss = self.critic_loss_fn(values, - experience.values, - experience.reward, - action_mask=experience.action_mask) - self.strategy.backward(critic_loss, self.critic, self.critic_optim) - self.strategy.optimizer_step(self.critic_optim) - self.critic_optim.zero_grad() - - return {'actor_loss': actor_loss.item(), 'critic_loss': critic_loss.item()} - - def _set_default_generate_kwargs(self, generate_kwargs: dict, actor: Actor) -> None: - origin_model = self.strategy._unwrap_actor(actor) - # use huggingface models method directly - if 'prepare_inputs_fn' not in generate_kwargs and hasattr(origin_model, 'prepare_inputs_for_generation'): - generate_kwargs['prepare_inputs_fn'] = origin_model.prepare_inputs_for_generation - - if 'update_model_kwargs_fn' not in generate_kwargs: - generate_kwargs['update_model_kwargs_fn'] = update_model_kwargs_fn diff --git a/applications/ChatGPT/chatgpt/trainer/rm.py b/applications/ChatGPT/chatgpt/trainer/rm.py deleted file mode 100644 index c07d65f84ca5..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/rm.py +++ /dev/null @@ -1,93 +0,0 @@ -from abc import ABC - -import loralib as lora -import torch -from chatgpt.dataset import RewardDataset -from chatgpt.models.loss import PairWiseLoss -from torch.optim import Adam, Optimizer -from torch.utils.data import DataLoader -from tqdm import tqdm - -from .strategies import Strategy -from .utils import is_rank_0 - - -class RewardModelTrainer(ABC): - """ - Trainer to use while training reward model. - - Args: - model (torch.nn.Module): the model to train - strategy (Strategy): the strategy to use for training - optim(Optimizer): the optimizer to use for training - train_dataset (RewardDataset): the dataset to use for training - eval_dataset (RewardDataset): the dataset to use for evaluation - batch_size (int, defaults to 1): the batch size while training - max_epochs (int, defaults to 2): the number of epochs to train - optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer - """ - - def __init__( - self, - model, - strategy: Strategy, - optim: Optimizer, - train_dataset: RewardDataset, - eval_dataset: RewardDataset, - batch_size: int = 1, - max_epochs: int = 2, - ) -> None: - super().__init__() - self.strategy = strategy - self.epochs = max_epochs - self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size) - self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size) - - self.model = strategy.setup_model(model) - if "DDP" in str(self.strategy): - self.model = self.model.module - self.loss_fn = PairWiseLoss() - self.optimizer = strategy.setup_optimizer(optim, self.model) - - def fit(self, use_lora): - epoch_bar = tqdm(range(self.epochs), desc='Train epoch', disable=not is_rank_0()) - for epoch in range(self.epochs): - step_bar = tqdm(range(self.train_dataloader.__len__()), - desc='Train step of epoch %d' % epoch, - disable=not is_rank_0()) - # train - self.model.train() - for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader: - chosen_ids = chosen_ids.squeeze(1).cuda() - c_mask = c_mask.squeeze(1).cuda() - reject_ids = reject_ids.squeeze(1).cuda() - r_mask = r_mask.squeeze(1).cuda() - chosen_reward = self.model(chosen_ids, attention_mask=c_mask) - reject_reward = self.model(reject_ids, attention_mask=r_mask) - loss = self.loss_fn(chosen_reward, reject_reward) - self.strategy.backward(loss, self.model, self.optimizer) - self.strategy.optimizer_step(self.optimizer) - self.optimizer.zero_grad() - step_bar.update() - step_bar.set_postfix({'loss': loss.item()}) - - # eval - self.model.eval() - with torch.no_grad(): - dist = 0 - loss_sum = 0 - for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader: - chosen_ids = chosen_ids.squeeze(1).cuda() - c_mask = c_mask.squeeze(1).cuda() - reject_ids = reject_ids.squeeze(1).cuda() - r_mask = r_mask.squeeze(1).cuda() - chosen_reward = self.model(chosen_ids, attention_mask=c_mask) - reject_reward = self.model(reject_ids, attention_mask=r_mask) - dist += (chosen_reward - reject_reward).mean().item() - loss = self.loss_fn(chosen_reward, reject_reward) - loss_sum += loss.item() - dist_mean = dist / self.eval_dataloader.__len__() - loss_mean = loss_sum / self.eval_dataloader.__len__() - epoch_bar.update() - step_bar.set_postfix({'loss': loss_mean, 'dist_mean': dist_mean}) - step_bar.close() diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/__init__.py b/applications/ChatGPT/chatgpt/trainer/strategies/__init__.py deleted file mode 100644 index f258c9b8a873..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/strategies/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .base import Strategy -from .colossalai import ColossalAIStrategy -from .ddp import DDPStrategy -from .naive import NaiveStrategy - -__all__ = ['Strategy', 'NaiveStrategy', 'DDPStrategy', 'ColossalAIStrategy'] diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/base.py b/applications/ChatGPT/chatgpt/trainer/strategies/base.py deleted file mode 100644 index 4347c08b4333..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/strategies/base.py +++ /dev/null @@ -1,131 +0,0 @@ -from abc import ABC, abstractmethod -from contextlib import nullcontext -from typing import Any, List, Tuple, Union - -import numpy as np -import torch -import torch.nn as nn -from chatgpt.models.base import Actor, Critic, RewardModel -from chatgpt.replay_buffer import ReplayBuffer -from torch.optim import Optimizer -from torch.utils.data import DataLoader - -from .sampler import DistributedSampler - -ModelOptimPair = Tuple[nn.Module, Optimizer] -ModelOrModelOptimPair = Union[nn.Module, ModelOptimPair] - - -class Strategy(ABC): - """ - Base class for training strategies. - """ - - def __init__(self) -> None: - super().__init__() - self.setup_distributed() - - @abstractmethod - def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: Optimizer, **kwargs) -> None: - pass - - @abstractmethod - def optimizer_step(self, optimizer: Optimizer, **kwargs) -> None: - pass - - @abstractmethod - def setup_distributed(self) -> None: - pass - - @abstractmethod - def setup_model(self, model: nn.Module) -> nn.Module: - pass - - @abstractmethod - def setup_optimizer(self, optimizer: Optimizer, model: nn.Module) -> Optimizer: - pass - - @abstractmethod - def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: - pass - - def model_init_context(self): - return nullcontext() - - def prepare( - self, *models_or_model_optim_pairs: ModelOrModelOptimPair - ) -> Union[List[ModelOrModelOptimPair], ModelOrModelOptimPair]: - """Prepare models or model-optimizer-pairs based on each strategy. - - Example:: - >>> # when fine-tuning actor and critic - >>> (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare((actor, actor_optim), (critic, critic_optim), reward_model, initial_model) - >>> # or when training reward model - >>> (reward_model, reward_model_optim) = strategy.prepare((reward_model, reward_model_optim)) - >>> # or just inference - >>> actor, critic = strategy.prepare(actor, critic) - - Returns: - Union[List[ModelOrModelOptimPair], ModelOrModelOptimPair]: Models or model-optimizer-pairs in the original order. - """ - - def prepare_model(model: nn.Module): - if isinstance(model, Actor): - return Actor(self.setup_model(self._unwrap_model(model))) - return self.setup_model(self._unwrap_model(model)) - - rets = [] - for arg in models_or_model_optim_pairs: - if isinstance(arg, tuple): - assert len(arg) == 2, f'Expect (model, optimizer) pair, got a tuple with size "{len(arg)}"' - model, optimizer = arg - model = prepare_model(model) - optimizer = self.setup_optimizer(optimizer, self._unwrap_model(model)) - rets.append((model, optimizer)) - elif isinstance(arg, nn.Module): - rets.append(prepare_model(arg)) - else: - raise RuntimeError(f'Expect model or (model, optimizer) pair, got {type(arg)}') - - if len(rets) == 1: - return rets[0] - return rets - - @staticmethod - def _unwrap_model(model: nn.Module) -> nn.Module: - """Useful for saving state dict. As actor is wrapped by Actor class again in `prepare()`, we should unwrap it before saving. - - Args: - model (nn.Module): an actor or a critic - """ - if isinstance(model, Actor): - return model.model - return model - - @staticmethod - def _unwrap_actor(actor: Actor) -> nn.Module: - """Get `actor.model` from a wrapped (by `prepare()`) actor. Useful for getting original huggingface model. - - Args: - actor (Actor): a wrapped actor - """ - return Strategy._unwrap_model(actor) - - @abstractmethod - def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None: - pass - - @abstractmethod - def load_model(self, model: nn.Module, path: str, map_location: Any = None, strict: bool = True) -> None: - pass - - @abstractmethod - def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None: - pass - - @abstractmethod - def load_optimizer(self, optimizer: Optimizer, path: str, map_location: Any = None) -> None: - pass - - def setup_sampler(self, dataset) -> DistributedSampler: - return DistributedSampler(dataset, 1, 0) diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py deleted file mode 100644 index b20b02d3d34d..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py +++ /dev/null @@ -1,171 +0,0 @@ -import warnings -from typing import Optional, Union - -import torch -import torch.distributed as dist -import torch.nn as nn -import torch.optim as optim -from chatgpt.models.base import Actor -from chatgpt.models.lora import LoraLinear -from torch.optim import Optimizer - -import colossalai -from colossalai.nn.optimizer import CPUAdam, HybridAdam -from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper, zero_optim_wrapper -from colossalai.nn.parallel.utils import get_static_torch_model -from colossalai.tensor import ProcessGroup, ShardSpec -from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext - -from .base import Strategy -from .ddp import DDPStrategy - - -class ColossalAIStrategy(DDPStrategy): - """ - The strategy for training with ColossalAI. - - Args: - stage(int): The stage to use in ZeRO. Choose in (1, 2, 3) - seed(int): The seed for the random number generator. - shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3. - This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future. - placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda') - If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU, - If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest. - pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3. - force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3. - search_range_mb(int): The search range in MB for the chunk size. Only for ZeRO-3. - hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3. - min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3. - gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3. - reduce_bugket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2. - overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2. - initial_scale(float): The initial scale for the optimizer. - growth_factor(float): The growth factor for the optimizer. - backoff_factor(float): The backoff factor for the optimizer. - growth_interval(int): The growth interval for the optimizer. - hysteresis(int): The hysteresis for the optimizer. - min_scale(float): The minimum scale for the optimizer. - max_scale(float): The maximum scale for the optimizer. - max_norm(float): The maximum norm for the optimizer. - norm_type(float): The norm type for the optimizer. - - """ - - def __init__( - self, - stage: int = 3, - seed: int = 42, - shard_init: bool = False, # only for stage 3 - placement_policy: str = 'cuda', - pin_memory: bool = True, # only for stage 3 - force_outputs_fp32: bool = False, # only for stage 3 - search_range_mb: int = 32, # only for stage 3 - hidden_dim: Optional[int] = None, # only for stage 3 - min_chunk_size_mb: float = 32, # only for stage 3 - gpu_margin_mem_ratio: float = 0.0, # only for stage 3 - reduce_bucket_size: int = 12 * 1024**2, # only for stage 1&2 - overlap_communication: bool = True, # only for stage 1&2 - initial_scale: float = 2**16, - growth_factor: float = 2, - backoff_factor: float = 0.5, - growth_interval: int = 1000, - hysteresis: int = 2, - min_scale: float = 1, - max_scale: float = 2**32, - max_norm: float = 0.0, - norm_type: float = 2.0) -> None: - super().__init__(seed) - assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"' - self.stage = stage - # TODO(ver217): support shard_init when using from_pretrained() - if shard_init: - warnings.warn( - f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()' - ) - self.shard_init = shard_init - self.gemini_config = dict(device=get_current_device(), - placement_policy=placement_policy, - pin_memory=pin_memory, - force_outputs_fp32=force_outputs_fp32, - strict_ddp_mode=shard_init, - search_range_mb=search_range_mb, - hidden_dim=hidden_dim, - min_chunk_size_mb=min_chunk_size_mb) - if stage == 3: - self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio) - else: - self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size, - overlap_communication=overlap_communication, - cpu_offload=(placement_policy == 'cpu')) - self.optim_kwargs = dict(initial_scale=initial_scale, - growth_factor=growth_factor, - backoff_factor=backoff_factor, - growth_interval=growth_interval, - hysteresis=hysteresis, - min_scale=min_scale, - max_scale=max_scale, - max_norm=max_norm, - norm_type=norm_type) - - def setup_distributed(self) -> None: - colossalai.launch_from_torch({}, seed=self.seed) - - def model_init_context(self): - if self.stage == 3: - world_size = dist.get_world_size() - shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None - default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None - return ColoInitContext(device=get_current_device(), - dtype=torch.half, - default_pg=shard_pg, - default_dist_spec=default_dist_spec) - return super().model_init_context() - - def setup_model(self, model: nn.Module) -> nn.Module: - return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config) - - def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer: - assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}' - return zero_optim_wrapper(model, optimizer, optim_config=self.zero_optim_config, **self.optim_kwargs) - - def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None: - optimizer.backward(loss) - - def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None: - optimizer.step() - - @staticmethod - def _unwrap_actor(actor: Actor) -> nn.Module: - model: Union[nn.Module, ZeroDDP] = Strategy._unwrap_actor(actor) - if isinstance(model, ZeroDDP): - return model.module - return model - - def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None: - unwrapped_model = self._unwrap_model(model) - # TODO : better way to get torch model from gemini model - # to get torch model from gemini model - if isinstance(unwrapped_model, ZeroDDP): - state_dict = unwrapped_model.state_dict() - unwrapped_model = get_static_torch_model(unwrapped_model) - if only_rank0 and dist.get_rank() != 0: - return - unwrapped_model.load_state_dict(state_dict) - # merge lora_weights into weights - for module in unwrapped_model.modules(): - if isinstance(module, LoraLinear): - module.merge_weights=True - module.eval() - # get state_dict and save - state_dict = unwrapped_model.state_dict() - if only_rank0 and dist.get_rank() != 0: - return - torch.save(state_dict, path) - - def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None: - if only_rank0: - raise RuntimeError( - f'Optimizer states are sharded when using ColossalAIStrategy. Only rank0 is not supported.') - torch.save(optimizer.state_dict(), path) diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/ddp.py b/applications/ChatGPT/chatgpt/trainer/strategies/ddp.py deleted file mode 100644 index c9f92c12fe0a..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/strategies/ddp.py +++ /dev/null @@ -1,93 +0,0 @@ -import os -import random - -import numpy as np -import torch -import torch.distributed as dist -import torch.nn as nn -from chatgpt.models.base import Actor -from chatgpt.models.lora import LoraLinear -from chatgpt.replay_buffer import ReplayBuffer -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.optim import Optimizer -from torch.utils.data import DataLoader - -from .base import Strategy -from .naive import NaiveStrategy -from .sampler import DistributedSampler - - -class DDPStrategy(NaiveStrategy): - """ - Strategy for distributed training using torch.distributed. - """ - - def __init__(self, seed: int = 42) -> None: - self.seed = seed - super().__init__() - - def setup_distributed(self) -> None: - try: - rank = int(os.environ['RANK']) - local_rank = int(os.environ['LOCAL_RANK']) - world_size = int(os.environ['WORLD_SIZE']) - host = os.environ['MASTER_ADDR'] - port = int(os.environ['MASTER_PORT']) - except KeyError as e: - raise RuntimeError( - f"Could not find {e} in the torch environment, visit https://www.colossalai.org/ for more information on launching with torch" - ) - dist.init_process_group('nccl', init_method=f'tcp://[{host}]:{port}', world_size=world_size, rank=rank) - self.set_seed(self.seed) - torch.cuda.set_device(local_rank) - - def set_seed(self, seed: int) -> None: - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - - def setup_model(self, model: nn.Module) -> nn.Module: - device = torch.cuda.current_device() - return DDP(model, device_ids=[device]) - - def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: - # DDP only mode, replay buffers on each rank are different. - # sampler = DistributedSampler(replay_buffer, - # num_replicas=dist.get_world_size(), - # rank=dist.get_rank(), - # shuffle=True, - # seed=self.seed, - # drop_last=True) - return DataLoader( - replay_buffer, - batch_size=replay_buffer.sample_batch_size, - # sampler=sampler, - shuffle=True, - drop_last=True, - pin_memory=pin_memory, - collate_fn=replay_buffer.collate_fn) - - @staticmethod - def _unwrap_actor(actor: Actor) -> nn.Module: - model: DDP = Strategy._unwrap_actor(actor) - return model.module - - def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None: - for module in model.modules(): - if isinstance(module, LoraLinear): - module.merge_weights=True - module.eval() - - if only_rank0 and dist.get_rank() != 0: - return - model = model.model.module - state_dict = model.state_dict() - torch.save(state_dict, path) - - def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None: - if only_rank0 and dist.get_rank() != 0: - return - super().save_optimizer(optimizer, path, only_rank0) - - def setup_sampler(self, dataset) -> DistributedSampler: - return DistributedSampler(dataset, dist.get_world_size(), dist.get_rank()) diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/naive.py b/applications/ChatGPT/chatgpt/trainer/strategies/naive.py deleted file mode 100644 index 99b8d6635394..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/strategies/naive.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Any - -import torch -import torch.nn as nn -import torch.optim as optim -from chatgpt.replay_buffer import ReplayBuffer -from torch.optim import Optimizer -from torch.utils.data import DataLoader - -from .base import Strategy - - -class NaiveStrategy(Strategy): - """ - Strategy for single GPU. No parallelism is used. - """ - - def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None: - loss.backward() - - def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None: - optimizer.step() - - def setup_distributed(self) -> None: - pass - - def setup_model(self, model: nn.Module) -> nn.Module: - return model - - def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer: - return optimizer - - def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: - return DataLoader(replay_buffer, - batch_size=replay_buffer.sample_batch_size, - shuffle=True, - drop_last=True, - pin_memory=pin_memory, - collate_fn=replay_buffer.collate_fn) - - def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None: - unwrapped_model = self._unwrap_model(model) - torch.save(unwrapped_model.state_dict(), path) - - def load_model(self, model: nn.Module, path: str, map_location: Any = None, strict: bool = True) -> None: - unwrapped_model = self._unwrap_model(model) - state_dict = torch.load(path, map_location=map_location) - unwrapped_model.load_state_dict(state_dict, strict=strict) - - def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None: - torch.save(optimizer.state_dict(), path) - - def load_optimizer(self, optimizer: Optimizer, path: str, map_location: Any = None) -> None: - state_dict = torch.load(path, map_location=map_location) - optimizer.load_state_dict(state_dict) diff --git a/applications/ChatGPT/chatgpt/trainer/utils.py b/applications/ChatGPT/chatgpt/trainer/utils.py deleted file mode 100644 index 6c9f7f085f8c..000000000000 --- a/applications/ChatGPT/chatgpt/trainer/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -import torch.distributed as dist - - -def is_rank_0() -> bool: - return not dist.is_initialized() or dist.get_rank() == 0 diff --git a/applications/ChatGPT/examples/README.md b/applications/ChatGPT/examples/README.md deleted file mode 100644 index 3876d20f02d7..000000000000 --- a/applications/ChatGPT/examples/README.md +++ /dev/null @@ -1,122 +0,0 @@ -# Examples - -## Install requirements - -```shell -pip install -r requirements.txt -``` - -## Train the reward model (Stage 2) -We use [rm-static](https://huggingface.co/datasets/Dahoas/rm-static) as dataset to train our reward model. It is a dataset of chosen & rejected response of the same prompt. - -You can download the dataset from huggingface automatically. - -Use these code to train your reward model. - -```shell -# Naive reward model training -python train_reward_model.py --pretrain --model --strategy naive -# use colossalai_zero2 -torchrun --standalone --nproc_per_node=2 train_reward_model.py --pretrain --model --strategy colossalai_zero2 -``` - -## Train with dummy prompt data (Stage 3) - -This script supports 3 strategies: - -- naive -- ddp -- colossalai - -It uses random generated prompt data. - -Naive strategy only support single GPU training: - -```shell -python train_dummy.py --strategy naive -# display cli help -python train_dummy.py -h -``` - -DDP strategy and ColossalAI strategy support multi GPUs training: - -```shell -# run DDP on 2 GPUs -torchrun --standalone --nproc_per_node=2 train_dummy.py --strategy ddp -# run ColossalAI on 2 GPUs -torchrun --standalone --nproc_per_node=2 train_dummy.py --strategy colossalai_zero2 -``` - -## Train with real prompt data (Stage 3) - -We use [awesome-chatgpt-prompts](https://huggingface.co/datasets/fka/awesome-chatgpt-prompts) as example dataset. It is a small dataset with hundreds of prompts. - -You should download `prompts.csv` first. - -This script also supports 3 strategies. - -```shell -# display cli help -python train_dummy.py -h -# run naive on 1 GPU -python train_prompts.py prompts.csv --strategy naive -# run DDP on 2 GPUs -torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy ddp -# run ColossalAI on 2 GPUs -torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy colossalai_zero2 -``` - -## Inference example(After Stage3) -We support naive inference demo after training. -```shell -# inference, using pretrain path to configure model -python inference.py --model_path --model --pretrain -# example -python inference.py --model_path ./actor_checkpoint_prompts.pt --pretrain bigscience/bloom-560m --model bloom -``` - - -#### data -- [x] [rm-static](https://huggingface.co/datasets/Dahoas/rm-static) -- [x] [hh-rlhf](https://huggingface.co/datasets/Anthropic/hh-rlhf) -- [ ] [openai/summarize_from_feedback](https://huggingface.co/datasets/openai/summarize_from_feedback) -- [ ] [openai/webgpt_comparisons](https://huggingface.co/datasets/openai/webgpt_comparisons) -- [ ] [Dahoas/instruct-synthetic-prompt-responses](https://huggingface.co/datasets/Dahoas/instruct-synthetic-prompt-responses) - -## Support Model - -### GPT -- [x] GPT2-S (s) -- [x] GPT2-M (m) -- [x] GPT2-L (l) -- [ ] GPT2-XL (xl) -- [x] GPT2-4B (4b) -- [ ] GPT2-6B (6b) -- [ ] GPT2-8B (8b) -- [ ] GPT2-10B (10b) -- [ ] GPT2-12B (12b) -- [ ] GPT2-15B (15b) -- [ ] GPT2-18B (18b) -- [ ] GPT2-20B (20b) -- [ ] GPT2-24B (24b) -- [ ] GPT2-28B (28b) -- [ ] GPT2-32B (32b) -- [ ] GPT2-36B (36b) -- [ ] GPT2-40B (40b) -- [ ] GPT3 (175b) - -### BLOOM -- [x] [BLOOM-560m](https://huggingface.co/bigscience/bloom-560m) -- [x] [BLOOM-1b1](https://huggingface.co/bigscience/bloom-1b1) -- [x] [BLOOM-3b](https://huggingface.co/bigscience/bloom-3b) -- [x] [BLOOM-7b](https://huggingface.co/bigscience/bloom-7b1) -- [ ] BLOOM-175b - -### OPT -- [x] [OPT-125M](https://huggingface.co/facebook/opt-125m) -- [x] [OPT-350M](https://huggingface.co/facebook/opt-350m) -- [ ] [OPT-1.3B](https://huggingface.co/facebook/opt-1.3b) -- [ ] [OPT-2.7B](https://huggingface.co/facebook/opt-2.7b) -- [ ] [OPT-6.7B](https://huggingface.co/facebook/opt-6.7b) -- [ ] [OPT-13B](https://huggingface.co/facebook/opt-13b) -- [ ] [OPT-30B](https://huggingface.co/facebook/opt-30b) diff --git a/applications/ChatGPT/examples/test_ci.sh b/applications/ChatGPT/examples/test_ci.sh deleted file mode 100755 index 0aa4a36fe514..000000000000 --- a/applications/ChatGPT/examples/test_ci.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env bash - -set -xue - -if [ -z "$PROMPT_PATH" ]; then - echo "Please set \$PROMPT_PATH to the path to prompts csv." - exit 1 -fi - -BASE=$(realpath $(dirname $0)) - -export OMP_NUM_THREADS=8 - -# install requirements -pip install -r ${BASE}/requirements.txt - -# train dummy -python ${BASE}/train_dummy.py --strategy naive --num_episodes 1 \ - --max_timesteps 2 --update_timesteps 2 \ - --max_epochs 1 --train_batch_size 2 --lora_rank 4 - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \ - --strategy colossalai_gemini --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'facebook/opt-350m' --model opt --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_dummy.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'facebook/opt-350m' --model opt - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \ - --strategy ddp --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'facebook/opt-350m' --model opt --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_dummy.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'facebook/opt-350m' --model opt - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py \ - --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'gpt2' --model gpt2 --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_dummy.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_dummy.pt --pretrain 'gpt2' --model gpt2 - -rm -rf ${BASE}/actor_checkpoint_dummy.pt - -# train prompts -python ${BASE}/train_prompts.py $PROMPT_PATH --strategy naive --num_episodes 1 \ - --max_timesteps 2 --update_timesteps 2 \ - --max_epochs 1 --train_batch_size 2 --lora_rank 4 - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \ - --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'facebook/opt-350m' --model opt --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_prompts.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'facebook/opt-350m' --model opt - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \ - --strategy ddp --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'gpt2' --model gpt2 --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_prompts.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'gpt2' --model gpt2 - -torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH \ - --strategy colossalai_gemini --num_episodes 1 --max_timesteps 2 \ - --update_timesteps 2 --max_epochs 1 --train_batch_size 2\ - --pretrain 'gpt2' --model gpt2 --lora_rank 4\ - --save_path ${BASE}/actor_checkpoint_prompts.pt -python ${BASE}/inference.py --model_path ${BASE}/actor_checkpoint_prompts.pt --pretrain 'gpt2' --model gpt2 - -rm -rf ${BASE}/actor_checkpoint_prompts.pt diff --git a/applications/ChatGPT/examples/train_dummy.py b/applications/ChatGPT/examples/train_dummy.py deleted file mode 100644 index c0ebf8f9b7b6..000000000000 --- a/applications/ChatGPT/examples/train_dummy.py +++ /dev/null @@ -1,148 +0,0 @@ -import argparse -from copy import deepcopy - -import torch -from chatgpt.models.base import RewardModel -from chatgpt.models.bloom import BLOOMActor, BLOOMCritic -from chatgpt.models.gpt import GPTActor, GPTCritic -from chatgpt.models.opt import OPTActor, OPTCritic -from chatgpt.trainer import PPOTrainer -from chatgpt.trainer.callbacks import SaveCheckpoint -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy -from torch.optim import Adam -from transformers import AutoTokenizer, BloomTokenizerFast -from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer - -from colossalai.nn.optimizer import HybridAdam - - -def preprocess_batch(samples): - input_ids = torch.stack(samples) - attention_mask = torch.ones_like(input_ids, dtype=torch.long) - return {'input_ids': input_ids, 'attention_mask': attention_mask} - - -def main(args): - # configure strategy - if args.strategy == 'naive': - strategy = NaiveStrategy() - elif args.strategy == 'ddp': - strategy = DDPStrategy() - elif args.strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) - elif args.strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') - else: - raise ValueError(f'Unsupported strategy "{args.strategy}"') - - # configure model - with strategy.model_init_context(): - if args.model == 'gpt2': - actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = GPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - elif args.model == 'bloom': - actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - elif args.model == 'opt': - actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = OPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - else: - raise ValueError(f'Unsupported model "{args.model}"') - - initial_model = deepcopy(actor).to(torch.cuda.current_device()) - reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device()) - - # configure optimizer - if args.strategy.startswith('colossalai'): - actor_optim = HybridAdam(actor.parameters(), lr=5e-6) - critic_optim = HybridAdam(critic.parameters(), lr=5e-6) - else: - actor_optim = Adam(actor.parameters(), lr=5e-6) - critic_optim = Adam(critic.parameters(), lr=5e-6) - - # configure tokenizer - if args.model == 'gpt2': - tokenizer = GPT2Tokenizer.from_pretrained('gpt2') - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'bloom': - tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'opt': - tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") - else: - raise ValueError(f'Unsupported model "{args.model}"') - - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare( - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model) - - callbacks = [] - if args.save_ckpt_path: - ckpt_callback = SaveCheckpoint( - args.save_ckpt_path, - args.save_ckpt_interval, - strategy, - actor, - critic, - actor_optim, - critic_optim, - ) - callbacks.append(ckpt_callback) - - # configure trainer - - trainer = PPOTrainer(strategy, - actor, - critic, - reward_model, - initial_model, - actor_optim, - critic_optim, - max_epochs=args.max_epochs, - train_batch_size=args.train_batch_size, - tokenizer=preprocess_batch, - max_length=128, - do_sample=True, - temperature=1.0, - top_k=50, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - callbacks=callbacks) - - random_prompts = torch.randint(tokenizer.vocab_size, (1000, 64), device=torch.cuda.current_device()) - trainer.fit(random_prompts, - num_episodes=args.num_episodes, - max_timesteps=args.max_timesteps, - update_timesteps=args.update_timesteps) - - # save model checkpoint after fitting - strategy.save_model(actor, args.save_path, only_rank0=True) - # save optimizer checkpoint on all ranks - if args.need_optim_ckpt: - strategy.save_optimizer(actor_optim, - 'actor_optim_checkpoint_dummy_%d.pt' % (torch.cuda.current_device()), - only_rank0=False) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--strategy', - choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'], - default='naive') - parser.add_argument('--model', type=str, default='gpt2', choices=['gpt2', 'bloom', 'opt']) - parser.add_argument('--pretrain', type=str, default=None) - parser.add_argument('--save_path', type=str, default='actor_checkpoint_dummy.pt') - parser.add_argument('--need_optim_ckpt', type=bool, default=False) - parser.add_argument('--num_episodes', type=int, default=50) - parser.add_argument('--max_timesteps', type=int, default=10) - parser.add_argument('--update_timesteps', type=int, default=10) - parser.add_argument('--max_epochs', type=int, default=5) - parser.add_argument('--train_batch_size', type=int, default=8) - parser.add_argument('--experience_batch_size', type=int, default=8) - parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") - parser.add_argument('--save_ckpt_path', - type=str, - default=None, - help="path to save checkpoint, None means not to save") - parser.add_argument('--save_ckpt_interval', type=int, default=1, help="the interval of episode to save checkpoint") - args = parser.parse_args() - main(args) diff --git a/applications/ChatGPT/examples/train_prompts.py b/applications/ChatGPT/examples/train_prompts.py deleted file mode 100644 index d4f31e61eb75..000000000000 --- a/applications/ChatGPT/examples/train_prompts.py +++ /dev/null @@ -1,131 +0,0 @@ -import argparse -from copy import deepcopy - -import pandas as pd -import torch -from chatgpt.models.base import RewardModel -from chatgpt.models.bloom import BLOOMActor, BLOOMCritic -from chatgpt.models.gpt import GPTActor, GPTCritic -from chatgpt.models.opt import OPTActor, OPTCritic -from chatgpt.trainer import PPOTrainer -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy -from torch.optim import Adam -from transformers import AutoTokenizer, BloomTokenizerFast -from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer - -from colossalai.nn.optimizer import HybridAdam - - -def main(args): - # configure strategy - if args.strategy == 'naive': - strategy = NaiveStrategy() - elif args.strategy == 'ddp': - strategy = DDPStrategy() - elif args.strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) - elif args.strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') - else: - raise ValueError(f'Unsupported strategy "{args.strategy}"') - - # configure model - with strategy.model_init_context(): - if args.model == 'gpt2': - actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = GPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - elif args.model == 'bloom': - actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - elif args.model == 'opt': - actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - critic = OPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device()) - else: - raise ValueError(f'Unsupported model "{args.model}"') - - initial_model = deepcopy(actor) - reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device()) - - - # configure optimizer - if args.strategy.startswith('colossalai'): - actor_optim = HybridAdam(actor.parameters(), lr=5e-6) - critic_optim = HybridAdam(critic.parameters(), lr=5e-6) - else: - actor_optim = Adam(actor.parameters(), lr=5e-6) - critic_optim = Adam(critic.parameters(), lr=5e-6) - - # configure tokenizer - if args.model == 'gpt2': - tokenizer = GPT2Tokenizer.from_pretrained('gpt2') - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'bloom': - tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'opt': - tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") - else: - raise ValueError(f'Unsupported model "{args.model}"') - - dataset = pd.read_csv(args.prompt_path)['prompt'] - - def tokenize_fn(texts): - batch = tokenizer(texts, return_tensors='pt', max_length=96, padding=True, truncation=True) - return {k: v.cuda() for k, v in batch.items()} - - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare( - (actor, actor_optim), (critic, critic_optim), reward_model, initial_model) - - # configure trainer - trainer = PPOTrainer( - strategy, - actor, - critic, - reward_model, - initial_model, - actor_optim, - critic_optim, - max_epochs=args.max_epochs, - train_batch_size=args.train_batch_size, - experience_batch_size=args.experience_batch_size, - tokenizer=tokenize_fn, - max_length=128, - do_sample=True, - temperature=1.0, - top_k=50, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - ) - - trainer.fit(dataset, - num_episodes=args.num_episodes, - max_timesteps=args.max_timesteps, - update_timesteps=args.update_timesteps) - # save model checkpoint after fitting - strategy.save_model(actor, args.save_path, only_rank0=True) - # save optimizer checkpoint on all ranks - if args.need_optim_ckpt: - strategy.save_optimizer(actor_optim, - 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()), - only_rank0=False) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('prompt_path') - parser.add_argument('--strategy', - choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'], - default='naive') - parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt']) - parser.add_argument('--pretrain', type=str, default=None) - parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts.pt') - parser.add_argument('--need_optim_ckpt', type=bool, default=False) - parser.add_argument('--num_episodes', type=int, default=10) - parser.add_argument('--max_timesteps', type=int, default=10) - parser.add_argument('--update_timesteps', type=int, default=10) - parser.add_argument('--max_epochs', type=int, default=5) - parser.add_argument('--train_batch_size', type=int, default=8) - parser.add_argument('--experience_batch_size', type=int, default=8) - parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") - args = parser.parse_args() - main(args) diff --git a/applications/ChatGPT/examples/train_prompts.sh b/applications/ChatGPT/examples/train_prompts.sh deleted file mode 100755 index db73ac8e8e85..000000000000 --- a/applications/ChatGPT/examples/train_prompts.sh +++ /dev/null @@ -1,18 +0,0 @@ -set_n_least_used_CUDA_VISIBLE_DEVICES() { - local n=${1:-"9999"} - echo "GPU Memory Usage:" - local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \ - | tail -n +2 \ - | nl -v 0 \ - | tee /dev/tty \ - | sort -g -k 2 \ - | awk '{print $1}' \ - | head -n $n) - export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') - echo "Now CUDA_VISIBLE_DEVICES is set to:" - echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" -} - -set_n_least_used_CUDA_VISIBLE_DEVICES 2 - -torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy colossalai_zero2 diff --git a/applications/ChatGPT/examples/train_reward_model.py b/applications/ChatGPT/examples/train_reward_model.py deleted file mode 100644 index 19b20b0847cc..000000000000 --- a/applications/ChatGPT/examples/train_reward_model.py +++ /dev/null @@ -1,101 +0,0 @@ -import argparse - -import loralib as lora -import torch -from chatgpt.dataset import RewardDataset -from chatgpt.models.base import RewardModel -from chatgpt.models.bloom import BLOOMRM -from chatgpt.models.gpt import GPTRM -from chatgpt.models.opt import OPTRM -from chatgpt.trainer import RewardModelTrainer -from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy -from datasets import load_dataset -from torch.optim import Adam -from transformers import AutoTokenizer, BloomTokenizerFast -from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer - -from colossalai.nn.optimizer import HybridAdam - - -def train(args): - # configure strategy - if args.strategy == 'naive': - strategy = NaiveStrategy() - elif args.strategy == 'ddp': - strategy = DDPStrategy() - elif args.strategy == 'colossalai_gemini': - strategy = ColossalAIStrategy(stage=3, placement_policy='cuda') - elif args.strategy == 'colossalai_zero2': - strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') - else: - raise ValueError(f'Unsupported strategy "{args.strategy}"') - - # configure model - with strategy.model_init_context(): - if args.model == 'bloom': - model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() - elif args.model == 'opt': - model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() - elif args.model == 'gpt2': - model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() - else: - raise ValueError(f'Unsupported model "{args.model}"') - - # configure tokenizer - if args.model == 'gpt2': - tokenizer = GPT2Tokenizer.from_pretrained('gpt2') - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'bloom': - tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) - tokenizer.pad_token = tokenizer.eos_token - elif args.model == 'opt': - tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") - else: - raise ValueError(f'Unsupported model "{args.model}"') - tokenizer.pad_token = tokenizer.eos_token - - max_len = 512 - - # configure optimizer - if args.strategy.startswith('colossalai'): - optim = HybridAdam(model.parameters(), lr=5e-5) - else: - optim = Adam(model.parameters(), lr=5e-5) - - # prepare for data and dataset - data = load_dataset(args.dataset) - train_data = data["train"] - eval_data = data['test'] - train_dataset = RewardDataset(train_data, tokenizer, max_len) - eval_dataset = RewardDataset(eval_data, tokenizer, max_len) - - trainer = RewardModelTrainer(model=model, - strategy=strategy, - optim=optim, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - batch_size=args.batch_size, - max_epochs=args.max_epochs) - - trainer.fit(use_lora=args.lora_rank) - - # save model checkpoint after fitting on only rank0 - strategy.save_model(model, 'rm_checkpoint.pt', only_rank0=True) - # save optimizer checkpoint on all ranks - strategy.save_optimizer(optim, 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), only_rank0=False) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--strategy', - choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'], - default='naive') - parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt'], default='bloom') - parser.add_argument('--pretrain', type=str, default=None) - parser.add_argument('--dataset', type=str, default='Dahoas/rm-static') - parser.add_argument('--save_path', type=str, default='rm_ckpt.pth') - parser.add_argument('--max_epochs', type=int, default=1) - parser.add_argument('--batch_size', type=int, default=4) - parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") - args = parser.parse_args() - train(args) diff --git a/applications/ChatGPT/examples/train_rm.sh b/applications/ChatGPT/examples/train_rm.sh deleted file mode 100755 index 6e11a148bfbe..000000000000 --- a/applications/ChatGPT/examples/train_rm.sh +++ /dev/null @@ -1,20 +0,0 @@ -set_n_least_used_CUDA_VISIBLE_DEVICES() { - local n=${1:-"9999"} - echo "GPU Memory Usage:" - local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \ - | tail -n +2 \ - | nl -v 0 \ - | tee /dev/tty \ - | sort -g -k 2 \ - | awk '{print $1}' \ - | head -n $n) - export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') - echo "Now CUDA_VISIBLE_DEVICES is set to:" - echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" -} - -set_n_least_used_CUDA_VISIBLE_DEVICES 2 - -# torchrun --standalone --nproc_per_node=2 train_reward_model.py --pretrain 'bigscience/bloomz-560m' --model 'bloom' --strategy colossalai_zero2 -torchrun --standalone --nproc_per_node=2 train_reward_model.py --model 'gpt2' --strategy colossalai_zero2 -# torchrun --standalone --nproc_per_node=2 train_reward_model.py --pretrain "facebook/opt-350m" --model 'opt' --strategy colossalai_zero2 diff --git a/applications/ChatGPT/requirements.txt b/applications/ChatGPT/requirements.txt deleted file mode 100644 index 15a960c2c650..000000000000 --- a/applications/ChatGPT/requirements.txt +++ /dev/null @@ -1,7 +0,0 @@ -transformers>=4.20.1 -tqdm -datasets -loralib -colossalai>=0.2.4 -torch -langchain diff --git a/applications/ChatGPT/version.txt b/applications/ChatGPT/version.txt deleted file mode 100644 index 6e8bf73aa550..000000000000 --- a/applications/ChatGPT/version.txt +++ /dev/null @@ -1 +0,0 @@ -0.1.0 diff --git a/applications/README.md b/applications/README.md new file mode 100644 index 000000000000..cd0435aae199 --- /dev/null +++ b/applications/README.md @@ -0,0 +1,12 @@ +# Applications + +This directory contains the applications that are powered by Colossal-AI. + +The list of applications include: + +- [X] [Chatbot](./Chat/README.md) +- [X] [FastFold](https://github.com/hpcaitech/FastFold): Optimizing AlphaFold (Biomedicine) Training and Inference on GPU Clusters + +> Please note that the `Chatbot` application is migrated from the original `ChatGPT` folder. + +You can find more example code for base models and functions in the [Examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples) directory. diff --git a/colossalai/_analyzer/__init__.py b/colossalai/_analyzer/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/colossalai/_analyzer/_subclasses/_meta_registration.py b/colossalai/_analyzer/_subclasses/_meta_registration.py index 20ab46054c8e..4049be79c70f 100644 --- a/colossalai/_analyzer/_subclasses/_meta_registration.py +++ b/colossalai/_analyzer/_subclasses/_meta_registration.py @@ -6,11 +6,15 @@ from typing import Callable, List, Optional, Tuple, Union import torch +from packaging import version from torch.utils._pytree import tree_map aten = torch.ops.aten -meta_lib = torch.library.Library("aten", "IMPL", "Meta") +try: + meta_lib = torch.library.Library("aten", "IMPL", "Meta") +except AttributeError: + meta_lib = None meta_table = {} @@ -50,432 +54,415 @@ def add_func(op): return wrapper -# ============================== Convolutions ====================================== -# https://github.com/pytorch/pytorch/pull/79834 -@register_meta(aten.convolution.default) -def meta_conv( - input_tensor: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - stride: List[int], - padding: List[int], - dilation: List[int], - is_transposed: bool, - output_padding: List[int], - groups: int, -): - - def _formula(ln: int, p: int, d: int, k: int, s: int) -> int: - """ - Formula to apply to calculate the length of some dimension of the output - See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html - Args: - ln: length of the dimension - p: padding in that dim - d: dilation in that dim - k: kernel size in that dim - s: stride in that dim - Returns: - The output length - """ - return (ln + 2 * p - d * (k - 1) - 1) // s + 1 - - def _formula_transposed(ln: int, p: int, d: int, k: int, s: int, op: int) -> int: - """ - Formula to apply to calculate the length of some dimension of the output - if transposed convolution is used. - See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html - Args: - ln: length of the dimension - p: padding in that dim - d: dilation in that dim - k: kernel size in that dim - s: stride in that dim - op: output padding in that dim - Returns: - The output length - """ - return (ln - 1) * s - 2 * p + d * (k - 1) + op + 1 - - def calc_conv_nd_return_shape( - dims: torch.Size, - kernel_size: torch.Size, - stride: Union[List[int], int], - padding: Union[List[int], int], - dilation: Union[List[int], int], - output_padding: Optional[Union[List[int], int]] = None, +if version.parse(torch.__version__) >= version.parse('1.12.0'): + # ============================== Convolutions ====================================== + # https://github.com/pytorch/pytorch/pull/79834 + @register_meta(aten.convolution.default) + def meta_conv( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: List[int], + padding: List[int], + dilation: List[int], + is_transposed: bool, + output_padding: List[int], + groups: int, ): - ret_shape = [] - if isinstance(stride, int): - stride = [stride] * len(dims) - elif len(stride) == 1: - stride = [stride[0]] * len(dims) - - if isinstance(padding, int): - padding = [padding] * len(dims) - elif len(padding) == 1: - padding = [padding[0]] * len(dims) - - if isinstance(dilation, int): - dilation = [dilation] * len(dims) - elif len(dilation) == 1: - dilation = [dilation[0]] * len(dims) - - output_padding_list: Optional[List[int]] = None - if output_padding: - if isinstance(output_padding, int): - output_padding_list = [output_padding] * len(dims) - elif len(output_padding) == 1: - output_padding_list = [output_padding[0]] * len(dims) - else: - output_padding_list = output_padding - - for i in range(len(dims)): - # If output_padding is present, we are dealing with a transposed convolution - if output_padding_list: - ret_shape.append( - _formula_transposed( - dims[i], - padding[i], - dilation[i], - kernel_size[i], - stride[i], - output_padding_list[i], - )) - else: - ret_shape.append(_formula(dims[i], padding[i], dilation[i], kernel_size[i], stride[i])) - return ret_shape - - def pick_memory_format(): - if input_tensor.is_contiguous(memory_format=torch.channels_last): - return torch.channels_last - elif input_tensor.is_contiguous(memory_format=torch.contiguous_format): - return torch.contiguous_format - elif input_tensor.is_contiguous(memory_format=torch.preserve_format): - return torch.preserve_format - - kernel_size = weight.shape[2:] - dims = input_tensor.shape[2:] - if is_transposed: - out_channels = groups * weight.shape[1] - - shape_out = calc_conv_nd_return_shape( - dims, - kernel_size, - stride, - padding, - dilation, - output_padding, - ) - - else: - out_channels = weight.shape[0] - if weight.shape[1] != input_tensor.shape[1] / groups: - raise RuntimeError("Invalid channel dimensions") - shape_out = calc_conv_nd_return_shape(dims, kernel_size, stride, padding, dilation) - out = input_tensor.new_empty((input_tensor.shape[0], out_channels, *shape_out)) - mem_fmt = pick_memory_format() - out = out.to(memory_format=mem_fmt) # type: ignore[call-overload] - return out - - -@register_meta(aten._convolution.default) -def meta__conv(input_tensor: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, stride: List[int], - padding: List[int], dilation: List[int], is_transposed: bool, output_padding: List[int], groups: int, - *extra_args): - out = meta_conv(input_tensor, weight, bias, stride, padding, dilation, is_transposed, output_padding, groups) - return out - - -@register_meta(aten.convolution_backward.default) -def meta_conv_backward(grad_output: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, bias_sizes, stride, - padding, dilation, transposed, output_padding, groups, output_mask): - return new_like(input), new_like(weight), new((bias_sizes)) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/AdaptiveAveragePooling.cpp -@register_meta(aten._adaptive_avg_pool2d_backward.default) -def meta_adaptive_avg_pool2d_backward( - grad_output: torch.Tensor, - input: torch.Tensor, -): - return new_like(input) - - -# ================================ RNN ============================================= -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/RNN.cpp -@register_meta(aten._cudnn_rnn.default) -def meta_cuda_rnn( - input, - weight, - weight_stride0, - weight_buf, - hx, - cx, - mode, - hidden_size, - proj_size, - num_layers, - batch_first, - dropout, - train, - bidirectional, - batch_sizes, - dropout_state, -): - - is_input_packed = len(batch_sizes) != 0 - if is_input_packed: - seq_length = len(batch_sizes) - mini_batch = batch_sizes[0] - batch_sizes_sum = input.shape[0] - else: - seq_length = input.shape[1] if batch_first else input.shape[0] - mini_batch = input.shape[0] if batch_first else input.shape[1] - batch_sizes_sum = -1 - - num_directions = 2 if bidirectional else 1 - out_size = proj_size if proj_size != 0 else hidden_size - if is_input_packed: - out_shape = [batch_sizes_sum, out_size * num_directions] - else: - out_shape = ([mini_batch, seq_length, out_size * - num_directions] if batch_first else [seq_length, mini_batch, out_size * num_directions]) - output = input.new_empty(out_shape) - - cell_shape = [num_layers * num_directions, mini_batch, hidden_size] - cy = new(0) if cx is None else cx.new_empty(cell_shape) - - hy = hx.new_empty([num_layers * num_directions, mini_batch, out_size]) - - # TODO: Query cudnnGetRNNTrainingReserveSize (expose to python) - reserve_shape = 0 if train else 0 - reserve = input.new_empty(reserve_shape, dtype=torch.uint8) - - return output, hy, cy, reserve, weight_buf - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/RNN.cpp -@register_meta(aten._cudnn_rnn_backward.default) -def meta_cudnn_rnn_backward(input: torch.Tensor, - weight: torch.Tensor, - weight_stride0: int, - hx: torch.Tensor, - cx: Optional[torch.Tensor] = None, - *args, - **kwargs): - return new_like(input), new_like(weight), new_like(hx), new_like(cx) if cx is not None else new( - ()) # (grad_input, grad_weight, grad_hx, grad_cx) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Activation.cpp -# ============================== Activations ======================================= -_unregistered_ewise = [ - aten.relu.default, - aten.prelu.default, - aten.hardswish.default, - aten.hardtanh.default, - aten.prelu_backward.default, - aten.hardswish_backward.default, - aten.hardtanh_backward.default, -] - - -@register_meta(_unregistered_ewise) -def meta_unregistered_ewise(input: torch.Tensor, *args): - return new_like(input) - - -# ============================== Normalization ===================================== -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp -@register_meta(aten.native_batch_norm.default) -def meta_bn(input: torch.Tensor, weight, bias, running_mean, running_var, training, momentum, eps): - n_input = input.size(1) - return new_like(input), new((n_input)), new((n_input)) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp -@register_meta(aten.native_batch_norm_backward.default) -def meta_bn_backward(dY: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, running_mean, running_var, save_mean, - save_invstd, train, eps, output_mask): - return new_like(input), new_like(weight), new_like(weight) # (dX, dgamma, dbeta) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp -@register_meta(aten.cudnn_batch_norm.default) -def meta_cudnn_bn(input: torch.Tensor, weight, bias, running_mean, running_var, training, momentum, eps): - n_input = input.size(1) - return new_like(input), new((n_input)), new((n_input)), new( - (0), dtype=torch.uint8) # (output, running_mean, running_var, reserve) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp -# NB: CuDNN only implements the backward algorithm for batchnorm -# in training mode (evaluation mode batchnorm has a different algorithm), -# which is why this doesn't accept a 'training' parameter. -@register_meta(aten.cudnn_batch_norm_backward.default) -def meta_cudnn_bn_backward(dY: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, running_mean, running_var, - save_mean, save_invstd, eps, reserve): - return new_like(input), new_like(weight), new_like(weight) # (dX, dgamma, dbeta) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/layer_norm.cpp -@register_meta(aten.native_layer_norm.default) -def meta_ln(input: torch.Tensor, normalized_shape, weight, bias, eps): - bs, n_input = input.size(0), input.size(1) - return new_like(input), new((bs, n_input, 1)), new((bs, n_input, 1)) # (output, running_mean, running_var) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/layer_norm.cpp -@register_meta(aten.native_layer_norm_backward.default) -def meta_ln_backward(dY: torch.Tensor, input: torch.Tensor, normalized_shape, mean, rstd, weight, bias, - grad_input_mask): - return new_like(input), new_like(weight), new_like(bias) # (dX, dgamma, dbeta) - - -# ================================== Misc ========================================== -# Maybe incorrect -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Im2Col.cpp -@register_meta(aten.im2col.default) -def meta_im2col(input: torch.Tensor, kernel_size, dilation, padding, stride): - return new_like(input) - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml -@register_meta(aten.eye.m_out) -def meta_eye(n: int, m: int, out: torch.Tensor): - return out - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml -@register_meta(aten.roll.default) -def meta_roll(input: torch.Tensor, shifts, dims): - return input - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Scalar.cpp -@register_meta(aten._local_scalar_dense.default) -def meta_local_scalar_dense(self: torch.Tensor): - return 0 - - -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorCompare.cpp -@register_meta(aten.where.self) -def meta_where_self(condition: torch.Tensor, self: torch.Tensor, other: torch.Tensor): - result_type = torch.result_type(self, other) - return new_like(condition + self + other, dtype=result_type) - - -@register_meta(aten.index.Tensor) -def meta_index_Tensor(self, indices): - assert indices, "at least one index must be provided" - # aten::index is the internal advanced indexing implementation - # checkIndexTensorTypes and expandTensors - result: List[Optional[torch.Tensor]] = [] - for i, index in enumerate(indices): - if index is not None: - assert index.dtype in [torch.long, torch.int8, torch.bool],\ - "tensors used as indices must be long, byte or bool tensors" - if index.dtype in [torch.int8, torch.bool]: - nonzero = index.nonzero() - k = len(result) - assert k + index.ndim <= self.ndim, f"too many indices for tensor of dimension {self.ndim}" - for j in range(index.ndim): - assert index.shape[j] == self.shape[ - k + - j], f"The shape of the mask {index.shape} at index {i} does not match the shape of the indexed tensor {self.shape} at index {k + j}" - result.append(nonzero.select(1, j)) - else: - result.append(index) - else: - result.append(index) - indices = result - assert len(indices) <= self.ndim, f"too many indices for tensor of dimension {self.ndim} (got {len(indices)})" - # expand_outplace - import torch._refs as refs - - indices = list(refs._maybe_broadcast(*indices)) - # add missing null tensors - while len(indices) < self.ndim: - indices.append(None) - - # hasContiguousSubspace - # true if all non-null tensors are adjacent - # See: - # https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing - # https://stackoverflow.com/questions/53841497/why-does-numpy-mixed-basic-advanced-indexing-depend-on-slice-adjacency - state = 0 - has_contiguous_subspace = False - for index in indices: - if state == 0: - if index is not None: - state = 1 - elif state == 1: - if index is None: - state = 2 - else: - if index is not None: - break - else: - has_contiguous_subspace = True - - # transposeToFront - # This is the logic that causes the newly inserted dimensions to show up - # at the beginning of the tensor, if they're not contiguous - if not has_contiguous_subspace: - dims = [] - transposed_indices = [] - for i, index in enumerate(indices): - if index is not None: - dims.append(i) - transposed_indices.append(index) - for i, index in enumerate(indices): - if index is None: - dims.append(i) - transposed_indices.append(index) - self = self.permute(dims) - indices = transposed_indices - - # AdvancedIndex::AdvancedIndex - # Now we can assume the indices have contiguous subspace - # This is simplified from AdvancedIndex which goes to more effort - # to put the input and indices in a form so that TensorIterator can - # take them. If we write a ref for this, probably that logic should - # get implemented - before_shape: List[int] = [] - after_shape: List[int] = [] - replacement_shape: List[int] = [] - for dim, index in enumerate(indices): - if index is None: - if replacement_shape: - after_shape.append(self.shape[dim]) - else: - before_shape.append(self.shape[dim]) - else: - replacement_shape = list(index.shape) - return self.new_empty(before_shape + replacement_shape + after_shape) - -# ============================== Embedding ========================================= -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Embedding.cpp -@register_meta(aten.embedding_dense_backward.default) -def meta_embedding_dense_backward(grad_output: torch.Tensor, indices: torch.Tensor, num_weights, padding_idx, - scale_grad_by_freq): - return new((num_weights, grad_output.size(-1)), - dtype=grad_output.dtype, - device=grad_output.device, - layout=grad_output.layout) - - -# ============================== Dropout =========================================== -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Dropout.cpp -@register_meta(aten.native_dropout.default) -def meta_native_dropout_default(input: torch.Tensor, p: float, train: bool = False): - # notice that mask is bool - return new_like(input), new_like(input, dtype=torch.bool) # (output, mask) + def _formula(ln: int, p: int, d: int, k: int, s: int) -> int: + """ + Formula to apply to calculate the length of some dimension of the output + See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html + Args: + ln: length of the dimension + p: padding in that dim + d: dilation in that dim + k: kernel size in that dim + s: stride in that dim + Returns: + The output length + """ + return (ln + 2 * p - d * (k - 1) - 1) // s + 1 + + def _formula_transposed(ln: int, p: int, d: int, k: int, s: int, op: int) -> int: + """ + Formula to apply to calculate the length of some dimension of the output + if transposed convolution is used. + See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html + Args: + ln: length of the dimension + p: padding in that dim + d: dilation in that dim + k: kernel size in that dim + s: stride in that dim + op: output padding in that dim + Returns: + The output length + """ + return (ln - 1) * s - 2 * p + d * (k - 1) + op + 1 + + def calc_conv_nd_return_shape( + dims: torch.Size, + kernel_size: torch.Size, + stride: Union[List[int], int], + padding: Union[List[int], int], + dilation: Union[List[int], int], + output_padding: Optional[Union[List[int], int]] = None, + ): + ret_shape = [] + if isinstance(stride, int): + stride = [stride] * len(dims) + elif len(stride) == 1: + stride = [stride[0]] * len(dims) + + if isinstance(padding, int): + padding = [padding] * len(dims) + elif len(padding) == 1: + padding = [padding[0]] * len(dims) + + if isinstance(dilation, int): + dilation = [dilation] * len(dims) + elif len(dilation) == 1: + dilation = [dilation[0]] * len(dims) + + output_padding_list: Optional[List[int]] = None + if output_padding: + if isinstance(output_padding, int): + output_padding_list = [output_padding] * len(dims) + elif len(output_padding) == 1: + output_padding_list = [output_padding[0]] * len(dims) + else: + output_padding_list = output_padding + + for i in range(len(dims)): + # If output_padding is present, we are dealing with a transposed convolution + if output_padding_list: + ret_shape.append( + _formula_transposed( + dims[i], + padding[i], + dilation[i], + kernel_size[i], + stride[i], + output_padding_list[i], + )) + else: + ret_shape.append(_formula(dims[i], padding[i], dilation[i], kernel_size[i], stride[i])) + return ret_shape + + def pick_memory_format(): + if input_tensor.is_contiguous(memory_format=torch.channels_last): + return torch.channels_last + elif input_tensor.is_contiguous(memory_format=torch.contiguous_format): + return torch.contiguous_format + elif input_tensor.is_contiguous(memory_format=torch.preserve_format): + return torch.preserve_format + + kernel_size = weight.shape[2:] + dims = input_tensor.shape[2:] + if is_transposed: + out_channels = groups * weight.shape[1] + + shape_out = calc_conv_nd_return_shape( + dims, + kernel_size, + stride, + padding, + dilation, + output_padding, + ) + else: + out_channels = weight.shape[0] + if weight.shape[1] != input_tensor.shape[1] / groups: + raise RuntimeError("Invalid channel dimensions") + shape_out = calc_conv_nd_return_shape(dims, kernel_size, stride, padding, dilation) + out = input_tensor.new_empty((input_tensor.shape[0], out_channels, *shape_out)) + mem_fmt = pick_memory_format() + out = out.to(memory_format=mem_fmt) # type: ignore[call-overload] + return out + + @register_meta(aten._convolution.default) + def meta__conv(input_tensor: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, stride: List[int], + padding: List[int], dilation: List[int], is_transposed: bool, output_padding: List[int], groups: int, + *extra_args): + out = meta_conv(input_tensor, weight, bias, stride, padding, dilation, is_transposed, output_padding, groups) + return out + + @register_meta(aten.convolution_backward.default) + def meta_conv_backward(grad_output: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, bias_sizes, stride, + padding, dilation, transposed, output_padding, groups, output_mask): + return new_like(input), new_like(weight), new((bias_sizes)) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/AdaptiveAveragePooling.cpp + @register_meta(aten._adaptive_avg_pool2d_backward.default) + def meta_adaptive_avg_pool2d_backward( + grad_output: torch.Tensor, + input: torch.Tensor, + ): + return new_like(input) + + # ================================ RNN ============================================= + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/RNN.cpp + @register_meta(aten._cudnn_rnn.default) + def meta_cuda_rnn( + input, + weight, + weight_stride0, + weight_buf, + hx, + cx, + mode, + hidden_size, + proj_size, + num_layers, + batch_first, + dropout, + train, + bidirectional, + batch_sizes, + dropout_state, + ): -# https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Dropout.cpp -@register_meta(aten.native_dropout_backward.default) -def meta_native_dropout_backward_default(grad: torch.Tensor, mask: torch.Tensor, scale: float): - return new_like(grad) # (grad_in) + is_input_packed = len(batch_sizes) != 0 + if is_input_packed: + seq_length = len(batch_sizes) + mini_batch = batch_sizes[0] + batch_sizes_sum = input.shape[0] + else: + seq_length = input.shape[1] if batch_first else input.shape[0] + mini_batch = input.shape[0] if batch_first else input.shape[1] + batch_sizes_sum = -1 + + num_directions = 2 if bidirectional else 1 + out_size = proj_size if proj_size != 0 else hidden_size + if is_input_packed: + out_shape = [batch_sizes_sum, out_size * num_directions] + else: + out_shape = ([mini_batch, seq_length, out_size * + num_directions] if batch_first else [seq_length, mini_batch, out_size * num_directions]) + output = input.new_empty(out_shape) + + cell_shape = [num_layers * num_directions, mini_batch, hidden_size] + cy = new(0) if cx is None else cx.new_empty(cell_shape) + + hy = hx.new_empty([num_layers * num_directions, mini_batch, out_size]) + + # TODO: Query cudnnGetRNNTrainingReserveSize (expose to python) + reserve_shape = 0 if train else 0 + reserve = input.new_empty(reserve_shape, dtype=torch.uint8) + + return output, hy, cy, reserve, weight_buf + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/RNN.cpp + @register_meta(aten._cudnn_rnn_backward.default) + def meta_cudnn_rnn_backward(input: torch.Tensor, + weight: torch.Tensor, + weight_stride0: int, + hx: torch.Tensor, + cx: Optional[torch.Tensor] = None, + *args, + **kwargs): + return new_like(input), new_like(weight), new_like(hx), new_like(cx) if cx is not None else new( + ()) # (grad_input, grad_weight, grad_hx, grad_cx) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Activation.cpp + # ============================== Activations ======================================= + _unregistered_ewise = [ + aten.relu.default, + aten.prelu.default, + aten.hardswish.default, + aten.hardtanh.default, + aten.hardswish_backward.default, + aten.hardtanh_backward.default, + ] + + if version.parse(torch.__version__) < version.parse('2.0.0'): + _unregistered_ewise += [ + aten.prelu_backward.default, + ] + + @register_meta(_unregistered_ewise) + def meta_unregistered_ewise(input: torch.Tensor, *args): + return new_like(input) + + # ============================== Normalization ===================================== + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp + @register_meta(aten.native_batch_norm.default) + def meta_bn(input: torch.Tensor, weight, bias, running_mean, running_var, training, momentum, eps): + n_input = input.size(1) + return new_like(input), new((n_input)), new((n_input)) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp + @register_meta(aten.native_batch_norm_backward.default) + def meta_bn_backward(dY: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, running_mean, running_var, + save_mean, save_invstd, train, eps, output_mask): + return new_like(input), new_like(weight), new_like(weight) # (dX, dgamma, dbeta) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp + @register_meta(aten.cudnn_batch_norm.default) + def meta_cudnn_bn(input: torch.Tensor, weight, bias, running_mean, running_var, training, momentum, eps): + n_input = input.size(1) + return new_like(input), new((n_input)), new((n_input)), new( + (0), dtype=torch.uint8) # (output, running_mean, running_var, reserve) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cudnn/BatchNorm.cpp + # NB: CuDNN only implements the backward algorithm for batchnorm + # in training mode (evaluation mode batchnorm has a different algorithm), + # which is why this doesn't accept a 'training' parameter. + @register_meta(aten.cudnn_batch_norm_backward.default) + def meta_cudnn_bn_backward(dY: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, running_mean, running_var, + save_mean, save_invstd, eps, reserve): + return new_like(input), new_like(weight), new_like(weight) # (dX, dgamma, dbeta) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/layer_norm.cpp + @register_meta(aten.native_layer_norm.default) + def meta_ln(input: torch.Tensor, normalized_shape, weight, bias, eps): + bs, n_input = input.size(0), input.size(1) + return new_like(input), new((bs, n_input, 1)), new((bs, n_input, 1)) # (output, running_mean, running_var) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/layer_norm.cpp + @register_meta(aten.native_layer_norm_backward.default) + def meta_ln_backward(dY: torch.Tensor, input: torch.Tensor, normalized_shape, mean, rstd, weight, bias, + grad_input_mask): + return new_like(input), new_like(weight), new_like(bias) # (dX, dgamma, dbeta) + + # ================================== Misc ========================================== + # Maybe incorrect + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Im2Col.cpp + @register_meta(aten.im2col.default) + def meta_im2col(input: torch.Tensor, kernel_size, dilation, padding, stride): + return new_like(input) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml + @register_meta(aten.roll.default) + def meta_roll(input: torch.Tensor, shifts, dims): + return input + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Scalar.cpp + @register_meta(aten._local_scalar_dense.default) + def meta_local_scalar_dense(self: torch.Tensor): + return 0 + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorCompare.cpp + @register_meta(aten.where.self) + def meta_where_self(condition: torch.Tensor, self: torch.Tensor, other: torch.Tensor): + result_type = torch.result_type(self, other) + return new_like(condition + self + other, dtype=result_type) + + # ============================== Embedding ========================================= + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Embedding.cpp + + @register_meta(aten.embedding_dense_backward.default) + def meta_embedding_dense_backward(grad_output: torch.Tensor, indices: torch.Tensor, num_weights, padding_idx, + scale_grad_by_freq): + return new((num_weights, grad_output.size(-1)), dtype=grad_output.dtype, layout=grad_output.layout) + + # ============================== Dropout =========================================== + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Dropout.cpp + @register_meta(aten.native_dropout.default) + def meta_native_dropout_default(input: torch.Tensor, p: float, train: bool = False): + # notice that mask is bool + return new_like(input), new_like(input, dtype=torch.bool) # (output, mask) + + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Dropout.cpp + @register_meta(aten.native_dropout_backward.default) + def meta_native_dropout_backward_default(grad: torch.Tensor, mask: torch.Tensor, scale: float): + return new_like(grad) # (grad_in) + + if version.parse(torch.__version__) < version.parse('1.13.0'): + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml + @register_meta(aten.eye.m_out) + def meta_eye(n: int, m: int, out: torch.Tensor): + return out + + @register_meta(aten.index.Tensor) + def meta_index_Tensor(self, indices): + assert indices, "at least one index must be provided" + # aten::index is the internal advanced indexing implementation + # checkIndexTensorTypes and expandTensors + result: List[Optional[torch.Tensor]] = [] + for i, index in enumerate(indices): + if index is not None: + assert index.dtype in [torch.long, torch.int8, torch.bool],\ + "tensors used as indices must be long, byte or bool tensors" + if index.dtype in [torch.int8, torch.bool]: + nonzero = index.nonzero() + k = len(result) + assert k + index.ndim <= self.ndim, f"too many indices for tensor of dimension {self.ndim}" + for j in range(index.ndim): + assert index.shape[j] == self.shape[ + k + + j], f"The shape of the mask {index.shape} at index {i} does not match the shape of the indexed tensor {self.shape} at index {k + j}" + result.append(nonzero.select(1, j)) + else: + result.append(index) + else: + result.append(index) + indices = result + assert len( + indices) <= self.ndim, f"too many indices for tensor of dimension {self.ndim} (got {len(indices)})" + # expand_outplace + import torch._refs as refs + + indices = list(refs._maybe_broadcast(*indices)) + # add missing null tensors + while len(indices) < self.ndim: + indices.append(None) + + # hasContiguousSubspace + # true if all non-null tensors are adjacent + # See: + # https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing + # https://stackoverflow.com/questions/53841497/why-does-numpy-mixed-basic-advanced-indexing-depend-on-slice-adjacency + state = 0 + has_contiguous_subspace = False + for index in indices: + if state == 0: + if index is not None: + state = 1 + elif state == 1: + if index is None: + state = 2 + else: + if index is not None: + break + else: + has_contiguous_subspace = True + + # transposeToFront + # This is the logic that causes the newly inserted dimensions to show up + # at the beginning of the tensor, if they're not contiguous + if not has_contiguous_subspace: + dims = [] + transposed_indices = [] + for i, index in enumerate(indices): + if index is not None: + dims.append(i) + transposed_indices.append(index) + for i, index in enumerate(indices): + if index is None: + dims.append(i) + transposed_indices.append(index) + self = self.permute(dims) + indices = transposed_indices + + # AdvancedIndex::AdvancedIndex + # Now we can assume the indices have contiguous subspace + # This is simplified from AdvancedIndex which goes to more effort + # to put the input and indices in a form so that TensorIterator can + # take them. If we write a ref for this, probably that logic should + # get implemented + before_shape: List[int] = [] + after_shape: List[int] = [] + replacement_shape: List[int] = [] + for dim, index in enumerate(indices): + if index is None: + if replacement_shape: + after_shape.append(self.shape[dim]) + else: + before_shape.append(self.shape[dim]) + else: + replacement_shape = list(index.shape) + return self.new_empty(before_shape + replacement_shape + after_shape) diff --git a/colossalai/_analyzer/_subclasses/_monkey_patch.py b/colossalai/_analyzer/_subclasses/_monkey_patch.py index 1c7b972ab2f6..b3ec98f0811f 100644 --- a/colossalai/_analyzer/_subclasses/_monkey_patch.py +++ b/colossalai/_analyzer/_subclasses/_monkey_patch.py @@ -1,7 +1,6 @@ import torch import torch.distributed as dist - -aten = torch.ops.aten +from packaging import version __all__ = [ "_TorchFactoryMethod", @@ -49,40 +48,46 @@ "scatter", ] -# TODO: dive deep here -# refer to https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorShape.cpp -_AliasATen = [ - aten.detach.default, - aten.detach_.default, - aten.t.default, - aten.transpose.int, - aten.view.default, - aten._unsafe_view.default, - aten._reshape_alias.default, -] +if version.parse(torch.__version__) >= version.parse('1.12.0'): + aten = torch.ops.aten + # TODO: dive deep here + # refer to https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorShape.cpp + _AliasATen = [ + aten.detach.default, + aten.detach_.default, + aten.t.default, + aten.transpose.int, + aten.view.default, + aten._unsafe_view.default, + aten._reshape_alias.default, + ] -_InplaceATen = [ - aten.add_.Tensor, - aten.add_.Scalar, - aten.sub_.Tensor, - aten.sub_.Scalar, - aten.mul_.Tensor, - aten.mul_.Scalar, - aten.div_.Tensor, - aten.div_.Scalar, - aten.pow_.Tensor, - aten.pow_.Scalar, -] + _InplaceATen = [ + aten.add_.Tensor, + aten.add_.Scalar, + aten.sub_.Tensor, + aten.sub_.Scalar, + aten.mul_.Tensor, + aten.mul_.Scalar, + aten.div_.Tensor, + aten.div_.Scalar, + aten.pow_.Tensor, + aten.pow_.Scalar, + ] -# use `MaybeInplace` because they call ``as_strided()`` or ``slice()`` -_MaybeInplaceATen = [ - aten.diagonal.default, - aten.expand.default, - aten.select.int, - aten.slice.Tensor, - aten.split.Tensor, - aten.squeeze.default, - aten.permute.default, - aten.unsqueeze.default, - aten.as_strided.default, -] + # use `MaybeInplace` because they call ``as_strided()`` or ``slice()`` + _MaybeInplaceATen = [ + aten.diagonal.default, + aten.expand.default, + aten.select.int, + aten.slice.Tensor, + aten.split.Tensor, + aten.squeeze.default, + aten.permute.default, + aten.unsqueeze.default, + aten.as_strided.default, + ] +else: + _AliasATen = [] + _InplaceATen = [] + _MaybeInplaceATen = [] diff --git a/colossalai/_analyzer/_subclasses/flop_tensor.py b/colossalai/_analyzer/_subclasses/flop_tensor.py index ab93551467b8..59991dc50912 100644 --- a/colossalai/_analyzer/_subclasses/flop_tensor.py +++ b/colossalai/_analyzer/_subclasses/flop_tensor.py @@ -11,6 +11,7 @@ from typing import Any, Callable, List, Optional, Union import torch +from packaging import version from torch.utils._pytree import tree_map from .meta_tensor import MetaTensor @@ -234,7 +235,28 @@ def matmul_flop_jit(inputs: List[Any], outputs: List[Any]) -> Number: # Inputs contains the shapes of two matrices. input_shapes = [v.shape for v in inputs] assert len(input_shapes) == 2, input_shapes - assert input_shapes[0][-1] == input_shapes[1][-2], input_shapes + + # There are three cases: 1) gemm, 2) gemv, 3) dot + if all(len(shape) == 2 for shape in input_shapes): + # gemm + assert input_shapes[0][-1] == input_shapes[1][-2], input_shapes + elif all(len(shape) == 1 for shape in input_shapes): + # dot + assert input_shapes[0][0] == input_shapes[1][0], input_shapes + + # expand shape + input_shapes[0] = torch.Size([1, input_shapes[0][0]]) + input_shapes[1] = torch.Size([input_shapes[1][0], 1]) + else: + # gemv + if len(input_shapes[0]) == 1: + assert input_shapes[0][0] == input_shapes[1][-2], input_shapes + input_shapes.reverse() + else: + assert input_shapes[1][0] == input_shapes[0][-1], input_shapes + + # expand the shape of the vector to [batch size, 1] + input_shapes[-1] = torch.Size([input_shapes[-1][-1], 1]) flops = reduce(operator.mul, input_shapes[0]) * input_shapes[-1][-1] return flops @@ -403,134 +425,139 @@ def zero_flop_jit(*args): return 0 -flop_mapping = { +if version.parse(torch.__version__) >= version.parse('1.12.0'): + flop_mapping = { # gemm - aten.mm.default: matmul_flop_jit, - aten.matmul.default: matmul_flop_jit, - aten.addmm.default: addmm_flop_jit, - aten.bmm.default: bmm_flop_jit, + aten.mm.default: matmul_flop_jit, + aten.matmul.default: matmul_flop_jit, + aten.addmm.default: addmm_flop_jit, + aten.bmm.default: bmm_flop_jit, # convolution - aten.convolution.default: conv_flop_jit, - aten._convolution.default: conv_flop_jit, - aten.convolution_backward.default: conv_backward_flop_jit, + aten.convolution.default: conv_flop_jit, + aten._convolution.default: conv_flop_jit, + aten.convolution_backward.default: conv_backward_flop_jit, # normalization - aten.native_batch_norm.default: batchnorm_flop_jit, - aten.native_batch_norm_backward.default: batchnorm_flop_jit, - aten.cudnn_batch_norm.default: batchnorm_flop_jit, - aten.cudnn_batch_norm_backward.default: partial(batchnorm_flop_jit, training=True), - aten.native_layer_norm.default: norm_flop_counter(2, 0), - aten.native_layer_norm_backward.default: norm_flop_counter(2, 0), + aten.native_batch_norm.default: batchnorm_flop_jit, + aten.native_batch_norm_backward.default: batchnorm_flop_jit, + aten.cudnn_batch_norm.default: batchnorm_flop_jit, + aten.cudnn_batch_norm_backward.default: partial(batchnorm_flop_jit, training=True), + aten.native_layer_norm.default: norm_flop_counter(2, 0), + aten.native_layer_norm_backward.default: norm_flop_counter(2, 0), # pooling - aten.avg_pool1d.default: ewise_flop_counter(1, 0), - aten.avg_pool2d.default: ewise_flop_counter(1, 0), - aten.avg_pool2d_backward.default: ewise_flop_counter(0, 1), - aten.avg_pool3d.default: ewise_flop_counter(1, 0), - aten.avg_pool3d_backward.default: ewise_flop_counter(0, 1), - aten.max_pool1d.default: ewise_flop_counter(1, 0), - aten.max_pool2d.default: ewise_flop_counter(1, 0), - aten.max_pool3d.default: ewise_flop_counter(1, 0), - aten.max_pool1d_with_indices.default: ewise_flop_counter(1, 0), - aten.max_pool2d_with_indices.default: ewise_flop_counter(1, 0), - aten.max_pool2d_with_indices_backward.default: ewise_flop_counter(0, 1), - aten.max_pool3d_with_indices.default: ewise_flop_counter(1, 0), - aten.max_pool3d_with_indices_backward.default: ewise_flop_counter(0, 1), - aten._adaptive_avg_pool2d.default: ewise_flop_counter(1, 0), - aten._adaptive_avg_pool2d_backward.default: ewise_flop_counter(0, 1), - aten._adaptive_avg_pool3d.default: ewise_flop_counter(1, 0), - aten._adaptive_avg_pool3d_backward.default: ewise_flop_counter(0, 1), - aten.embedding_dense_backward.default: ewise_flop_counter(0, 1), - aten.embedding.default: ewise_flop_counter(1, 0), -} - -ewise_flop_aten = [ + aten.avg_pool1d.default: ewise_flop_counter(1, 0), + aten.avg_pool2d.default: ewise_flop_counter(1, 0), + aten.avg_pool2d_backward.default: ewise_flop_counter(0, 1), + aten.avg_pool3d.default: ewise_flop_counter(1, 0), + aten.avg_pool3d_backward.default: ewise_flop_counter(0, 1), + aten.max_pool1d.default: ewise_flop_counter(1, 0), + aten.max_pool2d.default: ewise_flop_counter(1, 0), + aten.max_pool3d.default: ewise_flop_counter(1, 0), + aten.max_pool1d_with_indices.default: ewise_flop_counter(1, 0), + aten.max_pool2d_with_indices.default: ewise_flop_counter(1, 0), + aten.max_pool2d_with_indices_backward.default: ewise_flop_counter(0, 1), + aten.max_pool3d_with_indices.default: ewise_flop_counter(1, 0), + aten.max_pool3d_with_indices_backward.default: ewise_flop_counter(0, 1), + aten._adaptive_avg_pool2d.default: ewise_flop_counter(1, 0), + aten._adaptive_avg_pool2d_backward.default: ewise_flop_counter(0, 1), + aten._adaptive_avg_pool3d.default: ewise_flop_counter(1, 0), + aten._adaptive_avg_pool3d_backward.default: ewise_flop_counter(0, 1), + aten.embedding_dense_backward.default: ewise_flop_counter(0, 1), + aten.embedding.default: ewise_flop_counter(1, 0), + } + + ewise_flop_aten = [ # basic op - aten.add.Tensor, - aten.add_.Tensor, - aten.div.Tensor, - aten.div_.Tensor, - aten.div.Scalar, - aten.div_.Scalar, - aten.mul.Tensor, - aten.mul.Scalar, - aten.mul_.Tensor, - aten.neg.default, - aten.pow.Tensor_Scalar, - aten.rsub.Scalar, - aten.sum.default, - aten.sum.dim_IntList, - aten.mean.dim, + aten.add.Tensor, + aten.add_.Tensor, + aten.div.Tensor, + aten.div_.Tensor, + aten.div.Scalar, + aten.div_.Scalar, + aten.mul.Tensor, + aten.mul.Scalar, + aten.mul_.Tensor, + aten.neg.default, + aten.pow.Tensor_Scalar, + aten.rsub.Scalar, + aten.sum.default, + aten.sum.dim_IntList, + aten.mean.dim, # activation op - aten.hardswish.default, - aten.hardswish_.default, - aten.hardswish_backward.default, - aten.hardtanh.default, - aten.hardtanh_.default, - aten.hardtanh_backward.default, - aten.hardsigmoid_backward.default, - aten.hardsigmoid.default, - aten.gelu.default, - aten.gelu_backward.default, - aten.silu.default, - aten.silu_.default, - aten.silu_backward.default, - aten.sigmoid.default, - aten.sigmoid_backward.default, - aten._softmax.default, - aten._softmax_backward_data.default, - aten.relu_.default, - aten.relu.default, - aten.tanh.default, - aten.tanh_backward.default, - aten.threshold_backward.default, + aten.hardswish.default, + aten.hardswish_.default, + aten.hardswish_backward.default, + aten.hardtanh.default, + aten.hardtanh_.default, + aten.hardtanh_backward.default, + aten.hardsigmoid_backward.default, + aten.hardsigmoid.default, + aten.gelu.default, + aten.gelu_backward.default, + aten.silu.default, + aten.silu_.default, + aten.silu_backward.default, + aten.sigmoid.default, + aten.sigmoid_backward.default, + aten._softmax.default, + aten._softmax_backward_data.default, + aten.relu_.default, + aten.relu.default, + aten.tanh.default, + aten.tanh_backward.default, + aten.threshold_backward.default, # dropout - aten.native_dropout.default, - aten.native_dropout_backward.default, + aten.native_dropout.default, + aten.native_dropout_backward.default, # distribution - aten.bernoulli_.float, + aten.bernoulli_.float, # where - aten.where.self, -] -for op in ewise_flop_aten: - flop_mapping[op] = ewise_flop_counter(1, 0) - -# fix-me: this will be removed in future -zero_flop_aten = [ - aten.as_strided.default, - aten.as_strided_.default, - aten.cat.default, - aten.clone.default, - aten.copy_.default, - aten.detach.default, - aten.expand.default, - aten.empty_like.default, - aten.new_empty.default, - aten.new_empty_strided.default, - aten.ones_like.default, - aten._reshape_alias.default, - aten.select.int, - aten.select_backward.default, - aten.squeeze.dim, - aten.slice.Tensor, - aten.slice_backward.default, - aten.split.Tensor, - aten.permute.default, - aten.t.default, - aten.transpose.int, - aten._to_copy.default, - aten.unsqueeze.default, - aten.unbind.int, - aten._unsafe_view.default, - aten.view.default, - aten.zero_.default, - aten.zeros_like.default, -] - -for op in zero_flop_aten: - flop_mapping[op] = zero_flop_jit + aten.where.self, + ] + for op in ewise_flop_aten: + flop_mapping[op] = ewise_flop_counter(1, 0) + + # fix-me: this will be removed in future + zero_flop_aten = [ + aten.as_strided.default, + aten.as_strided_.default, + aten.cat.default, + aten.clone.default, + aten.copy_.default, + aten.detach.default, + aten.expand.default, + aten.empty_like.default, + aten.new_empty.default, + aten.new_empty_strided.default, + aten.ones_like.default, + aten._reshape_alias.default, + aten.select.int, + aten.select_backward.default, + aten.squeeze.dim, + aten.slice.Tensor, + aten.slice_backward.default, + aten.split.Tensor, + aten.permute.default, + aten.t.default, + aten.transpose.int, + aten._to_copy.default, + aten.unsqueeze.default, + aten.unbind.int, + aten._unsafe_view.default, + aten.view.default, + aten.zero_.default, + aten.zeros_like.default, + ] + + for op in zero_flop_aten: + flop_mapping[op] = zero_flop_jit +else: + flop_mapping = {} + elementwise_flop_aten = {} + zero_flop_aten = {} diff --git a/colossalai/_analyzer/fx/__init__.py b/colossalai/_analyzer/fx/__init__.py index 2e857b1b054b..aa01de0bbe6c 100644 --- a/colossalai/_analyzer/fx/__init__.py +++ b/colossalai/_analyzer/fx/__init__.py @@ -1,4 +1,3 @@ -from .bias_addition import * from .node_util import MetaInfo from .symbolic_profile import symbolic_profile -from .symbolic_trace import symbolic_trace +from .tracer.symbolic_trace import symbolic_trace diff --git a/colossalai/_analyzer/fx/codegen.py b/colossalai/_analyzer/fx/codegen.py index 1117c0103166..41d74f2e3719 100644 --- a/colossalai/_analyzer/fx/codegen.py +++ b/colossalai/_analyzer/fx/codegen.py @@ -1,8 +1,12 @@ from typing import Any, Callable, Dict, Iterable, List, Tuple import torch + +try: + from torch.fx.graph import CodeGen +except: + pass from torch.fx.graph import ( - CodeGen, PythonCode, _custom_builtins, _format_target, @@ -48,8 +52,8 @@ def _end_of_ckpt(node: Node, ckpt_level: int) -> bool: """ Check if the node could end the ckpt region at `ckpt_level` """ - if len(node.meta['info'].to_recompute) > ckpt_level: - return node.meta['info'].to_recompute[ckpt_level] is not None + if len(node.meta['info'].activation_checkpoint) > ckpt_level: + return node.meta['info'].activation_checkpoint[ckpt_level] is not None return True @@ -90,8 +94,8 @@ def _find_nested_ckpt_regions(node_list: List[Node], ckpt_level: int = 0): current_region = None for idx, node in enumerate(node_list): - if len(node.meta['info'].to_recompute) > ckpt_level: - act_ckpt_label = node.meta['info'].to_recompute[ckpt_level] + if len(node.meta['info'].activation_checkpoint) > ckpt_level: + act_ckpt_label = node.meta['info'].activation_checkpoint[ckpt_level] # this activation checkpoint label is not set yet # meaning this is the first node of the activation ckpt region @@ -134,7 +138,7 @@ def emit_ckpt_func(body, delete_unused_value_func, ckpt_level=0, in_ckpt=False): - """Emit ckpt fuction in nested way + """Emit ckpt function in nested way Args: body: forward code - in recursive calls, this part will be checkpoint @@ -152,12 +156,12 @@ def emit_ckpt_func(body, # label given by each layer, e.g. if you are currently at level (0, 1, 1) # the label will be '0_1_1' - label = "_".join([str(idx) for idx in node_list[0].meta['info'].to_recompute[:ckpt_level + 1]]) + label = "_".join([str(idx) for idx in node_list[0].meta['info'].activation_checkpoint[:ckpt_level + 1]]) ckpt_fn_def = _gen_ckpt_fn_def(label, inputs) ckpt_func.append(f'{ckpt_fn_def}\n') # if there is more level to fetch - if ckpt_level + 1 < max(map(lambda node: len(node.meta['info'].to_recompute), node_list)): + if ckpt_level + 1 < max(map(lambda node: len(node.meta['info'].activation_checkpoint), node_list)): ckpt_regions = _find_nested_ckpt_regions(node_list, ckpt_level + 1) start_idx = [item[0] for item in ckpt_regions] end_idx = [item[1] for item in ckpt_regions] @@ -215,7 +219,6 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func, ckpt_regions = _find_nested_ckpt_regions(nodes, 0) start_idx = [item[0] for item in ckpt_regions] end_idx = [item[1] for item in ckpt_regions] - node_list = list(nodes) node_idx = 0 diff --git a/colossalai/_analyzer/fx/graph_module.py b/colossalai/_analyzer/fx/graph_module.py index 779b42ebaafd..1fdedd758c01 100644 --- a/colossalai/_analyzer/fx/graph_module.py +++ b/colossalai/_analyzer/fx/graph_module.py @@ -1,4 +1,7 @@ +import linecache import os +import sys +import traceback import warnings from pathlib import Path from typing import Any, Dict, Optional, Union @@ -6,11 +9,74 @@ import torch import torch.fx import torch.nn as nn -from torch.fx.graph import PythonCode, _PyTreeCodeGen -from torch.fx.graph_module import _exec_with_source, _forward_from_src, _WrappedCall +from torch.fx.graph import PythonCode + +try: + from torch.fx.graph import _PyTreeCodeGen + SUPPORT_PT_CODEGEN = True +except ImportError: + SUPPORT_PT_CODEGEN = False + +from torch.fx.graph_module import _exec_with_source, _forward_from_src from torch.nn.modules.module import _addindent +# This is a copy of torch.fx.graph_module._WrappedCall. +# It should be removed when we stop supporting torch < 1.12.0. +class _WrappedCall: + + def __init__(self, cls, cls_call): + self.cls = cls + self.cls_call = cls_call + + # Previously, if an error occurred when valid + # symbolically-traced code was run with an invalid input, the + # user would see the source of the error as coming from + # `File "`, where N is some number. We use + # this function to generate a more informative error message. We + # return the traceback itself, a message explaining that the + # error occurred in a traced Module's generated forward + # function, and five lines of context surrounding the faulty + # line + @staticmethod + def _generate_error_message(frame_summary: traceback.FrameSummary) -> str: + # auxiliary variables (for readability) + err_lineno = frame_summary.lineno + assert err_lineno is not None + line = frame_summary.line + assert line is not None + err_line_len = len(line) + all_src_lines = linecache.getlines(frame_summary.filename) + + # constituent substrings of the error message + tb_repr = traceback.format_exc() + custom_msg = ("Call using an FX-traced Module, " + f"line {err_lineno} of the traced Module's " + "generated forward function:") + before_err = "".join(all_src_lines[err_lineno - 2:err_lineno]) + marker = "~" * err_line_len + "~~~ <--- HERE" + err_and_after_err = "\n".join(all_src_lines[err_lineno:err_lineno + 2]) + + # joined message + return "\n".join([tb_repr, custom_msg, before_err, marker, err_and_after_err]) + + def __call__(self, obj, *args, **kwargs): + try: + if self.cls_call is not None: + return self.cls_call(obj, *args, **kwargs) + else: + return super(self.cls, obj).__call__(*args, **kwargs) # type: ignore[misc] + except Exception as e: + assert e.__traceback__ + topmost_framesummary: traceback.FrameSummary = \ + traceback.StackSummary.extract(traceback.walk_tb(e.__traceback__))[-1] # type: ignore[arg-type] + if "eval_with_key" in topmost_framesummary.filename: + print(_WrappedCall._generate_error_message(topmost_framesummary), file=sys.stderr) + raise e.with_traceback(None) + else: + raise e + + class ColoGraphModule(torch.fx.GraphModule): """ ColoGraphGraphModule is an nn.Module generated from an fx.Graph. @@ -65,7 +131,7 @@ def recompile(self) -> PythonCode: called after editing the contained ``graph``, otherwise the generated code of this ``GraphModule`` will be out of date. """ - if isinstance(self._graph._codegen, _PyTreeCodeGen): + if SUPPORT_PT_CODEGEN and isinstance(self._graph._codegen, _PyTreeCodeGen): self._in_spec = self._graph._codegen.pytree_info.in_spec self._out_spec = self._graph._codegen.pytree_info.out_spec python_code = self._graph.python_code(root_module='self') diff --git a/colossalai/_analyzer/fx/node_util.py b/colossalai/_analyzer/fx/node_util.py index d06fa8b93fc6..fbe8400a437e 100644 --- a/colossalai/_analyzer/fx/node_util.py +++ b/colossalai/_analyzer/fx/node_util.py @@ -20,7 +20,7 @@ def union(a, b): return {**a, **b} -def compute_size_in_bytes(elem: torch.Tensor | Dict | List | Tuple | int) -> int: +def compute_size_in_bytes(elem: Union[torch.Tensor, Dict, List, Tuple, int]) -> int: """Compute the size of a tensor or a collection of tensors in bytes. Args: @@ -112,7 +112,7 @@ class MetaInfo: # should keep the same whenever manipulated # ============================= Invariant ================================== - to_recompute: Tuple[torch.Tensor] = () # (region_0, region_1, ...) support nested codegen + activation_checkpoint: Tuple[torch.Tensor] = () # (region_0, region_1, ...) support nested codegen to_offload: Optional[bool] = False sharding_spec: str = 'RR' @@ -195,8 +195,8 @@ def __repr__(self): s += f'\n\thas buffer of size {_format_memory(self.buffer_size)}' if self.output_size: s += f'\n\thas output activation of size {_format_memory(self.output_size)}' - if self.total_size: - s += f'\n\thas total activation of size {_format_memory(self.total_size)}' + # if self.total_size: + # s += f'\n\thas total activation of size {_format_memory(self.total_size)}' if self.temp_size: s += f'\n\thas temp activation of size {_format_memory(self.temp_size)}' if self.backward_size: diff --git a/colossalai/_analyzer/fx/passes/shape_prop.py b/colossalai/_analyzer/fx/passes/shape_prop.py index 3691497ed8cd..23e83013e02f 100644 --- a/colossalai/_analyzer/fx/passes/shape_prop.py +++ b/colossalai/_analyzer/fx/passes/shape_prop.py @@ -51,7 +51,10 @@ def _normalize_tuple(x): def _current_device(module): - return next(module.parameters()).device + try: + return next(module.parameters()).device + except StopIteration: + return torch.device('cpu') @compatibility(is_backward_compatible=False) @@ -111,7 +114,27 @@ def run_node(self, n: torch.fx.Node) -> Any: with self.global_hook: r = getattr(self, n.op)(n.target, args, kwargs) - unwrap_fn = lambda elem: elem._tensor if isinstance(elem, MetaTensor) else elem + def unwrap_fn(elem): + + def _convert_meta(t: torch.Tensor): + if t.device == 'meta': + return t + else: + return t.to('meta') + + if isinstance(elem, MetaTensor): + if getattr(self, '_is_param', False): + return torch.nn.Parameter(_convert_meta(elem._tensor)) + return _convert_meta(elem._tensor) + + elif isinstance(elem, torch.Tensor): + if isinstance(elem, torch.nn.Parameter): + return torch.nn.Parameter(_convert_meta(elem)) + return _convert_meta(elem) + + else: + return elem + is_pure_tensor = lambda elem: isinstance(elem, MetaTensor) and not isinstance(elem, torch.nn.Parameter) n_info = MetaInfo(n) n_info.outputs = _normalize_tuple(r) @@ -132,7 +155,11 @@ def run_node(self, n: torch.fx.Node) -> Any: n_info.inputs = tuple(v for v in args if is_pure_tensor(v)) + \ tuple(v for v in kwargs.values() if is_pure_tensor(v)) - n._meta_data = tree_map(unwrap_fn, _normalize_tuple(r)) # align with SPMD + # align with SPMD + if isinstance(r, (tuple, list)): + n._meta_data = tree_map(unwrap_fn, _normalize_tuple(r)) + else: + n._meta_data = unwrap_fn(r) n_info.global_ctx = self.global_hook.ctx n_info.curr_ctx = self.global_hook.ctx.copy() @@ -158,10 +185,48 @@ def call_function(self, target: 'Target', args: Tuple[Any, ...], kwargs: Dict[st Return Any: The value returned by the function invocation """ + convert_to_param = False + if target in (torch.transpose, torch.reshape) and isinstance(args[0], torch.nn.parameter.Parameter): + convert_to_param = True if target in self._custom_dispatch_func: - return self._custom_dispatch_func[target](*args, **kwargs) + res = self._custom_dispatch_func[target](*args, **kwargs) + else: + res = super().call_function(target, args, kwargs) + if convert_to_param: + return torch.nn.Parameter(res) + else: + return res + + def call_method(self, target: 'Target', args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Any: + """ + Execute a ``call_method`` node and return the result. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return + Any: The value returned by the method invocation + """ + # args[0] is the `self` object for this method call + self_obj, *args_tail = args + + target_method = getattr(self_obj.__class__, target) + + convert_to_parameter = False + if target_method in (torch.Tensor.view, torch.Tensor.transpose) and isinstance( + args[0], torch.nn.parameter.Parameter): + convert_to_parameter = True + # Execute the method and return the result + assert isinstance(target, str) + res = getattr(self_obj, target)(*args_tail, **kwargs) + if convert_to_parameter: + return torch.nn.Parameter(res) else: - return super().call_function(target, args, kwargs) + return res def propagate(self, *args, device=None): """ @@ -172,7 +237,14 @@ def propagate(self, *args, device=None): Returns: Any: The value returned from executing the Module """ - wrap_fn = lambda elem: MetaTensor(elem, device=device) + + # wrap_fn = lambda elem: MetaTensor(elem, device=device) + def wrap_fn(elem, device=device): + if isinstance(elem, torch.Tensor): + return MetaTensor(elem, device=device) + else: + return elem + with self._mode: return super().run(*tree_map(wrap_fn, args)) diff --git a/colossalai/_analyzer/fx/tracer/__init__.py b/colossalai/_analyzer/fx/tracer/__init__.py new file mode 100644 index 000000000000..6b1b2256aa44 --- /dev/null +++ b/colossalai/_analyzer/fx/tracer/__init__.py @@ -0,0 +1,2 @@ +from .bias_addition import * +from .custom_leaf_module import * diff --git a/colossalai/_analyzer/fx/bias_addition.py b/colossalai/_analyzer/fx/tracer/bias_addition.py similarity index 98% rename from colossalai/_analyzer/fx/bias_addition.py rename to colossalai/_analyzer/fx/tracer/bias_addition.py index 5359752d4cb4..1e75b47ca5b0 100644 --- a/colossalai/_analyzer/fx/bias_addition.py +++ b/colossalai/_analyzer/fx/tracer/bias_addition.py @@ -4,11 +4,10 @@ """ import torch -import torch.nn as nn import torch.nn.functional as F from torch.nn.modules.utils import _pair, _single, _triple -from .symbolic_trace import register_tracer_impl +from .tracer import register_tracer_impl __all__ = [] diff --git a/colossalai/_analyzer/fx/tracer/custom_leaf_module.py b/colossalai/_analyzer/fx/tracer/custom_leaf_module.py new file mode 100644 index 000000000000..112c7c9637d2 --- /dev/null +++ b/colossalai/_analyzer/fx/tracer/custom_leaf_module.py @@ -0,0 +1,29 @@ +import torch + +from .tracer import register_leaf_module, register_leaf_module_impl + +try: + import apex + register_leaf_module(apex.normalization.FusedLayerNorm) + register_leaf_module(apex.normalization.FusedRMSNorm) + register_leaf_module(apex.normalization.MixedFusedLayerNorm) + register_leaf_module(apex.normalization.MixedFusedRMSNorm) + + @register_leaf_module_impl(apex.normalization.FusedLayerNorm) + @register_leaf_module_impl(apex.normalization.FusedRMSNorm) + @register_leaf_module_impl(apex.normalization.MixedFusedLayerNorm) + @register_leaf_module_impl(apex.normalization.MixedFusedRMSNorm) + def torch_nn_normalize(self, input: torch.Tensor): + # check shape + if isinstance(self, torch.nn.BatchNorm1d): + assert input.dim() in [2, 3] + elif isinstance(self, torch.nn.BatchNorm2d): + assert input.dim() == 4 + elif isinstance(self, torch.nn.BatchNorm3d): + assert input.dim() == 5 + + # normalization maintain the same shape as the input + return input.clone() + +except (ImportError, AttributeError): + pass diff --git a/colossalai/_analyzer/fx/tracer/proxy.py b/colossalai/_analyzer/fx/tracer/proxy.py new file mode 100644 index 000000000000..ce379efdcf0d --- /dev/null +++ b/colossalai/_analyzer/fx/tracer/proxy.py @@ -0,0 +1,112 @@ +import operator +from typing import Any, Callable, Dict, Optional, Set, Union + +import torch +import torch.nn as nn +from torch.fx import Graph, Node, Proxy, Tracer +from torch.fx.graph import _Namespace +from torch.utils._pytree import tree_map + +from colossalai._analyzer._subclasses import MetaTensor + +Target = Union[Callable[..., Any], str] + + +class ColoProxy(Proxy): + _func_dispatch: Dict[Target, Callable[..., Any]] = {} + + def __init__(self, *args, data=None, **kwargs): + super().__init__(*args, **kwargs) + self._meta_data = data + + @property + def meta_data(self): + return self._meta_data + + @meta_data.setter + def meta_data(self, args): + wrap_fn = lambda x: MetaTensor(x) if isinstance(x, torch.Tensor) else x + self._meta_data = tree_map(wrap_fn, args) + + @classmethod + def __torch_function__(cls, orig_method, types, args=(), kwargs=None): + kwargs = {} if kwargs is None else kwargs + if orig_method in cls._func_dispatch: + impl = cls._func_dispatch.pop(orig_method) # avoid recursion + proxy = impl(*args, **kwargs) + cls._func_dispatch[orig_method] = impl + return proxy + else: + proxy = cls.from_torch_proxy(super().__torch_function__(orig_method, types, args, kwargs)) + unwrap_fn = lambda p: p.meta_data if isinstance(p, ColoProxy) else p + if proxy.meta_data is None: + proxy.meta_data = orig_method(*tree_map(unwrap_fn, args), **tree_map(unwrap_fn, kwargs)) + return proxy + + @classmethod + def from_torch_proxy(cls, proxy: Proxy): + return cls(proxy.node, proxy.tracer) + + def __repr__(self): + return f"ColoProxy({self.node.name}, meta_data={self.meta_data})" + + def __len__(self): + return len(self.meta_data) + + def __int__(self): + return int(self.meta_data) + + def __index__(self): + try: + return int(self.meta_data) + except: + return torch.zeros(self.meta_data.shape, dtype=torch.bool).numpy().__index__() + + def __float__(self): + return float(self.meta_data) + + def __bool__(self): + return self.meta_data + + def __getattr__(self, k): + return ColoAttribute(self, k, getattr(self._meta_data, k, None)) + + def __setitem__(self, key, value): + proxy = self.tracer.create_proxy('call_function', operator.setitem, (self, key, value), {}) + proxy.meta_data = self._meta_data + return proxy + + def __contains__(self, key): + if self.node.op == "placeholder": + # this is used to handle like + # if x in kwargs + # we don't handle this case for now + return False + return super().__contains__(key) + + def __isinstancecheck__(self, type): + return isinstance(self.meta_data, type) + + +class ColoAttribute(ColoProxy): + + def __init__(self, root, attr: str, data=None): + self.root = root + self.attr = attr + self.tracer = root.tracer + self._meta_data = data + self._node: Optional[Node] = None + + @property + def node(self): + # the node for attributes is added lazily, since most will just be method calls + # which do not rely on the getitem call + if self._node is None: + self._node = self.tracer.create_proxy('call_function', getattr, (self.root, self.attr), {}).node + return self._node + + def __call__(self, *args, **kwargs): + return self.tracer.create_proxy('call_method', self.attr, (self.root,) + args, kwargs) + + def __repr__(self): + return f"ColoAttribute({self.node.name}, attr={self.attr})" diff --git a/colossalai/_analyzer/fx/tracer/symbolic_trace.py b/colossalai/_analyzer/fx/tracer/symbolic_trace.py new file mode 100644 index 000000000000..2018863f6f5f --- /dev/null +++ b/colossalai/_analyzer/fx/tracer/symbolic_trace.py @@ -0,0 +1,157 @@ +from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Type, Union + +import torch +from torch.fx import Tracer +from torch.utils._pytree import tree_map + +from colossalai._analyzer._subclasses import MetaTensor + +try: + from ..codegen import ActivationCheckpointCodeGen + SUPPORT_ACTIVATION = True +except: + SUPPORT_ACTIVATION = False +from ..graph_module import ColoGraphModule +from .tracer import ColoTracer + + +def _default_device(): + return torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') + + +def _current_device(module: torch.nn.Module): + try: + return next(module.parameters()).device + except: + return _default_device() + + +def symbolic_trace( + root: Union[torch.nn.Module, Callable[..., Any]], + concrete_args: Optional[Dict[str, Any]] = None, + meta_args: Optional[Dict[str, Any]] = None, + trace_act_ckpt: bool = False, + bias_addition_split: bool = False, +) -> ColoGraphModule: + """ + Traces a ``torch.nn.Module`` or a function and returns a ``GraphModule`` with ``Node``s and ``MetaInfo`` + attached to the ``Node``s. + + Can be used to trace the usage of ``torch.utils.checkpoint`` and the path of module + (https://github.com/pytorch/examples/blob/main/fx/module_tracer.py). + + This tracer is able to trace basic control flow and for loops. + + It will split the bias addition into two parts if ``bias_addition_split`` is set to be ``True``. + (See ./bias_addition.py for more details). + + Examples: + 1. Tracing a ``torch.nn.Module`` with control flow. + + .. code-block:: python + + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def forward(self, x): + if x.size(0) > 1: + x = x.sum(dim=0) + return self.linear(x) + + traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}) + + # traced code like: + # def forward(self, x): + # linear_1 = self.linear(x) + # return linear_1 + + traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(2, 2, 2)}) + + # traced code like: + # def forward(self, x): + # sum = x.sum(dim=0); x = None + # linear = self.linear(sum); sum = None + # return linear + + 2. Tracing a ``torch.nn.Module`` with ``torch.utils.checkpoint``. + + .. code-block:: python + + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def forward(self, x): + def custom_forward(x): + return self.linear(x) + return torch.utils.checkpoint.checkpoint(custom_forward, x) + + traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}, trace_act_ckpt=True) + + # traced code like: + # def checkpoint_0(self, x): + # linear = self.linear(x); x = None + # return linear + # + # def forward(self, x): + # linear = torch.utils.checkpoint.checkpoint(checkpoint_0, x); x = None + # return linear + + 3. Tracing a ``torch.nn.Module`` with ``bias_addition_split``. + + .. code-block:: python + + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2, bias=True) + + def forward(self, x): + return self.linear(x) + + traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}, bias_addition_split=True) + + # traced code like: + # def forward(self, x): + # linear_bias = self.linear.bias + # linear_weight = self.linear.weight + # linear = torch._C._nn.linear(x, linear_weight); x = linear_weight = None + # add = linear + linear_bias; linear = linear_bias = None + # return add + + Args: + root (Union[torch.nn.Module, Callable[..., Any]]): The ``torch.nn.Module`` or function to be traced. + concrete_args (Optional[Dict[str, Any]], optional): Concrete arguments to be passed to the ``root``. + Defaults to {}. + meta_args (Optional[Dict[str, Any]], optional): Meta arguments to be passed to the ``root``. Mostly used + for tracing control flow. Defaults to {}. + trace_act_ckpt (bool, optional): Whether to trace the usage of ``torch.utils.checkpoint``. + Defaults to False. + bias_addition_split (bool, optional): Whether to split the bias addition into two parts. Defaults to False. + + Returns: + ColoGraphModule: A traced ``GraphModule`` that is ready for activation checkpoint ``CodeGen``. + + Remarks: + This part of ``symbolic_trace()`` is maintained by Colossal-AI team. If you encountered + any unexpected error during tracing, feel free to raise an issue on Colossal-AI GitHub + repo. We welcome any feedback and contributions to enhance the extensibility of + Colossal-AI. + """ + if meta_args: + device, orig_device = _default_device(), _current_device(root) + wrap_fn = lambda elem: MetaTensor(elem, device=device) if isinstance(elem, torch.Tensor) else elem + graph = ColoTracer(trace_act_ckpt=trace_act_ckpt, + bias_addition_split=bias_addition_split).trace(root.to(device), + concrete_args=concrete_args, + meta_args=tree_map(wrap_fn, meta_args)) + if trace_act_ckpt and SUPPORT_ACTIVATION: + graph.set_codegen(ActivationCheckpointCodeGen()) + root.to(orig_device) + else: + graph = Tracer().trace(root, concrete_args=concrete_args) + name = root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__ + return ColoGraphModule(root, graph, name) diff --git a/colossalai/_analyzer/fx/symbolic_trace.py b/colossalai/_analyzer/fx/tracer/tracer.py similarity index 53% rename from colossalai/_analyzer/fx/symbolic_trace.py rename to colossalai/_analyzer/fx/tracer/tracer.py index 5d858c87a3c8..6958a00a6a72 100644 --- a/colossalai/_analyzer/fx/symbolic_trace.py +++ b/colossalai/_analyzer/fx/tracer/tracer.py @@ -1,28 +1,19 @@ import functools import inspect -import operator from contextlib import contextmanager -from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Type, Union +from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Type, Union import torch import torch.nn as nn from torch.fx import Graph, Node, Proxy, Tracer -from torch.fx.graph import _Namespace from torch.utils._pytree import tree_map -from colossalai._analyzer._subclasses import MetaTensor, _TensorPropertyMethod, _TorchFactoryMethod +from colossalai._analyzer._subclasses import _TensorPropertyMethod, _TorchFactoryMethod -from .codegen import ActivationCheckpointCodeGen -from .graph_module import ColoGraphModule -from .node_util import MetaInfo +from ..node_util import MetaInfo +from .proxy import ColoProxy Target = Union[Callable[..., Any], str] -Argument = Optional[Union[Tuple[Any, ...], # actually Argument, but mypy can't represent recursive types - List[Any], # actually Argument - Dict[str, Any], # actually Argument - slice, # Slice[Argument, Argument, Argument], but slice is not a templated type in typing - 'Node',]] -zeros = torch.zeros def _truncate_suffix(s: str): @@ -32,17 +23,6 @@ def _truncate_suffix(s: str): return re.sub(r'_\d+$', '', s) -def _default_device(): - return torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') - - -def _current_device(module): - try: - return next(module.parameters()).device - except: - return _default_device() - - def register_tracer_impl(func: Callable[..., Any], name: Optional[str] = '_custom_impl'): def wrapper(impl): @@ -70,149 +50,6 @@ def register_non_leaf_module(module: nn.Module): ColoTracer._custom_non_leaf_module.add(module) -class ColoProxy(Proxy): - _func_dispatch: Dict[Target, Callable[..., Any]] = {} - - def __init__(self, *args, data=None, **kwargs): - super().__init__(*args, **kwargs) - self._meta_data = data - - @property - def meta_data(self): - return self._meta_data - - @meta_data.setter - def meta_data(self, args): - wrap_fn = lambda x: MetaTensor(x) if isinstance(x, torch.Tensor) else x - self._meta_data = tree_map(wrap_fn, args) - - @classmethod - def __torch_function__(cls, orig_method, types, args=(), kwargs=None): - kwargs = {} if kwargs is None else kwargs - if orig_method in cls._func_dispatch: - impl = cls._func_dispatch.pop(orig_method) # avoid recursion - proxy = impl(*args, **kwargs) - cls._func_dispatch[orig_method] = impl - return proxy - else: - proxy = cls.from_torch_proxy(super().__torch_function__(orig_method, types, args, kwargs)) - unwrap_fn = lambda p: p.meta_data if isinstance(p, ColoProxy) else p - if proxy.meta_data is None: - proxy.meta_data = orig_method(*tree_map(unwrap_fn, args), **tree_map(unwrap_fn, kwargs)) - return proxy - - @classmethod - def from_torch_proxy(cls, proxy: Proxy): - return cls(proxy.node, proxy.tracer) - - def __repr__(self): - return f"ColoProxy({self.node.name}, meta_data={self.meta_data})" - - def __len__(self): - return len(self.meta_data) - - def __int__(self): - return int(self.meta_data) - - def __index__(self): - try: - return int(self.meta_data) - except: - return zeros(self.meta_data.shape, dtype=torch.bool).numpy().__index__() - - def __float__(self): - return float(self.meta_data) - - def __bool__(self): - return self.meta_data - - def __getattr__(self, k): - return ColoAttribute(self, k, getattr(self._meta_data, k, None)) - - def __setitem__(self, key, value): - proxy = self.tracer.create_proxy('call_function', operator.setitem, (self, key, value), {}) - proxy.meta_data = self._meta_data - return proxy - - def __contains__(self, key): - if self.node.op == "placeholder": - # this is used to handle like - # if x in kwargs - # we don't handle this case for now - return False - return super().__contains__(key) - - def __isinstancecheck__(self, type): - return isinstance(self.meta_data, type) - - def size(self, dim=None): - if self._meta_data is None: - return self._meta_data.size(*[dim] if dim else []) - return self.tracer.create_proxy('call_method', 'size', (self, dim) if dim else (self,), {}) - - def dim(self): - if self._meta_data is not None: - return self._meta_data.dim() - return self.tracer.create_proxy('call_method', 'dim', (self,), {}) - - @property - def shape(self): - if self._meta_data is not None: - return self._meta_data.shape - return self.tracer.create_proxy('call_function', getattr, (self, 'shape'), {}) - - @property - def ndim(self): - if self._meta_data is not None: - return self._meta_data.ndim - return self.tracer.create_proxy('call_function', getattr, (self, 'ndim'), {}) - - @property - def device(self): - if self._meta_data is not None: - return self._meta_data.device - return self.tracer.create_proxy('call_function', getattr, (self, 'device'), {}) - - @property - def dtype(self): - if self._meta_data is not None: - return self._meta_data.dtype - return self.tracer.create_proxy('call_function', getattr, (self, 'dtype'), {}) - - def to(self, *args, **kwargs): - return self.tracer.create_proxy('call_method', 'to', (self, *args), {**kwargs}) - - def cpu(self, *args, **kwargs): - return self.tracer.create_proxy('call_method', 'cpu', (self, *args), {**kwargs}) - - def cuda(self, *args, **kwargs): - return self.tracer.create_proxy('call_method', 'cuda', (self, *args), {**kwargs}) - - -class ColoAttribute(ColoProxy): - - def __init__(self, root, attr: str, data=None): - self.root = root - self.attr = attr - self.tracer = root.tracer - self._meta_data = data - self._node: Optional[Node] = None - - @property - def node(self): - # the node for attributes is added lazily, since most will just be method calls - # which do not rely on the getitem call - if self._node is None: - self._node = self.tracer.create_proxy('call_function', getattr, (self.root, self.attr), {}).node - return self._node - - def __call__(self, *args, **kwargs): - return self.tracer.create_proxy('call_method', self.attr, (self.root,) + args, kwargs) - - def __repr__(self): - return f"ColoAttribute({self.node.name}, attr={self.attr})" - - class ColoTracer(Tracer): _custom_leaf_module: Set[Type[nn.Module]] = set() _custom_leaf_module_impl: Dict[Type[nn.Module], Callable[..., Any]] = {} @@ -249,7 +86,6 @@ def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: # we will enter the module and split the bias-addition ops if self.bias_addition_split and type(m) in self._bias_addition_module and m.bias is not None: return False - # user can specify which modules are leaf modules and which are not return (type(m) not in self._custom_non_leaf_module and (type(m) in self._custom_leaf_module or super().is_leaf_module(m, module_qualified_name))) @@ -306,29 +142,39 @@ def create_proxy(self, mod = self.root.get_submodule(target) self.disable_module_getattr = True try: - proxy.meta_data = self._custom_leaf_module_impl.get(type(mod), - mod.forward)(*tree_map(unwrap_fn, args), - **tree_map(unwrap_fn, kwargs)) + args = tree_map(unwrap_fn, args) + kwargs = tree_map(unwrap_fn, kwargs) + if type(mod) in self._custom_leaf_module: + target = self._custom_leaf_module_impl[type(mod)] + proxy.meta_data = target(mod, *args, **kwargs) + else: + proxy.meta_data = mod.forward(*args, **kwargs) finally: self.disable_module_getattr = False return proxy def create_node(self, *args, **kwargs) -> Node: node = super().create_node(*args, **kwargs) - n_info = MetaInfo(node, mod_dir=self.mod_dir, to_recompute=tuple(self.ckpt_regions)) + n_info = MetaInfo(node, mod_dir=self.mod_dir, activation_checkpoint=tuple(self.ckpt_regions)) return node def trace(self, root: torch.nn.Module, - concrete_args: Optional[Dict[str, torch.Tensor]] = {}, - meta_args: Optional[Dict[str, torch.Tensor]] = {}) -> Graph: + concrete_args: Optional[Dict[str, torch.Tensor]] = None, + meta_args: Optional[Dict[str, torch.Tensor]] = None) -> Graph: + + if meta_args is None: + meta_args = {} + + if concrete_args is None: + concrete_args = {} # check concrete and meta args have valid names sig = inspect.signature(root.forward) sig_names = set(sig.parameters.keys()) meta_arg_names = set(meta_args.keys()) concrete_arg_names = set(concrete_args.keys()) - + non_concrete_arg_names = sig_names - concrete_arg_names # update concrete args with default values for k, v in sig.parameters.items(): if k in sig_names - meta_arg_names and \ @@ -352,6 +198,34 @@ def _check_arg_name_valid(names: Iterable[str]): self.graph = super().trace(root, concrete_args=concrete_args) self.mod_dir = '' self.graph.lint() + + for node in self.graph.nodes: + if node.op == "placeholder": + # Removing default values for inputs as the forward pass will fail with them. + if node.target in non_concrete_arg_names: + node.args = () + # Without this, torch.jit.script fails because the inputs type is Optional[torch.Tensor]. + # It cannot infer on the attributes and methods the input should have, and fails. + node.type = torch.Tensor + # It is a concrete arg so it is not used and should be removed. + else: + if hasattr(torch.fx._symbolic_trace, "_assert_is_none"): + # Newer versions of torch.fx emit an assert statement + # for concrete arguments; delete those before we delete + # the concrete arg. + to_delete = [] + for user in node.users: + if user.target == torch.fx._symbolic_trace._assert_is_none: + to_delete.append(user) + for user in to_delete: + self.graph.erase_node(user) + + self.graph.erase_node(node) + + # TODO: solves GraphModule creation. + # Without this, return type annotation "Tuple" is causing code execution failure. + if node.op == "output": + node.type = None return self.graph @contextmanager @@ -454,7 +328,7 @@ def _post_check(self, non_concrete_arg_names: Set[str]): if node.op == "output": node.type = None self.graph.lint() - + def getattr(self, attr, attr_val, parameter_proxy_cache): return self._module_getattr(attr, attr_val, parameter_proxy_cache) @@ -487,134 +361,3 @@ def maybe_get_proxy_for_attr(attr_val, collection_to_search, parameter_proxy_cac return maybe_parameter_proxy return attr_val - - -def symbolic_trace( - root: Union[torch.nn.Module, Callable[..., Any]], - concrete_args: Optional[Dict[str, Any]] = {}, - meta_args: Optional[Dict[str, Any]] = {}, - trace_act_ckpt: bool = False, - bias_addition_split: bool = False, -) -> ColoGraphModule: - """ - Traces a ``torch.nn.Module`` or a function and returns a ``GraphModule`` with ``Node``s and ``MetaInfo`` - attached to the ``Node``s. - - Can be used to trace the usage of ``torch.utils.checkpoint`` and the path of module - (https://github.com/pytorch/examples/blob/main/fx/module_tracer.py). - - This tracer is able to trace basic control flow and for loops. - - It will split the bias addition into two parts if ``bias_addition_split`` is set to be ``True``. - (See ./bias_addition.py for more details). - - Examples: - 1. Tracing a ``torch.nn.Module`` with control flow. - - .. code-block:: python - - class MyModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(2, 2) - - def forward(self, x): - if x.size(0) > 1: - x = x.sum(dim=0) - return self.linear(x) - - traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}) - - # traced code like: - # def forward(self, x): - # linear_1 = self.linear(x) - # return linear_1 - - traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(2, 2, 2)}) - - # traced code like: - # def forward(self, x): - # sum = x.sum(dim=0); x = None - # linear = self.linear(sum); sum = None - # return linear - - 2. Tracing a ``torch.nn.Module`` with ``torch.utils.checkpoint``. - - .. code-block:: python - - class MyModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(2, 2) - - def forward(self, x): - def custom_forward(x): - return self.linear(x) - return torch.utils.checkpoint.checkpoint(custom_forward, x) - - traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}, trace_act_ckpt=True) - - # traced code like: - # def checkpoint_0(self, x): - # linear = self.linear(x); x = None - # return linear - # - # def forward(self, x): - # linear = torch.utils.checkpoint.checkpoint(checkpoint_0, x); x = None - # return linear - - 3. Tracing a ``torch.nn.Module`` with ``bias_addition_split``. - - .. code-block:: python - - class MyModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(2, 2, bias=True) - - def forward(self, x): - return self.linear(x) - - traced = symbolic_trace(MyModule(), meta_args={'x': torch.randn(1, 2, 2)}, bias_addition_split=True) - - # traced code like: - # def forward(self, x): - # linear_bias = self.linear.bias - # linear_weight = self.linear.weight - # linear = torch._C._nn.linear(x, linear_weight); x = linear_weight = None - # add = linear + linear_bias; linear = linear_bias = None - # return add - - Args: - root (Union[torch.nn.Module, Callable[..., Any]]): The ``torch.nn.Module`` or function to be traced. - concrete_args (Optional[Dict[str, Any]], optional): Concrete arguments to be passed to the ``root``. - Defaults to {}. - meta_args (Optional[Dict[str, Any]], optional): Meta arguments to be passed to the ``root``. Mostly used - for tracing control flow. Defaults to {}. - trace_act_ckpt (bool, optional): Whether to trace the usage of ``torch.utils.checkpoint``. - Defaults to False. - bias_addition_split (bool, optional): Whether to split the bias addition into two parts. Defaults to False. - - Returns: - ColoGraphModule: A traced ``GraphModule`` that is ready for activation checkpoint ``CodeGen``. - - Remarks: - This part of ``symbolic_trace()`` is maintained by Colossal-AI team. If you encountered - any unexpected error during tracing, feel free to raise an issue on Colossal-AI GitHub - repo. We welcome any feedback and contributions to enhance the extensibility of - Colossal-AI. - """ - if meta_args: - device, orig_device = _default_device(), _current_device(root) - wrap_fn = lambda elem: MetaTensor(elem, device=device) if isinstance(elem, torch.Tensor) else elem - graph = ColoTracer(trace_act_ckpt=trace_act_ckpt, - bias_addition_split=bias_addition_split).trace(root.to(device), - concrete_args=concrete_args, - meta_args=tree_map(wrap_fn, meta_args)) - if trace_act_ckpt: - graph.set_codegen(ActivationCheckpointCodeGen()) - root.to(orig_device) - else: - graph = Tracer().trace(root, concrete_args=concrete_args) - name = root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__ - return ColoGraphModule(root, graph, name) diff --git a/colossalai/amp/__init__.py b/colossalai/amp/__init__.py index 16da81f23898..963215476b6b 100644 --- a/colossalai/amp/__init__.py +++ b/colossalai/amp/__init__.py @@ -1,14 +1,16 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from .amp_type import AMP_TYPE -from colossalai.context import Config import torch.nn as nn -from torch.optim import Optimizer from torch.nn.modules.loss import _Loss -from .torch_amp import convert_to_torch_amp +from torch.optim import Optimizer + +from colossalai.context import Config + +from .amp_type import AMP_TYPE from .apex_amp import convert_to_apex_amp from .naive_amp import convert_to_naive_amp +from .torch_amp import convert_to_torch_amp __all__ = ['convert_to_amp', 'convert_to_naive_amp', 'convert_to_apex_amp', 'convert_to_torch_amp', 'AMP_TYPE'] diff --git a/colossalai/amp/naive_amp/mixed_precision_mixin/__init__.py b/colossalai/amp/naive_amp/mixed_precision_mixin/__init__.py new file mode 100644 index 000000000000..b0348e1477bb --- /dev/null +++ b/colossalai/amp/naive_amp/mixed_precision_mixin/__init__.py @@ -0,0 +1,9 @@ +from .base import MixedPrecisionMixin +from .bf16 import BF16MixedPrecisionMixin +from .fp16 import FP16MixedPrecisionMixin + +__all__ = [ + 'MixedPrecisionMixin', + 'FP16MixedPrecisionMixin', + 'BF16MixedPrecisionMixin', +] diff --git a/colossalai/amp/naive_amp/mixed_precision_mixin/base.py b/colossalai/amp/naive_amp/mixed_precision_mixin/base.py new file mode 100644 index 000000000000..a52a9747ad1e --- /dev/null +++ b/colossalai/amp/naive_amp/mixed_precision_mixin/base.py @@ -0,0 +1,91 @@ +from abc import ABC, abstractmethod + +import torch +from torch import Tensor + + +class MixedPrecisionMixin(ABC): + """A helper class for mixed precision training. This mixin is used in mixed precision optimizers. + + Attributes: + dtype (torc.dtype): The expected dtype of the gradients. + + Examples: + ```python + class MyMixedPrecisionOptimizer(OptimizerWrapper): + def __init__(self, optim: Optimizer): + super().__init__(optim) + self.mixed_precision = MixedPrecisionMixin() + + def backward(self, loss): + loss = self.mixed_precision.pre_backward(loss) + loss.backward() + + def backward_by_grad(self, tensor, grad): + grad = self.mixed_precision.pre_backward_by_grad(tensor, grad) + tensor.backward(grad) + + def step(self): + if self.mixed_precision.should_skip_step(): + self.zero_grad() + return + div_scale = self.mixed_precision.get_grad_div_scale() + # maybe clip grad here + # maybe scale grad here + self.optim.step() + + def zero_grad(self): + self.mixed_precision.pre_zero_grad() + return self.optim.zero_grad() + ``` + """ + dtype: torch.dtype + + @abstractmethod + def pre_backward(self, loss: Tensor) -> Tensor: + """Called before backward. + + Args: + loss (Tensor): Loss value. + + Returns: + Tensor: Loss value (possibly scaled). + """ + pass + + @abstractmethod + def pre_backward_by_grad(self, tensor: Tensor, grad: Tensor) -> Tensor: + """Called before backward by grad. This is helpful for pipeline parallelism. + + Args: + tensor (Tensor): Tensor to backward. + grad (Tensor): Gradient of the tensor. + + Returns: + Tensor: Gradient of the tensor (possibly scaled). + """ + pass + + @abstractmethod + def should_skip_step(self) -> bool: + """Called before step. + + Returns: + bool: Whether to skip the step. + """ + pass + + @abstractmethod + def pre_zero_grad(self) -> None: + """Called before zero_grad. + """ + pass + + @abstractmethod + def get_grad_div_scale(self) -> float: + """Called before step or clip_grad. To keep computation efficiency, this method does not (maybe) unscale grads. + + Returns: + float: A divisor for gradient clipping or step. + """ + pass diff --git a/colossalai/amp/naive_amp/mixed_precision_mixin/bf16.py b/colossalai/amp/naive_amp/mixed_precision_mixin/bf16.py new file mode 100644 index 000000000000..9454f6eb8413 --- /dev/null +++ b/colossalai/amp/naive_amp/mixed_precision_mixin/bf16.py @@ -0,0 +1,23 @@ +import torch +from torch import Tensor + +from .base import MixedPrecisionMixin + + +class BF16MixedPrecisionMixin(MixedPrecisionMixin): + dtype = torch.bfloat16 + + def pre_backward(self, loss: Tensor) -> Tensor: + return loss + + def pre_backward_by_grad(self, tensor: Tensor, grad: Tensor) -> Tensor: + return grad + + def should_skip_step(self) -> bool: + return False + + def pre_zero_grad(self) -> None: + pass + + def get_grad_div_scale(self) -> float: + return 1.0 diff --git a/colossalai/amp/naive_amp/mixed_precision_mixin/fp16.py b/colossalai/amp/naive_amp/mixed_precision_mixin/fp16.py new file mode 100644 index 000000000000..1ce8e42eb3ed --- /dev/null +++ b/colossalai/amp/naive_amp/mixed_precision_mixin/fp16.py @@ -0,0 +1,84 @@ +from abc import abstractmethod +from enum import Enum + +import torch +import torch.distributed as dist +from torch import Tensor + +from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler +from colossalai.utils import get_current_device + +from .base import MixedPrecisionMixin + + +class OptimState(Enum): + SCALED = 0 + UNSCALED = 1 + + +class FP16MixedPrecisionMixin(MixedPrecisionMixin): + dtype = torch.float16 + + def __init__(self, + initial_scale: float = 2**16, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32) -> None: + super().__init__() + self.grad_scaler = DynamicGradScaler(initial_scale=initial_scale, + min_scale=min_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + max_scale=max_scale) + self.optim_state = OptimState.UNSCALED + self.found_overflow = torch.zeros(1, dtype=torch.float, device=get_current_device()) + + @property + def loss_scale(self) -> float: + return self.grad_scaler.scale.item() + + @abstractmethod + def check_local_overflow(self) -> bool: + """Check whether there is overflow in the local process. This method should be implemented by subclasses. + + Returns: + bool: Whether there is overflow in the local process. + """ + pass + + def check_overflow(self) -> bool: + # clear previous overflow record + self.found_overflow.fill_(0.0) + if self.check_local_overflow(): + self.found_overflow.fill_(1.0) + dist.all_reduce(self.found_overflow, op=dist.ReduceOp.MAX) + return self.found_overflow.item() > 0 + + def pre_backward(self, loss: Tensor) -> Tensor: + loss = self.loss_scale * loss + self.optim_state = OptimState.SCALED + return loss + + def pre_backward_by_grad(self, tensor: Tensor, grad: Tensor) -> Tensor: + self.optim_state = OptimState.SCALED + return grad + + def should_skip_step(self) -> bool: + found_inf = self.check_overflow() + self.grad_scaler.update(found_inf) + if found_inf: + self.optim_state = OptimState.UNSCALED + return found_inf + + def pre_zero_grad(self) -> None: + pass + + def get_grad_div_scale(self) -> float: + assert self.optim_state == OptimState.SCALED, 'grads should be scaled before clipping' + self.optim_state = OptimState.UNSCALED + return self.loss_scale diff --git a/colossalai/amp/torch_amp/_grad_scaler.py b/colossalai/amp/torch_amp/_grad_scaler.py index 7b78998fb8c2..ed4b8e484436 100644 --- a/colossalai/amp/torch_amp/_grad_scaler.py +++ b/colossalai/amp/torch_amp/_grad_scaler.py @@ -240,7 +240,7 @@ def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16): for grads in per_dtype_grads.values(): torch._amp_foreach_non_finite_check_and_unscale_(grads, per_device_found_inf.get(device), per_device_inv_scale.get(device)) - # For tensor parallel paramters it should be all-reduced over tensor parallel process group + # For tensor parallel parameters it should be all-reduced over tensor parallel process group if gpc.is_initialized(ParallelMode.MODEL) and gpc.get_world_size(ParallelMode.MODEL) > 1: vals = [val for val in per_device_found_inf._per_device_tensors.values()] coalesced = _flatten_dense_tensors(vals) diff --git a/colossalai/auto_parallel/README.md b/colossalai/auto_parallel/README.md index 8e47e1bb0b4a..f011ec8ccbd7 100644 --- a/colossalai/auto_parallel/README.md +++ b/colossalai/auto_parallel/README.md @@ -16,8 +16,8 @@ A *symbolic profiler* for collecting computing and memory overhead related to st ### Solver **Solver** is designed to find the optimal execution plan for a given computation graph and cluster in two stages: -1) *Intra-op parallelism stage* is to find the plan with the minimum total execution time of all nodes with respect to the constraint of the memory budget. The optimaztion goal of intra-op parallelism solver is modified from Alpa 's intra-op parallelsim ILP solver. -2) *Activation checkpoint stage* is to search for the fastest execution plan that meets the memory budget on the computation graph after inserting the communication nodes by the intra-op parallelism stage. The algorithm to find optimial activation checkpoint is modified from Rotor . The reason we use two-stage optimization is that if the two tasks are formulated together, the solving time will be significantly increased, which will greatly affect the user experience of the system. On the contrary, solving in two hierarchical levels has many advantages. Firstly, compared with the computation graph with activation checkpointing, the original graph has fewer nodes, which can reduce the solving cost of intra-op parallelism solver. In addition, a more optimal solution can be found by adding the communication overhead into the activation checkpoint modeling. +1) *Intra-op parallelism stage* is to find the plan with the minimum total execution time of all nodes with respect to the constraint of the memory budget. The optimization goal of intra-op parallelism solver is modified from Alpa 's intra-op parallelism ILP solver. +2) *Activation checkpoint stage* is to search for the fastest execution plan that meets the memory budget on the computation graph after inserting the communication nodes by the intra-op parallelism stage. The algorithm to find optimal activation checkpoint is modified from Rotor . The reason we use two-stage optimization is that if the two tasks are formulated together, the solving time will be significantly increased, which will greatly affect the user experience of the system. On the contrary, solving in two hierarchical levels has many advantages. Firstly, compared with the computation graph with activation checkpointing, the original graph has fewer nodes, which can reduce the solving cost of intra-op parallelism solver. In addition, a more optimal solution can be found by adding the communication overhead into the activation checkpoint modeling. ### Generator **Generator** applies the searched execution plan to the computation graph and recompiles the computation graph to optimized PyTorch code. It has *a series compile pass* to insert a communication node or do the kernel substitution as the intra-op parallelism solver required. Additionally, we implement a *code generation* feature to recognize the annotation from the activation checkpoint solver and inject the activation checkpoint block following annotation instructions. diff --git a/colossalai/auto_parallel/meta_profiler/__init__.py b/colossalai/auto_parallel/meta_profiler/__init__.py index bfd36195149b..3741d8e5a8ad 100644 --- a/colossalai/auto_parallel/meta_profiler/__init__.py +++ b/colossalai/auto_parallel/meta_profiler/__init__.py @@ -1,3 +1,3 @@ from .meta_registry import * -from .metainfo import * from .registry import meta_register +from .shard_metainfo import * diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py index faeed9f29e61..0f2e9e44f91c 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import ewise_flop_counter as elementwise_flop_counter +from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import elementwise_flop_counter from ..registry import meta_register diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py b/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py index 281a92c0d4f1..e451748512b9 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..constants import BCAST_FUNC_OP, NO_SAVE_ACTIVATION from ..registry import meta_register @@ -17,7 +17,7 @@ def binary_elementwise_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, Train """Meta information generator for binary elementwise operations NOTE: Some of the binary elementwise operations will discard the input activation after computation, as they don't need those tensors for back propagation, for example, if there are two tensors being sent for `torch.add`, - they will be discarded right after add operation is done. We create a simple API in `MetaInfo` class to identify + they will be discarded right after add operation is done. We create a simple API in `ShardMetaInfo` class to identify this behavior, it is critical for better memory estimation. Returns: diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/conv.py b/colossalai/auto_parallel/meta_profiler/meta_registry/conv.py index d1bb6e7fa798..4336bf68363c 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/conv.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/conv.py @@ -2,6 +2,8 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( MemoryCost, OperationData, @@ -10,8 +12,6 @@ StrategiesVector, TrainCycleItem, ) -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from colossalai.tensor.sharding_spec import ShardingSpec from ..registry import meta_register @@ -110,18 +110,18 @@ def convnd_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L # calculate memory cost # TODO: use profiler to check conv temp memory # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_memory_cost = MemoryCost( - activation=activation_size([input_tensor, output_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]) if has_bias else activation_size(weight_tensor), - temp=0, - buffer=0) - - bwd_memory_cost = MemoryCost( - activation=activation_size([input_tensor, weight_tensor, bias_tensor]) - if has_bias else activation_size([input_tensor, weight_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]) if has_bias else activation_size(weight_tensor), - temp=0, - buffer=0) + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]) + if has_bias else compute_size_in_bytes(weight_tensor), + temp=0, + buffer=0) + + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor]) + if has_bias else compute_size_in_bytes([input_tensor, weight_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]) + if has_bias else compute_size_in_bytes(weight_tensor), + temp=0, + buffer=0) # total cost is the sum of forward and backward cost total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation, diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/embedding.py b/colossalai/auto_parallel/meta_profiler/meta_registry/embedding.py index 2997f31adff8..d5d80f5b3700 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/embedding.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/embedding.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..registry import meta_register @@ -34,11 +34,11 @@ def embedding_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem # NOTE: during the backward phase of torch.nn.Embedding, it seems when the input is large enough, it will # have a temp memory which is kind of weird and we don't know the reason yet, so currently we just assume # that there will be no temp memory, as the temp memory is significantly smaller than the gradient memory - fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]), + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]), parameter=0, temp=0, buffer=0) - bwd_memory_cost = MemoryCost(activation=activation_size([weight_tensor]), parameter=0, temp=0, buffer=0) + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([weight_tensor]), parameter=0, temp=0, buffer=0) total_memory_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/linear.py b/colossalai/auto_parallel/meta_profiler/meta_registry/linear.py index 617375721222..94dd9143e0ae 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/linear.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/linear.py @@ -3,6 +3,8 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( MemoryCost, OperationData, @@ -11,8 +13,6 @@ StrategiesVector, TrainCycleItem, ) -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from colossalai.tensor.sharding_spec import ShardingSpec from ..registry import meta_register @@ -112,14 +112,14 @@ def linear_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L # NOTE: Linear don't have buffer and temp in forward and backward phase # the forward activation cost is the size of output_tensor, parameter cost is the size of weight_tensor and bias_tensor # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), temp=0, buffer=0) # the backward activation cost is the size of input_tensor, weight_tensor and bias_tensor, parameter cost is 0 - bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor, bias_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), temp=0, buffer=0) @@ -148,14 +148,14 @@ def linear_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L # NOTE: Linear don't have buffer and temp in forward and backward phase # the forward activation cost is the size of output_tensor, parameter cost is the size of weight_tensor # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]), - parameter=activation_size(weight_tensor), + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]), + parameter=compute_size_in_bytes(weight_tensor), temp=0, buffer=0) # the backward activation cost is the size of input_tensor and weight_tensor, parameter cost is 0 - bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor]), - parameter=activation_size(weight_tensor), + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor]), + parameter=compute_size_in_bytes(weight_tensor), temp=0, buffer=0) @@ -210,48 +210,48 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L # Check dimension if all(len(tensor.shape) == 1 for tensor in input_tensors): # Dot - fwd_compute_cost = flop_mapping[torch.ops.aten.dot.default](input_tensors, output_tensors) + fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](input_tensors, output_tensors) bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor](input_tensors[0], output_tensors) * 2 - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0) elif len(input_tensors[0].shape) >= 2 and len(input_tensors[1].shape) == 1: # gemv case 1: matrix-vector multiplication # & # batched gemv case 1: batched matrix-vector multiplication - fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default]( + fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default]( [input_tensors[0].reshape(-1, input_tensors[0].shape[-1]), input_tensors[1]], output_tensors) # combine the dimensions of output bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor]( [output_tensors[0].reshape(-1), input_tensors[1]], output_tensors) + \ - flop_mapping[torch.ops.aten.mv.default]( + flop_mapping[torch.ops.aten.matmul.default]( [input_tensors[0].reshape(-1, input_tensors[0].shape[-1]).transpose(0, 1), output_tensors[0].reshape(-1)], output_tensors) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0) elif len(input_tensors[0].shape) == 1 and len(input_tensors[1].shape) == 2: # gemv case 2: vector-matrix multiplication - fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default](input_tensors, output_tensors) + fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](input_tensors, output_tensors) bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor]([output_tensors[0], input_tensors[0]], output_tensors) + \ - flop_mapping[torch.ops.aten.mv.default]([input_tensors[1], output_tensors[0]], output_tensors) + flop_mapping[torch.ops.aten.matmul.default]([input_tensors[1], output_tensors[0]], output_tensors) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, - temp=activation_size(input_tensors[1]), + temp=compute_size_in_bytes(input_tensors[1]), buffer=0) elif len(input_tensors[0].shape) == 1 and len(input_tensors[1].shape) >= 3: # batched gemv case 2: vector-batched matrix multiplication - fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default]( + fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default]( [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2]), input_tensors[0]], [output_tensors[0].reshape(-1)]) @@ -260,15 +260,15 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L [output_tensors[0].reshape(-1), input_tensors[0]], output_tensors ) + \ - flop_mapping[torch.ops.aten.mv.default]( + flop_mapping[torch.ops.aten.matmul.default]( [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2]).transpose(0, 1), output_tensors[0].reshape(-1)], output_tensors ) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors + [input_tensors[1]])) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors[0]), + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors + [input_tensors[1]])) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors[0]), parameter=0, - temp=activation_size(input_tensors[1]), + temp=compute_size_in_bytes(input_tensors[1]), buffer=0) elif len(input_tensors[0].shape) >= 2 and len(input_tensors[1].shape) == 2: @@ -287,8 +287,8 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L [input_tensors[0].reshape(-1, input_tensors[0].shape[-1])] ) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0) elif len(input_tensors[0].shape) == 2 and len(input_tensors[1].shape) >= 3: # batched gemm case 2: matrix-batched matrix multiplication @@ -306,11 +306,12 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2])] ) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors) + activation_size(input_tensors[1]), - temp=activation_size(output_tensors)) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors[0]), + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors) + + compute_size_in_bytes(input_tensors[1]), + temp=compute_size_in_bytes(output_tensors)) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors[0]), parameter=0, - temp=activation_size(input_tensors[1]) + activation_size(output_tensors)) + temp=compute_size_in_bytes(input_tensors[1]) + compute_size_in_bytes(output_tensors)) elif all(len(tensor.shape) >= 3 for tensor in input_tensors): # Batched matrix-batched matrix multiplication @@ -324,7 +325,7 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L else: _is_batch_dims_same = False - # retireve dimensions + # retrieve dimensions input_dim_00 = input_tensors[0].shape[-2] input_dim_01 = input_tensors[0].shape[-1] input_dim_10 = input_tensors[1].shape[-2] @@ -351,8 +352,8 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L [input_tensors[0].reshape(-1, input_dim_00, input_dim_01)] ) - fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors)) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors)) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors)) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors)) else: # Case 2: batch dimensions are different @@ -381,10 +382,10 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L ) fwd_mem_cost = MemoryCost( - activation=activation_size([output_tensors[0], extended_input_0, extended_input_1])) - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors) - - activation_size([extended_input_0, extended_input_1]), - temp=activation_size([extended_input_0, extended_input_1])) + activation=compute_size_in_bytes([output_tensors[0], extended_input_0, extended_input_1])) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors) - + compute_size_in_bytes([extended_input_0, extended_input_1]), + temp=compute_size_in_bytes([extended_input_0, extended_input_1])) # compute cost compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/non_spmd.py b/colossalai/auto_parallel/meta_profiler/meta_registry/non_spmd.py index 4634d3ccdcfd..12874810b13e 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/non_spmd.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/non_spmd.py @@ -4,8 +4,6 @@ import torch from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..registry import meta_register diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py b/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py index 3a1db396e188..b872fdc8bdcd 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/norm.py @@ -2,6 +2,8 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( MemoryCost, OperationData, @@ -10,8 +12,6 @@ StrategiesVector, TrainCycleItem, ) -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from colossalai.tensor.sharding_spec import ShardingSpec from ..registry import meta_register @@ -77,17 +77,18 @@ def batchnormnd_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleIt # calculate memory cost # the fwd activation cost is output plus saved mean and saved inv std # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, mean_tensor, var_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes( + [input_tensor, output_tensor, mean_tensor, var_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), temp=0, - buffer=activation_size([mean_tensor, var_tensor])) + buffer=compute_size_in_bytes([mean_tensor, var_tensor])) # the bwd memory cost is quite tricky here, BatchNorm will remove saved mean # and saved inv std during backward phase - bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), - temp=activation_size([mean_tensor, var_tensor]), - buffer=activation_size([mean_tensor, var_tensor])) + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), + temp=compute_size_in_bytes([mean_tensor, var_tensor]), + buffer=compute_size_in_bytes([mean_tensor, var_tensor])) # total cost is the sum of forward and backward cost total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation, @@ -131,15 +132,16 @@ def layernorm_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem # memory cost # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, weight_tensor, bias_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), + fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes( + [input_tensor, output_tensor, weight_tensor, bias_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), temp=0, - buffer=activation_size([running_mean, running_var])) + buffer=compute_size_in_bytes([running_mean, running_var])) - bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor, bias_tensor]), - parameter=activation_size([weight_tensor, bias_tensor]), - temp=activation_size([running_mean, running_var]), - buffer=activation_size([running_mean, running_var])) + bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor]), + parameter=compute_size_in_bytes([weight_tensor, bias_tensor]), + temp=compute_size_in_bytes([running_mean, running_var]), + buffer=compute_size_in_bytes([running_mean, running_var])) total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation, parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter, diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py b/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py index 21272ea09ac1..d785dfcca9ba 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..registry import meta_register @@ -52,8 +52,8 @@ def avgpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost) # calculate memory cost - fwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=activation_size(output_tensor)) - bwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=activation_size(input_tensor)) + fwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=compute_size_in_bytes(output_tensor)) + bwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=compute_size_in_bytes(input_tensor)) # total cost total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation) @@ -114,11 +114,11 @@ def maxpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, # calculate memory cost # NOTE: the index matrix will be discarded in backward phase # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_mem_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, index_matrix])) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor, index_matrix])) # temp memory for backward is the index matrix to be discarded - bwd_mem_cost = MemoryCost(activation=activation_size(input_tensor) - activation_size(index_matrix), - temp=activation_size(index_matrix)) + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensor) - compute_size_in_bytes(index_matrix), + temp=compute_size_in_bytes(index_matrix)) # total cost total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation, temp=bwd_mem_cost.temp) diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py index 332e649d2d7e..97fe3c6196f5 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..registry import meta_register @@ -35,11 +35,11 @@ def meta_func(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[tor # memory costs # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward - fwd_mem_cost = MemoryCost(activation=activation_size(outputs) * 2, parameter=0, temp=0, buffer=0) + fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(outputs) * 2, parameter=0, temp=0, buffer=0) - bwd_mem_cost = MemoryCost(activation=activation_size(outputs) * bwd_mem_out_factor, + bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(outputs) * bwd_mem_out_factor, parameter=0, - temp=activation_size(outputs) * bwd_mem_tmp_factor, + temp=compute_size_in_bytes(outputs) * bwd_mem_tmp_factor, buffer=0) total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation, diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/where.py b/colossalai/auto_parallel/meta_profiler/meta_registry/where.py index c67eb40bc80e..5cba1b5b6e2b 100644 --- a/colossalai/auto_parallel/meta_profiler/meta_registry/where.py +++ b/colossalai/auto_parallel/meta_profiler/meta_registry/where.py @@ -2,9 +2,9 @@ import torch +from colossalai._analyzer._subclasses.flop_tensor import flop_mapping +from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem -from colossalai.fx.profiler.memory_utils import activation_size -from colossalai.fx.profiler.opcount import flop_mapping from ..registry import meta_register diff --git a/colossalai/auto_parallel/meta_profiler/metainfo.py b/colossalai/auto_parallel/meta_profiler/shard_metainfo.py similarity index 70% rename from colossalai/auto_parallel/meta_profiler/metainfo.py rename to colossalai/auto_parallel/meta_profiler/shard_metainfo.py index 218187768a7b..0eee908b48b7 100644 --- a/colossalai/auto_parallel/meta_profiler/metainfo.py +++ b/colossalai/auto_parallel/meta_profiler/shard_metainfo.py @@ -15,11 +15,11 @@ from .constants import INPLACE_MODULE, INPLACE_OPS, NO_SAVE_ACTIVATION from .registry import meta_register -__all__ = ['MetaInfo'] +__all__ = ['ShardMetaInfo'] -class MetaInfo: - """MetaInfo class +class ShardMetaInfo: + """ShardMetaInfo class This class is used to store meta info based on sharding strategy and the given target function. """ @@ -46,9 +46,9 @@ def __init__(self, strategy: ShardingStrategy = None, target: Callable = None) - # target function self._target = target - # compute metainfo if possible + # compute shard_metainfo if possible if self._strategy is not None and self._target is not None: - self.compute_metainfo() + self.compute_shard_metainfo() @property def strategy(self) -> ShardingStrategy: @@ -62,24 +62,38 @@ def target(self) -> Callable: def strategy(self, strategy: ShardingStrategy) -> None: self._strategy = strategy if self._strategy is not None and self._target is not None: - self.compute_metainfo() + self.compute_shard_metainfo() @target.setter def target(self, target: Callable) -> None: self._target = target if self._strategy is not None and self._target is not None: - self.compute_metainfo() + self.compute_shard_metainfo() - def compute_sharded_opdata(self, operation_data: OperationData, sharding_spec: ShardingSpec) -> torch.Tensor: + def compute_sharded_opdata(self, operation_data: OperationData, sharding_spec: ShardingSpec): """ Compute sharded opdata based on the given data and sharding spec. """ - return OperationData(name=operation_data.name, - data=torch.zeros(sharding_spec.get_sharded_shape_per_device(), device="meta"), - type=operation_data.type, - logical_shape=operation_data.logical_shape) - def compute_metainfo(self): + if isinstance(sharding_spec, ShardingSpec): + op_data = OperationData(name=operation_data.name, + data=torch.zeros(sharding_spec.get_sharded_shape_per_device(), device="meta"), + type=operation_data.type, + logical_shape=operation_data.logical_shape) + elif isinstance(sharding_spec, (list, tuple)): + data = operation_data.data + assert isinstance(data, (list, tuple)), f"Data Should be list or tuple, but got {type(data)}." + assert len(data) == len(sharding_spec), f"Length of data and sharding spec should be the same." + sharded_data = [] + for d, s in zip(data, sharding_spec): + sharded_data.append(torch.zeros(s.get_sharded_shape_per_device(), device="meta")) + op_data = OperationData(name=operation_data.name, data=sharded_data, type=operation_data.type) + else: + raise ValueError(f"Sharding spec should be ShardingSpec or list, but got {type(sharding_spec)}.") + + return op_data + + def compute_shard_metainfo(self): """ Compute meta info based on sharding strategy and the given target function. """ diff --git a/colossalai/auto_parallel/offload/__init__.py b/colossalai/auto_parallel/offload/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/colossalai/auto_parallel/offload/amp_optimizer.py b/colossalai/auto_parallel/offload/amp_optimizer.py new file mode 100644 index 000000000000..a79e5006e7d2 --- /dev/null +++ b/colossalai/auto_parallel/offload/amp_optimizer.py @@ -0,0 +1,177 @@ +from typing import Dict, Tuple +from enum import Enum +import torch +from torch.optim import Optimizer + +from colossalai.logging import get_dist_logger +from colossalai.nn.optimizer import ColossalaiOptimizer +from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler +from colossalai.utils import get_current_device + +from .base_offload_module import BaseOffloadModule +from .region_manager import RegionManager +from .region import Region + + +class OptimState(Enum): + SCALED = 0 + UNSCALED = 1 + +class AMPOptimizer(ColossalaiOptimizer): + + """ + A wrapper for Optimizer. + Code reference: https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/optimizer/zero_optimizer.py + + Args: + optimizer (Optimizer): An Optimizer instance. + module (BaseOffloadModule): A ``BaseOffloadModule`` instance. + initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**16. + growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2. + backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5. + growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000. + hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2. + min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. + max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32. + norm_type (float, optional): norm_type used for `clip_grad_norm`. + """ + + def __init__(self, + optimizer: Optimizer, + module: BaseOffloadModule, + initial_scale: float = 2**16, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + min_scale: float = 1, + max_scale: float = 2**32, + clipping_norm: float = 0.0, + norm_type: float = 2.0): + + super().__init__(optimizer) + + self.module = module + self.optim_state = OptimState.UNSCALED + self.clipping_flag = clipping_norm > 0.0 + self.max_norm = clipping_norm + + self.region_manager: RegionManager = self.module.region_manager + self.param_to_range: Dict[torch.nn.Parameter, Tuple[int, int]] = dict() + self.param_to_region: Dict[torch.nn.Parameter, Region] = dict() + + self.fp32_to_fp16_params: Dict[torch.Tensor, torch.nn.Parameter] = dict() + + if self.clipping_flag: + assert norm_type == 2.0, "AMPOptimizer only supports L2 norm now" + + self.__init__optimizer() + + # Grad scaler + self.grad_scaler = DynamicGradScaler(initial_scale=initial_scale, + min_scale=min_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + max_scale=max_scale) + self._found_overflow: torch.Tensor = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + self._logger = get_dist_logger() + + def _set_grad_ptr(self): + for group in self.param_groups: + for fake_param in group['params']: + region = self.param_to_region[fake_param] + begin, end = self.param_to_range[fake_param] + + fake_param.data = region.cpu_grad[begin:end] + fake_param.grad = fake_param.data + fake_param.data = region.fp32_data[begin:end] + + def _update_fp16_params(self): + none_tensor = torch.empty([0]) + for group in self.param_groups: + for fake_param in group['params']: + assert fake_param.grad is None + fake_param.data = none_tensor + self.param_to_region[fake_param].cpu_grad = None + + def _check_overflow(self): + # clear previous overflow record + self._found_overflow.fill_(self.module.overflow_counter.item()) + return self._found_overflow.item() > 0 + + def _get_combined_scale(self): + loss_scale = 1 + + if self.optim_state == OptimState.SCALED: + loss_scale = self.loss_scale + self.optim_state = OptimState.UNSCALED + + combined_scale = loss_scale + + if combined_scale == 1: + return -1 + else: + return combined_scale + + @property + def loss_scale(self): + return self.grad_scaler.scale.item() + + def zero_grad(self, *args, **kwargs): + self.module.overflow_counter = torch.cuda.IntTensor([0]) + return self.optim.zero_grad(set_to_none=True) + + def step(self, *args, **kwargs): + # Copy gradients from model params to main params. + self._set_grad_ptr() + + found_inf = self._check_overflow() + if found_inf: + self.optim_state = OptimState.UNSCALED # no need to unscale grad + self.grad_scaler.update(found_inf) # update gradient scaler + self._logger.info(f'Found overflow. Skip step') + self.zero_grad() # reset all gradients + self._update_fp16_params() + return + + # get combined scale. combined scale = loss scale * clipping norm + # so that gradient = gradient / combined scale + combined_scale = self._get_combined_scale() + self.grad_scaler.update(found_inf) + + ret = self.optim.step(div_scale=combined_scale, *args, **kwargs) + self.zero_grad() + self._update_fp16_params() + return ret + + def clip_grad_norm(self, model: torch.nn.Module, max_norm: float, norm_type: float = 2.0): + raise NotImplementedError + + def backward(self, loss: torch.Tensor): + loss = self.loss_scale * loss + self.optim_state = OptimState.SCALED + self.module.backward(loss) + + def __init__optimizer(self): + + for group in self.optim.param_groups: + fake_params_list = list() + + for param in group['params']: + region = self.region_manager.get_region(param) + fake_param = torch.nn.Parameter(torch.empty([0])) + self.param_to_range[fake_param] = region.param_to_range[param] + self.param_to_region[fake_param] = region + fake_params_list.append(fake_param) + + # Reset existing state dict key to the new main param. + if param in self.optim.state: + self.optim.state[fake_param] = self.optim.state.pop(param) + + group['params'] = fake_params_list + + # Leverage state_dict() and load_state_dict() to + # recast preexisting per-param state tensors + self.optim.load_state_dict(self.optim.state_dict()) \ No newline at end of file diff --git a/colossalai/auto_parallel/offload/base_offload_module.py b/colossalai/auto_parallel/offload/base_offload_module.py new file mode 100644 index 000000000000..d0c328e134ff --- /dev/null +++ b/colossalai/auto_parallel/offload/base_offload_module.py @@ -0,0 +1,107 @@ +from functools import partial +from typing import Optional, Set + +import torch +import torch.nn as nn + +from colossalai.nn.parallel.data_parallel import _cast_float +from colossalai.zero.legacy.gemini.tensor_utils import free_storage + +from .region_manager import RegionManager +from .util import GlobalRuntimeInfo + + +class BaseOffloadModule: + """ + BaseOffloadModule: A model wrapper for parameter offloading. + + Args: + model (nn.Module): model to apply offloading. + region_manager (RegionManager): a ``RegionManager`` instance. + is_sync (bool): synchronous mode or not. + """ + + def __init__(self, model: nn.Module, region_manager: RegionManager, is_sync=True): + + self.model = model + self.region_manager = region_manager + self.grad_hook_list = [] + self.overflow_counter = torch.cuda.IntTensor([0]) + + self.grad_offload_stream = torch.cuda.current_stream() if is_sync else GlobalRuntimeInfo.d2h_stream + + self._cast_buffers() + + def register_grad_hook(self): + for p in self.model.parameters(): + if p.requires_grad: + self.grad_hook_list.append(p.register_hook(partial(self.grad_handle, p))) + + def remove_grad_hook(self): + for hook in self.grad_hook_list: + hook.remove() + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + def _pre_forward(self): + self.register_grad_hook() + for region in self.region_manager.region_list: + region.cpu_grad = None + + def forward(self, *args, **kwargs): + args, kwargs = _cast_float(args, torch.half), _cast_float(kwargs, torch.half) + self.model.zero_grad(set_to_none=True) + self._pre_forward() + outputs = self.model(*args, **kwargs) + return outputs + + def backward(self, loss): + loss.backward() + self._post_backward() + + def _post_backward(self): + torch.cuda.synchronize() + self.remove_grad_hook() + + for p in self.model.parameters(): + p.grad = None + + GlobalRuntimeInfo().fwd_prefetch_event_map.clear() + GlobalRuntimeInfo().bwd_prefetch_event_map.clear() + + def grad_handle(self, p, grad): + empty_grad = torch.empty_like(grad) + free_storage(empty_grad) + with torch._C.DisableTorchFunction(): + region = self.region_manager.get_region(p) + region.copy_grad_to_region_slice(p, grad) + if region.can_release: + self.overflow_counter += region.has_inf_or_nan + master_stream = torch.cuda.current_stream() + with torch.cuda.stream(self.grad_offload_stream): + GlobalRuntimeInfo().d2h_stream.wait_stream(master_stream) + region.move_grad_to_cpu() + return empty_grad + + def _cast_buffers(self): + for buffer in self.model.buffers(): + buffer.data = buffer.cuda() + + def parameters(self, recurse: bool = True): + return self.model.parameters(recurse) + + def named_parameters(self, prefix: str = '', recurse: bool = True): + return self.model.named_parameters(prefix, recurse) + + def named_buffers(self, prefix: str = '', recurse: bool = True): + return self.model.named_buffers(prefix, recurse) + + def named_children(self): + return self.model.named_children() + + def named_modules(self, + memo: Optional[Set[torch.nn.Module]] = None, + prefix: str = '', + remove_duplicate: bool = True): + return self.model.named_modules(memo, prefix, remove_duplicate) diff --git a/colossalai/auto_parallel/offload/mem_optimize.py b/colossalai/auto_parallel/offload/mem_optimize.py new file mode 100644 index 000000000000..d56166dea982 --- /dev/null +++ b/colossalai/auto_parallel/offload/mem_optimize.py @@ -0,0 +1,52 @@ +from typing import Dict + +import torch +import torch.fx +from torch.fx import GraphModule +from torch.utils._pytree import tree_map + +from colossalai.fx import ColoTracer, is_compatible_with_meta +from colossalai.fx.passes.meta_info_prop import MetaInfoProp + +from .base_offload_module import BaseOffloadModule +from .region_manager import RegionManager +from .runtime import runtime_asyn_offload_apply_pass, runtime_syn_offload_apply_pass +from .util import GlobalRuntimeInfo, compute_act_peak_mem, compute_max_param_mem, compute_total_param_mem + + +def memory_optimize(model: torch.nn.Module, + inps: Dict[str, torch.Tensor], + memory_budget: float = -1.0, + solver_name: str = 'asyn'): + + model = model.cpu().half() + tracer = ColoTracer() + assert is_compatible_with_meta() + wrap_fn = lambda x: x.to("meta") if isinstance(x, torch.Tensor) else x + meta_args = tree_map(wrap_fn, inps) + graph = tracer.trace(model, meta_args=meta_args) + gm = GraphModule(model, graph, model.__class__.__name__) + interp = MetaInfoProp(gm) + interp.propagate(*meta_args.values()) + + region_manager = RegionManager(graph, solver_name=solver_name, memory_budget=memory_budget) + region_manager._build_regions() + GlobalRuntimeInfo().region_list = region_manager.region_list + + act_peak_mem = compute_act_peak_mem(region_manager.region_list) / 1024**2 + max_param_mem = compute_max_param_mem(region_manager.region_list) / 1024**2 + total_param_mem = compute_total_param_mem(region_manager.region_list) / 1024**2 + print( + f"act_peak_mem={act_peak_mem:.3f} MB | max_param_mem={max_param_mem:.3f} MB | total_param_mem={total_param_mem:.3f}" + ) + + if solver_name == 'syn': + gm = runtime_syn_offload_apply_pass(gm, region_manager.region_list) + elif solver_name == 'asyn': + gm = runtime_asyn_offload_apply_pass(gm, region_manager.region_list) + else: + raise TypeError(f"Unknown solver name {solver_name}!") + + gm.recompile() + optimized_model = BaseOffloadModule(gm, region_manager, solver_name == 'syn') + return optimized_model diff --git a/colossalai/auto_parallel/offload/region.py b/colossalai/auto_parallel/offload/region.py new file mode 100644 index 000000000000..819ffbd96eb1 --- /dev/null +++ b/colossalai/auto_parallel/offload/region.py @@ -0,0 +1,145 @@ +from typing import Dict, List, Tuple + +import torch +from torch.fx import Node + +from colossalai.zero.legacy.gemini.tensor_utils import alloc_storage, free_storage + + +class Region: + """ + Region: A container owning a piece of contiguous nodes in the DNN computing graph. + + Args: + r_id (int): the index of the region in the computing graph. + """ + + def __init__(self, r_id: int = 0) -> None: + self.r_id: int = r_id + self.fp16_params: List[torch.nn.Parameter] = [] + self.param_size: int = 0 + self.shared_rid: int = self.r_id + + self.param_num: int = 0 + self.grad_num: int = 0 + self.fp16_data = None + self.fp32_data = None + self.cpu_grad = None + self.temp_fp32_data = None + self.param_to_range: Dict[torch.nn.Parameter, Tuple[int, int]] = dict() + + self.need_offload: bool = False + self.is_syn: bool = False + self.nodes: List[Node] = [] + self.fwd_prefetch_region = None + self.bwd_prefetch_region = None + + self.in_mem_pool_flag: bool = False + + @property + def can_release(self) -> bool: + """ + Check if the region can be released. + """ + return self.grad_num == self.param_num + + @property + def has_inf_or_nan(self) -> bool: + """ + Check if the grad of the region has inf or nan values on CUDA. + """ + return torch.isinf(self.fp16_data).any() | torch.isnan(self.fp16_data).any() + + def init_param_data(self, pre_alloc_tensor: torch.Tensor = None): + """ + Map the parameters in the region to a contiguous memory space. + """ + + self.fp16_data = torch.zeros(self.param_num, dtype=torch.half, device='cuda') + offset = 0 + for param in self.fp16_params: + param.data = param.data.cuda() + p_num = param.data.numel() + self.fp16_data[offset:offset + p_num].copy_(param.data.flatten()) + param.data = self.fp16_data[offset:offset + p_num].view(param.data.shape) + self.param_to_range[param] = (offset, offset + p_num) + offset += p_num + + self.fp32_data = self.fp16_data.float().cpu().pin_memory() + free_storage(self.fp16_data) + if self.in_mem_pool_flag and pre_alloc_tensor is not None: + self.fp16_data = pre_alloc_tensor + + def move_param_to_cuda(self): + """ + Move parameters from CPU to GPU. + It first moves float32 parameters to GPU and + then transforms float32 parameters to half-precision on the GPU. + The reason is that the performance of precision conversion on the CPU + is much slower than the data transfer overhead. + """ + + self.temp_fp32_data.copy_(self.fp32_data, non_blocking=True) + self.temp_fp32_data.record_stream(torch.cuda.current_stream()) + if not self.in_mem_pool_flag: + alloc_storage(self.fp16_data) + self.fp16_data[:self.param_num].copy_(self.temp_fp32_data) + self.fp16_data.record_stream(torch.cuda.current_stream()) + + self.__update_params_ptr() + + def move_grad_to_cpu(self): + """ + Move gradients from GPU to CPU. + """ + + self.cpu_grad = torch.empty(self.param_num, dtype=torch.half, pin_memory=True) + self.cpu_grad.copy_(self.fp16_data[:self.param_num], non_blocking=True) + self.fp16_data.record_stream(torch.cuda.current_stream()) + if not self.in_mem_pool_flag: + self.free_cuda_data() + + self.grad_num = 0 + + def free_cuda_data(self): + free_storage(self.fp16_data) + + # torch.cuda.empty_cache() + + def copy_grad_to_region_slice(self, param: torch.nn.Parameter, data_slice: torch.Tensor) -> None: + """ + Copy data slice to the memory space indexed by the input tensor in the region. + + Args: + param (torch.nn.Parameter): the param used to retrieve meta information + data_slice (torch.Tensor): the tensor to be copied to the region + """ + + begin, end = self.param_to_range[param] + self.fp16_data[begin:end].copy_(data_slice.data.flatten()) + param.data = self.fp16_data[begin:end].view(param.data.shape) + + self.grad_num += data_slice.numel() + + def split(self, cut_node_idx: int, cut_param_idx: int): + """ + Split the region into two and return the latter. + """ + new_reg = Region(r_id=self.r_id + 1) + new_reg.nodes = self.nodes[cut_node_idx:] + new_reg.fp16_params = self.fp16_params[cut_param_idx:] + for p in new_reg.fp16_params: + new_reg.param_size += p.data.numel() * p.data.element_size() + new_reg.param_num += p.data.numel() + + self.nodes = self.nodes[:cut_node_idx] + self.fp16_params = self.fp16_params[:cut_param_idx] + self.param_size -= new_reg.param_size + self.param_num -= new_reg.param_num + + return new_reg + + def __update_params_ptr(self) -> None: + for param in self.fp16_params: + begin, end = self.param_to_range[param] + param.data = self.fp16_data[begin:end].view(param.data.shape) diff --git a/colossalai/auto_parallel/offload/region_manager.py b/colossalai/auto_parallel/offload/region_manager.py new file mode 100644 index 000000000000..30bfaf00d493 --- /dev/null +++ b/colossalai/auto_parallel/offload/region_manager.py @@ -0,0 +1,526 @@ +from typing import List, Any, Dict, Tuple +import torch +from torch.fx import Graph, Node + +from .solver import SolverFactory +from .training_simulator import TrainingSimulator +from .region import Region +from .util import NodeInfo + + +class RegionManager: + """ + RegionManager is used to construct and manage the offload plan for the model execution. + + Args: + graph (Graph): a Graph object used for analysis and strategy generation. + solver_name (str): a solver name which specifies the preferences for plan searching. + memory_budget (float): the given memory budget. + cnode (List[str], optional): Common node List, should be the subset of input. + """ + + def __init__(self, + graph: Graph, + solver_name: str = 'asyn', + memory_budget: float = -1.0, + cnode: List[str] = None): + + self.graph = graph + assert graph.owning_module is not None, 'The given graph is not associated with a owning_module' + self.root_module = self.graph.owning_module + self.nodes = list(graph.nodes) + self.cnode = cnode + self.only_param_ops = [] + self.param_region_map: Dict[torch.nn.Parameter, Region] = dict() + self.shared_region_pairs: List[Tuple[Region, Region]] = list() + self.region_list: List[Region] = list() + self.rid_in_pool: List[int] = list() + self.mem_block_size: int = 0 + self.memory_budget = memory_budget + + self.solver_name = solver_name + self.require_pool: bool = solver_name == 'asyn' + + self.reg_to_block: Dict[int, int] = dict() + + def _build_regions(self): + """ + 1. Pre-processing, mainly contains linearized computing graph and + merge smaller regions into larger ones. + 2. Construct a solver to search for an efficient offload strategy. + 3. Post-processing, mainly contains early region placement if using asynchronous mode, + and initialize region data. + """ + + self._pre_process() + + solver_cls = SolverFactory.create(self.solver_name) + solver = solver_cls(self.region_list, self.memory_budget) + solver._call_solver() + + self._post_process(solver.best_ts) + + def _pre_process(self): + + init_region_list = self._linearize_graph() + + if len(self.shared_region_pairs) > 1: + raise NotImplementedError( + 'The current version only considers at most one pair of parameter sharing.') + + elif len(self.shared_region_pairs) == 1: + shared_regs = self.shared_region_pairs[0] + assert shared_regs[0].shared_rid == shared_regs[1].r_id \ + and shared_regs[1].shared_rid == shared_regs[0].r_id + fst_id = shared_regs[0].r_id + lst_id = shared_regs[1].r_id + regs_left_out = init_region_list[:fst_id + 1] + regs_right_out = init_region_list[lst_id:] + hold_regs = init_region_list[fst_id + 1:lst_id] + else: + regs_left_out = [] + regs_right_out = [] + hold_regs = init_region_list + + self.mem_block_size = self._search_block_size(hold_regs) + hold_regs = self._merge_small_regions(hold_regs) + + if self.require_pool: + for reg in hold_regs: + reg.in_mem_pool_flag = True + self.rid_in_pool.append(reg.r_id) + + self.region_list.extend(regs_left_out) + self.region_list.extend(hold_regs) + + for reg in regs_right_out: + reg.r_id = self.region_list[-1].r_id + 1 + self.region_list[reg.shared_rid].shared_rid = reg.r_id + self.region_list.append(reg) + + self._process_shared_region() + + self.max_param_num = max([reg.param_num for reg in self.region_list]) + self.memory_budget -= self.max_param_num * torch.tensor([], dtype=torch.float32).element_size() + + def _post_process(self, ts: TrainingSimulator = None): + if self.require_pool: + self._early_region_placement(ts) + self._init_region_data() + + def _early_region_placement(self, ts: TrainingSimulator): + """ + Implemented the early region placement strategy to avoid GPU memory fragmentation. + It maps all region data into a contiguous memory space and + reuses the same memory space for regions that do not coexist. + + Args: + ts (TrainingSimulator): the best training simulator, which records region execution flow. + + Raises: + NotImplementedError: due to the naive implementation, + it may not find a suitable region placement strategy for the given execution flow. + """ + + reg_flow = torch.cat( + [ts.fwd_reg_flow, ts.bwd_reg_flow], dim=0) + mem_block_num = torch.max( + torch.sum(reg_flow[:, self.rid_in_pool], dim=1)) + coexist_matrix = torch.logical_or( + ts.fwd_reg_flow, ts.bwd_reg_flow) + + block_to_regs = {} + for block_idx in range(mem_block_num): + block_to_regs[block_idx] = [] + for reg in self.region_list: + if reg.r_id in self.rid_in_pool: + cur_reg_appears = coexist_matrix[:, reg.r_id] + cur_reg_coexists = torch.sum( + coexist_matrix[cur_reg_appears], dim=0).bool() + for block_idx in range(mem_block_num): + if not any(cur_reg_coexists[block_to_regs[block_idx]]): + block_to_regs[block_idx].append(reg.r_id) + self.reg_to_block[reg.r_id] = block_idx + break + + if reg.r_id not in self.reg_to_block: + raise NotImplementedError( + f'can not find a block from the memory pool to store parameters of the region') + self.memory_pool = torch.chunk(torch.zeros(int( + mem_block_num * self.mem_block_size / 2), dtype=torch.half, device='cuda'), chunks=int(mem_block_num)) + + def _merge_small_regions(self, orig_reg_list: List[Region]) -> List[Region]: + """ + Merge smaller regions into larger ones for better bandwidth utilization and easier management. + It is inspired by Gemini. + + Args: + orig_reg_list (List[Region]): original region list. + + Returns: + List[Region]: region list after merging. + """ + + r_id = orig_reg_list[0].r_id + region = Region(r_id=r_id) + region_list = [region] + + for orig_reg in orig_reg_list: + if region_list[-1].param_size + orig_reg.param_size > self.mem_block_size: + r_id += 1 + region = Region(r_id=r_id) + region_list.append(region) + region.param_size += orig_reg.param_size + region.param_num += orig_reg.param_num + region.nodes.extend(orig_reg.nodes) + region.fp16_params.extend(orig_reg.fp16_params) + self.__update_param_region_map(orig_reg.fp16_params, region) + + return region_list + + def _search_block_size(self, + region_list: List[Region], + search_interval_byte: int = 1024, + search_range_byte: int = 128 * 1024 ** 2) -> int: + """ + Search for a suitable memory block size. + + Args: + region_list (List[Region]): region list. + search_interval_byte (int): searching interval in byte. + search_range_byte (int): searching range in byte. + + Returns: + int: the best memory block size. + """ + + def _get_wasted_mem(size_list: List[int], blk_size: int): + """ + Get wasted byte for a certain block size. + """ + acc_wasted = 0 + left = 0 + for s in size_list: + if left + s > blk_size: + acc_wasted += blk_size - left + left = s + left += s + acc_wasted += blk_size - left + return acc_wasted + + param_size_list = [ + region.param_size for region in region_list if region.r_id == region.shared_rid] + + start_size = max(param_size_list) + min_mem_waste = float('+inf') + best_block_size = start_size + + for block_size in range(start_size, start_size + search_range_byte + 1, search_interval_byte): + temp_waste = 0 + temp_waste += _get_wasted_mem(param_size_list, block_size) + if temp_waste < min_mem_waste: + min_mem_waste = temp_waste + best_block_size = block_size + + return best_block_size + + def _init_region_data(self): + """ + Initialize region data, which maps the parameters in the region to a contiguous memory space. + """ + + self.temp_fp32_data = torch.zeros(self.max_param_num, device='cuda', dtype=torch.float32) + + for region in self.region_list: + pre_alloc_tensor = None + if self.require_pool and region.r_id in self.rid_in_pool: + block_idx = self.reg_to_block[region.r_id] + pre_alloc_tensor = self.memory_pool[block_idx] + + if region.r_id <= region.shared_rid: + region.init_param_data(pre_alloc_tensor) + else: + shared_region = self.region_list[region.shared_rid] + region.fp16_data = shared_region.fp16_data + region.fp32_data = shared_region.fp32_data + region.param_to_range = shared_region.param_to_range + region.temp_fp32_data = self.temp_fp32_data[:region.param_num].detach( + ) + + torch.cuda.empty_cache() + + def _process_shared_region(self): + """ + Special processing for the shared region, which uses GPT2 and Bert case as a priori knowledge. + """ + + if len(self.shared_region_pairs): + assert len(self.shared_region_pairs) <= 1 + former_reg, latter_reg = self.shared_region_pairs[0] + assert latter_reg.param_num >= former_reg.param_num + embedding_node = former_reg.nodes[-1] + assert embedding_node.op == 'call_module' and isinstance( + self.root_module.get_submodule(embedding_node.target), torch.nn.Embedding) + if latter_reg.param_num > former_reg.param_num: + for idx, n in enumerate(latter_reg.nodes): + if (n.op == 'call_module' and isinstance(self.root_module.get_submodule(n.target), + torch.nn.Linear)) or \ + (n.op == 'call_function' and n.target is torch.nn.functional.linear): + cut_node_idx = idx + 1 + break + assert len(latter_reg.fp16_params) == 2 + new_reg = latter_reg.split(cut_node_idx, 1) + for p in new_reg.fp16_params: + self.param_region_map[p] = new_reg + self.region_list.insert(new_reg.r_id, new_reg) + for reg in self.region_list[new_reg.r_id + 1:]: + reg.r_id += 1 + latter_reg.shared_rid = former_reg.r_id + former_reg.shared_rid = latter_reg.r_id + + def _linearize_graph(self) -> List[Region]: + """Linearizing the graph + + Args: + graph (Graph): The computing graph to be optimized. + + Returns: + List[Region]: each region contains the actual 'node' in linearized manner. + + Remarks: + Do merge the inplace ops and shape-consistency ops into the previous node. + """ + + # List of target name that could be seen as common node + common_ops = ["getattr", "getitem", "size"] + + def _is_cop(target: Any) -> bool: + """Check if an op could be seen as common node + + Args: + target (Any): node target + + Returns: + bool + """ + + if isinstance(target, str): + return target in common_ops + else: + return target.__name__ in common_ops + + def _is_act(data: Any) -> bool: + """Check if an op could be seen as parameter computation start + + Args: + data (Any): meta_data + + Returns: + bool + """ + + label = False + if isinstance(data, torch.Tensor): + return True + elif isinstance(data, (tuple, list)): + for d in data: + label = label or _is_act(d) + return label + + def _maybe_param_comp_start() -> bool: + """Check if an op could be seen as parameter computation start + + Args: + n (Node): node + + Returns: + bool + """ + + label = False + if n.op == "get_attr": + label = True + elif n.op == "call_module": + target = n.target + submod = self.root_module.get_submodule(target) + if ( + len(list(submod.named_parameters(recurse=False))) != 0 + or len(list(submod.named_buffers(recurse=False))) != 0 + ): + label = True + + return label and not sum([v for _, v in param_op_deps.items()]) + + def _is_param_comp_end() -> bool: + """Check if an op could be seen as parameter computation end + + Args: + n (Node): node + + Returns: + bool + """ + + def _is_inplace(n: Node): + """Get the inplace argument from ``torch.fx.Node`` + """ + inplace = False + if n.op == "call_function": + inplace = n.kwargs.get("inplace", False) + elif n.op == "call_module": + inplace = getattr(n.graph.owning_module.get_submodule( + n.target), "inplace", False) + return inplace + + label = False + + if n.op == "call_module": + target = n.target + submod = self.root_module.get_submodule(target) + if ( + len(list(submod.named_parameters(recurse=False))) != 0 + or len(list(submod.named_buffers(recurse=False))) != 0 + ): + label = True + + elif n.op == "call_function": + label = any(map(lambda x: x.name in self.only_param_ops, n.all_input_nodes)) and any( + map(lambda x: x.name not in self.only_param_ops and not _is_cop(n.target), n.all_input_nodes)) + + return label and not sum([v for _, v in param_op_deps.items()]) and not any(map(_is_inplace, n.users)) + + def _exception_node_handling(): + # TODO meta info prop bug + if n.name.__contains__("transpose") and n.meta['fwd_out'][0].dim() <= 2: + n.meta['fwd_out'] = [] + + # make sure that item in cnode is valid + if self.cnode: + for name in self.cnode: + try: + assert next(node for node in self.graph.nodes if node.name == name).op == "placeholder", \ + f"Common node {name} is not an input of the model." + except StopIteration: + raise ValueError(f"Common node name {name} not in graph.") + else: + self.cnode = [] + + node_id = 0 + region_id = 0 + + param_op_deps = {} + + deps = {} + region_list = [] + region = Region(r_id=region_id) + + act_n = None + + for n in self.graph.nodes: + if n.op != "placeholder" and n.op != "output": + for n_par in n.all_input_nodes: + if n_par.op != "placeholder" and n_par.name not in self.cnode: + deps[n_par] -= 1 + if n_par.op != "placeholder" and n_par.name in self.only_param_ops: + param_op_deps[n_par] -= 1 + + if act_n in region.nodes and _maybe_param_comp_start(): + ns = [] + border_n_idx = region.nodes.index(act_n) + if border_n_idx < len(region.nodes): + ns = region.nodes[border_n_idx + 1:] + region.nodes = region.nodes[:border_n_idx + 1] + region_list.append(region) + region_id += 1 + region = Region(r_id=region_id) + region.nodes = ns + + _exception_node_handling() + region.nodes.append(n) + self._set_node_and_region_info(node_id, n, region) + node_id += 1 + + # if the node could free all dependencies in graph + # we could begin a new region + if _is_param_comp_end(): + region_list.append(region) + region_id += 1 + region = Region(r_id=region_id) + + # propagate common node attr if possible + if len(n.all_input_nodes) == len([node for node in n.all_input_nodes if node.name in self.cnode + ]) or _is_cop(n.target): + self.cnode.append(n.name) + else: + deps[n] = len( + [user for user in n.users if user.op != "output"]) + + # propagate param node attr if possible + if len(n.all_input_nodes) == len([node for node in n.all_input_nodes if node.name in self.only_param_ops + ]) or n.op == "get_attr": + self.only_param_ops.append(n.name) + param_op_deps[n] = len( + [user for user in n.users if user.op != "output"]) + + # record last activation node + if _is_act(n._meta_data): + act_n = n + + if len(region.nodes): + region_list.append(region) + + return region_list + + def _set_node_and_region_info(self, node_id: int, cur_n: Node, cur_reg: Region): + + cur_n.node_info = NodeInfo(node_id) + + if cur_n.op == 'call_module': + target = cur_n.target + submod = self.root_module.get_submodule(target) + for p in list(submod.parameters(recurse=False)): + + if p in self.param_region_map: + cur_reg.shared_rid = self.param_region_map[p].r_id + self.param_region_map[p].shared_rid = cur_reg.r_id + self.shared_region_pairs.append( + (self.param_region_map[p], cur_reg)) + else: + self.param_region_map[p] = cur_reg + + cur_reg.fp16_params.append(p) + cur_reg.param_num += p.data.numel() + cur_reg.param_size += p.data.numel() * p.data.element_size() + + elif cur_n.op == "get_attr": + attr_itr = self.root_module + atoms = cur_n.target.split(".") + for atom in atoms: + attr_itr = getattr(attr_itr, atom) + + if isinstance(attr_itr, torch.nn.Parameter): + + if attr_itr in self.param_region_map: + cur_reg.shared_rid = self.param_region_map[attr_itr].r_id + self.param_region_map[attr_itr].shared_rid = cur_reg.r_id + self.shared_region_pairs.append( + (self.param_region_map[attr_itr], cur_reg)) + else: + self.param_region_map[attr_itr] = cur_reg + + cur_reg.fp16_params.append(attr_itr) + cur_reg.param_num += attr_itr.data.numel() + cur_reg.param_size += attr_itr.data.numel() * attr_itr.data.element_size() + + def get_region(self, param: torch.nn.Parameter) -> Region: + """ + Return the region owning the parameter. + + Args: + param (torch.nn.Parameter): a torch parameter object + """ + return self.param_region_map[param] + + def __update_param_region_map(self, params: List[torch.nn.Parameter], region: Region): + for p in params: + self.param_region_map[p] = region diff --git a/colossalai/auto_parallel/offload/runtime.py b/colossalai/auto_parallel/offload/runtime.py new file mode 100644 index 000000000000..764ac608826b --- /dev/null +++ b/colossalai/auto_parallel/offload/runtime.py @@ -0,0 +1,256 @@ +from typing import List + +import torch +from torch.fx.node import Node + +from .region import Region +from .util import GlobalRuntimeInfo, requires_upload_p_in_fwd + + +class SynPreFwdPostBwdOP(torch.autograd.Function): + """ + A customized prefetch and offload operation. + + Args: + input_: input tensor. + fwd_info: information dict, which contains region indices + that need to be uploaded or freed during forward pass. + bwd_info: information dict, which contains region indices + that need to be uploaded during backward pass. + """ + + @staticmethod + def forward(ctx, input_, fwd_info, bwd_info): + ctx.bwd_info = bwd_info + d2h_rid = fwd_info.get('d2h_rid', None) + if d2h_rid is not None: + free_region = GlobalRuntimeInfo().region_list[d2h_rid] + assert isinstance(free_region, Region) + free_region.free_cuda_data() + + h2d_rid = fwd_info.get('h2d_rid', None) + if h2d_rid is not None: + h2d_region = GlobalRuntimeInfo().region_list[h2d_rid] + assert isinstance(h2d_region, Region) + h2d_region.move_param_to_cuda() + + return input_ + + @staticmethod + def backward(ctx, grad_output): + + h2d_rid = ctx.bwd_info.get('h2d_rid', None) + if h2d_rid is not None: + pref_region = GlobalRuntimeInfo().region_list[h2d_rid] + assert isinstance(pref_region, Region) + pref_region.move_param_to_cuda() + + return grad_output, None, None + + +class AsynPreFwdPostBwdOP(torch.autograd.Function): + """ + A customized prefetch and offload operation. + + Args: + input_: input tensor. + fwd_info: information dict, which contains region indices + that need to be prefetched, waited, or freed during forward pass. + bwd_info: information dict, which contains region indices + that need to be prefetched or waited during backward pass. + """ + + @staticmethod + def forward(ctx, input_, fwd_info, bwd_info): + ctx.bwd_info = bwd_info + + sync_rid = fwd_info.get('sync_rid', None) + if sync_rid is not None: + prefetch_event = GlobalRuntimeInfo().fwd_prefetch_event_map.get(sync_rid, None) + if prefetch_event: + prefetch_event.wait() + + h2d_rid = fwd_info.get('h2d_rid', None) + if h2d_rid is not None: + pref_region = GlobalRuntimeInfo().region_list[h2d_rid] + assert isinstance(pref_region, Region) + master_stream = torch.cuda.current_stream() + with torch.cuda.stream(GlobalRuntimeInfo().h2d_stream): + GlobalRuntimeInfo().h2d_stream.wait_stream(master_stream) + pref_region.move_param_to_cuda() + + prefetch_event = torch.cuda.Event() + prefetch_event.record(GlobalRuntimeInfo().h2d_stream) + GlobalRuntimeInfo().fwd_prefetch_event_map[h2d_rid] = prefetch_event + + return input_ + + @staticmethod + def backward(ctx, grad_output): + + sync_rid = ctx.bwd_info.get('sync_rid', None) + if sync_rid is not None: + wait_region = GlobalRuntimeInfo().region_list[sync_rid] + assert isinstance(wait_region, Region) + prefetch_event = GlobalRuntimeInfo().bwd_prefetch_event_map.get(sync_rid, None) + if prefetch_event: + prefetch_event.wait() + else: + wait_region.move_param_to_cuda() + + h2d_rid = ctx.bwd_info.get('h2d_rid', None) + if h2d_rid is not None: + pref_region = GlobalRuntimeInfo().region_list[h2d_rid] + assert isinstance(pref_region, Region) + master_stream = torch.cuda.current_stream() + with torch.cuda.stream(GlobalRuntimeInfo().h2d_stream): + GlobalRuntimeInfo().h2d_stream.wait_stream(master_stream) + pref_region.move_param_to_cuda() + + prefetch_event = torch.cuda.Event() + prefetch_event.record(GlobalRuntimeInfo().h2d_stream) + GlobalRuntimeInfo().bwd_prefetch_event_map[h2d_rid] = prefetch_event + return grad_output, None, None + + +def convert_fwd_upload_bwd_offload_to_action(tensor, fwd_info, bwd_info): + ''' + Convert Upload and Offload operation into runtime action. + + Argument: + tensor(torch.Tensor): input tensor. + fwd_info(dict): information dict, which contains region indices + that need to be uploaded, or freed during forward pass. + bwd_info(dict): information dict, which contains region indices + that need to be uploaded during backward pass. + ''' + with torch._C.DisableTorchFunction(): + ret = SynPreFwdPostBwdOP.apply(tensor, fwd_info, bwd_info) + return ret + + +def convert_fwd_prefetch_bwd_offload_to_action(tensor, fwd_info, bwd_info): + ''' + Convert Prefetch and Offload operation into runtime action. + + Argument: + tensor(torch.Tensor): input tensor. + fwd_info(dict): information dict, which contains region indices + that need to be prefetched, waited, or freed during forward pass. + bwd_info(dict): information dict, which contains region indices + that need to be prefetched or waited during backward pass. + ''' + with torch._C.DisableTorchFunction(): + ret = AsynPreFwdPostBwdOP.apply(tensor, fwd_info, bwd_info) + return ret + + +def replace_node_users(orig_node: Node, inserted_node: Node, rep_user_nodes: List[Node] = None): + user_list = list(orig_node.users.keys()) + if rep_user_nodes is not None: + user_list = rep_user_nodes + for user in user_list: + if user == inserted_node: + continue + new_args = list(user.args) + new_kwargs = dict(user.kwargs) + # the origin node may be a positional argument or key word argument of user node + if orig_node in new_args: + # substitute the origin node with offload_apply_node + new_args[new_args.index(orig_node)] = inserted_node + user.args = tuple(new_args) + elif str(orig_node) in new_kwargs: + # substitute the origin node with offload_apply_node + new_kwargs[str(orig_node)] = inserted_node + user.kwargs = new_kwargs + + +def runtime_syn_offload_apply_pass(gm: torch.fx.GraphModule, region_list: List[Region]): + """ + This pass is used to add the synchronous upload and offload spec apply node to the origin graph. + """ + mod_graph = gm.graph + last_inp_node = tuple(mod_graph.nodes)[0] + + for r_idx, region in enumerate(region_list): + # forward upload + fwd_info = {} + if requires_upload_p_in_fwd(region_list[region.shared_rid]): + fwd_info['h2d_rid'] = region.r_id + + # forward offload + if r_idx > 0 and region_list[r_idx - 1].need_offload: + fwd_info['d2h_rid'] = r_idx - 1 + + bwd_info = {} + # backward upload + if r_idx > 0 and region_list[r_idx - 1].need_offload: + bwd_info['h2d_rid'] = region_list[r_idx - 1].r_id + + if fwd_info or bwd_info: + with mod_graph.inserting_after(last_inp_node): + new_node = mod_graph.create_node('call_function', + convert_fwd_upload_bwd_offload_to_action, + args=(last_inp_node, fwd_info, bwd_info)) + replace_node_users(last_inp_node, new_node) + + last_inp_node = region.nodes[-1] + + return gm + + +def runtime_asyn_offload_apply_pass(gm: torch.fx.GraphModule, region_list: List[Region]): + """ + This pass is used to add the asynchronous prefetch and offload spec apply node to the origin graph. + """ + mod_graph = gm.graph + + # upload parameters of the first region + last_inp_node = tuple(mod_graph.nodes)[0] + first_region_with_p = [region for region in region_list if region.param_size][0] + fwd_info = {"h2d_rid": first_region_with_p.r_id} + with mod_graph.inserting_after(last_inp_node): + upload_apply_node = mod_graph.create_node('call_function', + convert_fwd_upload_bwd_offload_to_action, + args=(last_inp_node, fwd_info, {})) + replace_node_users(last_inp_node, upload_apply_node) + last_inp_node = upload_apply_node + + for r_idx, region in enumerate(region_list): + # forward prefetch + fwd_info = {} + if region.param_size: + fwd_info['sync_rid'] = region.r_id + fwd_prefetch_region = region.fwd_prefetch_region + if fwd_prefetch_region and requires_upload_p_in_fwd(region_list[fwd_prefetch_region.shared_rid]): + fwd_info['h2d_rid'] = fwd_prefetch_region.r_id + + # forward offload + if r_idx > 0 and region_list[r_idx - 1].need_offload: + fwd_info['d2h_rid'] = r_idx - 1 + + bwd_info = {} + # backward prefetch + if r_idx > 0 and region_list[r_idx - 1].need_offload: + bwd_info['sync_rid'] = r_idx - 1 + if r_idx > 0 and region_list[r_idx - 1].bwd_prefetch_region: + bwd_info['h2d_rid'] = region_list[r_idx - 1].bwd_prefetch_region.r_id + + if fwd_info or bwd_info: + with mod_graph.inserting_after(last_inp_node): + new_node = mod_graph.create_node('call_function', + convert_fwd_prefetch_bwd_offload_to_action, + args=(last_inp_node, fwd_info, bwd_info)) + replace_node_users(last_inp_node, new_node) + + last_inp_node = region.nodes[-1] + + if region.bwd_prefetch_region: + bwd_info = {'h2d_rid': region.bwd_prefetch_region.r_id} + with mod_graph.inserting_after(last_inp_node): + new_node = mod_graph.create_node('call_function', + convert_fwd_prefetch_bwd_offload_to_action, + args=(last_inp_node, {}, bwd_info)) + replace_node_users(last_inp_node, new_node) + # gm.graph.print_tabular() + return gm diff --git a/colossalai/auto_parallel/offload/solver.py b/colossalai/auto_parallel/offload/solver.py new file mode 100644 index 000000000000..161f7ff86898 --- /dev/null +++ b/colossalai/auto_parallel/offload/solver.py @@ -0,0 +1,523 @@ +import time +from typing import List, Dict, Type +from abc import ABC, abstractmethod + +NOT_NVML = False +try: + from pynvml import * +except: + NOT_NVML = True + +import torch +from torch.fx.node import Node +from colossalai.utils.cuda import get_current_device + +from .training_simulator import TrainingSimulator, SynTrainingSimulator, AsynTrainingSimulator +from .region import Region +from .util import NodeInfo, NvDevicePower + + +def benchmark_func(func, number=1, repeat=1, warmup=3): + """ + benchmark data transfer cost. + """ + + for i in range(warmup): + func() + + costs = [] + + for i in range(repeat): + torch.cuda.synchronize() + begin = time.time() + for i in range(number): + func() + torch.cuda.synchronize() + costs.append((time.time() - begin) / number) + + return sum(costs) / len(costs) + + +class Solver(ABC): + """ + The parameter offload solver. + + Args: + region_list (List[Region]): represents the linearized DNN computing graph. + memory_budget (float): the given memory budget. + error_factor (float): the error factor. + It is used to reduce the memory budget. Due to some errors in the estimation of peak memory and execution time. + """ + + def __init__(self, + region_list: List[Region], + memory_budget: float = -1.0, + error_factor: float = 0.95) -> None: + + self.region_list = region_list + + self.error_factor: float = error_factor + if memory_budget > 0: + self.memory_budget = memory_budget * self.error_factor + else: + self.memory_budget = torch.cuda.get_device_properties( + get_current_device()).total_memory * self.error_factor + + self.link_to_bandwidth: Dict[str, Dict[float, float]] = self._profile_bandwidth() + self.comp_power: float = self._extract_computing_power() + + @abstractmethod + def _call_solver(self): + raise NotImplementedError + + @abstractmethod + def _try_to_offload(self, *args): + raise NotImplementedError + + @abstractmethod + def _eval_one_choice(self, *args): + raise NotImplementedError + + def _compute_offload_profit(self, total_mem_saving: float, peak_mem_saving: float, extra_cost: float): + """ + Compute the profits of the offload strategies, + which packages the memory savings information for subsequent comparisons. + + Args: + total_mem_saving (float): the total memory saving of the offload strategy. + peak_mem_saving (float): the peak memory saving of the offload strategy. + extra_cost (float): extra data transfer cost. + + Returns: + tuple: profit information, the first term represents memory savings per unit of time. + """ + + if extra_cost == 0: + # means data transfer overhead can be completely overlapped + return (float('inf'), total_mem_saving, peak_mem_saving) + return (total_mem_saving / extra_cost, total_mem_saving, peak_mem_saving) + + def _compare_profit(self, profit_a: tuple, profit_b: tuple) -> bool: + """ + Compare the profits of the two offload strategies using the dictionary order algorithm. + + Args: + profit_a (tuple): the profit of a offload strategy. + profit_b (tuple): the profit of another offload strategy. + + Returns: + bool: whether profit_a is greater than profit_b. + """ + + for val1, val2 in zip(profit_a, profit_b): + if val1 != val2: + return val1 > val2 + return False + + def _update_state(self, best_ts: TrainingSimulator): + """ + Update the solver state. + """ + + self.best_ts = best_ts + self._update_node_mem_info(best_ts.fwd_node_mem, best_ts.bwd_node_mem) + + def _update_node_mem_info(self, + fwd_mem_info: Dict[Node, float], + bwd_mem_info: Dict[Node, float]): + """ + Update the runtime memory information of the node. + + Args: + fwd_mem_info (Dict[Node, float]): the runtime memory of each node in forward pass. + bwd_mem_info (Dict[Node, float]): the runtime memory of each node in backward pass. + """ + + for node, mem in fwd_mem_info.items(): + assert hasattr(node, 'node_info') and isinstance( + node.node_info, NodeInfo) + node.node_info.runtime_fwd_mem = mem + for node, mem in bwd_mem_info.items(): + assert hasattr(node, 'node_info') and isinstance( + node.node_info, NodeInfo) + node.node_info.runtime_bwd_mem = mem + + def _extract_computing_power(self): + """ + return the FP16 computing performance of the current NVIDIA GPU. + + Raises: + TypeError: Unknown NVIDIA GPU device. + """ + + nvmlInit() + handle = nvmlDeviceGetHandleByIndex(0) + device_name = nvmlDeviceGetName(handle) + units = 1e12 + + if device_name.__contains__("RTX 3080"): + return NvDevicePower.RTX3080_FP16 * units + elif device_name.__contains__("RTX 3090"): + return NvDevicePower.RTX3090_FP16 * units + elif device_name.__contains__('V100'): + return NvDevicePower.V100_FP16 * units + elif device_name.__contains__("A100"): + return NvDevicePower.A100_FP16 * units + else: + raise TypeError(f'Unknown NVIDIA GPU device name {device_name}') + + def _profile_bandwidth(self): + """ + Profile the bidirectional communication bandwidth between CPU and GPU + using data volumes ranging from 1KB to 1GB. + """ + + print('profiling bandwidth ......') + link_to_bandwidth = {} + links = ['h2d', 'd2h'] + + for link in links: + t_size = 1024 + size_to_bandwidth = {} + + # from 1KB to 1GB + for i in range(21): + if link == 'h2d': + src_tensor = torch.ones( + int(t_size), dtype=torch.int8, pin_memory=True) + dst_tensor = torch.ones( + (int(t_size)), dtype=torch.int8, device='cuda') + elif link == 'd2h': + src_tensor = torch.ones( + int(t_size), dtype=torch.int8, device='cuda') + dst_tensor = torch.ones( + (int(t_size)), dtype=torch.int8, pin_memory=True) + + def func(): + dst_tensor.copy_(src_tensor) + + size_to_bandwidth[t_size] = t_size / benchmark_func(func, number=5, repeat=3) + print(f'size: {t_size / 1024 ** 2:.3f} MB, ' + f'{src_tensor.device.type}-to-{dst_tensor.device.type} ' + f'bandwidth: {size_to_bandwidth[t_size] / 1024 ** 3:.3f} GB/s') + + t_size *= 2 + + link_to_bandwidth[link] = size_to_bandwidth + return link_to_bandwidth + + +class SynGreedySolver(Solver): + + def __init__(self, + region_list: List[Region], + memory_budget: float = -1.0) -> None: + super().__init__(region_list, memory_budget) + + self.best_ts: SynTrainingSimulator = None + self._init_state() + + def _init_state(self): + """ + Initialize the solver state when without offloading. + """ + + ts = SynTrainingSimulator(self.region_list, self.comp_power, self.link_to_bandwidth) + ts.execute() + self._update_state(ts) + + def _call_solver(self): + """ + Call the solver to search an efficient parameter offloading strategy for the linearized graph. + The solver adopts greedy algorithm. + + Raises: + NotImplementedError: Unable to find a solution for the given memory budget. + """ + + print("search offloading strategy ......") + while self.best_ts.peak_mem > self.memory_budget: + offload_region = None + best_ts = None + max_profit = (0,) + + # search which region should be offloaded, + # the last region does not need to be offloaded. + for region in self.region_list[:-1]: + if region.param_size and not region.need_offload: + temp_ts, profit = self._try_to_offload(region) + if self._compare_profit(profit, max_profit): + offload_region = region + max_profit = profit + best_ts = temp_ts + + if offload_region is not None and best_ts is not None: + offload_region.need_offload = True + offload_region.is_syn = True + self._update_state(best_ts) + else: + raise NotImplementedError( + f"can't find the offload strategy met the memory budget {self.memory_budget / 1024 ** 2} MB, " + f"it needs {self.best_ts.peak_mem / 1024 ** 2:.3f} MB at least!") + + def _call_solver_l2l(self): + """ + The layer-wise offload strategy. + """ + + for region in self.region_list[:-1]: + region.need_offload = True + region.is_syn = True + + def _try_to_offload(self, offload_region: Region): + + # record previous information + orig_need_offload = offload_region.need_offload + assert not orig_need_offload + offload_region.need_offload = True + + ts, profit = self._eval_one_choice(offload_region) + + # restore previous information + offload_region.need_offload = orig_need_offload + return ts, profit + + def _eval_one_choice(self, offload_region: Region): + """ + Evaluate the profit of a strategy choice. + + Args: + offload_region (Region): the offload region of current choice. + + Returns: + SynTrainingSimulator: the training simulator corresponding to the current strategy. + tuple: contains memory saving and cost information of the current strategy. + """ + + ts = SynTrainingSimulator(self.region_list, self.comp_power, self.link_to_bandwidth) + ts.execute() + + extra_comm_cost = 2.0 * \ + ts._get_communication_overhead('h2d', offload_region.param_size) + # the shared region needs to be moved twice + if offload_region.r_id < offload_region.shared_rid: + extra_comm_cost *= 2.0 + profit = self._compute_offload_profit( + ts.total_mem_saving, self.best_ts.peak_mem - ts.peak_mem, extra_comm_cost) + + return ts, profit + + +class AsynGreedySolver(Solver): + + def __init__(self, + region_list: List[Region], + memory_budget: float = -1.0, + search_window_size: int = 3): + super().__init__(region_list, memory_budget) + + self.search_window_size = search_window_size + # Records the prefetch execution location of the offloaded region + self.region_to_region_map = {} + self.best_ts: AsynTrainingSimulator = None + + self._init_state() + + def _init_state(self): + """ + Initialize the solver state when without offloading. + """ + + ts = AsynTrainingSimulator(self.region_list, self.comp_power, self.link_to_bandwidth) + ts.execute() + self._update_state(ts) + print("init peak memory", self.best_ts.peak_mem / 1024 ** 2, "MB") + + def _call_solver(self): + """ + Call the solver to search an efficient parameter offloading strategy for the linearized graph. + The solver adopts greedy algorithm. + + Raises: + NotImplementedError: Unable to find a solution for the given memory budget. + """ + + print("search for offloading strategy ......") + # Records the prefetch execution location of the offloaded region + region_to_region_map = {} + while self.best_ts.peak_mem > self.memory_budget: + region_to_offload = None + max_offload_profit = (0,) + best_offl_ts = None + + # search which region should be offloaded, + # the last region does not need to be offloaded + for region in self.region_list[:-1]: + if region.param_size and not region.need_offload: + max_prefetch_profit = (0,) + best_pref_ts = None + + # search when to prefetch the region offloaded + for host_region in self.region_list[region.r_id + 1:region.r_id + 1 + self.search_window_size]: + if host_region.bwd_prefetch_region is not None: + continue + + temp_ts, profit = self._try_to_offload( + host_region, region) + + if self._compare_profit(profit, max_prefetch_profit): + region_to_region_map[region.r_id] = host_region + max_prefetch_profit = profit + best_pref_ts = temp_ts + if profit[0] == float('inf'): + break + + if self._compare_profit(max_prefetch_profit, max_offload_profit): + region_to_offload = region + max_offload_profit = max_prefetch_profit + best_offl_ts = best_pref_ts + + if (region_to_offload is not None) and (best_offl_ts is not None): + region_to_offload.need_offload = True + if region_to_region_map[region_to_offload.r_id] == region_to_offload: + region_to_offload.is_syn = True + else: + region_to_region_map[region_to_offload.r_id].bwd_prefetch_region = region_to_offload + self.region_to_region_map[region_to_offload.r_id] = region_to_region_map[region_to_offload.r_id] + + self._update_state(best_offl_ts) + + elif self.region_to_region_map.__len__() > 0: + self._repair_strategy() + else: + raise NotImplementedError( + f"can't find the offload strategy met the memory budget {self.memory_budget / 1024 ** 2} MB, " + f"it needs {self.best_ts.peak_mem / 1024 ** 2:.3f} MB at least!") + + region_to_region_map.clear() + + def _try_to_offload(self, host_region: Region, offload_region: Region): + """ + Attempts to offload the region and prefetch it in backward pass. + """ + + # record previous information + orig_prefetch = host_region.bwd_prefetch_region + orig_is_syn = offload_region.is_syn + orig_need_offload = offload_region.need_offload + + if host_region == offload_region: + offload_region.is_syn = True + else: + host_region.bwd_prefetch_region = offload_region + offload_region.need_offload = True + + ts, profit = self._eval_one_choice() + + # restore previous information + host_region.bwd_prefetch_region = orig_prefetch + offload_region.is_syn = orig_is_syn + offload_region.need_offload = orig_need_offload + + return ts, profit + + def _try_convert_to_syn_upload(self, host_region: Region, offload_region: Region): + """ + Attempts to convert asynchronous prefetch into synchronous upload operations. + """ + + # record previous information + orig_prefetch = host_region.bwd_prefetch_region + orig_is_syn = offload_region.is_syn + assert orig_prefetch is not None and not orig_is_syn + + host_region.bwd_prefetch_region = None + offload_region.is_syn = True + + ts, profit = self._eval_one_choice() + + # restore previous information + host_region.bwd_prefetch_region = orig_prefetch + offload_region.is_syn = orig_is_syn + + return ts, profit + + def _repair_strategy(self): + """ + Repair offload strategy. + It attempts to convert asynchronous prefetch into synchronous upload operations and selects the best one. + The repair process does not end until peak memory is reduced or there is no asynchronous prefetch operation. + """ + print("repair strategy ......") + + peak_mem_saving = 0 + while len(self.region_to_region_map) and peak_mem_saving <= 0: + + max_profit = (0,) + best_ts = None + undo_host_region = None + undo_offload_region = None + + for offload_region_id, host_region in self.region_to_region_map.items(): + offload_region = self.region_list[offload_region_id] + assert host_region.bwd_prefetch_region == offload_region + assert offload_region.need_offload + assert not offload_region.is_syn + + ts, profit = self._try_convert_to_syn_upload(host_region, + offload_region) + + if self._compare_profit(profit, max_profit): + undo_host_region = host_region + undo_offload_region = offload_region + max_profit = profit + best_ts = ts + + if best_ts is None: + raise NotImplementedError('repair error!') + + assert not undo_offload_region.is_syn + undo_offload_region.is_syn = True + undo_host_region.bwd_prefetch_region = None + + peak_mem_saving = self.best_ts.peak_mem - best_ts.peak_mem + + self._update_state(best_ts) + self.region_to_region_map.pop(undo_offload_region.r_id) + + return best_ts + + def _eval_one_choice(self): + """ + Evaluate the profit of a strategy choice. + + Returns: + AsynTrainingSimulator: the training simulator corresponding to the current strategy. + tuple: contains memory saving and cost information of the current strategy. + """ + + ts = AsynTrainingSimulator(self.region_list, self.comp_power, self.link_to_bandwidth) + ts.execute() + + extra_comm_cost = max(ts.iter_end_time - self.best_ts.iter_end_time, 0) + profit = self._compute_offload_profit( + ts.total_mem_saving, self.best_ts.peak_mem - ts.peak_mem, extra_comm_cost) + + return ts, profit + + +class SolverFactory: + solvers: Dict[str, Type[Solver]] = { + 'syn': SynGreedySolver, + 'asyn': AsynGreedySolver + } + + @staticmethod + def create(solver_name: str) -> Type[Solver]: + if solver_name not in SolverFactory.solvers: + raise TypeError(f"Unknown parameter offload policy {solver_name}") + return SolverFactory.solvers[solver_name] + + @staticmethod + def get_solver_names(): + return tuple(SolverFactory.solvers.keys()) diff --git a/colossalai/auto_parallel/offload/training_simulator.py b/colossalai/auto_parallel/offload/training_simulator.py new file mode 100644 index 000000000000..de58023ec2d6 --- /dev/null +++ b/colossalai/auto_parallel/offload/training_simulator.py @@ -0,0 +1,458 @@ +import bisect +from typing import List, Dict +from collections import OrderedDict +from abc import ABC, abstractmethod + +from torch.fx.node import Node + +from .region import Region +from .util import * + + +@dataclass +class ExecutionPeriod: + start_time: float = 0 + end_time: float = 0 + + +class TrainingSimulator(ABC): + """ + The Training Simulator is used to simulate the training process. + It records computation, communication, and runtime memory during forward and backward passes. + + Args: + region_list (List[Region]): represents the linearized DNN computing graph. + comp_power (float): the NVIDIA GPU FP16 computing power. + link_to_bw (Dict[str, Dict[float, float]]): communication links and the corresponding bandwidth. + """ + + def __init__(self, + region_list: List[Region], + comp_power: float, + link_to_bw: Dict[str, Dict[float, float]]) -> None: + self.region_list = region_list + self.region_num = len(region_list) + + self.runtime_mem: int = 0 + self.peak_mem: int = 0 + self.total_mem_saving: int = 0 + + self.fwd_node_mem: Dict[Node, float] = {} + self.bwd_node_mem: Dict[Node, float] = {} + + # Node dependencies in backward pass + self.bwd_node_deps: Dict[Node, int] = {} + + self.comp_power: float = comp_power + self.link_to_bandwidth: Dict[str, Dict[float, float]] = link_to_bw + + @abstractmethod + def execute(self): + raise NotImplementedError + + @abstractmethod + def _eval_fwd_mem_per_region(self, region: Region): + raise NotImplementedError + + @abstractmethod + def _eval_bwd_mem_per_region(self, region: Region): + raise NotImplementedError + + def _get_bandwidth(self, link: str, comm_volumn: float) -> float: + """ + Get the data transfer bandwidth. + + Args: + link (str): the data transfer link. + comm_volumn (float): the amount of data transferred. + + Returns: + float: the data transfer bandwidth. + """ + + assert len(self.link_to_bandwidth) + if link not in self.link_to_bandwidth: + raise TypeError(f"Unknown data transfer link {link}") + + # size_list = sorted(list(map(float, self.link_to_bandwidth[link].keys()))) + size_list = sorted(self.link_to_bandwidth[link].keys()) + d_idx = bisect.bisect_left(size_list, comm_volumn) + return self.link_to_bandwidth[link][size_list[d_idx]] + + def _get_communication_overhead(self, link: str, comm_volumn: float) -> float: + return comm_volumn / self._get_bandwidth(link, comm_volumn) + + def _get_computing_overhead(self, flop: float) -> float: + return flop / self.comp_power + + +class SynTrainingSimulator(TrainingSimulator): + + def __init__(self, + region_list: List[Region], + comp_power: float, + link_to_bw: Dict[str, Dict[float, float]]) -> None: + super().__init__(region_list, comp_power, link_to_bw) + + def execute(self): + """ + Simulate synchronous training process. + """ + + for reg in self.region_list: + self._eval_fwd_mem_per_region(reg) + + for reg in self.region_list.__reversed__(): + self._eval_bwd_mem_per_region(reg) + + def _eval_fwd_mem_per_region(self, region: Region): + """ + Evaluate the runtime and peak memory when the forward execution reaches the current region. + """ + + # upload parameters of the current region + if requires_upload_p_in_fwd(self.region_list[region.shared_rid]): + self.runtime_mem += region.param_size + + for node in region.nodes: + self.runtime_mem += calculate_fwd_tmp(node) + \ + calculate_fwd_out(node) + self.fwd_node_mem[node] = self.runtime_mem + self.peak_mem = max(self.runtime_mem, self.peak_mem) + self.total_mem_saving += node.node_info.runtime_fwd_mem - self.runtime_mem + + if region.need_offload: + self.runtime_mem -= region.param_size + + def _eval_bwd_mem_per_region(self, region: Region): + """ + Evaluate the runtime and peak memory when the backward execution reaches the current region. + """ + + # upload parameters of the current region + if region.need_offload: + self.runtime_mem += region.param_size + + # add the gradient of the parameter + if region.r_id < region.shared_rid: + # gradient accumulation is required for shared parameters + self.runtime_mem += 2.0 * region.param_size + else: + self.runtime_mem += region.param_size + + for node in region.nodes.__reversed__(): + + self.runtime_mem -= calculate_fwd_out(node) + self.runtime_mem += node.meta['bwd_mem_tmp'] + \ + node.meta['bwd_mem_out'] + self.peak_mem = max(self.runtime_mem, self.peak_mem) + + # The memory savings of a node may be negative due to parameter prefetch. + self.total_mem_saving += node.node_info.runtime_bwd_mem - self.runtime_mem + self.bwd_node_mem[node] = self.runtime_mem + + self.runtime_mem -= (node.meta['bwd_mem_tmp'] + + calculate_fwd_tmp(node)) + + # free bwd_mem_out + self.bwd_node_deps[node] = len(node.all_input_nodes) + for user_node in node.users: + if user_node in self.bwd_node_deps: + self.bwd_node_deps[user_node] -= 1 + if self.bwd_node_deps[user_node] <= 0: + self.runtime_mem -= user_node.meta['bwd_mem_out'] + + if self.runtime_mem < 0: + raise ValueError(f"region id: {region.r_id}, node name: {node.name}, " + f"runtime_mem: {self.runtime_mem / 1024 ** 2:.3f}MB ---" + f"runtime memory computed less than 0, which is miscalculated!") + + # release parameter and offload gradient in region + if region.r_id == region.shared_rid: + self.runtime_mem -= 2.0 * region.param_size + elif region.r_id < region.shared_rid: + self.runtime_mem -= 3.0 * region.param_size + elif self.region_list[region.shared_rid].need_offload: + self.runtime_mem -= region.param_size + + +class AsynTrainingSimulator(TrainingSimulator): + + def __init__(self, + region_list: List[Region], + comp_power: float, + link_to_bw: Dict[str, Dict[float, float]]) -> None: + super().__init__(region_list, comp_power, link_to_bw) + + self.iter_end_time: int = 0 + # the last computation execution period + self.last_comp: ExecutionPeriod = ExecutionPeriod( + start_time=0, end_time=0) + # the last parameter prefetch execution period + self.last_h2d: ExecutionPeriod = ExecutionPeriod( + start_time=0, end_time=0) + # the last gradient offload execution period + self.last_d2h: ExecutionPeriod = ExecutionPeriod( + start_time=0, end_time=0) + # the forward computation execution period of the region + self.fwd_reg_to_comp: OrderedDict[int, ExecutionPeriod] = OrderedDict() + # the forward parameter prefetch execution period of the region + self.fwd_reg_to_pref: OrderedDict[int, ExecutionPeriod] = OrderedDict() + # the backward computation execution period of the region + self.bwd_reg_to_comp: OrderedDict[int, ExecutionPeriod] = OrderedDict() + # the backward parameter prefetch execution period of the region + self.bwd_reg_to_pref: OrderedDict[int, ExecutionPeriod] = OrderedDict() + # the gradient offload execution period of the region + # which is divided into those that are waiting and those that have been released + self.bwd_reg_to_offl_waiting: OrderedDict[int, + ExecutionPeriod] = OrderedDict() + self.bwd_reg_to_offl_freed: OrderedDict[int, + ExecutionPeriod] = OrderedDict() + # the region buffer, which records regions that are offloaded but not released + self.reg_buffer_to_free: List[int] = [] + + # node dependencies in backward pass + self.bwd_node_deps: Dict[Node, int] = {} + + # the region execution flow, + # where fwd_reg_flow[i,j] denotes whether the parameters of j-th region are in the GPU + # when the execution reaches the i-th region. + self.fwd_reg_flow = torch.zeros( + (self.region_num, self.region_num)).bool() + self.bwd_reg_flow = torch.zeros( + (self.region_num, self.region_num)).bool() + + def execute(self): + """ + Simulate asynchronous training process. + In forward pass, parameter prefetching is advanced by one region. + In backward pass, parameter prefetching is executed at the specified location, + and gradient offloading is urgent. + """ + + for reg in self.region_list: + if reg.param_size and reg.r_id < self.region_num - 1: + for nr in self.region_list[reg.r_id + 1:]: + if nr.param_size and requires_upload_p_in_fwd(self.region_list[nr.shared_rid]): + reg.fwd_prefetch_region = nr + break + self._eval_fwd_cost_per_region(reg) + self._eval_fwd_mem_per_region(reg) + + for reg in self.region_list.__reversed__(): + self._eval_bwd_cost_per_region(reg) + self._eval_bwd_mem_per_region(reg) + + # release remaining grads + for reg_id, offl_exec in self.bwd_reg_to_offl_waiting.items(): + self.bwd_reg_to_offl_freed[reg_id] = offl_exec + self.runtime_mem -= self.region_list[reg_id].param_size + self.bwd_reg_to_offl_waiting.clear() + + self.iter_end_time = max( + self.last_comp.end_time, self.last_d2h.end_time) + + def _insert_h2d_exec(self, region: Region, is_fwd: bool = True): + """ + Insert parameter prefetch execution period of the current region to the end of the h2d stream + """ + + pref_start_time = max(self.last_h2d.end_time, self.last_comp.end_time) + pref_end_time = pref_start_time + \ + 2.0 * self._get_communication_overhead('h2d', region.param_size) + pref_ep = ExecutionPeriod( + start_time=pref_start_time, end_time=pref_end_time) + if is_fwd: + self.fwd_reg_to_pref[region.r_id] = pref_ep + else: + self.bwd_reg_to_pref[region.r_id] = pref_ep + self.last_h2d = pref_ep + + def _insert_comp_exec(self, region: Region, is_fwd: bool = True): + """ + Insert computation execution period of the current region to the end of the computing stream + """ + + if is_fwd: + reg_to_comp = self.fwd_reg_to_comp + reg_to_pref = self.fwd_reg_to_pref + flop_key = 'fwd_flop' + else: + reg_to_comp = self.bwd_reg_to_comp + reg_to_pref = self.bwd_reg_to_pref + flop_key = 'bwd_flop' + comp_start_time = max(self.last_comp.end_time, reg_to_pref.get( + region.r_id, ExecutionPeriod(0, 0)).end_time) + comp_end_time = comp_start_time + \ + sum([self._get_computing_overhead(node.meta.get(flop_key, 0)) + for node in region.nodes]) + comp_ep = ExecutionPeriod( + start_time=comp_start_time, end_time=comp_end_time) + reg_to_comp[region.r_id] = comp_ep + self.last_comp = comp_ep + + def _insert_d2h_exec(self, region: Region): + """ + Insert gradient offload execution period of the current region to the end of the d2h stream + """ + + offl_start_time = max(self.last_d2h.end_time, self.last_comp.end_time) + offl_end_time = offl_start_time + \ + self._get_communication_overhead('d2h', region.param_size) + offl_ep = ExecutionPeriod( + start_time=offl_start_time, end_time=offl_end_time) + self.bwd_reg_to_offl_waiting[region.r_id] = offl_ep + self.last_d2h = offl_ep + + def _eval_fwd_cost_per_region(self, region: Region): + """ + Evaluate computation and communication execution period of the region in forward pass. + """ + + # upload parameters of the first region + if region.r_id == 0: + self._insert_h2d_exec(region) + + # prefetch parameters of the next region + fwd_prefetch_region = region.fwd_prefetch_region + if fwd_prefetch_region and requires_upload_p_in_fwd(self.region_list[fwd_prefetch_region.shared_rid]): + self._insert_h2d_exec(fwd_prefetch_region) + + # execute computation + self._insert_comp_exec(region) + + def _eval_fwd_mem_per_region(self, region: Region): + """ + Evaluate the runtime and peak memory when the forward execution reaches the current region. + """ + + # upload parameters of the current region + if region.r_id <= 0: + self.runtime_mem += region.param_size + self.fwd_reg_flow[region.r_id, region.r_id] = True + else: + self.fwd_reg_flow[region.r_id] = self.fwd_reg_flow[region.r_id - 1] + self.fwd_reg_flow[region.r_id, + self.reg_buffer_to_free] = False + self.reg_buffer_to_free.clear() + + # prefetch parameters of the next region + fwd_prefetch_region = region.fwd_prefetch_region + if fwd_prefetch_region and requires_upload_p_in_fwd(self.region_list[fwd_prefetch_region.shared_rid]): + self.runtime_mem += fwd_prefetch_region.param_size + self.fwd_reg_flow[region.r_id, + fwd_prefetch_region.r_id] = True + + for node in region.nodes: + self.runtime_mem += calculate_fwd_tmp(node) + \ + calculate_fwd_out(node) + self.peak_mem = max(self.runtime_mem, self.peak_mem) + + self.total_mem_saving += node.node_info.runtime_fwd_mem - self.runtime_mem + self.fwd_node_mem[node] = self.runtime_mem + + if region.need_offload: + self.runtime_mem -= region.param_size + + assert len( + self.reg_buffer_to_free) <= 1, f'{len(self.reg_buffer_to_free)}' + self.reg_buffer_to_free.append(region.r_id) + + def _eval_bwd_cost_per_region(self, region: Region): + """ + Evaluate computation and communication execution period of the region in backward pass. + """ + + # upload parameters of the current region + if region.is_syn: + assert region.need_offload + self._insert_h2d_exec(region, is_fwd=False) + + # prefetch parameters of the region choiced, which is parallel to computation + if region.bwd_prefetch_region is not None: + self._insert_h2d_exec(region.bwd_prefetch_region, is_fwd=False) + + # execute computation + self._insert_comp_exec(region, is_fwd=False) + + # offload gradient + if requires_offload_g_in_bwd(region): + self._insert_d2h_exec(region) + + assert len(self.reg_buffer_to_free) == 0 + for reg_id, offl_exec in self.bwd_reg_to_offl_waiting.items(): + if offl_exec.end_time >= self.last_comp.start_time: + break + self.reg_buffer_to_free.append(reg_id) + self.bwd_reg_to_offl_freed[reg_id] = offl_exec + + for reg_id in self.reg_buffer_to_free: + self.bwd_reg_to_offl_waiting.pop(reg_id) + + def _eval_bwd_mem_per_region(self, region: Region): + """ + Evaluate the runtime and peak memory when the backward execution reaches the current region. + """ + + if region.r_id + 1 < self.region_num: + self.bwd_reg_flow[region.r_id] = self.bwd_reg_flow[region.r_id + 1] + else: + self.bwd_reg_flow[region.r_id] = self.fwd_reg_flow[-1] + self.bwd_reg_flow[region.r_id, + self.reg_buffer_to_free] = False + + # free gradients in the buffer + while len(self.reg_buffer_to_free): + reg_id = self.reg_buffer_to_free.pop(0) + self.runtime_mem -= self.region_list[reg_id].param_size + + # upload parameters of the current region + if region.is_syn: + self.runtime_mem += region.param_size + self.bwd_reg_flow[region.r_id, region.r_id] = True + + # prefetch parameters of the region choiced + bwd_prefetch_region = region.bwd_prefetch_region + if bwd_prefetch_region: + self.runtime_mem += bwd_prefetch_region.param_size + self.bwd_reg_flow[region.r_id, + bwd_prefetch_region.r_id] = True + + # add the gradient of the parameter + if region.r_id < region.shared_rid: + # gradient accumulation is required for shared parameters + self.runtime_mem += 2.0 * region.param_size + else: + self.runtime_mem += region.param_size + + for node in region.nodes.__reversed__(): + + self.runtime_mem -= calculate_fwd_out(node) + self.runtime_mem += node.meta['bwd_mem_tmp'] + \ + node.meta['bwd_mem_out'] + self.peak_mem = max(self.runtime_mem, self.peak_mem) + + # The memory savings of a node may be negative due to parameter prefetch. + self.total_mem_saving += node.node_info.runtime_bwd_mem - self.runtime_mem + + self.bwd_node_mem[node] = self.runtime_mem + + self.runtime_mem -= (node.meta['bwd_mem_tmp'] + + calculate_fwd_tmp(node)) + + # free bwd_mem_out + self.bwd_node_deps[node] = len(node.all_input_nodes) + for user_node in node.users: + if user_node in self.bwd_node_deps: + self.bwd_node_deps[user_node] -= 1 + if self.bwd_node_deps[user_node] <= 0: + self.runtime_mem -= user_node.meta['bwd_mem_out'] + + if self.runtime_mem < 0: + raise ValueError(f"region id: {region.r_id}, node name: {node.name}, " + f"runtime_mem: {self.runtime_mem / 1024 ** 2:.3f}MB ---" + f"runtime memory computed less than 0, which is miscalculated!") + + # release parameters of the region + if requires_release_p_in_bwd(self.region_list[region.shared_rid]): + self.runtime_mem -= region.param_size diff --git a/colossalai/auto_parallel/offload/util.py b/colossalai/auto_parallel/offload/util.py new file mode 100644 index 000000000000..6b010512cc9c --- /dev/null +++ b/colossalai/auto_parallel/offload/util.py @@ -0,0 +1,99 @@ +from dataclasses import dataclass +from typing import List + +import torch + +from colossalai.context.singleton_meta import SingletonMeta +from colossalai.fx.profiler import calculate_fwd_out, calculate_fwd_tmp + +from .region import Region + + +@dataclass +class NodeInfo: + node_id: int = 0 + runtime_fwd_mem: float = 0 + runtime_bwd_mem: float = 0 + + +class NvDevicePower: + """ + NVIDIA GPU computing performance (TFLOPs). + """ + + RTX3080_FP16 = 70 + RTX3080_FP32 = 34.1 + + RTX3090_FP16 = 71 + RTX3090_FP32 = 35.7 + + V100_FP16 = 31.4 + V100_FP32 = 15.7 + + A100_FP16 = 78 + A100_FP32 = 19.5 + + +class GlobalRuntimeInfo(metaclass=SingletonMeta): + + def __init__(self): + self.h2d_stream = torch.cuda.Stream() + self.d2h_stream = torch.cuda.Stream() + self.fwd_prefetch_event_map = {} + self.bwd_prefetch_event_map = {} + self.region_list = [] + + +def compute_act_peak_mem(region_list: List[Region]) -> float: + act_peak_mem = 0 + runtime_mem = 0 + # forward + for region in region_list: + for node in region.nodes: + runtime_mem = runtime_mem + \ + calculate_fwd_tmp(node) + calculate_fwd_out(node) + act_peak_mem = max(runtime_mem, act_peak_mem) + # backward + bwd_deps = {} + for region in region_list.__reversed__(): + for node in region.nodes.__reversed__(): + runtime_mem -= calculate_fwd_out(node) + runtime_mem = runtime_mem + \ + node.meta['bwd_mem_tmp'] + node.meta['bwd_mem_out'] + + act_peak_mem = max(runtime_mem, act_peak_mem) + + runtime_mem = runtime_mem - \ + node.meta['bwd_mem_tmp'] - calculate_fwd_tmp(node) + + # free bwd_mem_out + bwd_deps[node] = len(node.all_input_nodes) + for user_node in node.users: + if user_node in bwd_deps: + bwd_deps[user_node] -= 1 + if bwd_deps[user_node] <= 0: + runtime_mem -= user_node.meta['bwd_mem_out'] + + return act_peak_mem + + +def compute_max_param_mem(region_list: List[Region]) -> float: + return max(region.param_size for region in region_list) + + +def compute_total_param_mem(region_list: List[Region]) -> float: + return sum(region.param_size for region in region_list if region.r_id <= region.shared_rid) + + +def requires_upload_p_in_fwd(shared_reg: Region): + return (shared_reg.r_id >= shared_reg.shared_rid) or (shared_reg.r_id < shared_reg.shared_rid + and shared_reg.need_offload) + + +def requires_release_p_in_bwd(shared_reg: Region): + return (shared_reg.r_id >= shared_reg.shared_rid) or (shared_reg.r_id < shared_reg.shared_rid + and shared_reg.need_offload) + + +def requires_offload_g_in_bwd(region: Region): + return region.param_size and (region.r_id <= region.shared_rid) diff --git a/colossalai/auto_parallel/passes/comm_metainfo_pass.py b/colossalai/auto_parallel/passes/comm_metainfo_pass.py index ab3acb0563ff..ffda58e0689f 100644 --- a/colossalai/auto_parallel/passes/comm_metainfo_pass.py +++ b/colossalai/auto_parallel/passes/comm_metainfo_pass.py @@ -4,7 +4,7 @@ from torch.fx import GraphModule from torch.fx.node import Node -from colossalai.auto_parallel.meta_profiler import MetaInfo +from colossalai.auto_parallel.meta_profiler import ShardMetaInfo from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply, runtime_comm_spec_apply from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, TrainCycleItem from colossalai.tensor.comm_spec import CommSpec @@ -14,15 +14,15 @@ shape_consistency_manager = ShapeConsistencyManager() -def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec, - target_sharding_spec: ShardingSpec) -> MetaInfo: +def _construct_shard_meta_info(node: Node, origin_sharding_spec: ShardingSpec, + target_sharding_spec: ShardingSpec) -> ShardMetaInfo: # get comm_action_sequence and total_cost from shape_consistency_manager _, comm_action_sequence, total_cost = shape_consistency_manager.shape_consistency( origin_sharding_spec, target_sharding_spec) - meta_info = MetaInfo() + meta_info = ShardMetaInfo() # NOTE: the cost in shape_consistency_manager.mem_cost is the count in number of numel - # get mem cost for MetaInfo + # get mem cost for ShardMetaInfo mem_cost = shape_consistency_manager.mem_cost(comm_action_sequence) # extract user that has _meta_data and extract element length input_node = next(n for n in node._input_nodes if hasattr(n, '_meta_data')) @@ -36,12 +36,12 @@ def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec, meta_info.memory_cost = mem_cost - # get computation cost for MetaInfo + # get computation cost for ShardMetaInfo meta_info.compute_cost = TrainCycleItem(total_cost['forward'] * element_length, total_cost['backward'] * element_length, total_cost['total'] * element_length) - # get tensor shape for MetaInfo + # get tensor shape for ShardMetaInfo origin_sharding_spec: ShardingSpec target_sharding_spec: ShardingSpec input_shape = origin_sharding_spec.get_sharded_shape_per_device() @@ -54,7 +54,7 @@ def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec, return meta_info -def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -> MetaInfo: +def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -> ShardMetaInfo: """ This method is used to construct `MetaInto` for shape consistency node """ @@ -65,17 +65,17 @@ def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) - origin_sharding_spec, target_sharding_spec = origin_spec_dict[node_index], sharding_spec_dict[node_index][ user_node_index] - return _construct_meta_info(node, origin_sharding_spec, target_sharding_spec) + return _construct_shard_meta_info(node, origin_sharding_spec, target_sharding_spec) -def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> MetaInfo: +def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> ShardMetaInfo: # extract node_index and op_data_name node_index, op_data_name = node.args[2], node.args[3] comm_action = comm_actions_dict[node_index][op_data_name] if isinstance(comm_action.comm_spec, CommSpec): # this case is for all_reduce, there will be no memory cost - meta_info = MetaInfo() + meta_info = ShardMetaInfo() meta_info.memory_cost = TrainCycleItem(MemoryCost(), MemoryCost(), MemoryCost) output_node = next(n for n in node.users if hasattr(n, '_meta_data')) element_length = output_node._meta_data.element_size() @@ -93,7 +93,7 @@ def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> M # this case will be handled by shape consistency manager origin_sharding_spec, target_sharding_spec = comm_action.comm_spec['src_spec'], comm_action.comm_spec[ 'tgt_spec'] - meta_info = _construct_meta_info(node, origin_sharding_spec, target_sharding_spec) + meta_info = _construct_shard_meta_info(node, origin_sharding_spec, target_sharding_spec) return meta_info @@ -105,9 +105,9 @@ def comm_metainfo_pass(gm: GraphModule, sharding_spec_dict: Dict, origin_spec_di """ for node in gm.graph.nodes: if node.target == runtime_apply: - setattr(node, 'best_metainfo', _runtime_apply_meta_info(node, origin_spec_dict, sharding_spec_dict)) + setattr(node, 'best_strategy_info', _runtime_apply_meta_info(node, origin_spec_dict, sharding_spec_dict)) elif node.target == runtime_comm_spec_apply: - setattr(node, 'best_metainfo', _runtime_comm_spec_apply_meta_info(node, comm_actions_dict)) + setattr(node, 'best_strategy_info', _runtime_comm_spec_apply_meta_info(node, comm_actions_dict)) else: pass return gm diff --git a/colossalai/auto_parallel/passes/meta_info_prop.py b/colossalai/auto_parallel/passes/meta_info_prop.py index f7e07ef1ec18..0673b767de7b 100644 --- a/colossalai/auto_parallel/passes/meta_info_prop.py +++ b/colossalai/auto_parallel/passes/meta_info_prop.py @@ -7,7 +7,7 @@ from torch.fx import GraphModule from torch.fx.node import Node -from colossalai.auto_parallel.meta_profiler import MetaInfo +from colossalai.auto_parallel.meta_profiler import ShardMetaInfo from colossalai.auto_parallel.passes.constants import OUTPUT_SAVED_MOD, OUTPUT_SAVED_OPS from colossalai.fx._compatibility import compatibility from colossalai.fx.profiler import GraphInfo @@ -96,12 +96,12 @@ def node_handler(self, node: Node) -> None: """ Handle other kind of nodes """ - assert hasattr(node, 'best_metainfo'), f"Cannot find best_metainfo in node {node}, {node.op}" + assert hasattr(node, 'best_strategy_info'), f"Cannot find best_strategy_info in node {node}, {node.op}" graph_info = GraphInfo() - meta_info = node.best_metainfo - meta_info: MetaInfo + meta_info = node.best_strategy_info + meta_info: ShardMetaInfo - # set data_ptr for input_tensor in MetaInfo class + # set data_ptr for input_tensor in ShardMetaInfo class input_tensors: List[torch.Tensor] = meta_info.fwd_in buffer_tensors: List[torch.Tensor] = meta_info.fwd_buffer output_tensors: List[torch.Tensor] = meta_info.fwd_out @@ -148,7 +148,7 @@ def node_handler(self, node: Node) -> None: graph_info.fwd_tmp = buffer_tensors graph_info.fwd_out = output_tensors - # fetch other memory informations + # fetch other memory information memory_cost = meta_info.memory_cost graph_info.fwd_mem_tmp = memory_cost.fwd.temp graph_info.fwd_mem_out = memory_cost.fwd.activation diff --git a/colossalai/auto_parallel/passes/runtime_apply_pass.py b/colossalai/auto_parallel/passes/runtime_apply_pass.py index 9d83f105748b..2049a06187d2 100644 --- a/colossalai/auto_parallel/passes/runtime_apply_pass.py +++ b/colossalai/auto_parallel/passes/runtime_apply_pass.py @@ -4,7 +4,7 @@ import torch from torch.fx.node import Node -from colossalai.auto_parallel.meta_profiler import MetaInfo +from colossalai._analyzer.fx.node_util import MetaInfo from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( CommAction, CommType, @@ -128,9 +128,10 @@ def _shape_consistency_apply(gm: torch.fx.GraphModule): runtime_apply, args=(node, origin_dict_node, input_dict_node, node_to_index_dict[node], user_node_index)) - if 'activation_checkpoint' in user_node.meta: - shape_consistency_node.meta['activation_checkpoint'] = user_node.meta['activation_checkpoint'] - + if hasattr(user_node.meta['info'], 'activation_checkpoint'): + MetaInfo(shape_consistency_node, + mod_dir=user_node.meta['info'].mod_dir, + activation_checkpoint=tuple(user_node.meta['info'].activation_checkpoint)) new_args = list(user_node.args) new_kwargs = dict(user_node.kwargs) # the origin node may be a positional argument or key word argument of user node @@ -210,14 +211,15 @@ def _comm_spec_apply(gm: torch.fx.GraphModule): # substitute the origin node with comm_spec_apply_node new_kwargs[str(node)] = comm_spec_apply_node user.kwargs = new_kwargs - - if 'activation_checkpoint' in node.meta: - comm_spec_apply_node.meta['activation_checkpoint'] = node.meta['activation_checkpoint'] + if hasattr(node.meta['info'], 'activation_checkpoint'): + MetaInfo(comm_spec_apply_node, + mod_dir=node.meta['info'].mod_dir, + activation_checkpoint=tuple(node.meta['info'].activation_checkpoint)) return gm -def _act_annotataion_pass(gm: torch.fx.GraphModule): +def _act_annotation_pass(gm: torch.fx.GraphModule): """ This pass is used to add the act annotation to the new inserted nodes. """ diff --git a/colossalai/auto_parallel/passes/runtime_preparation_pass.py b/colossalai/auto_parallel/passes/runtime_preparation_pass.py index e63bfdfe730c..9a2314826448 100644 --- a/colossalai/auto_parallel/passes/runtime_preparation_pass.py +++ b/colossalai/auto_parallel/passes/runtime_preparation_pass.py @@ -6,6 +6,7 @@ from torch.fx import symbolic_trace from torch.fx.node import Node +from colossalai._analyzer.fx.node_util import MetaInfo from colossalai.auto_parallel.tensor_shard.constants import RESHAPE_FUNC_OP from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( CommAction, @@ -53,7 +54,7 @@ def size_processing(size: Union[int, torch.Size], return size -def solution_annotatation_pass(gm: torch.fx.GraphModule, solution: List[int], +def solution_annotation_pass(gm: torch.fx.GraphModule, solution: List[int], strategies_constructor: StrategiesConstructor): """ This method is used to stick the solution strategy to the nodes and add the information @@ -74,9 +75,9 @@ def solution_annotatation_pass(gm: torch.fx.GraphModule, solution: List[int], origin_node_sharding_spec_dict[node_index] = strategies_vector[strategy_index].get_sharding_spec_by_name( str(node)) - # attach the corresponding metainfo if node has the attribute `metainfo_vector` - if hasattr(node, 'metainfo_vector'): - setattr(node, 'best_metainfo', node.metainfo_vector[strategy_index]) + # attach the corresponding metainfo if node has the attribute `strategies_info` + if hasattr(node, 'strategies_info'): + setattr(node, 'best_strategy_info', node.strategies_info[strategy_index]) # the dict to get input sharding specs of user node sharding_spec_convert_dict = {} @@ -148,7 +149,7 @@ def size_value_converting_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMesh def _extract_target_dim(node): ''' - A helper function to etract the target dimension from size node. + A helper function to extract the target dimension from size node. There are two usages of torch.Tensor.size: 1. tensor.size() 2. tensor.size(dim) @@ -168,12 +169,15 @@ def _post_processing(node, size_processing_node): This function is used to process the dependency between the size node and its users after inserting the size_process_node. ''' - # store original node and processing node pair in node_pairs dictioanry + # store original node and processing node pair in node_pairs dictionary # It will be used to replace the original node with processing node in slice object node_pairs[node] = size_processing_node size_processing_node._meta_data = node._meta_data - if 'activation_checkpoint' in node.meta: - size_processing_node.meta['activation_checkpoint'] = node.meta['activation_checkpoint'] + + if hasattr(node.meta['info'], 'activation_checkpoint'): + MetaInfo(size_processing_node, + mod_dir=node.meta['info'].mod_dir, + activation_checkpoint=tuple(node.meta['info'].activation_checkpoint)) user_list = list(node.users.keys()) for user in user_list: @@ -384,15 +388,16 @@ def module_params_sharding_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMes """ mod_graph = gm.graph nodes = tuple(mod_graph.nodes) - # This stream is created for overlaping the communication and computation. + # This stream is created for overlapping the communication and computation. reduction_stream = torch.cuda.Stream() - def _add_hook_for_grad_communication(node, param): + def _add_hook_for_grad_communication(node, param, name=None): comm_actions = node.best_strategy.communication_actions - def _filter_param_to_hook(node, op_data, comm_action): - if node.op == 'call_module' and op_data.type == OperationDataType.PARAM and op_data.name == param.name and comm_action.comm_type == CommType.HOOK: + def _filter_param_to_hook(node, op_data, comm_action, name): + + if node.op == 'call_module' and op_data.type == OperationDataType.PARAM and op_data.name == name and comm_action.comm_type == CommType.HOOK: return True if node.op == 'get_attr' and isinstance( node._meta_data, torch.nn.parameter.Parameter) and comm_action.comm_type == CommType.HOOK: @@ -402,7 +407,7 @@ def _filter_param_to_hook(node, op_data, comm_action): for operation_data, comm_action in comm_actions.items(): comm_spec_to_use = comm_action.comm_spec # register hook to the parameters - if _filter_param_to_hook(node, operation_data, comm_action): + if _filter_param_to_hook(node, operation_data, comm_action, name=name): def wrapper(param, comm_spec, stream, overlap): @@ -422,7 +427,7 @@ def _shard_param(param, target_sharding_spec): if target_sharding_spec.dim_partition_dict != {}: origin_sharding_spec = ShardingSpec(device_mesh, param.shape, {}) setattr(param, 'sharding_spec', origin_sharding_spec) - # TODO: build a ColoParamter class to manager the distributed parameters + # TODO: build a ColoParameter class to manager the distributed parameters # we could use .data here, because all the operations just happen before the real training # loop, so we don't need to track these operations in the autograd graph. param = torch.nn.Parameter( @@ -442,7 +447,7 @@ def _shard_param(param, target_sharding_spec): param = _shard_param(param, target_sharding_spec) setattr(target_module, name, param) - _add_hook_for_grad_communication(node, param) + _add_hook_for_grad_communication(node, param, name) sharded_buffer_dict = {} # apply the sharding spec of buffers @@ -491,7 +496,7 @@ def runtime_preparation_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMesh, strategies_constructor: StrategiesConstructor, overlap=False): - gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict = solution_annotatation_pass( + gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict = solution_annotation_pass( gm, solution, strategies_constructor) gm = size_value_converting_pass(gm, device_mesh) gm = node_args_converting_pass(gm, device_mesh) diff --git a/colossalai/auto_parallel/tensor_shard/initialize.py b/colossalai/auto_parallel/tensor_shard/initialize.py index 60472eee52ca..b406ca6fb7e0 100644 --- a/colossalai/auto_parallel/tensor_shard/initialize.py +++ b/colossalai/auto_parallel/tensor_shard/initialize.py @@ -6,6 +6,10 @@ from torch.fx import GraphModule from torch.fx.graph import Graph +from colossalai._analyzer.fx.codegen import ActivationCheckpointCodeGen +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass from colossalai.auto_parallel.tensor_shard.options import DataloaderOption, ShardOption, SolverOptions, SolverPerference @@ -13,8 +17,6 @@ from colossalai.auto_parallel.tensor_shard.solver import CostGraph, GraphAnalyser, Solver, StrategiesConstructor from colossalai.device.alpha_beta_profiler import AlphaBetaProfiler from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.graph_module import ColoGraphModule -from colossalai.fx.tracer import ColoTracer from colossalai.tensor.sharding_spec import ShardingSpec @@ -126,6 +128,7 @@ def solve_solution(gm: ColoGraphModule, strategy_constructor: StrategiesConstruc def transform_to_sharded_model(gm: ColoGraphModule, + meta_args: Dict, solution: List[int], device_mesh: DeviceMesh, strategies_constructor: StrategiesConstructor, @@ -142,6 +145,7 @@ def transform_to_sharded_model(gm: ColoGraphModule, strategies_constructor, overlap=overlap) gm = runtime_apply_pass(gm) + shape_prop_pass(gm, *meta_args.values(), sharding_spec_dict, origin_spec_dict, comm_actions_dict) gm.recompile() sharding_spec_dicts = (sharding_spec_dict, origin_spec_dict, comm_actions_dict) @@ -243,10 +247,13 @@ def initialize_model(model: nn.Module, solution will be used to debug or help to analyze the sharding result. Therefore, we will not just return a series of integers, but return the best strategies. ''' - tracer = ColoTracer(trace_act_ckpt=True) + tracer = ColoTracer(trace_act_ckpt=True, bias_addition_split=True) graph = tracer.trace(root=model, meta_args=meta_args) + graph.set_codegen(ActivationCheckpointCodeGen()) gm = ColoGraphModule(model, graph, model.__class__.__name__) + + shape_prop_pass(gm, *meta_args.values()) gm.recompile() strategies_constructor = build_strategy_constructor(graph, @@ -261,7 +268,9 @@ def initialize_model(model: nn.Module, if save_solver_solution: torch.save(solution, solution_path) - gm, sharding_spec_dicts = transform_to_sharded_model(gm, solution, device_mesh, strategies_constructor, overlap) + gm, sharding_spec_dicts = transform_to_sharded_model(gm, meta_args, solution, device_mesh, strategies_constructor, + overlap) + model_to_return = ModuleWrapper(gm, *sharding_spec_dicts) if return_solution: diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/batch_norm_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/batch_norm_handler.py index 57b623b0122c..cb1bb36b7879 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/batch_norm_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/batch_norm_handler.py @@ -2,8 +2,6 @@ import torch -from colossalai.auto_parallel.meta_profiler.metainfo import MetaInfo - from ..sharding_strategy import OperationData, OperationDataType, StrategiesVector from .node_handler import MetaInfoModuleHandler, ModuleHandler from .registry import operator_registry diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/bmm_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/bmm_handler.py index 9e1d958e15ab..da2b733c9f7a 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/bmm_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/bmm_handler.py @@ -81,7 +81,10 @@ def get_operation_data_mapping(self) -> Dict[str, OperationData]: def get_strategy_generator(self) -> List[StrategyGenerator]: op_data_mapping = self.get_operation_data_mapping() generators = [] - generators.append(BatchedMatMulStrategyGenerator(op_data_mapping, self.device_mesh)) + generator = BatchedMatMulStrategyGenerator(op_data_mapping, self.device_mesh) + # addbmm will shrink the first batch dim + generator.squeeze_batch_dim = True + generators.append(generator) return generators def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, List[ShardingStrategy]]: diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/embedding_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/embedding_handler.py index e154105b672d..112ee194b4ec 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/embedding_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/embedding_handler.py @@ -155,7 +155,7 @@ def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, Li Convert the sharding spec from the logical shape to the physical shape. """ # create multiple sharding strategies for the inputs - # as input can be multi-dimensinal and the partition dim is only 2D, + # as input can be multi-dimensional and the partition dim is only 2D, # we need to map the partition at logical dim 0 to one of the first few dimensions of the input and output strategies = _convert_logical_sharding_to_physical_sharding_spec_for_embedding(strategy=strategy, input_name=str( @@ -221,7 +221,7 @@ def post_process(self, strategy: ShardingStrategy): Convert the sharding spec from the logical shape to the physical shape. """ # create multiple sharding strategies for the inputs - # as input can be multi-dimensinal and the partition dim is only 2D, + # as input can be multi-dimensional and the partition dim is only 2D, # we need to map the partition at logical dim 0 to one of the first few dimensions of the input and output strategies = _convert_logical_sharding_to_physical_sharding_spec_for_embedding(strategy=strategy, input_name=str( diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/linear_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/linear_handler.py index 59091dab519f..ea541e434009 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/linear_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/linear_handler.py @@ -23,7 +23,7 @@ def _update_sharding_spec_for_transposed_weight_for_linear(strategy: ShardingStr weight_name: str) -> ShardingStrategy: """ This function is a helper function used by both module node handler and function node handler. This function will - convert the sharding spec for the transposed weight to the correct partititon spec. + convert the sharding spec for the transposed weight to the correct partition spec. Args: strategy (ShardingStrategy): the strategy generated by the strategy generator. @@ -197,7 +197,7 @@ def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, Li strategy = _update_sharding_spec_for_transposed_weight_for_linear(strategy=strategy, weight_name='weight') # create multiple sharding strategies for the inputs - # as input can be multi-dimensinal and the partition dim is only 2D, + # as input can be multi-dimensional and the partition dim is only 2D, # we need to map the partition at dim 0 to one of the first few dimensions of the input strategies = _convert_logical_sharding_to_physical_sharding_spec_for_linear(strategy=strategy, input_name=str(self.node.args[0]), @@ -267,7 +267,7 @@ def post_process(self, strategy: ShardingStrategy): strategy = _update_sharding_spec_for_transposed_weight_for_linear(strategy=strategy, weight_name=str(self.node.args[1])) # create multiple sharding strategies for the inputs - # as input can be multi-dimensinal and the partition dim is only 2D, + # as input can be multi-dimensional and the partition dim is only 2D, # we need to map the partition at dim 0 to one of the first few dimensions of the input strategies = _convert_logical_sharding_to_physical_sharding_spec_for_linear(strategy=strategy, input_name=str(self.node.args[0]), diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/matmul_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/matmul_handler.py index f3c9d0cbf826..fa51114a5c94 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/matmul_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/matmul_handler.py @@ -48,8 +48,8 @@ def get_matmul_type(input_dim: int, other_dim: int): Determine which type of matmul operation should be executed for the given tensor dimensions. Args: - input_dim (int): the number of dimensions for the input tenosr - other_dim (int): the number of dimensions for the other tenosr + input_dim (int): the number of dimensions for the input tensor + other_dim (int): the number of dimensions for the other tensor """ if input_dim == 1 and other_dim == 1: matmul_type = MatMulType.DOT @@ -206,7 +206,7 @@ def _remove_sharding_on_broadcast_dim(key, strategy): # e.g. [1, 2, 4] x [4, 4, 8] -> [4, 2, 8] # the dim 0 of [1, 2, 4] is multiplied to 4 tensor_shape[dim_idx] = 1 - elif broadcast_type == BroadcastType.PADDDING: + elif broadcast_type == BroadcastType.PADDING: # if the dim is padded # we remove its sharding tensor_shape[dim_idx] = None @@ -268,13 +268,13 @@ def _update_sharding_spec(key, strategy, physical_batch_dim): dim_partition_dict = sharding_spec.dim_partition_dict entire_shape = sharding_spec.entire_shape - # upddate the dimension index for the matrix dimensions + # update the dimension index for the matrix dimensions if 2 in dim_partition_dict: dim_partition_dict[len(self.batch_dims_before_view) + 1] = dim_partition_dict.pop(2) if 1 in dim_partition_dict: dim_partition_dict[len(self.batch_dims_before_view)] = dim_partition_dict.pop(1) - # map the logical batch dim to phyiscal batch dim + # map the logical batch dim to physical batch dim if 0 in dim_partition_dict: batch_dim_shard = dim_partition_dict.pop(0) dim_partition_dict[physical_batch_dim] = batch_dim_shard @@ -414,7 +414,7 @@ def _get_logical_shape_for_dot(self): def _get_logical_shape_for_mm(self): """ - We need to handle the input tensor for a matrix-matrix multiplcation as the input + We need to handle the input tensor for a matrix-matrix multiplication as the input tensor can be a 1D or 2D tensor. If it is a 1D tensor, 1 will be prepended to its shape (e.g. [4] -> [1, 4]). """ diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py b/colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py index 136e57c5e0f5..b4b7b0e794d1 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py @@ -4,7 +4,7 @@ import torch from torch.fx.node import Node -from colossalai.auto_parallel.meta_profiler.metainfo import MetaInfo, meta_register +from colossalai.auto_parallel.meta_profiler.shard_metainfo import ShardMetaInfo, meta_register from colossalai.auto_parallel.tensor_shard.options import ShardOption, SolverPerference from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( OperationData, @@ -75,7 +75,7 @@ def update_resharding_cost(self, strategy: ShardingStrategy) -> None: prev_strategy.get_sharding_spec_by_name(node_name) for prev_strategy in prev_strategy_vector ] - # create data structrure to store costs + # create data structure to store costs if node not in resharding_costs: resharding_costs[node] = [] @@ -188,7 +188,7 @@ def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesV remove_strategy_list = [] for strategy in self.strategies_vector: shard_axis_list = [] - last_axis = len(self.device_mesh.mesh_shape) - 1 + last_axis = len(self.device_mesh.shape) - 1 for op_data, sharding_spec in strategy.sharding_specs.items(): if op_data.data is not None and isinstance(op_data.data, torch.Tensor): for dim, shard_axes in sharding_spec.dim_partition_dict.items(): @@ -212,7 +212,7 @@ def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesV return self.strategies_vector def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, List[ShardingStrategy]]: - # tranform the strategy generated + # transform the strategy generated # e.g. to process the sharding strategy for the transposed weights return strategy @@ -258,7 +258,7 @@ class MetaInfoNodeHandler(NodeHandler): def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesVector: """ This method is inherited from NodeHandler. It will register the strategies first, - and rewrite the memory_cost and compute_cost of the strategy using the MetaInfo class. + and rewrite the memory_cost and compute_cost of the strategy using the ShardMetaInfo class. """ super().register_strategy(compute_resharding_cost=compute_resharding_cost) target = self.get_target_function() @@ -266,15 +266,15 @@ def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesV # is not patched, we will use the default cost model to compute the cost. # TODO: patch all torch functions and modules to make it clean if meta_register.has(target.__class__) or meta_register.has(target): - metainfo_vector = [] + strategies_info = [] for strategy in self.strategies_vector: - metainfo = MetaInfo(strategy, target) + metainfo = ShardMetaInfo(strategy, target) strategy.compute_cost = metainfo.compute_cost strategy.memory_cost = metainfo.memory_cost - metainfo_vector.append(metainfo) + strategies_info.append(metainfo) # attach metainfos to the handler - setattr(self, "metainfo_vector", metainfo_vector) + setattr(self, "strategies_info", strategies_info) else: logger = get_dist_logger() @@ -313,7 +313,7 @@ class MetaInfoModuleHandler(ModuleHandler): def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesVector: """ This method is inherited from NodeHandler. It will register the strategies first, - and rewrite the memory_cost and compute_cost of the strategy using the MetaInfo class. + and rewrite the memory_cost and compute_cost of the strategy using the ShardMetaInfo class. """ super().register_strategy(compute_resharding_cost=compute_resharding_cost) target = self.get_target_function() @@ -321,15 +321,15 @@ def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesV # is not patched, we will use the default cost model to compute the cost. # TODO: patch all torch functions and modules to make it clean if meta_register.has(target.__class__) or meta_register.has(target): - metainfo_vector = [] + strategies_info = [] for strategy in self.strategies_vector: - metainfo = MetaInfo(strategy, target) + metainfo = ShardMetaInfo(strategy, target) strategy.compute_cost = metainfo.compute_cost strategy.memory_cost = metainfo.memory_cost - metainfo_vector.append(metainfo) + strategies_info.append(metainfo) # attach metainfos to the handler - setattr(self, "metainfo_vector", metainfo_vector) + setattr(self, "strategies_info", strategies_info) else: logger = get_dist_logger() diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/batch_norm_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/batch_norm_generator.py index 1f3812429fc2..416dc9c29cad 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/batch_norm_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/batch_norm_generator.py @@ -24,7 +24,7 @@ class BatchNormStrategyGenerator(StrategyGenerator): To keep the math consistency, there are two way to do BatchNorm if the input shards on batch dimension: 1. We gather the input partitions through batch dimension, then do the normal BatchNorm. - 2. We do the SyncBatchNorm on the each input partition seperately, the SyncBN op will help + 2. We do the SyncBatchNorm on the each input partition separately, the SyncBN op will help us to keep the computing correctness. In this generator, both methods will be considered. """ @@ -44,7 +44,7 @@ def update_compute_cost(self, strategy: ShardingStrategy): ''' Compute the computation cost per device with this specific strategy. - Note: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size. ''' # TODO: a constant coefficient need to be added. # 1D: (L) * N * Cin @@ -212,7 +212,7 @@ def split_input_batch(self, mesh_dim_0): # set communication action # For SyncBN case, we don't need to do communication for weight and bias. - # TODO: the communication happens interally at SyncBN operation. We need to replace the BN operation + # TODO: the communication happens internally at SyncBN operation. We need to replace the BN operation # to SyncBN operation instead of inserting a communication node. output_comm_action = self.get_communication_action( sharding_spec=sharding_spec_mapping["output"], @@ -250,7 +250,7 @@ def split_input_batch_1d(self, mesh_dim_0, mesh_dim_1): # set communication action # For SyncBN case, we don't need to do communication for gradients of weight and bias. - # TODO: the communication happens interally at SyncBN operation. We need to replace the BN operation + # TODO: the communication happens internally at SyncBN operation. We need to replace the BN operation # to SyncBN operation instead of inserting a communication node. output_comm_action = self.get_communication_action( sharding_spec=sharding_spec_mapping["output"], @@ -298,7 +298,7 @@ def split_input_both_dim(self, mesh_dim_0, mesh_dim_1): # set communication action # For SyncBN case, we don't need to do communication for gradients of weight and bias. - # TODO: the communication happens interally at SyncBN operation. We need to replace the BN operation + # TODO: the communication happens internally at SyncBN operation. We need to replace the BN operation # to SyncBN operation instead of inserting a communication node. output_comm_action = self.get_communication_action( sharding_spec=sharding_spec_mapping["output"], diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/binary_elementwise_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/binary_elementwise_generator.py index fd7f811c8972..d27cc046eaf3 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/binary_elementwise_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/binary_elementwise_generator.py @@ -51,7 +51,7 @@ def update_memory_cost(self, strategy: ShardingStrategy) -> ShardingStrategy: # compute fwd memory cost in bytes # as the elementwise ops are not memory-intensive - # we approximate the fwd memroy cost to be the output + # we approximate the fwd memory cost to be the output # and the backward memory cost to be grad of input and other input_bytes = self._compute_size_in_bytes(strategy, 'input') other_bytes = self._compute_size_in_bytes(strategy, 'other') diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/conv_strategy_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/conv_strategy_generator.py index c2154b3104d3..e605a68a326b 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/conv_strategy_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/conv_strategy_generator.py @@ -38,9 +38,9 @@ def update_compute_cost(self, strategy: ShardingStrategy): ''' Compute the computation cost per device with this specific strategy. - Note: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size. ''' - # TODO: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + # TODO: compute_cost need to be divided by TFLOPS, now it just shows the computation size. # 1D: (L) * N * Cout * Cin * kernel # 2D: (H * W) * N * Cout * Cin * kernel # 3D: (H * W * D) * N * Cout * Cin * kernel diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/layer_norm_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/layer_norm_generator.py index fbb6070f7e82..65b173bbf65d 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/layer_norm_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/layer_norm_generator.py @@ -34,9 +34,9 @@ def update_compute_cost(self, strategy: ShardingStrategy): ''' Compute the computation cost per device with this specific strategy. - Note: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size. ''' - # TODO: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + # TODO: compute_cost need to be divided by TFLOPS, now it just shows the computation size. # TODO: a constant coefficient need to be added. sharded_input_shape = strategy.sharding_specs[self.op_data['input']].get_sharded_shape_per_device() diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py index 5d70e131d1e9..aa1581b99e0f 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py @@ -776,10 +776,6 @@ def validate(self) -> bool: bias_op_data = self.op_data['bias'] assert bias_op_data.data.dim() < 3 and len(bias_op_data.logical_shape) == 2 - if self.op_data['output'].data.dim() == 2: - # addbmm will shrink the first batch dim - self.squeeze_batch_dim = True - def update_compute_cost(self, strategy: ShardingStrategy) -> ShardingStrategy: fwd_compute_cost = self.op_data['input'].data.shape[-1] * reduce(operator.mul, self.op_data['output'].data.shape) @@ -988,7 +984,7 @@ def split_batch_dim_both_contract(self, mesh_dim_0, mesh_dim_1): def collate_strategies(self) -> List[ShardingStrategy]: strategy_list = [] device_mesh_is_1d = True - if len(self.device_mesh.mesh_shape) == 2 and 1 not in self.device_mesh.mesh_shape: + if len(self.device_mesh.shape) == 2 and 1 not in self.device_mesh.shape: device_mesh_is_1d = False if device_mesh_is_1d: @@ -996,10 +992,10 @@ def collate_strategies(self) -> List[ShardingStrategy]: # Sb = Sb x Sb # can be None as it is only for 1D device mesh # only for 1D device mesh - if len(self.device_mesh.mesh_shape) == 1: + if len(self.device_mesh.shape) == 1: mesh_dim = 0 else: - mesh_dim = self.device_mesh.mesh_shape.index(1) + mesh_dim = self.device_mesh.shape.index(1) strategy_list.append(self.split_one_batch_dim(mesh_dim)) else: # for 2D device mesh diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/normal_pooling_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/normal_pooling_generator.py index 9df6d2fbfa12..b7db42f8f67e 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/normal_pooling_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/normal_pooling_generator.py @@ -17,7 +17,7 @@ class NormalPoolStrategyGenerator(StrategyGenerator): """ NormalPoolStrategyGenerator is a generic class to generate strategies for pool operation like MaxPoolxd. The reason we call this normal pool is AvgPoolxd and MaxPoolxd are taking the kernel size element from image, - and reduce them depening on the operation type. + and reduce them depending on the operation type. """ def validate(self) -> bool: @@ -35,9 +35,9 @@ def update_compute_cost(self, strategy: ShardingStrategy) -> TrainCycleItem: ''' Compute the computation cost per device with this specific strategy. - Note: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size. ''' - # TODO: compute_cost need to be devided by TFLOPS, now it just shows the computation size. + # TODO: compute_cost need to be divided by TFLOPS, now it just shows the computation size. # 1D: (Lout) * N * C * kernel # 2D: (H * W) * N * Cout * Cin * kernel # 3D: (H * W * D) * N * Cout * Cin * kernel diff --git a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/strategy_generator.py b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/strategy_generator.py index 6d68521aaea7..d42429745c61 100644 --- a/colossalai/auto_parallel/tensor_shard/node_handler/strategy/strategy_generator.py +++ b/colossalai/auto_parallel/tensor_shard/node_handler/strategy/strategy_generator.py @@ -225,7 +225,7 @@ def _compute_size_in_bytes_helper(sharding_spec, meta_data): if isinstance(meta_data, torch.Tensor): element_bytes = _compute_size_in_bytes_helper(sharding_spec, meta_data) else: - # if meta_data is not a tensor, we count the memroy as 0 + # if meta_data is not a tensor, we count the memory as 0 element_bytes = 0 total_bytes += element_bytes @@ -233,7 +233,7 @@ def _compute_size_in_bytes_helper(sharding_spec, meta_data): if isinstance(op_data.data, torch.Tensor): total_bytes = _compute_size_in_bytes_helper(strategy.sharding_specs[op_data], op_data.data) else: - # if op_data.data is not a tensor, we count the memroy as 0 + # if op_data.data is not a tensor, we count the memory as 0 total_bytes = 0 return total_bytes diff --git a/colossalai/auto_parallel/tensor_shard/solver/cost_graph.py b/colossalai/auto_parallel/tensor_shard/solver/cost_graph.py index 74290453ca0c..1b2d3ad57407 100644 --- a/colossalai/auto_parallel/tensor_shard/solver/cost_graph.py +++ b/colossalai/auto_parallel/tensor_shard/solver/cost_graph.py @@ -9,7 +9,7 @@ class CostGraph: 1. To feed the quadratic resharding costs into solver, we need to linearize it. We build edge_cost in CostGraph, and it stored every combinations of strategies for a src-dst node pair in an 1D list. 2. To reduce the searching space, we merge computationally-trivial operators, such as - element-wise operators, transpose, and reduction, into their following nodes. The merging infomation will + element-wise operators, transpose, and reduction, into their following nodes. The merging information will be given by the StrategiesVector depending on the type of target node and following nodes. Argument: @@ -90,7 +90,7 @@ def _check_tensor_in_node(data): if self.simplify and strategies_vector.check_merge(): for followed_node in strategies_vector.predecessor_nodes: # we only merge node pairs which src node has a tensor element inside. - # This is necessay because the node without a tensor element inside will not + # This is necessary because the node without a tensor element inside will not # be assigned any strategy. if _check_tensor_in_node(followed_node._meta_data): self.merge_pair.append((followed_node, dst_node)) diff --git a/colossalai/auto_parallel/tensor_shard/solver/graph_analysis.py b/colossalai/auto_parallel/tensor_shard/solver/graph_analysis.py index be39a74cb237..171aa8b3399f 100644 --- a/colossalai/auto_parallel/tensor_shard/solver/graph_analysis.py +++ b/colossalai/auto_parallel/tensor_shard/solver/graph_analysis.py @@ -83,7 +83,7 @@ def graph(self) -> Graph: def liveness_analysis(self) -> List[LiveStage]: """ - Analyse the graph to obtain the variable liveness information. This function returns + Analyses the graph to obtain the variable liveness information. This function returns an ordered dictionary where the key is the compute stage ID and the value is a LivenessStage object. """ compute_nodes = self.graph.nodes @@ -91,7 +91,7 @@ def liveness_analysis(self) -> List[LiveStage]: # checked: record all variables created since the first stage # all: record the live variables only exist until the current stage. - # this can be different from the `checked list`` as some varialbes may be destroyed prior to this stage. + # this can be different from the `checked list`` as some variables may be destroyed prior to this stage. # unique: record the unique live variables only exist until the current stage. # this is different from `all list` as some variables are duplicated. checked_variables = LiveVariableVector() @@ -103,7 +103,7 @@ def liveness_analysis(self) -> List[LiveStage]: # find new living variables # ############################# # detect whether the current op is an in-place op - # if it is an in-place op, we would deem it as a duplciate var + # if it is an in-place op, we would deem it as a duplicate var is_inplace = False if node.op == 'call_function': # check if this is an inplace op such as torch.nn.functional.relu(x, inplace=True) diff --git a/colossalai/auto_parallel/tensor_shard/solver/solver.py b/colossalai/auto_parallel/tensor_shard/solver/solver.py index f5c6663dce80..564c5f09220c 100644 --- a/colossalai/auto_parallel/tensor_shard/solver/solver.py +++ b/colossalai/auto_parallel/tensor_shard/solver/solver.py @@ -44,7 +44,7 @@ def __init__(self, graph: The computing graph to be optimized. strategies_constructor: It will provide all the possible strategies for each node in the computing graph. cost_graph: A graph data structure to simplify the edge cost graph. - graph_analyser: graph_analyser will analyse the graph to obtain the variable liveness information, which will be used to generate memory constraints. + graph_analyser: graph_analyser will analyses the graph to obtain the variable liveness information, which will be used to generate memory constraints. memory_budget: Memory constraint for the solution. solution_numbers: If solution_numbers is larger than one, solver will us a serious of solutions based on different memory budget. memory_increasing_coefficient: If solution_numbers is larger than one, we will use this coefficient to generate new memory budget. diff --git a/colossalai/auto_parallel/tensor_shard/solver/strategies_constructor.py b/colossalai/auto_parallel/tensor_shard/solver/strategies_constructor.py index 59ead1ca8fac..044a8ac847ea 100644 --- a/colossalai/auto_parallel/tensor_shard/solver/strategies_constructor.py +++ b/colossalai/auto_parallel/tensor_shard/solver/strategies_constructor.py @@ -137,9 +137,9 @@ def _check_no_strategy_for_data(data): shard_option=self.solver_options.shard_option, solver_perference=self.solver_options.solver_perference) handler.register_strategy() - # attach metainfo_vector to node - if hasattr(handler, 'metainfo_vector'): - setattr(node, 'metainfo_vector', handler.metainfo_vector) + # attach strategies_info to node + if hasattr(handler, 'strategies_info'): + setattr(node, 'strategies_info', handler.strategies_info) # call_function node elif node.op == 'call_function': @@ -150,9 +150,9 @@ def _check_no_strategy_for_data(data): shard_option=self.solver_options.shard_option, solver_perference=self.solver_options.solver_perference) handler.register_strategy() - # attach metainfo_vector to node - if hasattr(handler, 'metainfo_vector'): - setattr(node, 'metainfo_vector', handler.metainfo_vector) + # attach strategies_info to node + if hasattr(handler, 'strategies_info'): + setattr(node, 'strategies_info', handler.strategies_info) # call_method node elif node.op == 'call_method': @@ -163,9 +163,9 @@ def _check_no_strategy_for_data(data): shard_option=self.solver_options.shard_option, solver_perference=self.solver_options.solver_perference) handler.register_strategy() - # attach metainfo_vector to node - if hasattr(handler, 'metainfo_vector'): - setattr(node, 'metainfo_vector', handler.metainfo_vector) + # attach strategies_info to node + if hasattr(handler, 'strategies_info'): + setattr(node, 'strategies_info', handler.strategies_info) # output node elif node.op == 'output': diff --git a/colossalai/auto_parallel/tensor_shard/utils/broadcast.py b/colossalai/auto_parallel/tensor_shard/utils/broadcast.py index 28aa551328d7..307348ea1eaf 100644 --- a/colossalai/auto_parallel/tensor_shard/utils/broadcast.py +++ b/colossalai/auto_parallel/tensor_shard/utils/broadcast.py @@ -21,7 +21,7 @@ class BroadcastType(Enum): EQUAL = auto() - PADDDING = auto() + PADDING = auto() MULTIPLE = auto() @@ -69,18 +69,18 @@ def get_broadcast_dim_info(logical_shape, physical_shape): for i in range(logical_num_dims): # get the trailing dim size logical_dim_idx = logical_num_dims - i - 1 - phyiscal_dim_idx = physical_num_dims - i - 1 + physical_dim_idx = physical_num_dims - i - 1 logical_dim_size = logical_shape[logical_dim_idx] - if phyiscal_dim_idx >= 0: - physical_dim_size = physical_shape[phyiscal_dim_idx] + if physical_dim_idx >= 0: + physical_dim_size = physical_shape[physical_dim_idx] if physical_dim_size == logical_dim_size: logical_dim_broadcast_info[logical_dim_idx] = BroadcastType.EQUAL elif physical_dim_size == 1 and physical_dim_size != logical_dim_size: logical_dim_broadcast_info[logical_dim_idx] = BroadcastType.MULTIPLE else: - logical_dim_broadcast_info[logical_dim_idx] = BroadcastType.PADDDING + logical_dim_broadcast_info[logical_dim_idx] = BroadcastType.PADDING return logical_dim_broadcast_info @@ -117,7 +117,7 @@ def recover_sharding_spec_for_broadcast_shape(logical_sharding_spec: ShardingSpe for shape_dim, mesh_dim in logical_dim_partition.items(): logical_broadcast_type = logical_dim_broadcast_info[shape_dim] - if logical_broadcast_type == BroadcastType.PADDDING or logical_broadcast_type == BroadcastType.MULTIPLE: + if logical_broadcast_type == BroadcastType.PADDING or logical_broadcast_type == BroadcastType.MULTIPLE: removed_dims.extend(mesh_dim) else: # get the corresponding physical dim diff --git a/colossalai/auto_parallel/tensor_shard/utils/factory.py b/colossalai/auto_parallel/tensor_shard/utils/factory.py index 05331e560001..347c10aa102d 100644 --- a/colossalai/auto_parallel/tensor_shard/utils/factory.py +++ b/colossalai/auto_parallel/tensor_shard/utils/factory.py @@ -30,7 +30,7 @@ def generate_sharding_spec(input_: Union[Node, torch.Tensor], device_mesh: Devic """ if isinstance(input_, Node): - assert hasattr(input_, '_meta_data'), f'The given node has no attribte _meta_data' + assert hasattr(input_, '_meta_data'), f'The given node has no attribute _meta_data' meta_tensor = input_._meta_data assert meta_tensor is not None, "The given node's _meta_data attribute is None" shape = meta_tensor.shape diff --git a/colossalai/auto_parallel/tensor_shard/utils/misc.py b/colossalai/auto_parallel/tensor_shard/utils/misc.py index 9e402dab7578..475e95fc4326 100644 --- a/colossalai/auto_parallel/tensor_shard/utils/misc.py +++ b/colossalai/auto_parallel/tensor_shard/utils/misc.py @@ -46,8 +46,8 @@ def check_sharding_spec_validity(sharding_spec: ShardingSpec, tensor: torch.Tens # make sure all dims are covered in sharding spec sharding_len = len(sharding_spec.sharding_sequence) tensor_num_dim = tensor.dim() - num_devices_in_col = sharding_spec.device_mesh.mesh_shape[0] - num_devices_in_row = sharding_spec.device_mesh.mesh_shape[1] + num_devices_in_col = sharding_spec.device_mesh.shape[0] + num_devices_in_row = sharding_spec.device_mesh.shape[1] assert sharding_len == tensor_num_dim, \ f'The ShardingSpec ({sharding_spec.sharding_sequence}) is created for {sharding_len}-dimension tensor, but the given tensor is {tensor_num_dim}-dimension ({tensor.shape}).' diff --git a/colossalai/auto_parallel/tensor_shard/utils/reshape.py b/colossalai/auto_parallel/tensor_shard/utils/reshape.py index a32a14bf7d57..d0ebbd7e8b1b 100644 --- a/colossalai/auto_parallel/tensor_shard/utils/reshape.py +++ b/colossalai/auto_parallel/tensor_shard/utils/reshape.py @@ -6,12 +6,12 @@ class PreviousStatus(Enum): """ - This class shows the status of previous comparision. + This class shows the status of previous comparison. """ RESET = 0 - # ORIGIN means the dimension size of original tensor is larger in the previous comparision. + # ORIGIN means the dimension size of original tensor is larger in the previous comparison. ORIGIN = 1 - # TGT means the dimension size of target tensor is larger in the previous comparision. + # TGT means the dimension size of target tensor is larger in the previous comparison. TGT = 2 @@ -91,7 +91,7 @@ def detect_reshape_mapping(origin_shape: torch.Size, tgt_shape: torch.Size) -> D tgt_index += 1 if previous_label == PreviousStatus.TGT: - # if the target dimension size is larger in the previous comparision, which means + # if the target dimension size is larger in the previous comparison, which means # the origin dimension size has already accumulated larger than target dimension size, so # we need to offload the origin dims and tgt dims into the reshape_mapping_dict. reshape_mapping_dict[tuple(origin_dims)] = tuple(tgt_dims) @@ -111,7 +111,7 @@ def detect_reshape_mapping(origin_shape: torch.Size, tgt_shape: torch.Size) -> D origin_index += 1 if previous_label == PreviousStatus.ORIGIN: - # if the origin element is larger in the previous comparision, which means + # if the origin element is larger in the previous comparison, which means # the target element has already accumulated larger than origin element, so # we need to offload the origin dims and tgt dims into the reshape_mapping_dict. reshape_mapping_dict[tuple(origin_dims)] = tuple(tgt_dims) @@ -139,7 +139,7 @@ def check_keep_sharding_status(input_dim_partition_dict: Dict[int, List[int]], Rule: For a sharded dimension of input tensor, if it is not the minimum element of the input tuple, the function will return false. - To illustrate this issue, there are two cases to analyse: + To illustrate this issue, there are two cases to analyze: 1. no sharded dims in the input tuple: we could do the reshape operation safely just as the normal operation without distributed tensor. 2. sharded dims in the input tuple: the sharded dim must be the minimum element, then during shape diff --git a/colossalai/autochunk/autochunk_codegen.py b/colossalai/autochunk/autochunk_codegen.py index 2cbc6c9221aa..cc98c1570b4a 100644 --- a/colossalai/autochunk/autochunk_codegen.py +++ b/colossalai/autochunk/autochunk_codegen.py @@ -40,7 +40,7 @@ def _gen_chunk_slice_dim(chunk_dim: int, chunk_indice_name: str, shape: List) -> return new_shape -def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_ouput_dim: int, chunk_size=2) -> str: +def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_output_dim: int, chunk_size=2) -> str: """ Generate chunk loop start @@ -52,7 +52,7 @@ def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_oup Args: chunk_input (List[Node]): chunk input node chunk_output (Node): chunk output node - chunk_ouput_dim (int): chunk output node chunk dim + chunk_output_dim (int): chunk output node chunk dim chunk_size (int): chunk size. Defaults to 2. Returns: @@ -74,7 +74,7 @@ def _gen_loop_start(chunk_input: List[Node], chunk_output: List[Node], chunk_oup input_node.name, input_node.name) out_shape = get_node_shape(chunk_output[0]) - chunk_shape = out_shape[chunk_ouput_dim[0]] + chunk_shape = out_shape[chunk_output_dim[0]] context += "chunk_size = %d\nfor chunk_idx in range(0, %d, chunk_size):\n" % (chunk_size, chunk_shape) return context @@ -287,7 +287,7 @@ def emit_code_with_chunk(body: List[str], body = _replace_new_tensor_like_shape(search_chunk, chunk_infos, region_idx, node_idx, node, body) # new tensor body = _replace_new_tensor_shape(search_chunk, chunk_infos, region_idx, node_idx, node, body) - # reassgin reshape size + # reassign reshape size body[-1] = _replace_reshape_size(body[-1], node.name, chunk_infos[region_idx]["reshape_size"]) body[-1] = " " + body[-1] delete_unused_value_func(node, body, chunk_inputs_names) diff --git a/colossalai/autochunk/estimate_memory.py b/colossalai/autochunk/estimate_memory.py index 08a55f9aa04a..77bc2ef17bc3 100644 --- a/colossalai/autochunk/estimate_memory.py +++ b/colossalai/autochunk/estimate_memory.py @@ -153,7 +153,7 @@ def estimate_chunk_inference_mem(self, node_list: List, chunk_infos: Dict = None Returns: act_memory_peak_log (List): peak memory of every node - act_memory_after_node_log (List): memory after excuting every node + act_memory_after_node_log (List): memory after executing every node active_node_list_log (List): active nodes of every node. active nodes refer to nodes generated but not deleted. """ diff --git a/colossalai/autochunk/search_chunk.py b/colossalai/autochunk/search_chunk.py index 326445ee9f12..59645c80e808 100644 --- a/colossalai/autochunk/search_chunk.py +++ b/colossalai/autochunk/search_chunk.py @@ -16,7 +16,7 @@ class SearchChunk(object): This is the core class for AutoChunk. It defines the framework of the strategy of AutoChunk. - Chunks will be selected one by one utill search stops. + Chunks will be selected one by one until search stops. The chunk search is as follows: 1. find the peak memory node @@ -73,7 +73,7 @@ def _init_trace(self) -> None: def _find_peak_region(self, mem_peak: List) -> int: """ - find peak node, along with its neighbour nodes exceeds max mem + find peak node, along with its neighbor nodes exceeds max mem """ max_value = max(mem_peak) max_idx = mem_peak.index(max_value) @@ -118,7 +118,7 @@ def _search_max_chunk_region(self, active_node: List, peak_region: int, chunk_re chunk_region_start (int) chunk_region_end (int) """ - # check if peak node already in chunkinfo + # check if peak node already in chunk info if chunk_regions is not None: for i in chunk_regions: if i["region"][0] < peak_region[0] <= i["region"][1] or \ diff --git a/colossalai/autochunk/trace_flow.py b/colossalai/autochunk/trace_flow.py index 16815215f52b..a1080fda1541 100644 --- a/colossalai/autochunk/trace_flow.py +++ b/colossalai/autochunk/trace_flow.py @@ -64,7 +64,7 @@ def check_index_compute(self, start_idx, end_dim, end_node, end_idx): return False return True - def _assgin_single_node_flow( + def _assign_single_node_flow( self, arg_node: Node, start_idx: int, @@ -177,7 +177,7 @@ def _get_all_node_info(self, end_dim, start_idx, end_idx): if get_node_shape(arg) is None: continue arg_list.append(arg) - flow_flag = self._assgin_single_node_flow( + flow_flag = self._assign_single_node_flow( arg, start_idx, end_idx, @@ -315,7 +315,7 @@ def _get_prepose_nodes(self, all_node_info: Dict, start_idx: int, end_idx: int, chunk_info["args"]["prepose_nodes"] = prepose_nodes def _get_non_chunk_inputs(self, chunk_info, start_idx, end_idx): - # we need to log input nodes to avoid deleteing them in the loop + # we need to log input nodes to avoid deleting them in the loop chunk_node_list = self.node_mgr.get_node_slice_by_idx(start_idx, end_idx + 1) # also need to get some prepose node's arg out of non_chunk_inputs for n in chunk_info["args"]["prepose_nodes"]: @@ -366,8 +366,8 @@ def flow_search(self, start_idx, start_dim, end_idx, end_dim): # find non chunk inputs chunk_info = self._get_non_chunk_inputs(chunk_info, start_idx, end_idx) - # reassgin reshape size, some size may have changed due to chunk - chunk_info = self._reassgin_reshape_size(chunk_info) + # reassign reshape size, some size may have changed due to chunk + chunk_info = self._reassign_reshape_size(chunk_info) return chunk_info @@ -428,10 +428,10 @@ def _update_chunk_info(self, chunk_info: Dict, new_all_node_info: Dict, output: chunk_info["outputs_dim"].append(output_dim) return True - def _reassgin_reshape_size(self, chunk_info): + def _reassign_reshape_size(self, chunk_info): """ Some shape args in reshape may have changed due to chunk - reassgin those changed shape + reassign those changed shape """ chunk_region = chunk_info["region"] reshape_size = {} @@ -479,7 +479,7 @@ def check_region_start_end(self, start_node: Node, start_dim: int, start_idx: in # check index source align if not self.check_index_source(start_dim, start_node, start_idx, end_dim, end_node): return False - # check index copmute + # check index compute if not self.check_index_compute(start_idx, end_dim, end_node, end_idx): return False return True diff --git a/colossalai/autochunk/trace_indice.py b/colossalai/autochunk/trace_indice.py index 307f4de326d7..fbe0741b8827 100644 --- a/colossalai/autochunk/trace_indice.py +++ b/colossalai/autochunk/trace_indice.py @@ -8,7 +8,7 @@ class TraceIndice(object): """ - Trace all indice infomation for every node. + Trace all indice information for every node. Indice is a logical concept. Equal dims can been treated as one indice. eg. dim(x1) = [a, b, c] @@ -18,7 +18,7 @@ class TraceIndice(object): dim(x1)=dim(x2)=dim(x3)=[a, b, c] This class will record every node's dims' indice, compute and source. - Attibutes: + Attributes: node_list (List) indice_trace_list (List): [{"indice": [...], "compute": [...], "source": [...]}, {...}] indice_view_list (Dict): not used for now @@ -153,7 +153,7 @@ def _inherit_all_indice(self, node_from: Node, node_to: Node) -> None: def _inherit_more_indice_from_node_with_exclude(self, node_from: Node, node_to: Node, exclude: List = None) -> None: """ - inheirt indice from node without init + inherit indice from node without init """ if exclude == None: exclude = [] @@ -301,7 +301,7 @@ def _assign_permute_indice(self, node: Node, node_idx: int) -> None: def _assign_linear_indice(self, node: Node, node_idx: int) -> None: """ Assign indice for linear op. - 1. copy trace from input node and change last indice accroding to weight + 1. copy trace from input node and change last indice according to weight 2. mark equal for input node last indice, weight first dim and bias dim. 3. inherit input's computation, mark computation for last dim. @@ -360,7 +360,7 @@ def _assign_baddbmm_indice(self, node: Node, node_idx: int) -> None: def _assign_matmul_indice(self, node: Node, node_idx: int) -> None: """ Assign indice for matmul op. - 1. copy trace from matmul_left and change last indice accroding to matmul_right. (assert they have same length) + 1. copy trace from matmul_left and change last indice according to matmul_right. (assert they have same length) 2. mark equal for input matmul_left -1 indice and matmul_right -2 dim. 3. inherit matmul_left and matmul_right computation, mark computation for last dim. @@ -397,7 +397,7 @@ def _assign_conv2d_indice(self, node: Node, node_idx: int) -> None: input_node = node.args[0] assert len(get_node_shape(input_node)) == 4 - # assgin index + # assign index self._assign_indice_as_input(node, node_idx, input_node) self._del_dim(node_idx, 1) self._add_dim(node_idx, 1) @@ -415,7 +415,7 @@ def _assign_interpolate_indice(self, node: Node, node_idx: int) -> None: assert node.kwargs['size'] is None assert len(get_node_shape(node)) == 4 - # assgin index + # assign index self._assign_indice_as_input(node, node_idx) self._mark_computation(node, node_idx, [-1, -2]) @@ -461,7 +461,7 @@ def _assign_elementwise_indice(self, node, idx): nodes_in.append(node_in) self._inherit_more_indice_from_node_with_exclude(node_in, node) - def _assgin_no_change_indice(self, node, idx): + def _assign_no_change_indice(self, node, idx): self._assign_indice_as_input(node, idx) for node_in in node.args: if type(node_in) == type(node): @@ -720,11 +720,11 @@ def _assign_view_reshape_indice(self, node: Node, node_idx: int) -> None: Assign indice for view and reshape op. 1. get origin shape and target shape by meta info. 2. compute the real value of -1 in target shape. - 3. determine changed dim, and assgin indice for generated dim. + 3. determine changed dim, and assign indice for generated dim. 4. log changed dim and generated dim for restore 5. inherit computation. 6. look into view list to see whether the view is associated with other, - if so assgin equal dim according to previous view. + if so assign equal dim according to previous view. Args: node (node) @@ -792,7 +792,7 @@ def _assign_view_reshape_indice(self, node: Node, node_idx: int) -> None: self._add_dim(node_idx, i) dim_from.reverse() - # inheirt indice from current node + # inherit indice from current node if len(dim_from) != 0 and len(dim_to) != 0: if dim_diff == 1: if origin_shape[dim_from[0]] == 1: @@ -852,7 +852,7 @@ def trace_indice(self) -> None: elif "split" == node_name: self._assign_split_indice(node, idx) elif any(i == node_name for i in ["to", "contiguous", "clone", "type", "float"]): - self._assgin_no_change_indice(node, idx) + self._assign_no_change_indice(node, idx) elif "new_ones" == node_name: self._assign_all_indice(node, idx) elif "flatten" == node_name: @@ -914,7 +914,7 @@ def trace_indice(self) -> None: elif "conv2d" == node_name: self._assign_conv2d_indice(node, idx) elif "identity" == node_name: - self._assgin_no_change_indice(node, idx) + self._assign_no_change_indice(node, idx) elif any(n == node_name for n in ["sigmoid", "dropout", "relu", "silu", "gelu"]): self._assign_elementwise_indice(node, idx) else: diff --git a/colossalai/booster/__init__.py b/colossalai/booster/__init__.py index 3b3f45bb0fe2..841054a9c672 100644 --- a/colossalai/booster/__init__.py +++ b/colossalai/booster/__init__.py @@ -1,4 +1,3 @@ from .accelerator import Accelerator from .booster import Booster -from .environment_table import EnvironmentTable from .plugin import Plugin diff --git a/colossalai/booster/accelerator.py b/colossalai/booster/accelerator.py index 63ba193e3e4f..fc2c4a40068b 100644 --- a/colossalai/booster/accelerator.py +++ b/colossalai/booster/accelerator.py @@ -3,12 +3,52 @@ __all__ = ['Accelerator'] +_supported_devices = [ + 'cpu', + 'cuda', + + # To be supported + # 'xpu', + # 'npu', + # 'tpu', +] + class Accelerator: + """ + Accelerator is an abstraction for the hardware device that is used to run the model. + + Args: + device (str): The device to be used. Currently only support 'cpu' and 'gpu'. + """ - def __init__(self, device: torch.device): + def __init__(self, device: str): self.device = device - def setup_model(self, model: nn.Module) -> nn.Module: - # TODO: implement this method - pass + assert self.device in _supported_devices, f"Device {self.device} is not supported yet, supported devices include {_supported_devices}" + + def bind(self): + """ + Set the default device for the current process. + """ + if self.device == 'cpu': + pass + elif self.device == 'cuda': + # TODO(FrankLeeeee): use global environment to check if it is a dist job + # if is_distributed: + # local_rank = EnvTable().get_local_rank() + # torch.cuda.set_device(torch.device(f'cuda:{local_rank}')) + torch.cuda.set_device(torch.device('cuda')) + pass + else: + raise ValueError(f"Device {self.device} is not supported yet") + + def configure_model(self, model: nn.Module) -> nn.Module: + """ + Move the model to the device. + + Args: + model (nn.Module): The model to be moved. + """ + model = model.to(torch.device(self.device)) + return model diff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py index 7b351ae343d2..cee547b33b0c 100644 --- a/colossalai/booster/booster.py +++ b/colossalai/booster/booster.py @@ -1,13 +1,17 @@ +import warnings from contextlib import contextmanager -from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union +from typing import Callable, Iterator, List, Optional, Tuple, Union import torch import torch.nn as nn -from torch import Tensor from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler as LRScheduler from torch.utils.data import DataLoader +from colossalai.checkpoint_io import GeneralCheckpointIO +from colossalai.interface import ModelWrapper + +from .accelerator import Accelerator from .mixed_precision import MixedPrecision, mixed_precision_factory from .plugin import Plugin @@ -17,30 +21,31 @@ class Booster: """ Booster is a high-level API for training neural networks. It provides a unified interface for - training with different precisio, accelerator, and plugin. + training with different precision, accelerator, and plugin. Examples: - >>> colossalai.launch(...) - >>> plugin = GeminiPlugin(stage=3, ...) - >>> booster = Booster(precision='fp16', plugin=plugin) - >>> - >>> model = GPT2() - >>> optimizer = Adam(model.parameters()) - >>> dataloader = Dataloader(Dataset) - >>> lr_scheduler = LinearWarmupScheduler() - >>> criterion = GPTLMLoss() - >>> - >>> model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader) - >>> - >>> for epoch in range(max_epochs): - >>> for input_ids, attention_mask in dataloader: - >>> outputs = model(input_ids, attention_mask) - >>> loss = criterion(outputs.logits, input_ids) - >>> booster.backward(loss, optimizer) - >>> optimizer.step() - >>> lr_scheduler.step() - >>> optimizer.zero_grad() + ```python + colossalai.launch(...) + plugin = GeminiPlugin(...) + booster = Booster(precision='fp16', plugin=plugin) + + model = GPT2() + optimizer = HybridAdam(model.parameters()) + dataloader = Dataloader(Dataset) + lr_scheduler = LinearWarmupScheduler() + criterion = GPTLMLoss() + model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader) + + for epoch in range(max_epochs): + for input_ids, attention_mask in dataloader: + outputs = model(input_ids, attention_mask) + loss = criterion(outputs.logits, input_ids) + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + ``` Args: device (str or torch.device): The device to run the training. Default: 'cuda'. @@ -51,23 +56,53 @@ class Booster: """ def __init__(self, - device: Union[str, torch.device] = 'cuda', + device: str = 'cuda', mixed_precision: Union[MixedPrecision, str] = None, plugin: Optional[Plugin] = None) -> None: - # validate and set precision - if isinstance(MixedPrecision, str): - # the user will take the default arguments for amp training - self.mixed_precision = mixed_precision_factory(mixed_precision) - elif isinstance(mixed_precision, MixedPrecision): - # the user can customize the arguments by passing the precision object - self.mixed_precision = mixed_precision + if plugin is not None: + assert isinstance( + plugin, Plugin), f'Expected the argument plugin to be an instance of Plugin, but got {type(plugin)}.' + self.plugin = plugin + + # set accelerator + if self.plugin and self.plugin.control_device(): + self.accelerator = None + warnings.warn('The plugin will control the accelerator, so the device argument will be ignored.') else: - raise ValueError( - f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.' - ) + self.accelerator = Accelerator(device) - def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler, - dataloader: DataLoader) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]: + # set precision + if self.plugin and self.plugin.control_precision(): + warnings.warn('The plugin will control the precision, so the mixed_precision argument will be ignored.') + self.mixed_precision = None + elif mixed_precision is None: + self.mixed_precision = None + else: + # validate and set precision + if isinstance(mixed_precision, str): + # the user will take the default arguments for amp training + self.mixed_precision = mixed_precision_factory(mixed_precision) + elif isinstance(mixed_precision, MixedPrecision): + # the user can customize the arguments by passing the precision object + self.mixed_precision = mixed_precision + else: + raise ValueError( + f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.' + ) + + if self.plugin is not None and self.plugin.control_checkpoint_io(): + self.checkpoint_io = self.plugin.get_checkpoint_io() + else: + self.checkpoint_io = GeneralCheckpointIO() + + def boost( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]: """ Boost the model, optimizer, criterion, lr_scheduler, and dataloader. @@ -75,19 +110,34 @@ def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_ model (nn.Module): The model to be boosted. optimizer (Optimizer): The optimizer to be boosted. criterion (Callable): The criterion to be boosted. - lr_scheduler (LRScheduler): The lr_scheduler to be boosted. dataloader (DataLoader): The dataloader to be boosted. + lr_scheduler (LRScheduler): The lr_scheduler to be boosted. """ # TODO(FrankLeeeee): consider multi-model and multi-optimizer case - # TODO(lsg): Add plugin control logic - # e.g. - # if self.plugin is not None and self.plugin.control_boost: - # ... + # TODO(FrankLeeeee): consider multi-dataloader case # transform model for mixed precision - model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion) - return model, optimizer, criterion, lr_scheduler, dataloader + if self.plugin: + model, optimizer, criterion, dataloader, lr_scheduler = self.plugin.configure( + model, optimizer, criterion, dataloader, lr_scheduler) + + if self.plugin and not self.plugin.control_device(): + # transform model for accelerator + model = self.accelerator.configure(model) + + if self.mixed_precision and (self.plugin is None or self.plugin and not self.plugin.control_precision()): + # transform model for mixed precision + # when mixed_precision is specified and the plugin is not given or does not control the precision + model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion) + + return model, optimizer, criterion, dataloader, lr_scheduler def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None: + """Backward pass. + + Args: + loss (torch.Tensor): The loss to be backpropagated. + optimizer (Optimizer): The optimizer to be updated. + """ # TODO: implement this method with plugin optimizer.backward(loss) @@ -104,21 +154,111 @@ def execute_pipeline(self, pass def no_sync(self, model: nn.Module) -> contextmanager: - # TODO: implement this method - pass + """Context manager to disable gradient synchronization across DP process groups. - def save(self, - obj: Union[nn.Module, Optimizer, LRScheduler], - path_like: str, - plan: str = 'torch', - **kwargs) -> None: - # TODO: implement this method - pass + Args: + model (nn.Module): The model to be disabled gradient synchronization. - def load(self, - obj: Union[nn.Module, Optimizer, LRScheduler], - path_like: str, - plan: str = 'torch', - **kwargs) -> None: - # TODO: implement this method - pass + Returns: + contextmanager: Context to disable gradient synchronization. + """ + assert self.plugin is not None, f'no_sync is only enabled when a plugin is provided and the plugin supports no_sync.' + assert self.plugin.support_no_sync, f'The plugin {self.plugin.__class__.__name__} does not support no_sync.' + return self.plugin.no_sync(model) + + def load_model(self, model: Union[nn.Module, ModelWrapper], checkpoint: str, strict: bool = True): + """Load model from checkpoint. + + Args: + model (nn.Module or ModelWrapper): A model boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local path. + It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path. + strict (bool, optional): whether to strictly enforce that the keys + in :attr:`state_dict` match the keys returned by this module's + :meth:`~torch.nn.Module.state_dict` function. Defaults to True. + """ + self.checkpoint_io.load_model(model, checkpoint, strict) + + def save_model(self, + model: Union[nn.Module, ModelWrapper], + checkpoint: str, + shard: bool = False, + gather_dtensor: bool = True, + prefix: Optional[str] = None, + size_per_shard: int = 1024, + use_safetensors: bool = False): + """Save model to checkpoint. + + Args: + model (nn.Module or ModelWrapper): A model boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local path. + It is a file path if ``shard=False``. Otherwise, it is a directory path. + shard (bool, optional): Whether to save checkpoint a sharded way. + If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False. + gather_dtensor (bool, optional): whether to gather the distributed tensor to the first device. Default: True. + prefix (str, optional): A prefix added to parameter and buffer + names to compose the keys in state_dict. Defaults to None. + size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024. + use_safetensors (bool, optional): whether to use safe tensors. Default: False. If set to True, the checkpoint will be saved. + """ + self.checkpoint_io.save_model(model, + checkpoint=checkpoint, + shard=shard, + gather_dtensor=gather_dtensor, + prefix=prefix, + size_per_shard=size_per_shard, + use_safetensors=use_safetensors) + + def load_optimizer(self, optimizer: Optimizer, checkpoint: str): + """Load optimizer from checkpoint. + + Args: + optimizer (Optimizer): An optimizer boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local path. + It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path. + prefix (str, optional): A prefix added to parameter and buffer + names to compose the keys in state_dict. Defaults to None. + size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024. + """ + self.checkpoint_io.load_optimizer(optimizer, checkpoint) + + def save_optimizer(self, + optimizer: Optimizer, + checkpoint: str, + shard: bool = False, + gather_dtensor: bool = True, + prefix: Optional[str] = None, + size_per_shard: int = 1024): + """ + Save optimizer to checkpoint. + + Args: + optimizer (Optimizer): An optimizer boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local path. + It is a file path if ``shard=False``. Otherwise, it is a directory path. + shard (bool, optional): Whether to save checkpoint a sharded way. + If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. Default: True. + prefix (str, optional): A prefix added to parameter and buffer + names to compose the keys in state_dict. Defaults to None. + size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024. + """ + self.checkpoint_io.save_optimizer(optimizer, checkpoint, shard, gather_dtensor, prefix, size_per_shard) + + def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """Save lr scheduler to checkpoint. + + Args: + lr_scheduler (LRScheduler): A lr scheduler boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local file path. + """ + self.checkpoint_io.save_lr_scheduler(lr_scheduler, checkpoint) + + def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """Load lr scheduler from checkpoint. + + Args: + lr_scheduler (LRScheduler): A lr scheduler boosted by Booster. + checkpoint (str): Path to the checkpoint. It must be a local file path. + """ + self.checkpoint_io.load_lr_scheduler(lr_scheduler, checkpoint) diff --git a/colossalai/booster/environment_table.py b/colossalai/booster/environment_table.py deleted file mode 100644 index 4b16f120c1b9..000000000000 --- a/colossalai/booster/environment_table.py +++ /dev/null @@ -1,18 +0,0 @@ -from typing import List - -__all__ = ['EnvironmentTable'] - - -class EnvironmentTable: - - def __init__(self, intra_op_world_sizes: List[int]): - # TODO: implement this method - pass - - @property - def is_master(self) -> bool: - # TODO: implement this method - pass - - # TODO: implement more utility methods as given in - # https://github.com/hpcaitech/ColossalAI/issues/3051 diff --git a/colossalai/booster/interface/__init__.py b/colossalai/booster/interface/__init__.py deleted file mode 100644 index 8892a13e1814..000000000000 --- a/colossalai/booster/interface/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .optimizer import OptimizerWrapper - -__all__ = ['OptimizerWrapper'] diff --git a/colossalai/booster/mixed_precision/__init__.py b/colossalai/booster/mixed_precision/__init__.py index 3cf0ad28cdbe..0df9d84159f9 100644 --- a/colossalai/booster/mixed_precision/__init__.py +++ b/colossalai/booster/mixed_precision/__init__.py @@ -1,17 +1,19 @@ from .bf16 import BF16MixedPrecision from .fp8 import FP8MixedPrecision from .fp16_apex import FP16ApexMixedPrecision +from .fp16_naive import FP16NaiveMixedPrecision from .fp16_torch import FP16TorchMixedPrecision from .mixed_precision_base import MixedPrecision __all__ = [ 'MixedPrecision', 'mixed_precision_factory', 'FP16_Apex_MixedPrecision', 'FP16_Torch_MixedPrecision', - 'FP32_MixedPrecision', 'BF16_MixedPrecision', 'FP8_MixedPrecision' + 'FP32_MixedPrecision', 'BF16_MixedPrecision', 'FP8_MixedPrecision', 'FP16NaiveMixedPrecision' ] _mixed_precision_mapping = { 'fp16': FP16TorchMixedPrecision, 'fp16_apex': FP16ApexMixedPrecision, + 'fp16_naive': FP16NaiveMixedPrecision, 'bf16': BF16MixedPrecision, 'fp8': FP8MixedPrecision } diff --git a/colossalai/booster/mixed_precision/fp16_apex.py b/colossalai/booster/mixed_precision/fp16_apex.py index 266a750734b1..e184271e932a 100644 --- a/colossalai/booster/mixed_precision/fp16_apex.py +++ b/colossalai/booster/mixed_precision/fp16_apex.py @@ -1,5 +1,38 @@ +from typing import Any, Optional, Union + +import torch + from .mixed_precision_base import MixedPrecision class FP16ApexMixedPrecision(MixedPrecision): - pass + """ + Precision for mixed precision training in FP16 using apex AMP. + + Args: + opt_level(str, optional, default="O1" ): Pure or mixed precision optimization level. Accepted values are “O0”, “O1”, “O2”, and “O3”, explained in detail above Apex AMP Documentation. + cast_model_type (torch.dtype, optional, default=None): Casts your model’s parameters and buffers to the desired type. + patch_torch_functions (bool, optional, default=None): Patch all Torch functions and Tensor methods to perform Tensor Core-friendly ops like GEMMs and convolutions in FP16, and any ops that benefit from FP32 precision in FP32. + keep_batchnorm_fp32 (bool or str, optional, default=None): To enhance precision and enable cudnn batchnorm (which improves performance), it’s often beneficial to keep batchnorm weights in FP32 even if the rest of the model is FP16. + master_weights (bool, optional, default=None): Maintain FP32 master weights to accompany any FP16 model weights. FP32 master weights are stepped by the optimizer to enhance precision and capture small gradients. + loss_scale (float or str, optional, default=None): If loss_scale is a float value, use this value as the static (fixed) loss scale. If loss_scale is the string "dynamic", adaptively adjust the loss scale over time. Dynamic loss scale adjustments are performed by Amp automatically. + cast_model_outputs (torch.dpython:type, optional, default=None): Option to ensure that the outputs of your model(s) are always cast to a particular type regardless of opt_level. + num_losses(int, optional, default=1): Option to tell AMP in advance how many losses/backward passes you plan to use. When used in conjunction with the loss_id argument to `amp.scale_loss`, enables Amp to use a different loss scale per loss/backward pass, which can improve stability. If num_losses is left to 1, Amp will still support multiple losses/backward passes, but use a single global loss scale for all of them. + verbosity(int, default=1): Set to 0 to suppress Amp-related output. + min_loss_scale(float, default=None): Sets a floor for the loss scale values that can be chosen by dynamic loss scaling. The default value of None means that no floor is imposed. If dynamic loss scaling is not used, min_loss_scale is ignored. + max_loss_scale(float, default=2.**24 ): Sets a ceiling for the loss scale values that can be chosen by dynamic loss scaling. If dynamic loss scaling is not used, max_loss_scale is ignored. + """ + + def __init__(self, + opt_level: Optional[str] = "O1", + cast_model_type: torch.dtype = None, + patch_torch_functions: bool = None, + keep_batchnorm_fp32: Union[bool, str] = None, + master_weights: bool = None, + loss_scale: Union[float, str] = None, + cast_model_outputs: Any = None, + num_losses: Optional[int] = 1, + verbosity: int = 1, + min_loss_scale: float = None, + max_loss_scale: float = 2.**24) -> None: + pass diff --git a/colossalai/booster/mixed_precision/fp16_naive.py b/colossalai/booster/mixed_precision/fp16_naive.py new file mode 100644 index 000000000000..5d0d815257f3 --- /dev/null +++ b/colossalai/booster/mixed_precision/fp16_naive.py @@ -0,0 +1,26 @@ +from .mixed_precision_base import MixedPrecision + + +class FP16NaiveMixedPrecision(MixedPrecision): + """ + Precision for mixed precision training in FP16 using naive AMP. + + Args: + log_num_zeros_in_grad(bool): return number of zeros in the gradients. + initial_scale(int): initial scale of gradient scaler. + growth_factor(int): the growth rate of loss scale. + backoff_factor(float): the decrease rate of loss scale. + hysteresis(int): delay shift in dynamic loss scaling. + max_scale(int): maximum loss scale allowed. + verbose(bool): if set to `True`, will print debug info. + """ + + def __init__(self, + log_num_zeros_in_grad: bool, + initial_scale: int, + growth_factor: int, + backoff_factor: float, + hysteresis: int, + max_scale: int, + verbose: bool = None) -> None: + pass diff --git a/colossalai/booster/mixed_precision/fp16_torch.py b/colossalai/booster/mixed_precision/fp16_torch.py index 054f78d2e226..26fd92bd50b8 100644 --- a/colossalai/booster/mixed_precision/fp16_torch.py +++ b/colossalai/booster/mixed_precision/fp16_torch.py @@ -5,7 +5,8 @@ from torch import Tensor from torch.optim import Optimizer -from ..interface import OptimizerWrapper +from colossalai.interface import ModelWrapper, OptimizerWrapper + from .mixed_precision_base import MixedPrecision __all__ = ['FP16_Torch_MixedPrecision', 'TorchAMPOptimizer', 'TorchAMPModule'] @@ -45,7 +46,9 @@ def backward(self, loss: Tensor, *args, **kwargs) -> None: scaled_loss.backward(*args, **kwargs) def step(self, *args, **kwargs) -> Optional[float]: - return self.scaler.step(self.optim, *args, **kwargs) + out = self.scaler.step(self.optim, *args, **kwargs) + self.scaler.update() + return out def scale_loss(self, loss: Tensor) -> Tensor: return self.scaler.scale(loss) @@ -67,7 +70,7 @@ def clip_grad_by_norm(self, super().clip_grad_by_norm(max_norm, norm_type, error_if_nonfinite, *args, **kwargs) -class TorchAMPModule(nn.Module): +class TorchAMPModule(ModelWrapper): """ Module wrapper for mixed precision training in FP16 using PyTorch AMP. @@ -76,8 +79,7 @@ class TorchAMPModule(nn.Module): """ def __init__(self, module: nn.Module): - super().__init__() - self.module = module + super().__init__(module) def forward(self, *args, **kwargs): with torch.cuda.amp.autocast(): @@ -113,10 +115,12 @@ def __init__(self, def configure(self, model: nn.Module, - optimizer: Optimizer, - criterion: Callable = None) -> Tuple[nn.Module, OptimizerWrapper, Callable]: + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable]: model = TorchAMPModule(model) - optimizer = TorchAMPOptimizer(optimizer, **self.torch_amp_kwargs) + if optimizer is not None: + optimizer = TorchAMPOptimizer(optimizer, **self.torch_amp_kwargs) if criterion is not None: criterion = TorchAMPModule(criterion) return model, optimizer, criterion diff --git a/colossalai/booster/mixed_precision/mixed_precision_base.py b/colossalai/booster/mixed_precision/mixed_precision_base.py index d1e8acc82cc6..8caa34e505e1 100644 --- a/colossalai/booster/mixed_precision/mixed_precision_base.py +++ b/colossalai/booster/mixed_precision/mixed_precision_base.py @@ -1,10 +1,10 @@ from abc import ABC, abstractmethod -from typing import Callable, Tuple +from typing import Callable, Optional, Tuple import torch.nn as nn from torch.optim import Optimizer -from ..interface import OptimizerWrapper +from colossalai.interface import OptimizerWrapper class MixedPrecision(ABC): @@ -15,7 +15,8 @@ class MixedPrecision(ABC): @abstractmethod def configure(self, model: nn.Module, - optimizer: Optimizer, - criterion: Callable = None) -> Tuple[nn.Module, OptimizerWrapper, Callable]: + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable]: # TODO: implement this method pass diff --git a/colossalai/booster/plugin.py b/colossalai/booster/plugin.py deleted file mode 100644 index 32e0a7bde3f7..000000000000 --- a/colossalai/booster/plugin.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import List, Tuple - -import torch -import torch.nn as nn -from torch.optim import Optimizer -from torch.utils.data import DataLoader - -from colossalai.device.device_mesh import DeviceMesh - -__all__ = ['Plugin'] - - -class Plugin: - - @property - def supported_devices(self) -> List[torch.device]: - pass - - @property - def supported_precisions(self) -> List[str]: - pass - - @property - def control_precision(self) -> bool: - pass - - @property - def control_device(self) -> bool: - pass - - @property - def support_no_sync(self) -> bool: - pass - - def setup_model(self, model: nn.Module, device_mesh_pool: DeviceMesh) -> nn.Module: - pass - - def setup_optimizer(self, optimizer: Optimizer) -> Optimizer: - pass - - def setup_dataloader(self, dataloader: DataLoader) -> DataLoader: - pass - - @property - def device_mesh_shape(self) -> List[Tuple[int, ...]]: - pass diff --git a/colossalai/booster/plugin/__init__.py b/colossalai/booster/plugin/__init__.py new file mode 100644 index 000000000000..a3b87b5f11d3 --- /dev/null +++ b/colossalai/booster/plugin/__init__.py @@ -0,0 +1,13 @@ +from .gemini_plugin import GeminiPlugin +from .low_level_zero_plugin import LowLevelZeroPlugin +from .plugin_base import Plugin +from .torch_ddp_plugin import TorchDDPPlugin + +__all__ = ['Plugin', 'TorchDDPPlugin', 'GeminiPlugin', 'LowLevelZeroPlugin'] + +import torch +from packaging import version + +if version.parse(torch.__version__) >= version.parse('1.12.0'): + from .torch_fsdp_plugin import TorchFSDPPlugin + __all__.append('TorchFSDPPlugin') diff --git a/colossalai/booster/plugin/dp_plugin_base.py b/colossalai/booster/plugin/dp_plugin_base.py new file mode 100644 index 000000000000..d5da5938bfd9 --- /dev/null +++ b/colossalai/booster/plugin/dp_plugin_base.py @@ -0,0 +1,70 @@ +import random + +import numpy as np +import torch +import torch.distributed as dist +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from .plugin_base import Plugin + + +class DPPluginBase(Plugin): + """This is a base class for all DP plugins. It sets up world size and rank, and provides data loader creation. + """ + + def __init__(self) -> None: + super().__init__() + assert dist.is_initialized( + ), 'torch.distributed is not initialized, please use colossalai.launch to create the distributed environment' + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + + def prepare_dataloader(self, + dataset, + batch_size, + shuffle=False, + seed=1024, + drop_last=False, + pin_memory=False, + num_workers=0, + **kwargs): + r""" + Prepare a dataloader for distributed training. The dataloader will be wrapped by + `torch.utils.data.DataLoader` and `torch.utils.data.DistributedSampler`. + + + Args: + dataset (`torch.utils.data.Dataset`): The dataset to be loaded. + shuffle (bool, optional): Whether to shuffle the dataset. Defaults to False. + seed (int, optional): Random worker seed for sampling, defaults to 1024. + add_sampler: Whether to add ``DistributedDataParallelSampler`` to the dataset. Defaults to True. + drop_last (bool, optional): Set to True to drop the last incomplete batch, if the dataset size + is not divisible by the batch size. If False and the size of dataset is not divisible by + the batch size, then the last batch will be smaller, defaults to False. + pin_memory (bool, optional): Whether to pin memory address in CPU memory. Defaults to False. + num_workers (int, optional): Number of worker threads for this dataloader. Defaults to 0. + kwargs (dict): optional parameters for ``torch.utils.data.DataLoader``, more details could be found in + `DataLoader `_. + + Returns: + :class:`torch.utils.data.DataLoader`: A DataLoader used for training or testing. + """ + _kwargs = kwargs.copy() + sampler = DistributedSampler(dataset, num_replicas=self.world_size, rank=self.rank, shuffle=shuffle) + + # Deterministic dataloader + def seed_worker(worker_id): + worker_seed = seed + np.random.seed(worker_seed) + torch.manual_seed(worker_seed) + random.seed(worker_seed) + + return DataLoader(dataset, + batch_size=batch_size, + sampler=sampler, + worker_init_fn=seed_worker, + drop_last=drop_last, + pin_memory=pin_memory, + num_workers=num_workers, + **_kwargs) diff --git a/colossalai/booster/plugin/gemini_plugin.py b/colossalai/booster/plugin/gemini_plugin.py new file mode 100644 index 000000000000..7b6e17337d36 --- /dev/null +++ b/colossalai/booster/plugin/gemini_plugin.py @@ -0,0 +1,412 @@ +import gc +import logging +import os +import warnings +from pathlib import Path +from typing import Callable, Iterator, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader + +from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO +from colossalai.checkpoint_io.utils import ( + get_model_base_filenames, + get_optimizer_base_filenames, + get_shard_filename, + load_shard_state_dict, + save_state_dict, + save_state_dict_shards, +) +from colossalai.cluster import DistCoordinator +from colossalai.interface import ModelWrapper, OptimizerWrapper +from colossalai.utils import get_current_device +from colossalai.zero import GeminiDDP, zero_model_wrapper, zero_optim_wrapper +from colossalai.zero.gemini import ZeroOptimizer +from colossalai.zero.gemini.memory_tracer import MemStats + +from .dp_plugin_base import DPPluginBase + +__all__ = ['GeminiPlugin'] + +SUPPORTED_PRECISION = ['fp16', 'bf16'] +PRECISION_STR_TO_DTYPE = {'fp16': torch.half, 'bf16': torch.bfloat16} + + +class GeminiCheckpointIO(GeneralCheckpointIO): + + def __init__(self) -> None: + super().__init__() + self.coordinator = DistCoordinator() + + def save_unsharded_model(self, model: GeminiDDP, checkpoint: str, gather_dtensor: bool, use_safetensors: bool): + """ + Save sharded model to checkpoint but only on master process. + The model should be unwrapped in self.load_model via ModelWrapper.unwrap. + As there is communication when getting state dict, model.state_dict() must be called on all processes. + """ + state_dict = model.state_dict(only_rank_0=True) + if self.coordinator.is_master(): + save_state_dict(state_dict, checkpoint, use_safetensors) + + def load_unsharded_model(self, model: GeminiDDP, checkpoint: str, strict: bool = True): + """ + Load model from checkpoint with automatic unwrapping. + The model should be unwrapped in self.load_model via ModelWrapper.unwrap. + """ + super().load_unsharded_model(model, checkpoint, strict=strict) + + def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool): + """ + Save unsharded optimizer state dict to checkpoint. + After calling optimizer.state_dict(), the complete optimizer states will be collected on master rank. + As there is communication when getting state dict, optimizer.state_dict() must be called on all processes. + The saving process will only be executed by master rank. + """ + state_dict = optimizer.state_dict() + if self.coordinator.is_master(): + save_state_dict(state_dict, checkpoint, use_safetensors=False) + + def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str): + """ + Loading unsharded optimizer from checkpoint file. + For each process, only loading optimizer states of parameters it controls. + """ + super().load_unsharded_optimizer(optimizer, checkpoint) + + def save_sharded_model(self, + model: GeminiDDP, + checkpoint_path: str, + gather_dtensor: bool = False, + prefix: Optional[str] = None, + max_shard_size: int = 1024, + use_safetensors: bool = False): + """ + Save sharded model. + As there is communication when getting state dict, model.state_dict() must be called on all processes. + """ + if os.path.isfile(checkpoint_path): + logging.error(f"Provided path ({checkpoint_path}) should be a directory, not a file") + return + + Path(checkpoint_path).mkdir(parents=True, exist_ok=True) + + state_dict_shard = model.state_dict_shard(max_shard_size=max_shard_size, only_rank_0=True, dtype=torch.float32) + weights_name, save_index_file = get_model_base_filenames(prefix, use_safetensors) + index_file = CheckpointIndexFile(checkpoint_path) + + # Save shards of optimizer states. + is_master = self.coordinator.is_master() + total_size = save_state_dict_shards(sharded_state_dict=state_dict_shard, + checkpoint=checkpoint_path, + index_file=index_file, + base_filename=weights_name, + is_master=is_master, + use_safetensors=use_safetensors) + + # only save the index file on the master rank + if self.coordinator.is_master(): + index_file.append_meta_data("total_size", total_size) + index_file.write_index_file(save_index_file) + logging.info(f"The model is split into checkpoint shards. " + f"You can find where each parameters has been saved in the " + f"index located at {save_index_file}.") + + def load_sharded_model(self, + model: GeminiDDP, + checkpoint_index_file: Path, + strict: bool = False, + use_safetensors: bool = False): + """ + Load shard model, load model from multiple files. + """ + return super().load_sharded_model(model, checkpoint_index_file, strict, use_safetensors, load_sub_module=False) + + def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: Path, gather_dtensor: bool, prefix: str, + size_per_shard: int): + """ + Save sharded optimizer state dict to checkpoint folder. + As there is communication when getting state dict, this must be called on all processes. + """ + + # If optimizer is wrapped, unwrap it. + if isinstance(optimizer, OptimizerWrapper): + optimizer = optimizer.unwrap() + + assert isinstance(optimizer, ZeroOptimizer) + + if os.path.isfile(checkpoint): + logging.error(f"Provided path ({checkpoint}) should be a directory, not a file") + return + + Path(checkpoint).mkdir(parents=True, exist_ok=True) + + # Preparing file paths and index file. + states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix) + index_file = CheckpointIndexFile(checkpoint) + + # Store the information of param groups to param_group_file. + index_file.append_meta_data("param_groups", param_group_file) + group_file_path = os.path.join(checkpoint, param_group_file) + param_groups = optimizer.get_param_groups_for_saving() + torch.save(param_groups, group_file_path) + + # States are broken into shards within max_shard_size. + state_dict_shard = optimizer.state_shard(prefix=prefix, max_shard_size=size_per_shard, only_rank_0=True) + + # Save shards of optimizer states. + is_master = self.coordinator.is_master() + total_size = save_state_dict_shards(sharded_state_dict=state_dict_shard, + checkpoint=checkpoint, + index_file=index_file, + base_filename=states_name, + is_master=is_master, + use_safetensors=False) + + # Wrap up index file. Only save it on master rank. + if self.coordinator.is_master(): + index_file.append_meta_data("total_size", total_size) + index_file.write_index_file(save_index_file) + logging.info(f"The optimizer is going to be split to checkpoint shards. " + f"You can find where each parameters has been saved in the " + f"index located at {save_index_file}.") + + def load_sharded_optimizer(self, optimizer: Optimizer, checkpoint_index_file: Path, prefix: str): + """ + Loading sharded optimizer from checkpoint folder, with index file given. + For each process, only loading optimizer states of parameters it controls. + """ + + if not os.path.isfile(checkpoint_index_file): + logging.error(f"Provided path ({checkpoint_index_file}) should be a file") + + # If optimizer is wrapped, unwrap it. + if isinstance(optimizer, OptimizerWrapper): + optimizer = optimizer.unwrap() + + assert isinstance(optimizer, ZeroOptimizer) + + # Read checkpoint index file. + ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) + + # Load param_groups. + param_group_path = ckpt_index_file.get_param_group_filename() + if param_group_path is None: + raise RuntimeError(f'Invalid index file path {checkpoint_index_file} for an optimizer. \ + Lacking param group file under current directory.') + saved_param_groups = torch.load(param_group_path) + optimizer.load_param_groups(saved_param_groups) + + checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames() + + # Load optimizer states from shard files under checkpoint path. + # For each file, only load the states managed by current process. + for shard_file in checkpoint_files: + state_dict_shard = load_shard_state_dict(Path(shard_file), use_safetensors=False) + optimizer.load_param_states(state_dict_shard) + del state_dict_shard + gc.collect() + + optimizer.optimizer_loading_epilogue() + + def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """ + Save model to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_lr_scheduler(lr_scheduler, checkpoint) + + +class GeminiModel(ModelWrapper): + + def __init__(self, module: nn.Module, gemini_config: dict, verbose: bool = False) -> None: + super().__init__(module) + self.module = zero_model_wrapper(module, zero_stage=3, gemini_config=gemini_config, verbose=verbose) + + def unwrap(self): + # as save/load state dict is coupled with the GeminiDDP, we only return GeminiDDP model + return self.module + + +class GeminiOptimizer(OptimizerWrapper): + + def __init__(self, + module: GeminiDDP, + optimizer: Optimizer, + zero_optim_config: dict, + optim_kwargs: dict, + verbose: bool = False) -> None: + optimizer = zero_optim_wrapper(module, + optimizer, + optim_config=zero_optim_config, + **optim_kwargs, + verbose=verbose) + super().__init__(optimizer) + + def backward(self, loss: Tensor, *args, **kwargs): + self.optim.backward(loss) + + def clip_grad_by_norm(self, + max_norm: Union[float, int], + norm_type: Union[float, int] = 2, + error_if_nonfinite: bool = False, + *args, + **kwargs) -> Tensor: + warnings.warn(f'Gemini controls grad clipping by itself, so you should not use clip_grad_by_norm') + + def clip_grad_by_value(self, clip_value: float, *args, **kwargs) -> None: + raise NotImplementedError('Gemini does not support clip_grad_by_value') + + +class GeminiPlugin(DPPluginBase): + """ + Plugin for Gemini. + + Example: + >>> from colossalai.booster import Booster + >>> from colossalai.booster.plugin import GeminiPlugin + >>> + >>> model, train_dataset, optimizer, criterion = ... + >>> plugin = GeminiPlugin() + + >>> train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8) + >>> booster = Booster(plugin=plugin) + >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion) + + Args: + device (torch.device): device to place the model. + placement_policy (str, optional): "cpu", "cuda", "auto". Defaults to "cpu". + precision (str, optional): precision. Support 'fp16' and 'bf16'. Defaults to 'fp16'. + pin_memory (bool, optional): use pin memory on CPU. Defaults to False. + force_outputs_fp32 (bool, optional): force outputs are fp32. Defaults to False. + strict_ddp_mode (bool, optional): use strict ddp mode (only use dp without other parallelism). Defaults to False. + search_range_m (int, optional): chunk size searching range divided by 2^20. Defaults to 32. + hidden_dim (int, optional): the hidden dimension of DNN. + Users can provide this argument to speed up searching. + If users do not know this argument before training, it is ok. We will use a default value 1024. + min_chunk_size_m (float, optional): the minimum chunk size divided by 2^20. + If the aggregate size of parameters is still smaller than the minimum chunk size, + all parameters will be compacted into one small chunk. + memstats (MemStats, optional) the memory statistics collector by a runtime memory tracer. + gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward) + which will be used when using hybrid CPU optimizer. + This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto". + Defaults to 0.0. + initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**16. + min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. + growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2. + backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5. + growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000. + hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2. + max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32. + max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do + clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm. + norm_type (float, optional): norm_type used for `clip_grad_norm`. + verbose (bool, optional): verbose mode. Debug info including chunk search result will be printed. Defaults to False. + """ + + def __init__( + self, + device: Optional[torch.device] = None, + placement_policy: str = "cpu", + precision: str = "fp16", + pin_memory: bool = False, + force_outputs_fp32: bool = False, + strict_ddp_mode: bool = False, + search_range_m: int = 32, + hidden_dim: Optional[int] = None, + min_chunk_size_m: float = 32, + memstats: Optional[MemStats] = None, + gpu_margin_mem_ratio: float = 0.0, + initial_scale: float = 2**16, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32, + max_norm: float = 0.0, + norm_type: float = 2.0, + verbose: bool = False, + ) -> None: + super().__init__() + assert precision in SUPPORTED_PRECISION, f'precision {precision} is not supported' + self.gemini_config = dict( + device=(device or get_current_device()), + placement_policy=placement_policy, + pin_memory=pin_memory, + force_outputs_fp32=force_outputs_fp32, + strict_ddp_mode=strict_ddp_mode, + search_range_m=search_range_m, + hidden_dim=hidden_dim, + min_chunk_size_m=min_chunk_size_m, + memstats=memstats, + mixed_precision=PRECISION_STR_TO_DTYPE[precision], + ) + self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio,) + self.optim_kwargs = dict(initial_scale=initial_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + min_scale=min_scale, + max_scale=max_scale, + max_norm=max_norm, + norm_type=norm_type) + self.verbose = verbose + + def support_no_sync(self) -> bool: + return False + + def control_precision(self) -> bool: + return True + + def supported_precisions(self) -> List[str]: + return SUPPORTED_PRECISION + + def control_device(self) -> bool: + return True + + def supported_devices(self) -> List[str]: + return ['cuda'] + + def configure( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]: + + if not isinstance(model, ModelWrapper): + # convert model to sync bn + # FIXME(ver217): gemini does not support sync bn + # In torch/nn/modules/_functions.py, line 22, ``mean, invstd = torch.batch_norm_stats(input, eps)`` will get fp32 mean and invstd even though the input is fp16. + # This inconsistency of dtype will cause the error. + # We have two possible solutions: + # 1. keep batch norm always in fp32. This is hard for gemini, as it use chunks. + # 2. patch sync bn or write a new on. This is relatively easy, but we need to test it. + # model = nn.SyncBatchNorm.convert_sync_batchnorm(model, None) + + # wrap the model with Gemini + model = GeminiModel(model, self.gemini_config, self.verbose) + + if optimizer is not None and \ + not isinstance(optimizer, OptimizerWrapper): + optimizer = GeminiOptimizer(model.unwrap(), optimizer, self.zero_optim_config, self.optim_kwargs, + self.verbose) + + return model, optimizer, criterion, dataloader, lr_scheduler + + def control_checkpoint_io(self) -> bool: + return True + + def get_checkpoint_io(self) -> CheckpointIO: + return GeminiCheckpointIO() + + def no_sync(self, model: nn.Module) -> Iterator[None]: + raise NotImplementedError diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py new file mode 100644 index 000000000000..94d722080367 --- /dev/null +++ b/colossalai/booster/plugin/low_level_zero_plugin.py @@ -0,0 +1,226 @@ +import warnings +from functools import partial +from typing import Callable, Iterator, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils._pytree import tree_map +from torch.utils.data import DataLoader + +from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO +from colossalai.interface import ModelWrapper, OptimizerWrapper +from colossalai.utils import get_current_device +from colossalai.zero import zero_model_wrapper, zero_optim_wrapper + +from .dp_plugin_base import DPPluginBase +from .torch_ddp_plugin import TorchDDPCheckpointIO + +__all__ = ['LowLevelZeroPlugin'] + + +def _convert_floating_point(x, dtype: torch.dtype = torch.float16): + if isinstance(x, torch.Tensor) and torch.is_floating_point(x): + return x.to(dtype) + return x + + +SUPPORTED_PRECISION = ['fp16', 'bf16', 'fp32'] + + +class LowLevelZeroCheckpointIO(TorchDDPCheckpointIO): + + def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool): + """ + Save optimizer to checkpoint but only on master process. + """ + # TODO(ver217): optimizer state dict is sharded, and cannot get full state dict now + warnings.warn( + 'LowLevelZeroPlugin does not support save full optimizer checkpoint now. Save it on every process.') + checkpoint = f'{checkpoint}.rank{self.coordinator.rank}' + GeneralCheckpointIO.save_unsharded_optimizer(self, optimizer, checkpoint, gather_dtensor) + + def load_optimizer(self, optimizer: Optimizer, checkpoint: str): + warnings.warn( + 'LowLevelZeroPlugin can only load optimizer checkpoint saved by itself with the same number of processes.') + checkpoint = f'{checkpoint}.rank{self.coordinator.rank}' + super().load_optimizer(optimizer, checkpoint) + + +class LowLevelZeroModel(ModelWrapper): + + def __init__(self, module: nn.Module, stage: int, precision: str) -> None: + super().__init__(module) + self.dtype = None + if precision == 'fp16': + self.dtype = torch.float16 + elif precision == 'bf16': + self.dtype = torch.bfloat16 + module = zero_model_wrapper(module, zero_stage=stage) + if self.dtype is not None: + module = module.to(self.dtype) + module = module.to(get_current_device()) + self.module = module + self.convert_fn = None + if self.dtype is not None: + self.convert_fn = partial(_convert_floating_point, dtype=self.dtype) + + def forward(self, *args, **kwargs): + if self.convert_fn is not None: + args = tree_map(self.convert_fn, args) + kwargs = tree_map(self.convert_fn, kwargs) + return super().forward(*args, **kwargs) + + +class LowLevelZeroOptimizer(OptimizerWrapper): + + def __init__(self, + module: nn.Module, + optimizer: Optimizer, + zero_optim_config: dict, + optim_kwargs: dict, + verbose: bool = False) -> None: + optimizer = zero_optim_wrapper(module, + optimizer, + optim_config=zero_optim_config, + **optim_kwargs, + verbose=verbose) + super().__init__(optimizer) + + def backward(self, loss: Tensor, *args, **kwargs): + self.optim.backward(loss) + + def clip_grad_by_norm(self, + max_norm: Union[float, int], + norm_type: Union[float, int] = 2, + error_if_nonfinite: bool = False, + *args, + **kwargs) -> Tensor: + warnings.warn(f'LowLevelZero controls grad clipping by itself, so you should not use clip_grad_by_norm') + + def clip_grad_by_value(self, clip_value: float, *args, **kwargs) -> None: + raise NotImplementedError('LowLevelZero does not support clip_grad_by_value') + + +class LowLevelZeroPlugin(DPPluginBase): + """ + Plugin for low level zero. + + Example: + >>> from colossalai.booster import Booster + >>> from colossalai.booster.plugin import LowLevelZeroPlugin + >>> + >>> model, train_dataset, optimizer, criterion = ... + >>> plugin = LowLevelZeroPlugin() + + >>> train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8) + >>> booster = Booster(plugin=plugin) + >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion) + + Args: + strage (int, optional): ZeRO stage. Defaults to 1. + precision (str, optional): precision. Support 'fp16', 'bf16' and 'fp32'. Defaults to 'fp16'. + initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32. + min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. + growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2. + backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5. + growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000. + hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2. + max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32. + max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do + clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm. + norm_type (float, optional): norm_type used for `clip_grad_norm`. + reduce_bucket_size_in_m (int, optional): grad reduce bucket size in M. Defaults to 12. + communication_dtype (torch.dtype, optional): communication dtype. If not specified, the dtype of param will be used. Defaults to None. + overlap_communication (bool, optional): whether to overlap communication and computation. Defaults to True. + cpu_offload (bool, optional): whether to offload grad, master weight and optimizer state to cpu. Defaults to False. + verbose (bool, optional): verbose mode. Debug info including grad overflow will be printed. Defaults to False. + """ + + def __init__( + self, + stage: int = 1, + precision: str = 'fp16', + initial_scale: float = 2**32, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32, + max_norm: float = 0.0, + norm_type: float = 2.0, + reduce_bucket_size_in_m: int = 12, + communication_dtype: Optional[torch.dtype] = None, + overlap_communication: bool = True, + cpu_offload: bool = False, + verbose: bool = False, + ) -> None: + super().__init__() + assert stage in (1, 2), f'LowLevelZeroPlugin only supports stage 1/2 training' + assert precision in SUPPORTED_PRECISION, f'LowLevelZeroPlugin only supports {SUPPORTED_PRECISION} training' + + self.stage = stage + self.precision = precision + self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size_in_m * 1024 * 1024, + communication_dtype=communication_dtype, + overlap_communication=overlap_communication, + cpu_offload=cpu_offload) + self.optim_kwargs = dict(initial_scale=initial_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + min_scale=min_scale, + max_scale=max_scale, + max_norm=max_norm, + norm_type=norm_type) + self.verbose = verbose + + def support_no_sync(self) -> bool: + return False + + def control_precision(self) -> bool: + return True + + def supported_precisions(self) -> List[str]: + return SUPPORTED_PRECISION + + def control_device(self) -> bool: + return True + + def supported_devices(self) -> List[str]: + return ['cuda'] + + def configure( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]: + + if not isinstance(model, ModelWrapper): + model = LowLevelZeroModel(model, self.stage, self.precision) + + if optimizer is not None and \ + not isinstance(optimizer, OptimizerWrapper): + optimizer = LowLevelZeroOptimizer(model.unwrap(), + optimizer, + self.zero_optim_config, + self.optim_kwargs, + self.verbose) + + return model, optimizer, criterion, dataloader, lr_scheduler + + def control_checkpoint_io(self) -> bool: + return True + + def get_checkpoint_io(self) -> CheckpointIO: + return LowLevelZeroCheckpointIO() + + def no_sync(self, model: nn.Module) -> Iterator[None]: + raise NotImplementedError diff --git a/colossalai/booster/plugin/plugin_base.py b/colossalai/booster/plugin/plugin_base.py new file mode 100644 index 000000000000..aa78f6827003 --- /dev/null +++ b/colossalai/booster/plugin/plugin_base.py @@ -0,0 +1,83 @@ +from abc import ABC, abstractmethod +from typing import Callable, Iterator, List, Optional, Tuple, Union + +import torch.nn as nn +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader, Dataset + +from colossalai.checkpoint_io import CheckpointIO +from colossalai.interface import OptimizerWrapper + +__all__ = ['Plugin'] + + +class Plugin(ABC): + + @abstractmethod + def supported_devices(self) -> List[str]: + pass + + @abstractmethod + def supported_precisions(self) -> List[str]: + pass + + @abstractmethod + def control_precision(self) -> bool: + pass + + @abstractmethod + def control_device(self) -> bool: + pass + + @abstractmethod + def support_no_sync(self) -> bool: + pass + + @abstractmethod + def configure( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]: + # implement this method + pass + + @abstractmethod + def control_checkpoint_io(self) -> bool: + """ + Whether the plugin controls the checkpoint io + """ + pass + + @abstractmethod + def get_checkpoint_io(self) -> CheckpointIO: + """ + Get checkpoint io object for this plugin, only invoked when control_checkpoint_io is True. + """ + pass + + @abstractmethod + def no_sync(self, model: nn.Module) -> Iterator[None]: + """ + Context manager to disable gradient synchronization. + """ + pass + + @abstractmethod + def prepare_dataloader(self, + dataset: Dataset, + batch_size: int, + shuffle: bool = False, + seed: int = 1024, + drop_last: bool = False, + pin_memory: bool = False, + num_workers: int = 0, + **kwargs): + """Prepare a dataloader for distributed training. The dataloader will be wrapped by + `torch.utils.data.DataLoader` + """ + pass diff --git a/colossalai/booster/plugin/torch_ddp_plugin.py b/colossalai/booster/plugin/torch_ddp_plugin.py new file mode 100644 index 000000000000..71b435155503 --- /dev/null +++ b/colossalai/booster/plugin/torch_ddp_plugin.py @@ -0,0 +1,173 @@ +from typing import Callable, Iterator, List, Optional, Tuple, Union + +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader + +from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO +from colossalai.cluster import DistCoordinator +from colossalai.interface import ModelWrapper, OptimizerWrapper + +from .dp_plugin_base import DPPluginBase + +__all__ = ['TorchDDPPlugin'] + + +class TorchDDPCheckpointIO(GeneralCheckpointIO): + + def __init__(self) -> None: + super().__init__() + self.coordinator = DistCoordinator() + + def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool = True): + """ + Load model from checkpoint with automatic unwrapping. + """ + # the model should be unwrapped in self.load_model via ModelWrapper.unwrap + return super().load_unsharded_model(model, checkpoint, strict=strict) + + def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool): + """ + Save model to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors) + + def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool): + """ + Save optimizer to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor) + + def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """ + Save model to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_lr_scheduler(lr_scheduler, checkpoint) + + def save_sharded_model(self, + model: nn.Module, + checkpoint_path: str, + gather_dtensor: bool = True, + prefix: Optional[str] = None, + max_shard_size: int = 1024, + use_safetensors: bool = False): + """ + Save model to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_sharded_model(model, checkpoint_path, gather_dtensor, prefix, max_shard_size, use_safetensors) + + def save_sharded_optimizer(self, + optimizer: Optimizer, + checkpoint: str, + gather_dtensor: bool = True, + prefix: Optional[str] = None, + size_per_shard: int = 1024): + """ + Save optimizer to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard) + + +class TorchDDPModel(ModelWrapper): + + def __init__(self, module: nn.Module, *args, **kwargs) -> None: + super().__init__(module) + self.module = DDP(module, *args, **kwargs) + + def unwrap(self): + return self.module.module + + +class TorchDDPPlugin(DPPluginBase): + """ + Plugin for PyTorch DDP. + + Example: + >>> from colossalai.booster import Booster + >>> from colossalai.booster.plugin import TorchDDPPlugin + >>> + >>> model, train_dataset, optimizer, criterion = ... + >>> plugin = TorchDDPPlugin() + + >>> train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8) + >>> booster = Booster(plugin=plugin) + >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion) + + Args: + broadcast_buffers (bool, optional): Whether to broadcast buffers in the beginning of training. Defaults to True. + bucket_cap_mb (int, optional): The bucket size in MB. Defaults to 25. + find_unused_parameters (bool, optional): Whether to find unused parameters. Defaults to False. + check_reduction (bool, optional): Whether to check reduction. Defaults to False. + gradient_as_bucket_view (bool, optional): Whether to use gradient as bucket view. Defaults to False. + static_graph (bool, optional): Whether to use static graph. Defaults to False. + """ + + def __init__(self, + broadcast_buffers: bool = True, + bucket_cap_mb: int = 25, + find_unused_parameters: bool = False, + check_reduction: bool = False, + gradient_as_bucket_view: bool = False, + static_graph: bool = False) -> None: + super().__init__() + self.ddp_kwargs = dict(broadcast_buffers=broadcast_buffers, + bucket_cap_mb=bucket_cap_mb, + find_unused_parameters=find_unused_parameters, + check_reduction=check_reduction, + gradient_as_bucket_view=gradient_as_bucket_view, + static_graph=static_graph) + + def support_no_sync(self) -> bool: + return True + + def control_precision(self) -> bool: + return False + + def supported_precisions(self) -> List[str]: + return ['fp16', 'fp16_apex', 'bf16', 'fp8'] + + def control_device(self) -> bool: + return True + + def supported_devices(self) -> List[str]: + return ['cuda'] + + def configure( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]: + # cast model to cuda + model = model.cuda() + + # convert model to sync bn + model = nn.SyncBatchNorm.convert_sync_batchnorm(model, None) + + # wrap the model with PyTorch DDP + model = TorchDDPModel(model, **self.ddp_kwargs) + + if optimizer is not None and \ + not isinstance(optimizer, OptimizerWrapper): + optimizer = OptimizerWrapper(optimizer) + + return model, optimizer, criterion, dataloader, lr_scheduler + + def control_checkpoint_io(self) -> bool: + return True + + def get_checkpoint_io(self) -> CheckpointIO: + return TorchDDPCheckpointIO() + + def no_sync(self, model: nn.Module) -> Iterator[None]: + assert isinstance(model, TorchDDPModel), 'Model is not boosted by TorchDDPPlugin.' + return model.module.no_sync() diff --git a/colossalai/booster/plugin/torch_fsdp_plugin.py b/colossalai/booster/plugin/torch_fsdp_plugin.py new file mode 100644 index 000000000000..abfffa9b099e --- /dev/null +++ b/colossalai/booster/plugin/torch_fsdp_plugin.py @@ -0,0 +1,223 @@ +import warnings +from pathlib import Path +from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from packaging import version +from torch.distributed import ProcessGroup + +if version.parse(torch.__version__) >= version.parse('1.12.0'): + from torch.distributed.fsdp import FullStateDictConfig + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp import StateDictType + from torch.distributed.fsdp.fully_sharded_data_parallel import ( + BackwardPrefetch, + CPUOffload, + FullStateDictConfig, + MixedPrecision, + ShardingStrategy, + ) +else: + raise RuntimeError("FSDP is not supported while torch version under 1.12.0.") + +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader + +from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO, utils +from colossalai.cluster import DistCoordinator +from colossalai.interface import ModelWrapper, OptimizerWrapper + +from .dp_plugin_base import DPPluginBase + +__all__ = ['TorchFSDPPlugin'] + + +class TorchFSDPCheckpointIO(GeneralCheckpointIO): + + def __init__(self) -> None: + super().__init__() + self.coordinator = DistCoordinator() + + def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool): + checkpoint = utils.load_state_dict(checkpoint) + model.load_state_dict(checkpoint) + + def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path): + checkpoint = utils.load_state_dict(checkpoint) + fsdp_model = optimizer.unwrap_model() + sharded_osd = FSDP.scatter_full_optim_state_dict(checkpoint, fsdp_model) + optimizer.load_state_dict(sharded_osd) + + def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool): + """ + Save model to checkpoint but only on master process. + """ + # the model should be unwrapped in self.load_model via ModelWrapper.unwrap + cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg): + full_model_state = model.state_dict() + utils.save_state_dict(full_model_state, checkpoint_file_path=checkpoint, use_safetensors=use_safetensors) + + def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool): + """ + Save optimizer to checkpoint but only on master process. + """ + assert isinstance(optimizer, FSDPOptimizerWrapper) + fsdp_model = optimizer.unwrap_model() + full_optimizer_state = FSDP.full_optim_state_dict(fsdp_model, optim=optimizer, rank0_only=True) + utils.save_state_dict(full_optimizer_state, checkpoint_file_path=checkpoint, use_safetensors=False) + + def save_sharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, prefix: Optional[str], + size_per_shard: int, use_safetensors: bool): + """ + Save model to checkpoint but only on master process. + """ + raise NotImplementedError("Sharded model checkpoint is not supported yet.") + + def load_sharded_model(self, + model: nn.Module, + checkpoint_index_file: Path, + strict: bool = False, + use_safetensors: bool = False, + load_sub_module: bool = True): + """ + Load model to checkpoint but only on master process. + """ + raise NotImplementedError("Sharded model checkpoint is not supported yet.") + + def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool, prefix: str, + size_per_shard: int): + """ + Save optimizer to checkpoint but only on master process. + """ + raise NotImplementedError("Sharded optimizer checkpoint is not supported yet.") + + def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, size_per_shard: int): + """ + Load optimizer to checkpoint but only on master process. + """ + raise NotImplementedError("Sharded optimizer checkpoint is not supported yet.") + + def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """ + Save model to checkpoint but only on master process. + """ + if self.coordinator.is_master(): + super().save_lr_scheduler(lr_scheduler, checkpoint) + + +class TorchFSDPModel(ModelWrapper): + + def __init__(self, module: nn.Module, *args, **kwargs) -> None: + super().__init__(module) + self.module = FSDP(module, *args, **kwargs) + + def unwrap(self): + return self.module + + +class FSDPOptimizerWrapper(OptimizerWrapper): + + def __init__(self, optimizer: Optimizer, model: nn.Module): + self.model = model + super().__init__(optimizer) + + def unwrap_model(self) -> nn.Module: + return self.model + + +class TorchFSDPPlugin(DPPluginBase): + """ + Plugin for PyTorch FSDP. + + Example: + >>> from colossalai.booster import Booster + >>> from colossalai.booster.plugin import TorchFSDPPlugin + >>> + >>> model, train_dataset, optimizer, criterion = ... + >>> plugin = TorchFSDPPlugin() + + >>> train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8) + >>> booster = Booster(plugin=plugin) + >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion) + + Args: + See https://pytorch.org/docs/stable/fsdp.html for details. + """ + + if version.parse(torch.__version__) >= version.parse('1.12.0'): + + def __init__( + self, + process_group: Optional[ProcessGroup] = None, + sharding_strategy: Optional[ShardingStrategy] = None, + cpu_offload: Optional[CPUOffload] = None, + auto_wrap_policy: Optional[Callable] = None, + backward_prefetch: Optional[BackwardPrefetch] = None, + mixed_precision: Optional[MixedPrecision] = None, + ignored_modules: Optional[Iterable[torch.nn.Module]] = None, + param_init_fn: Optional[Callable[[nn.Module], None]] = None, + sync_module_states: bool = False, + ): + super().__init__() + self.fsdp_kwargs = dict(process_group=process_group, + sharding_strategy=sharding_strategy, + cpu_offload=cpu_offload, + auto_wrap_policy=auto_wrap_policy, + backward_prefetch=backward_prefetch, + mixed_precision=mixed_precision, + ignored_modules=ignored_modules, + param_init_fn=param_init_fn, + sync_module_states=sync_module_states) + else: + raise RuntimeError("FSDP is not supported while torch version under 1.12.0.") + + def support_no_sync(self) -> bool: + False + + def no_sync(self, model: nn.Module) -> Iterator[None]: + raise NotImplementedError("Torch fsdp no_sync func not supported yet.") + + def control_precision(self) -> bool: + return True + + def supported_precisions(self) -> List[str]: + return ['fp16', 'bf16'] + + def control_device(self) -> bool: + return True + + def supported_devices(self) -> List[str]: + return ['cuda'] + + def configure( + self, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + criterion: Optional[Callable] = None, + dataloader: Optional[DataLoader] = None, + lr_scheduler: Optional[LRScheduler] = None, + ) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]: + + # wrap the model with PyTorch FSDP + fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs) + + if optimizer is not None: + if len(optimizer.param_groups) > 1: + warnings.warn( + 'TorchFSDPPlugin does not support optimizer that use multi param groups. The results may not be as expected if used.' + ) + optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults) + + if not isinstance(optimizer, FSDPOptimizerWrapper): + optimizer = FSDPOptimizerWrapper(optimizer, fsdp_model) + + return fsdp_model, optimizer, criterion, dataloader, lr_scheduler + + def control_checkpoint_io(self) -> bool: + return True + + def get_checkpoint_io(self) -> CheckpointIO: + return TorchFSDPCheckpointIO() diff --git a/colossalai/checkpoint_io/__init__.py b/colossalai/checkpoint_io/__init__.py new file mode 100644 index 000000000000..c25048e25754 --- /dev/null +++ b/colossalai/checkpoint_io/__init__.py @@ -0,0 +1,5 @@ +from .checkpoint_io_base import CheckpointIO +from .general_checkpoint_io import GeneralCheckpointIO +from .index_file import CheckpointIndexFile + +__all__ = ['CheckpointIO', 'CheckpointIndexFile', 'GeneralCheckpointIO'] diff --git a/colossalai/checkpoint_io/checkpoint_io_base.py b/colossalai/checkpoint_io/checkpoint_io_base.py new file mode 100644 index 000000000000..baff24e1cb25 --- /dev/null +++ b/colossalai/checkpoint_io/checkpoint_io_base.py @@ -0,0 +1,332 @@ +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Optional, Union + +import torch +import torch.nn as nn +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler + +from colossalai.interface import ModelWrapper + +from .utils import has_index_file + +__all__ = ['CheckpointIO'] + + +class CheckpointIO(ABC): + """ + CheckpointIO is the base class for all checkpoint IO classes. It defines the interface for checkpoint IO. + + + Examples: + >>> from colossalai.checkpoint_io import GeneralCheckpointIO + >>> checkpoint_io = CheckpointIO() + >>> + >>> # load model from checkpoint + >>> model = checkpoint_io.load_model(model, 'model.pt') + >>> + >>> # save model to checkpoint, any distributed tensor is gathered by default + >>> checkpoint_io.save_model(model, 'model.pt') + >>> + >>> # if the model contains distributed tensor, and you don't want to gather it + >>> # each rank will save its own shard of the distributed tensor + >>> checkpoint_io.save_model(model, 'model.pt', gather_dtensor=False) + >>> + >>> # save model to sharded checkpoints + >>> checkpoint_io.save_model(model, './checkpoints/', shard=True) + >>> + >>> # save model to sharded and assume we don't want to gather distributed tensors + >>> checkpoint_io.save_model(model, './checkpoints/', shard=True, gather_dtensor=False) + >>> + >>> # Note: + >>> # 1. we don't support loading from distributed tensors, conversion from distributed tensors + >>> # checkpoints to full tensor checkpoint should be done offline via our CLI + >>> # 2. you don't have to specify whether the model is sharded or not when loading the model + >>> # as it will be automatically detected + >>> + >>> # load model from sharded checkpoints + >>> model = checkpoint_io.load_model(model, './checkpoints/') + >>> + >>> # load model from unsharded checkpoints + >>> model = checkpoint_io.load_model(model, './checkpoints/') + >>> + >>> # load optimizer from checkpoint + >>> optimizer = checkpoint_io.load_optimizer(optimizer, 'optimizer.pt') + >>> + >>> # save optimizer to checkpoint + >>> checkpoint_io.save_optimizer(optimizer, 'optimizer.pt') + """ + + # ====================================== + # Public methods + # ====================================== + def load_model(self, + model: Union[nn.Module, ModelWrapper], + checkpoint: str, + strict: bool = True) -> Union[nn.Module, ModelWrapper]: + """ + Load model from checkpoint. + + Args: + model (nn.Module): model to be loaded. + checkpoint (str): checkpoint path. This value is made compatibility with the model checkpoints in the + mainstream model zoos such as Hugging Face and TIMM. The checkpoint path can be: + 1. a file path, e.g. 'model.pt' + 2. a path to a json file which defines the index to the sharded checkpoint + 3. a path to a folder containing a unique .index.json file for sharded checkpoint + Distributed tensors cannot be loaded directly unless gathered offline via our CLI. + strict (bool): whether to strictly enforce that the param name in + the checkpoint match the keys returned by this module's. + """ + # since we only support loaded sharded and unsharded weight format + # containing no distributed tensors, dtensor -> full tensor conversion + # should be done offline via our CLI + # the existence of index file means it is a sharded checkpoint + index_file_exists, index_file_path = has_index_file(checkpoint) + + # return the origin model instead of the unwrapped model + origin_model = model + + if isinstance(model, ModelWrapper): + model = model.unwrap() + + if index_file_exists: + self.load_sharded_model(model, index_file_path, strict) + else: + self.load_unsharded_model(model, checkpoint, strict) + + return origin_model + + def save_model(self, + model: Union[nn.Module, ModelWrapper], + checkpoint: str, + shard: bool = False, + gather_dtensor: bool = True, + prefix: str = None, + size_per_shard: int = 1024, + use_safetensors: bool = False): + """ + Save model to checkpoint. + + Examples: + >>> from colossalai.checkpoint_io import GeneralCheckpointIO + >>> checkpoint_io = CheckpointIO() + >>> + >>> # save model to a single file + >>> save_model(model, 'model.pt') + >>> + >>> # save model to a sharded checkpoint + >>> save_model(model, './checkpoints/', shard=True) + + Args: + model (nn.Module): model to be saved. + checkpoint (str): checkpoint path. The checkpoint path can be : + 1. a file path, e.g. 'model.pt' + 2. a directory path to save the sharded checkpoint, e.g. './checkpoints/' when shard = True. + shard (bool): whether to shard the checkpoint. Default: False. If set to True, the checkpoint will be sharded into + multiple files. The model shards will be specified by a `model.index.json` file. When shard = True, please ensure + that the checkpoint path is a directory path instead of a file path. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. Default: True. + prefix (str): If specified, weights are saved in the format pytorch_model..bin. Default: None. + size_per_shard (int): size per shard in MB. Default: 1024. This value is only used when shard = True. + use_safetensors (bool): whether to use safe tensors. Default: False. If set to True, the checkpoint will be saved + """ + + if isinstance(model, ModelWrapper): + model = model.unwrap() + + if shard: + self.save_sharded_model(model, checkpoint, gather_dtensor, prefix, size_per_shard, use_safetensors) + else: + self.save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors) + + def load_optimizer(self, optimizer: Optimizer, checkpoint: str, prefix: str = None, size_per_shard: int = 1024): + """ + Load optimizer from checkpoint. + + Args: + optimizer (Optimizer): optimizer to be loaded. + checkpoint (str): checkpoint path. This value is made compatibility with the model checkpoints in the + prefix (str, optional): A prefix added to parameter and buffer + names to compose the keys in state_dict. Defaults to None. + size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024. + """ + + index_file_exists, index_file_path = has_index_file(checkpoint) + + if Path(checkpoint).is_dir() and not index_file_exists: + # if the checkpoint is a directory and there is no index file, raise error + raise ValueError(f'Cannot find index file in {checkpoint}') + + if index_file_exists: + # the existence of index file means it is a sharded checkpoint + self.load_sharded_optimizer(optimizer, index_file_path, prefix) + else: + self.load_unsharded_optimizer(optimizer, checkpoint) + + def save_optimizer(self, + optimizer: Optimizer, + checkpoint: str, + shard: bool = False, + gather_dtensor=True, + prefix: str = None, + size_per_shard: int = 1024): + """ + Save optimizer to checkpoint. Optimizer states saving is not compatible with safetensors. + + Args: + optimizer (Optimizer): optimizer to be saved. + checkpoint (str): checkpoint path. The checkpoint path can be : + 1. a file path, e.g. 'model.pt' + 2. a path to a json file which defines the index to the sharded checkpoint for the optimizer + 3. a path to a folder containing a unique .index.json file for sharded checkpoint + shard (bool): whether to shard the checkpoint. Default: False. If set to True, the checkpoint will be sharded into + multiple files. The optimizer shards will be specified by a `optimizer.index.json` file. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. Default: True. + prefix (str): prefix for the optimizer checkpoint when shard = True. Default: None. + size_per_shard (int): size per shard in MB. Default: 1024. This value is only used when shard is set to True. + """ + + if shard: + self.save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard) + else: + self.save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor) + + # ======================================================== + # Abstract methods for model loading/saving implementation + # ======================================================== + @abstractmethod + def load_sharded_model(self, model: nn.Module, index_file_path: str, strict: bool): + """ + Load model from sharded checkpoint. + + Args: + model (nn.Module): model to be loaded. + index_file_path (str): checkpoint path. It should be path to the .index.json file or a path to a directory which contains a .index.json file. + strict (bool): whether to strictly enforce that the param name in + the checkpoint match the keys returned by this module's. + """ + pass + + @abstractmethod + def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool): + """ + Load model from unsharded checkpoint. + + Args: + model (nn.Module): model to be loaded. + checkpoint (str): checkpoint path. It should be a single file path pointing to a model weight binary. + strict (bool): whether to strictly enforce that the param name in + the checkpoint match the keys returned by this module's. + """ + pass + + @abstractmethod + def save_sharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, prefix: Optional[str], + size_per_shard: int, use_safetensors: bool): + """ + Save model to sharded checkpoint. + + Args: + model (nn.Module): model to be saved. + checkpoint (str): checkpoint path. It should be a directory path. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. + prefix (str): prefix for the model checkpoint. + size_per_shard (int): size per shard in MB. + use_safetensors (bool): whether to use safe tensors. + """ + pass + + @abstractmethod + def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool): + """ + Save model to unsharded checkpoint. + + Args: + model (nn.Module): model to be saved. + checkpoint (str): checkpoint path. It should be a single file path pointing to a model weight binary. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. + use_safetensors (bool): whether to use safe tensors. + """ + pass + + # ======================================================== + # Abstract methods for optimizer loading/saving implementation + # ======================================================== + + @abstractmethod + def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, prefix: str): + """ + Load optimizer from sharded checkpoint. + + Args: + optimizer (Optimizer): optimizer to be loaded. + index_file_path (str): checkpoint path. It should be path to the .index.json file or a path to a directory which contains a .index.json file. + prefix (str): prefix for the optimizer checkpoint. + """ + pass + + @abstractmethod + def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path): + """ + Load optimizer from unsharded checkpoint. + + Args: + optimizer (Optimizer): optimizer to be loaded. + checkpoint (str): checkpoint path. It should be a single file path pointing to a model weight binary. + """ + pass + + @abstractmethod + def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: Path, gather_dtensor: bool, prefix: str, + size_per_shard: int): + """ + Save optimizer to sharded checkpoint. + + Args: + optimizer (Optimizer): optimizer to be saved. + checkpoint (Path): checkpoint path. It should be a directory path. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. + prefix (str): prefix for the optimizer checkpoint. + size_per_shard (int): size per shard in MB. + """ + pass + + @abstractmethod + def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path, gather_dtensor: bool): + """ + Save optimizer to unsharded checkpoint. + + Args: + optimizer (Optimizer): optimizer to be saved. + checkpoint (str): checkpoint path. It should be a single file path pointing to a model weight binary. + gather_dtensor (bool): whether to gather the distributed tensor to the first device. + """ + pass + + # ============================================ + # methods for loading and saving lr scheduler + # as this is quite standard, there is no need + # to make them abstract + # ============================================ + def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """ + Save lr scheduler to checkpoint. + + Args: + lr_scheduler (LRScheduler): lr scheduler to be saved. + checkpoint: checkpoint path. The checkpoint path can only be a file path. + """ + torch.save(lr_scheduler.state_dict(), checkpoint) + + def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str): + """ + Load lr scheduler from checkpoint. + + Args: + lr_scheduler (LRScheduler): lr scheduler to be loaded. + checkpoint (str): the path for a single checkpoint file. + """ + state_dict = torch.load(checkpoint) + lr_scheduler.load_state_dict(state_dict) diff --git a/colossalai/checkpoint_io/general_checkpoint_io.py b/colossalai/checkpoint_io/general_checkpoint_io.py new file mode 100644 index 000000000000..83e4bdcc863b --- /dev/null +++ b/colossalai/checkpoint_io/general_checkpoint_io.py @@ -0,0 +1,225 @@ +import gc +import logging +import os +from functools import reduce +from pathlib import Path +from typing import Iterator, Optional, OrderedDict, Tuple + +import torch.distributed as dist +import torch.nn as nn +from torch.optim import Optimizer + +from colossalai.interface import OptimizerWrapper + +from .checkpoint_io_base import CheckpointIO +from .index_file import CheckpointIndexFile +from .utils import ( + get_model_base_filenames, + get_optimizer_base_filenames, + get_shard_filename, + is_safetensors_available, + load_param_groups_into_optimizer, + load_shard_state_dict, + load_state_dict, + load_state_dict_into_model, + load_states_into_optimizer, + save_param_groups, + save_state_dict, + save_state_dict_shards, + shard_model_checkpoint, + shard_optimizer_checkpoint, + sharded_optimizer_loading_epilogue, + unwrap_optimizer, +) + +__all__ = ['GeneralCheckpointIO'] + + +class GeneralCheckpointIO(CheckpointIO): + """ + Checkpoint IO + """ + + def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool): + checkpoint = load_state_dict(checkpoint) + model.load_state_dict(checkpoint, strict=strict) + + def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool): + state_dict = model.state_dict() + + # TODO(FrankLeeeee): add support for gather_dtensor + if gather_dtensor: + pass + + # save the checkpoint + save_state_dict(state_dict, checkpoint, use_safetensors) + + def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, prefix: str): + """ + Load sharded optimizer with the given path to index file. + """ + + # If optimizer is wrapped, unwrap it. + if isinstance(optimizer, OptimizerWrapper): + optimizer = unwrap_optimizer(optimizer) + + # Read checkpoint index file. + ckpt_index_file = CheckpointIndexFile.from_file(index_file_path) + + # Load param_groups + param_group_path = ckpt_index_file.get_param_group_filename() + if param_group_path is None: + raise RuntimeError(f'Invalid index file path {index_file_path} for an optimizer. \ + Lacking param group file under current directory.') + id_map = load_param_groups_into_optimizer(optimizer, param_group_path) + + checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames() + + for shard_file in checkpoint_files: + state_dict = load_shard_state_dict(Path(shard_file), use_safetensors=False) + load_states_into_optimizer(optimizer, state_dict, id_map) + del state_dict + gc.collect() + + sharded_optimizer_loading_epilogue(optimizer) + + def save_sharded_optimizer( + self, + optimizer: Optimizer, + checkpoint: Path, + gather_dtensor: bool, + prefix: str, + size_per_shard: int, + ): + """ + Save sharded optimizer checkpoint under the given checkpointing path. + The following files will be created under the path: + - An index file (pytorch_optim.bin.index.json) containing a map between optimizer states and file names + - A group file (pytorch_optim_group.bin) recording information of param_groups + - Multiple files (pytorch_optim-000XX.bin) that store state tensors of optimizer in a sharding way + """ + + # If optimizer is wrapped, unwrap it. + if isinstance(optimizer, OptimizerWrapper): + optimizer = unwrap_optimizer(optimizer) + + if os.path.isfile(checkpoint): + logging.error(f"Provided path ({checkpoint}) should be a directory, not a file") + return + + Path(checkpoint).mkdir(parents=True, exist_ok=True) + + # Offload optimizer states. States are broken into shards within max_shard_size. + state_dict = optimizer.state_dict() + sharded_state = shard_optimizer_checkpoint(state_dict, max_shard_size=size_per_shard) + + # Preparing file paths and index file. + states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix) + index_file = CheckpointIndexFile(checkpoint) + + # Store the information of param groups to param_group_file. + index_file.append_meta_data("param_groups", param_group_file) + group_file_path = os.path.join(checkpoint, param_group_file) + save_param_groups(state_dict, group_file_path) + + # Save shards of optimizer states. + # In general cases, is_master is set to True to get the right behavior. + total_size = save_state_dict_shards(sharded_state_dict=sharded_state, + checkpoint=checkpoint, + index_file=index_file, + base_filename=states_name, + is_master=True, + use_safetensors=False) + + # Wrap up index file. + index_file.append_meta_data("total_size", total_size) + index_file.write_index_file(save_index_file) + logging.info(f"The optimizer is going to be split to checkpoint shards. " + f"You can find where each parameters has been saved in the " + f"index located at {save_index_file}.") + + def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path): + checkpoint = load_state_dict(checkpoint) + optimizer.load_state_dict(checkpoint) + + def save_unsharded_optimizer( + self, + optimizer: Optimizer, + checkpoint: Path, + gather_dtensor: bool, + ): + # TODO(FrankLeeeee): handle distributed tensors + save_state_dict(optimizer.state_dict(), checkpoint, use_safetensors=False) + + def save_sharded_model(self, + model: nn.Module, + checkpoint_path: str, + gather_dtensor: bool = False, + prefix: Optional[str] = None, + max_shard_size: int = 1024, + use_safetensors: bool = False): + """ + implement this method as it can be supported by Huggingface model, + save shard model, save model to multiple files + """ + if os.path.isfile(checkpoint_path): + logging.error(f"Provided path ({checkpoint_path}) should be a directory, not a file") + return + + Path(checkpoint_path).mkdir(parents=True, exist_ok=True) + + # shard checkpoint + state_dict = model.state_dict() + state_dict_shard = shard_model_checkpoint(state_dict, max_shard_size=max_shard_size) + weights_name, save_index_file = get_model_base_filenames(prefix, use_safetensors) + index_file = CheckpointIndexFile(checkpoint_path) + + # Save shards of optimizer states. + # In general cases, is_master is set to True to get the right behavior. + total_size = save_state_dict_shards(sharded_state_dict=state_dict_shard, + checkpoint=checkpoint_path, + index_file=index_file, + base_filename=weights_name, + is_master=True, + use_safetensors=use_safetensors) + + index_file.append_meta_data("total_size", total_size) + index_file.write_index_file(save_index_file) + logging.info(f"The model is going to be split to checkpoint shards. " + f"You can find where each parameters has been saved in the " + f"index located at {save_index_file}.") + + def load_sharded_model(self, + model: nn.Module, + checkpoint_index_file: Path, + strict: bool = False, + use_safetensors: bool = False, + load_sub_module: bool = True): + """ + load shard model, load model from multiple files + """ + use_safetensors = False + if "safetensors" in checkpoint_index_file.name: + use_safetensors = True + + if use_safetensors and not is_safetensors_available(): + raise ImportError("`safe_serialization` requires the `safetensors` library: `pip install safetensors`.") + + # read checkpoint index file + ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) + checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames() + missing_keys = [] + + for shard_file in checkpoint_files: + state_dict = load_shard_state_dict(Path(shard_file), use_safetensors) + load_state_dict_into_model(model, state_dict, missing_keys, strict, load_sub_module) + del state_dict + gc.collect() + + if strict: + remain_keys = reduce(lambda a, b: a & b, map(set, missing_keys)) + if len(remain_keys) > 0: + error_msgs = 'Missing key(s) in state_dict: {}. '.format(', '.join( + '"{}"'.format(k) for k in missing_keys)) + raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( + self.__class__.__name__, "\n\t".join(error_msgs))) diff --git a/colossalai/checkpoint_io/index_file.py b/colossalai/checkpoint_io/index_file.py new file mode 100644 index 000000000000..388cf3fbe9bb --- /dev/null +++ b/colossalai/checkpoint_io/index_file.py @@ -0,0 +1,182 @@ +import json +import os +from collections import OrderedDict +from pathlib import Path +from typing import Any, Dict, List, Union + +from .utils import is_dtensor_checkpoint + +__all__ = ['CheckpointIndexFile'] + + +class CheckpointIndexFile: + """ + This class is a data structure to keep the content in the index.json file for sharded checkpoint. + + Example: + >>> index = CheckpointIndexFile.from_file('model.index.json') + >>> index.append_metadata('model_type', 'bert') + >>> index.append_weight_map('bert.embeddings.word_embeddings.weight', 'model_0001-of-0002.bin') + >>> index.export('new_index.json') + """ + + def __init__(self, root_path=None) -> None: + self.root_path = root_path + + # use ordered dict to preserve the tensor checkpoint order + self.metadata: Dict = OrderedDict() + self.weight_map: Dict = OrderedDict() + + @staticmethod + def from_file(index_path: Union[str, Path]): + """ + Create a CheckpointIndexFile object from a json file. + + Args: + index_path (str): path to the json file. + + Returns: + CheckpointIndexFile: CheckpointIndexFile object. + """ + index = CheckpointIndexFile() + index.load(index_path) + return index + + def load(self, json_path: str): + """ + Load the index file from a json file. + + Args: + json_path (str): path to the json file. + """ + # load the json file + with open(json_path, 'r') as f: + index = json.load(f) + + # assign attributes if exists + if "metadata" in index: + self.metadata = index["metadata"] + if "weight_map" in index: + self.weight_map = index["weight_map"] + + # assign the root directory for the index file + self.root_path = Path(json_path).absolute().parent + + def export(self, json_path: str): + """ + Export the index file to a json file. + + Args: + json_path (str): path to the json file. + """ + # create the index file + index = dict() + index["metadata"] = self.metadata + index["weight_map"] = self.weight_map + + # export the index file + with open(json_path, 'w') as f: + json.dump(index, f, indent=4) + + def append_weight_map(self, param_name: str, shard_file: str): + """ + Append a weight map entry to the index file. + + Args: + param_name (str): name of the parameter. + shard_file (str): name of the shard file. + """ + self.weight_map[param_name] = shard_file + + def append_meta_data(self, name: str, val: Any): + """ + Append a metadata entry to the index file. + + Args: + name (str): name of the metadata. + val (Any): value of the metadata. + """ + self.metadata[name] = val + + def contains_dtensor(self): + """ + Check if the index file contains any distributed tensor. The distributed tensors will be stored in + `dtensor/module.linear.weight.*.bin` or `dtensor/module.linear.weight.*.safetensors` in the weight map. + + Returns: + bool: True if the index file contains any distributed tensor, False otherwise. + """ + for value in self.weight_map.values(): + if value.endswith(".*.bin") or value.endswith(".*.safetensors"): + return True + return False + + def get_checkpoint_filenames(self) -> List[str]: + """ + Get the set of checkpoint filenames in the weight map. + + Returns: + list: checkpoint shard filenames. + """ + # read the checkpoint file list from the json file and get a list of unique file names + checkpoint_files = sorted(list(set(self.weight_map.values()))) + + # get the absolute paths for all checkpoint files + checkpoint_files = [str(self.root_path.joinpath(f)) for f in checkpoint_files] + + dtensor_list = [] + checkpoint_list = [] + + for ckpt_file in checkpoint_files: + if is_dtensor_checkpoint(ckpt_file): + dtensor_list.append(ckpt_file) + else: + checkpoint_list.append(ckpt_file) + + return checkpoint_list, dtensor_list + + def assert_no_dtensor_checkpoint(self): + for val in self.weight_map.values(): + if is_dtensor_checkpoint(val): + raise ValueError(f"Checkpoint file {val} contains distributed tensor") + + def get_checkpoint_file(self, param_name: str) -> str: + """ + Get the checkpoint file name for a parameter. + + Args: + param_name (str): name of the parameter. + + Returns: + str: checkpoint file name. + """ + ckpt_path = self.weight_map[param_name] + return ckpt_path + + def get_all_param_names(self): + """ + Get all the weight keys. + """ + return list(self.weight_map.keys()) + + def get_param_group_filename(self) -> Union[str, None]: + """ + Get the file name of param_group file if this is a checkpoint for optimizer. + Returns: + str: param_group file name + """ + filename = self.metadata.get("param_groups", None) + if filename: + return str(self.root_path.joinpath(filename)) + else: + return None + + def write_index_file(self, save_index_file): + """ + Write index file. + """ + save_index_file = os.path.join(self.root_path, save_index_file) + index = {"metadata": self.metadata, "weight_map": self.weight_map} + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2) + "\n" + f.write(content) diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py new file mode 100644 index 000000000000..8837776aee4d --- /dev/null +++ b/colossalai/checkpoint_io/utils.py @@ -0,0 +1,658 @@ +# coding=utf-8 +import os +import re +from collections import abc as container_abcs +from collections import defaultdict +from itertools import chain +from pathlib import Path +from typing import Iterator, List, Mapping, Optional, OrderedDict, Tuple + +import torch +import torch.nn as nn +from torch.optim import Optimizer + +from colossalai.interface import OptimizerWrapper +from colossalai.nn.optimizer import ColossalaiOptimizer +from colossalai.tensor.d_tensor import is_distributed_tensor + +SAFE_WEIGHTS_NAME = "model.safetensors" +WEIGHTS_NAME = "pytorch_model.bin" +STATES_NAME = "pytorch_optim.bin" +SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" +WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json" +STATES_INDEX_NAME = "pytorch_optim.bin.index.json" +GROUP_FILE_NAME = "pytorch_optim_group.bin" + +# ====================================== +# General helper functions +# ====================================== + + +def calculate_tensor_size(tensor: torch.Tensor) -> float: + """ + Calculate the size of a parameter in MB. Used to compute whether a group of params exceed the shard size. + If so, a new shard should be created. + + Args: + tensor (torch.Tensor): the tensor to calculate size for. + + Returns: + float: size of the tensor in MB. + """ + return tensor.numel() * tensor.element_size() / 1024 / 1024 + + +def is_safetensors_available() -> bool: + """ + Check whether safetensors is available. + + Returns: + bool: whether safetensors is available. + """ + try: + import safetensors + return True + except ImportError: + return False + + +def is_dtensor_checkpoint(checkpoint_file_path: str) -> bool: + """ + Check whether the checkpoint file is a dtensor checkpoint. + + Args: + checkpoint_file_path (str): path to the checkpoint file. + + Returns: + bool: whether the checkpoint file is a dtensor checkpoint. + """ + if checkpoint_file_path.endswith('.*.safetensors') or checkpoint_file_path.endswith('.*.bin'): + return True + else: + return False + + +def is_safetensor_checkpoint(checkpoint_file_path: str) -> bool: + """ + Check whether the checkpoint file is a safetensor checkpoint. + + Args: + checkpoint_file_path (str): path to the checkpoint file. + + Returns: + bool: whether the checkpoint file is a safetensor checkpoint. + """ + if checkpoint_file_path.endswith('.safetensors'): + return True + else: + return False + + +# ====================================== +# Helper functions for saving shard file +# ====================================== +def unwrap_optimizer(optimizer: OptimizerWrapper): + ''' + Unwrap a wrapped optimizer. + This method should be used before saving/loading it to/from sharded checkpoints. + ''' + + # TODO(Baizhou): ColossalaiOptimizer will be replaced with OptimizerWrapper in the future + unwrapped_optim = optimizer.optim + if isinstance(unwrapped_optim, ColossalaiOptimizer): + unwrapped_optim = unwrapped_optim.optim + return unwrapped_optim + + +def save_state_dict_shards(sharded_state_dict: Iterator[Tuple[OrderedDict, int]], + checkpoint: str, + index_file: "CheckpointIndexFile", + base_filename: str, + is_master: bool, + use_safetensors: bool = False) -> int: + ''' + Save sharded state dict only on master rank, this method can be used by both model and optimizer states. + Args: + sharded_state_dict (Iterator[Tuple[OrderedDict, int]]): a generator of shards, each shard contains state dict and shard size. + checkpoint (str): The path of checkpoint directory as string. + index_file (CheckpointIndexFile): The index file object to be updated. + base_filename (str): Decides the prefix of filenames of shards. + is_master (bool): Whether current rank is master. + use_safetensors (bool): Whether to use safetensors to save checkpoint. + + Returns: + int: the total size of shards + ''' + + total_size = 0 + for idx, shard_pair in enumerate(sharded_state_dict): + if not is_master: + continue + shard, current_size = shard_pair + shard_file = get_shard_filename(base_filename, idx) + total_size = total_size + current_size + for key in shard.keys(): + index_file.append_weight_map(key, shard_file) + checkpoint_file_path = os.path.join(checkpoint, shard_file) + + # Only save on master rank. + save_state_dict(shard, checkpoint_file_path, use_safetensors=use_safetensors) + + return total_size + + +def shard_model_checkpoint(state_dict: torch.Tensor, max_shard_size: int = 1024) -> Iterator[Tuple[OrderedDict, int]]: + """ + Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a + given size. + """ + current_block = {} + current_block_size = 0 + + for key, weight in state_dict.items(): + ret_block = None + ret_block_size = 0 + if not is_distributed_tensor(weight): + weight_size = calculate_tensor_size(weight) + + # If this weight is going to tip up over the maximal size, we split. + if current_block_size + weight_size > max_shard_size and current_block_size > 0: + ret_block = current_block + ret_block_size = current_block_size + current_block = {} + current_block_size = 0 + current_block[key] = weight + current_block_size += weight_size + + if ret_block != None: + yield ret_block, ret_block_size + + yield current_block, current_block_size + + +def shard_optimizer_checkpoint(state_dict: dict, max_shard_size: int = 1024) -> Iterator[Tuple[OrderedDict, int]]: + """ + Splits an optimizer state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a + given size. + """ + + # Only split state_dict['state']; state_dict['param_group'] is not considered in this function. + states = state_dict['state'] + + current_block = {} + current_block_size = 0 + + for param_id, state in states.items(): + + ret_block = None + ret_block_size = 0 + + # A state might contain more than one tensors. + # e.g. each Adam state includes: 'step', 'exp_avg', 'exp_avg_sq' + state_size = 0 + isDTensor = False + for state_tensor in state.values(): + + # When state_tensor is not of Tensor class, + # e.g., a SGD optimizer with momentum set to 0 can have None as state + # The calculation of tensor size should be skipped to avoid error. + if not isinstance(state_tensor, torch.Tensor): + continue + + # If the states are stored as DTensors, mark isDTensor as true. + if is_distributed_tensor(state_tensor): + isDTensor = True + state_size += calculate_tensor_size(state_tensor) + + if not isDTensor: + + if current_block_size + state_size > max_shard_size and current_block_size > 0: + ret_block = current_block + ret_block_size = current_block_size + current_block = {} + current_block_size = 0 + + current_block[param_id] = state + current_block_size += state_size + + if ret_block != None: + yield ret_block, ret_block_size + + yield current_block, current_block_size + + +def load_shard_state_dict(checkpoint_file: Path, use_safetensors: bool = False): + """ + load shard state dict into model + """ + if use_safetensors and not checkpoint_file.suffix == ".safetensors": + raise Exception("load the model using `safetensors`, but no file endwith .safetensors") + if use_safetensors: + from safetensors.torch import load_file as safe_load_file + from safetensors.torch import safe_open + with safe_open(checkpoint_file, framework="pt") as f: + metadata = f.metadata() + if metadata["format"] != "pt": + raise NotImplementedError( + f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet.") + return safe_load_file(checkpoint_file) + else: + return torch.load(checkpoint_file) + + +def load_state_dict_into_model(model: nn.Module, + state_dict: torch.Tensor, + missing_keys: List, + strict: bool = False, + load_sub_module: bool = True): + r"""Copies parameters and buffers from :attr:`state_dict` into + this module and its descendants. + + Args: + state_dict (dict): a dict containing parameters and + persistent buffers. + """ + if not isinstance(state_dict, Mapping): + raise TypeError("Expected state_dict to be dict-like, got {}.".format(type(state_dict))) + + unexpected_keys: List[str] = [] + sub_missing_keys: List[str] = [] + error_msgs: List[str] = [] + + # copy state_dict so _load_from_state_dict can modify it + metadata = getattr(state_dict, '_metadata', None) + state_dict = OrderedDict(state_dict) + if metadata is not None: + state_dict._metadata = metadata + + def load(module: nn.Module, state_dict, prefix="", load_sub_module: bool = True): + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + args = (state_dict, prefix, local_metadata, True, sub_missing_keys, [], error_msgs) + # Parameters of module and children will start with prefix. We can exit early if there are none in this + # state_dict + if len([key for key in state_dict if key.startswith(prefix)]) > 0: + module._load_from_state_dict(*args) + if load_sub_module: + for name, child in module._modules.items(): + if child is not None: + load(child, state_dict, prefix + name + ".") + + load(model, state_dict, "", load_sub_module) + del load + + missing_keys = missing_keys.append(sub_missing_keys) + + if strict: + if len(unexpected_keys) > 0: + error_msgs = 'Unexpected key(s) in state_dict: {}. '.format(', '.join( + '"{}"'.format(k) for k in unexpected_keys)) + raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( + model.__class__.__name__, "\n\t".join(error_msgs))) + + +def load_param_groups_into_optimizer(optimizer: Optimizer, param_group_path: str) -> dict: + """ + Load information of param_groups into an initialized optimizer. + """ + + # Load list of param_groups from given file path. + # The params in saved_groups are in the form of integer indices. + saved_groups = torch.load(param_group_path) + if not isinstance(saved_groups, List): + raise ValueError(f'The param_groups saved at {param_group_path} is not of List type') + + # The params in param_groups are in the form of pytorch tensors. + # For more details, please view source code of Optimizer class in pytorch. + param_groups = optimizer.param_groups + + # Check the compatibility of saved_groups and param_groups. + if len(param_groups) != len(saved_groups): + raise ValueError("loaded state dict has a different number of original parameter groups") + param_lens = (len(g['params']) for g in param_groups) + saved_lens = (len(g['params']) for g in saved_groups) + if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)): + raise ValueError("loaded state dict contains a parameter group " + "that doesn't match the size of optimizer's group") + + # Creating mapping from id to parameters. + id_map = { + old_id: p for old_id, p in zip(chain.from_iterable((g['params'] for g in saved_groups + )), chain.from_iterable((g['params'] for g in param_groups))) + } + + # Update parameter groups, setting their 'params' value. + def update_group(group, new_group): + new_group['params'] = group['params'] + return new_group + + updated_groups = [update_group(g, ng) for g, ng in zip(param_groups, saved_groups)] + + optimizer.__dict__.update({'param_groups': updated_groups}) + return id_map + + +def load_states_into_optimizer(optimizer: Optimizer, state_dict: dict, id_map: dict): + r"""Copies states from `state_dict` into an Optimizer object. + + Args: + optimizer(Optimizer): An initialized Optimizer object to be loaded + state_dict(dict): a mapping from tensor index (an integer) + to its states to be loaded (a mapping from state name to a tensor). + id_map(dict): a mapping from tensor index (an integer) + to its corresponding parameter (a tensor) whose states will be updated. + """ + + def cast(param, value, key=None): + r"""Make a deep copy of value, casting all tensors to device of param.""" + if isinstance(value, torch.Tensor): + # Floating-point types are a bit special here. They are the only ones + # that are assumed to always match the type of params. + # Make sure state['step'] is not casted https://github.com/pytorch/pytorch/issues/74424 + if (key != "step"): + if param.is_floating_point(): + value = value.to(param.dtype) + value = value.to(param.device) + return value + elif isinstance(value, dict): + return {k: cast(param, v, key=k) for k, v in value.items()} + elif isinstance(value, container_abcs.Iterable): + return type(value)(cast(param, v) for v in value) + else: + return value + + # Copy state assigned to params (and cast tensors to appropriate types). + # State that is not assigned to params is copied as is (needed for + # backward compatibility). + new_states = defaultdict(dict) + for k, v in state_dict.items(): + if k in id_map: + param = id_map[k] + new_states[param] = cast(param, v) + else: + new_states[k] = v + + optimizer.state.update(new_states) + + +def sharded_optimizer_loading_epilogue(optimizer: Optimizer): + r"""Do the cleaning up work after state_dict has been loaded into optimizer + + Args: + optimizer(Optimizer): An optimizer object whose state has just been loaded. + """ + + # Do the cleaning up as in src code of Pytorch. + optimizer._hook_for_profile() # To support multiprocessing pickle/unpickle. + optimizer.defaults.setdefault('differentiable', False) + + +# ====================================== +# Helper functions for saving state dict +# ====================================== + + +def save_state_dict(state_dict: dict, checkpoint_file_path: str, use_safetensors: bool) -> None: + """ + Save state dict to checkpoint. + + Args: + state_dict (dict): state dict. + checkpoint_file_path (str): path to the checkpoint file. + use_safetensors (bool): whether to use safetensors to save the checkpoint. + """ + if use_safetensors: + assert is_safetensors_available(), "safetensors is not available." + assert checkpoint_file_path.endswith('.safetensors'), \ + "safetensors only supports .safetensors suffix for checkpoint file." + from safetensors.torch import save_file as safe_save_file + safe_save_file(state_dict, checkpoint_file_path, metadata={"format": "pt"}) + else: + torch.save(state_dict, checkpoint_file_path) + + +def save_param_groups(state_dict: dict, group_file_path: str) -> None: + """ + Save information of param_groups to given file path. + + Args: + state_dict (dict): state dict. + group_file_path (str): path to the group file. + """ + param_groups = state_dict["param_groups"] + torch.save(param_groups, group_file_path) + + +def save_dtensor(name: str, tensor: torch.Tensor, index_file: "CheckpointIndexFile", use_safetensors: bool) -> None: + """ + Save distributed tensor to checkpoint. This checkpoint will be a dictionary which contains + only one tensor. + + Args: + tensor (Tensor): tensor to be saved. + index_file (CheckpointIndexFile): path to the checkpoint file. + size_per_shard (int): size per shard in MB. + """ + root_path = index_file.root_path + output_root_path = root_path.joinpath('dtensor') + + # create directory + output_root_path.mkdir(exist_ok=True) + + # save tensor to this directory + # TODO(YuliangLiu): get index of the tensor shard + # e.g. index = + index = 0 + + # save tensor to file + ckpt_file_name = generate_dtensor_file_name(name, index, use_safetensors) + ckpt_file_path = output_root_path.joinpath(ckpt_file_name) + + # dtensor ckpt file always contains only one tensor + state_dict = {name: tensor} + save_state_dict(state_dict, str(ckpt_file_path), use_safetensors) + + # update the weight map + # * means all shards + ckpt_file_name_in_weight_map = 'dtensor/' + generate_dtensor_file_name(name, '*', use_safetensors) + index_file.append_weight_map(name, ckpt_file_name_in_weight_map) + + +def get_checkpoint_file_suffix(use_safetensors: bool) -> str: + """ + Get checkpoint file suffix. + + Args: + use_safetensors (bool): whether to use safetensors to save the checkpoint. + + Returns: + str: checkpoint file suffix. + """ + if use_safetensors: + return '.safetensors' + else: + return '.bin' + + +def generate_checkpoint_shard_file_name(index: int, + total_number: int, + use_safetensors: bool, + prefix: str = None) -> str: + """ + Generate checkpoint shard file name. + + Args: + index (int): index of the shard. + total_number (int): total number of shards. + use_safetensors (bool): whether to use safetensors to save the checkpoint. + prefix (str): prefix of the shard file name. Default: None. + + Returns: + str: checkpoint shard file name. + """ + suffix = get_checkpoint_file_suffix(use_safetensors) + + if prefix is None: + return f"{index:05d}-of-{total_number:05d}.{suffix}" + else: + return f"{prefix}-{index:05d}-of-{total_number:05d}.{suffix}" + + +def generate_dtensor_file_name(param_name: str, index: int, use_safetensors: bool) -> str: + """ + Generate dtensor file name. + + Args: + param_name (str): name of the distributed parameter. + index (int): index of the shard. + use_safetensors (bool): whether to use safetensors to save the checkpoint. + + Returns: + str: dtensor file name. + """ + suffix = get_checkpoint_file_suffix(use_safetensors) + return f'{param_name}.{index}.{suffix}' + + +def save_state_dict_as_shard( + state_dict: dict, + checkpoint_path: str, + index: int, + total_number: int, + use_safetensors: bool, + prefix: str = None, +) -> None: + """ + Save state dict as shard. + + Args: + state_dict (dict): state dict. + checkpoint_path (str): path to the checkpoint file. + index (int): index of the shard. + total_number (int): total number of shards. + prefix (str): prefix of the shard file name. + use_safetensors (bool): whether to use safetensors to save the checkpoint. + """ + # generate the shard name + shard_file_name = generate_checkpoint_shard_file_name(index, total_number, use_safetensors, prefix) + shard_file_path = Path(checkpoint_path).joinpath(shard_file_name).absolute() + + # save the shard + save_state_dict(state_dict, str(shard_file_path), use_safetensors) + + +# ======================================== +# Helper functions for loading state dict +# ======================================== + + +def has_index_file(checkpoint_path: str) -> Tuple[bool, Optional[Path]]: + """ + Check whether the checkpoint has an index file. + + Args: + checkpoint_path (str): path to the checkpoint. + + Returns: + Tuple[bool, Optional[Path]]: a tuple of (has_index_file, index_file_path) + """ + checkpoint_path = Path(checkpoint_path) + if checkpoint_path.is_file(): + # check if it is .index.json + reg = re.compile("(.*?).index((\..*)?).json") + if reg.fullmatch(checkpoint_path.name) is not None: + return True, checkpoint_path + else: + return False, None + elif checkpoint_path.is_dir(): + # check if there is only one a file ending with .index.json in this directory + index_files = list(checkpoint_path.glob('*.index.*json')) + + # if we found a .index.json file, make sure there is only one + if len(index_files) > 0: + assert len( + index_files + ) == 1, f'Expected to find one .index.json file in {checkpoint_path}, but found {len(index_files)}' + + if len(index_files) == 1: + return True, index_files[0] + else: + return False, None + else: + raise RuntimeError(f'Invalid checkpoint path {checkpoint_path}. Expected a file or a directory.') + + +def load_state_dict(checkpoint_file_path: Path): + """ + Load state dict from checkpoint. + + Args: + checkpoint_file_path (Path): path to the checkpoint file. + + Returns: + dict: state dict. + """ + + assert not is_dtensor_checkpoint(checkpoint_file_path), \ + f'Cannot load state dict from dtensor checkpoint {checkpoint_file_path}, you should convert the distributed tensors to gathered tensors with our CLI offline.' + + if is_safetensor_checkpoint(checkpoint_file_path): + assert is_safetensors_available(), \ + f'Cannot load state dict from safetensor checkpoint {checkpoint_file_path}, because safetensors is not available. Please install safetensors first with pip install safetensors.' + # load with safetensors + from safetensors import safe_open + state_dict = {} + with safe_open(checkpoint_file_path, framework="pt", device="cpu") as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + return state_dict + + else: + # load with torch + return torch.load(checkpoint_file_path) + + +def add_prefix(weights_name: str, prefix: Optional[str] = None) -> str: + if prefix is not None and len(prefix) > 0: + splits = weights_name.split(".") + splits = splits[:-1] + [prefix] + splits[-1:] + weights_name = ".".join(splits) + + return weights_name + + +def get_model_base_filenames(prefix: str = None, use_safetensors: bool = False): + """ + generate base model weight filenames + """ + weights_name = SAFE_WEIGHTS_NAME if use_safetensors else WEIGHTS_NAME + weights_name = add_prefix(weights_name, prefix) + + save_index_file = SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME + save_index_file = add_prefix(save_index_file, prefix) + + return weights_name, save_index_file + + +def get_optimizer_base_filenames(prefix: str = None): + """ + generate base optimizer state filenames + """ + states_name = STATES_NAME + states_name = add_prefix(states_name, prefix) + + save_index_file = STATES_INDEX_NAME + save_index_file = add_prefix(save_index_file, prefix) + + param_group_file = GROUP_FILE_NAME + param_group_file = add_prefix(param_group_file, prefix) + + return states_name, save_index_file, param_group_file + + +def get_shard_filename(weights_name: str, idx: int): + """ + get shard file name + """ + shard_file = weights_name.replace(".bin", f"-{idx+1:05d}.bin") + shard_file = shard_file.replace(".safetensors", f"-{idx + 1:05d}.safetensors") + return shard_file diff --git a/colossalai/cli/benchmark/benchmark.py b/colossalai/cli/benchmark/benchmark.py index f40f8f2f995e..97a9f45722dd 100644 --- a/colossalai/cli/benchmark/benchmark.py +++ b/colossalai/cli/benchmark/benchmark.py @@ -10,7 +10,8 @@ from colossalai.context.random import reset_seeds from colossalai.core import global_context as gpc from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.utils import MultiTimer, free_port +from colossalai.testing import free_port +from colossalai.utils import MultiTimer from .models import MLP diff --git a/colossalai/cli/benchmark/models.py b/colossalai/cli/benchmark/models.py index 38ea54188b8c..f8fd1c41a059 100644 --- a/colossalai/cli/benchmark/models.py +++ b/colossalai/cli/benchmark/models.py @@ -1,4 +1,5 @@ import torch + import colossalai.nn as col_nn diff --git a/colossalai/cli/check/check_installation.py b/colossalai/cli/check/check_installation.py index 44d7840700ef..4a481f3bd122 100644 --- a/colossalai/cli/check/check_installation.py +++ b/colossalai/cli/check/check_installation.py @@ -31,7 +31,7 @@ def check_installation(): found_aot_cuda_ext = _check_aot_built_cuda_extension_installed() cuda_version = _check_cuda_version() torch_version, torch_cuda_version = _check_torch_version() - colossalai_verison, prebuilt_torch_version_required, prebuilt_cuda_version_required = _parse_colossalai_version() + colossalai_version, prebuilt_torch_version_required, prebuilt_cuda_version_required = _parse_colossalai_version() # if cuda_version is None, that means either # CUDA_HOME is not found, thus cannot compare the version compatibility @@ -57,7 +57,7 @@ def check_installation(): click.echo(f'#### Installation Report ####') click.echo(f'\n------------ Environment ------------') - click.echo(f"Colossal-AI version: {to_click_output(colossalai_verison)}") + click.echo(f"Colossal-AI version: {to_click_output(colossalai_version)}") click.echo(f"PyTorch version: {to_click_output(torch_version)}") click.echo(f"System CUDA version: {to_click_output(cuda_version)}") click.echo(f"CUDA version required by PyTorch: {to_click_output(torch_cuda_version)}") @@ -76,7 +76,7 @@ def check_installation(): click.echo("") click.echo(f"Note:") click.echo( - f"1. AOT (ahead-of-time) compilation of the CUDA kernels occurs during installation when the environment varialbe CUDA_EXT=1 is set" + f"1. AOT (ahead-of-time) compilation of the CUDA kernels occurs during installation when the environment variable CUDA_EXT=1 is set" ) click.echo(f"2. If AOT compilation is not enabled, stay calm as the CUDA kernels can still be built during runtime") @@ -88,7 +88,7 @@ def check_installation(): click.echo(f"Note:") click.echo(f"1. The table above checks the version compatibility of the libraries/tools in the current environment") click.echo( - f" - PyTorch version mistach: whether the PyTorch version in the current environment is compatible with the PyTorch version used for AOT compilation" + f" - PyTorch version mismatch: whether the PyTorch version in the current environment is compatible with the PyTorch version used for AOT compilation" ) click.echo( f" - System and PyTorch CUDA version match: whether the CUDA version in the current environment is compatible with the CUDA version required by PyTorch" @@ -137,7 +137,7 @@ def _parse_colossalai_version(): # 1. X.X.X+torchX.XXcuXX.X (when colossalai is installed with CUDA extensions) # 2. X.X.X (when colossalai is not installed with CUDA extensions) # where X represents an integer. - colossalai_verison = colossalai.__version__.split('+')[0] + colossalai_version = colossalai.__version__.split('+')[0] try: torch_version_for_aot_build = colossalai.__version__.split('torch')[1].split('cu')[0] @@ -145,7 +145,7 @@ def _parse_colossalai_version(): except: torch_version_for_aot_build = None cuda_version_for_aot_build = None - return colossalai_verison, torch_version_for_aot_build, cuda_version_for_aot_build + return colossalai_version, torch_version_for_aot_build, cuda_version_for_aot_build def _check_aot_built_cuda_extension_installed(): diff --git a/colossalai/cli/launcher/__init__.py b/colossalai/cli/launcher/__init__.py index 8d9ec147d401..808e4e84574f 100644 --- a/colossalai/cli/launcher/__init__.py +++ b/colossalai/cli/launcher/__init__.py @@ -28,7 +28,7 @@ type=str, default=None, help= - "Specify computing devices to NOT use during execution. Mutually exclusive with --include. Formatting is the same as --includ," + "Specify computing devices to NOT use during execution. Mutually exclusive with --include. Formatting is the same as --include," " only effective when used with --hostfile.") @click.option("--num_nodes", type=int, diff --git a/colossalai/cli/launcher/hostinfo.py b/colossalai/cli/launcher/hostinfo.py index 065cbc37101f..d1b88b229fb8 100644 --- a/colossalai/cli/launcher/hostinfo.py +++ b/colossalai/cli/launcher/hostinfo.py @@ -38,7 +38,7 @@ def is_host_localhost(hostname: str, port: str = None) -> None: # socket.getfqdn("127.0.0.1") does not return localhost # on some users' machines - # thus, we directly return True if hostname is locahost, 127.0.0.1 or 0.0.0.0 + # thus, we directly return True if hostname is localhost, 127.0.0.1 or 0.0.0.0 if hostname in ("localhost", "127.0.0.1", "0.0.0.0"): return True diff --git a/colossalai/cli/launcher/multinode_runner.py b/colossalai/cli/launcher/multinode_runner.py index a51e1e371f13..85b241e96292 100644 --- a/colossalai/cli/launcher/multinode_runner.py +++ b/colossalai/cli/launcher/multinode_runner.py @@ -114,7 +114,7 @@ def recv_from_all(self) -> dict: Receive messages from all hosts Returns: - msg_from_node (dict): a dictionry which contains messages from each node + msg_from_node (dict): a dictionary which contains messages from each node """ msg_from_node = dict() diff --git a/colossalai/cli/launcher/run.py b/colossalai/cli/launcher/run.py index 6411b4302e95..5e74c2c4f5b8 100644 --- a/colossalai/cli/launcher/run.py +++ b/colossalai/cli/launcher/run.py @@ -154,7 +154,7 @@ def _arg_dict_to_list(arg_dict): extra_launch_args = dict() torch_version = version.parse(torch.__version__) - assert torch_version.major == 1 + assert torch_version.major >= 1 if torch_version.minor < 9: cmd = [ @@ -164,9 +164,7 @@ def _arg_dict_to_list(arg_dict): ] else: # extra launch args for torch distributed launcher with torch >= 1.9 - default_torchrun_rdzv_args = dict(rdzv_backend="c10d", - rdzv_endpoint=f"{master_addr}:{master_port}", - rdzv_id="colossalai-default-job") + default_torchrun_rdzv_args = dict(master_addr=master_addr, master_port=master_port) # update rdzv arguments for key in default_torchrun_rdzv_args.keys(): @@ -298,7 +296,7 @@ def launch_multi_processes(args: Config) -> None: # receive the stop status msg_from_node = runner.recv_from_all() - # printe node status + # print node status click.echo("\n====== Stopping All Nodes =====") for hostname, msg in msg_from_node.items(): click.echo(f"{hostname}: {msg}") diff --git a/colossalai/cluster/__init__.py b/colossalai/cluster/__init__.py new file mode 100644 index 000000000000..2fbdfd3cc999 --- /dev/null +++ b/colossalai/cluster/__init__.py @@ -0,0 +1,5 @@ +from .device_mesh_manager import DeviceMeshManager +from .dist_coordinator import DistCoordinator +from .process_group_manager import ProcessGroupManager + +__all__ = ['DistCoordinator', 'ProcessGroupManager', 'DeviceMeshManager'] diff --git a/colossalai/cluster/device_mesh_manager.py b/colossalai/cluster/device_mesh_manager.py new file mode 100644 index 000000000000..8754baa19792 --- /dev/null +++ b/colossalai/cluster/device_mesh_manager.py @@ -0,0 +1,117 @@ +from dataclasses import dataclass +from typing import Dict, List, Tuple, Union + +import torch +import torch.distributed as dist + +from colossalai.device.alpha_beta_profiler import AlphaBetaProfiler +from colossalai.device.device_mesh import DeviceMesh + + +@dataclass +class DeviceMeshInfo: + ''' + This class is used to store the information used to initialize the device mesh. + + Args: + physical_ids (List[int]): The physical ids of the current booster. For example, if we have the last 4 GPUs on a 8-devices cluster, then the physical ids should be [4, 5, 6, 7]. + mesh_shapes (List[Union[torch.Size, List[int], Tuple[int]]]): The shape of the mesh. For example, if we have 4 GPUs and we want to use 2D mesh with mesh shape [2, 2], then the mesh shape should be [2, 2]. + ''' + physical_ids: List[int] + mesh_shape: Union[torch.Size, List[int], Tuple[int]] = None + + def __post_init__(self): + if self.mesh_shape is not None: + world_size = len(self.physical_ids) + mesh_shape_numel = torch.Size(self.mesh_shape).numel() + assert world_size == mesh_shape_numel, f'the numel of mesh_shape should be equal to world size, but got {world_size} != {mesh_shape_numel}' + + +def initialize_device_mesh(device_mesh_info: DeviceMeshInfo): + ''' + This method is used to initialize the device mesh. + + Args: + device_mesh_info (DeviceMeshInfo): The information used to initialize device mesh. + ''' + # parse the device mesh info + physical_devices = device_mesh_info.physical_ids + physical_mesh = torch.tensor(physical_devices) + logical_mesh_shape = device_mesh_info.mesh_shape + + if logical_mesh_shape is None: + ab_profiler = AlphaBetaProfiler(physical_devices) + # search for the best logical mesh shape + logical_mesh_id = ab_profiler.search_best_logical_mesh() + logical_mesh_id = torch.Tensor(logical_mesh_id).to(torch.int) + + else: + logical_mesh_id = physical_mesh.reshape(logical_mesh_shape) + + device_mesh = DeviceMesh(physical_mesh_id=physical_mesh, logical_mesh_id=logical_mesh_id, init_process_group=True) + return device_mesh + + +class DeviceMeshManager: + """ + Device mesh manager is responsible for creating and managing device meshes. + """ + + def __init__(self): + self.device_mesh_store: Dict[str, DeviceMesh] = dict() + + def create_device_mesh(self, name, device_mesh_info: DeviceMeshInfo) -> DeviceMesh: + """ + Create a device mesh and store it in the manager. + + Args: + name (str): name of the device mesh + device_mesh_info (DeviceMeshInfo): the information used to initialize the device mesh + """ + if name not in self.device_mesh_store: + device_mesh = initialize_device_mesh(device_mesh_info) + self.device_mesh_store[name] = device_mesh + return device_mesh + else: + raise ValueError(f'Device mesh {name} already exists.') + + def get(self, name: str) -> DeviceMesh: + """ + Get a device mesh by name. + + Args: + name (str): name of the device mesh + + Returns: + DeviceMesh: the device mesh + """ + if name in self.device_mesh_store: + return self.device_mesh_store[name] + else: + raise ValueError(f'Device mesh {name} does not exist.') + + def destroy(self, name: str) -> None: + """ + Destroy a device mesh by name. + + Args: + name (str): name of the device mesh + """ + if name in self.device_mesh_store: + for pgs in self.device_mesh_store[name].process_groups_dict.values(): + for pg in pgs: + dist.destroy_process_group(pg) + del self.device_mesh_store[name] + else: + raise ValueError(f'Device mesh {name} does not exist.') + + def destroy_all(self): + """ + Destroy all device meshes. + """ + for name in self.device_mesh_store: + for pgs in self.device_mesh_store[name].process_groups_dict.values(): + for pg in pgs: + dist.destroy_process_group(pg) + + self.device_mesh_store.clear() diff --git a/colossalai/cluster/dist_coordinator.py b/colossalai/cluster/dist_coordinator.py new file mode 100644 index 000000000000..3ee364ec3364 --- /dev/null +++ b/colossalai/cluster/dist_coordinator.py @@ -0,0 +1,194 @@ +import functools +import os +from contextlib import contextmanager + +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from colossalai.context.singleton_meta import SingletonMeta + + +class DistCoordinator(metaclass=SingletonMeta): + """ + This class is used to coordinate distributed training. It is a singleton class, which means that there is only one instance of this + class in the whole program. + + There are some terms that are used in this class: + - rank: the rank of the current process + - world size: the total number of processes + - local rank: the rank of the current process on the current node + - master: the process with rank 0 + - node master: the process with local rank 0 on the current node + + Example: + >>> from colossalai.cluster.dist_coordinator import DistCoordinator + >>> coordinator = DistCoordinator() + >>> + >>> if coordinator.is_master(): + >>> do_something() + >>> + >>> coordinator.print_on_master('hello world') + + Attributes: + rank (int): the rank of the current process + world_size (int): the total number of processes + local_rank (int): the rank of the current process on the current node + """ + + def __init__(self): + assert dist.is_initialized( + ), 'Distributed is not initialized. Please call `torch.distributed.init_process_group` or `colossalai.launch` first.' + self._rank = dist.get_rank() + self._world_size = dist.get_world_size() + # this is often passed by launchers such as torchrun + self._local_rank = os.environ.get('LOCAL_RANK', -1) + + @property + def rank(self) -> int: + return self._rank + + @property + def world_size(self) -> int: + return self._world_size + + @property + def local_rank(self) -> int: + return self._local_rank + + def _assert_local_rank_set(self): + """ + Assert that the local rank is set. This is often passed by launchers such as torchrun. + """ + assert self.local_rank >= 0, 'The environment variable LOCAL_RANK is not set, thus the coordinator is not aware of the local rank of the current process.' + + def is_master(self, process_group: ProcessGroup = None) -> bool: + """ + Check if the current process is the master process (rank is 0). It can accept a sub process group to check the rank 0 with respect to the process. + + Args: + process_group (ProcessGroup, optional): process group to use for the rank 0 check. Defaults to None, which refers to the default process group. + + Returns: + bool: True if the current process is the master process, False otherwise + """ + rank = dist.get_rank(group=process_group) + return rank == 0 + + def is_node_master(self) -> bool: + """ + Check if the current process is the master process on the current node (local rank is 0). + + Returns: + bool: True if the current process is the master process on the current node, False otherwise + """ + self._assert_local_rank_set() + return self.local_rank == 0 + + def is_last_process(self, process_group: ProcessGroup = None) -> bool: + """ + Check if the current process is the last process (rank is world size - 1). It can accept a sub process group to check the last rank with respect to the process. + + Args: + process_group (ProcessGroup, optional): process group to use for the last rank check. Defaults to None, which refers to the default process group. + + Returns: + bool: True if the current process is the last process, False otherwise + """ + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + return rank == world_size - 1 + + def print_on_master(self, msg: str, process_group: ProcessGroup = None): + """ + Print message only from rank 0. + + Args: + msg (str): message to print + process_group (ProcessGroup, optional): process group to use for the rank 0 check. Defaults to None, which refers to the default process group. + """ + rank = dist.get_rank(group=process_group) + if rank == 0: + print(msg) + + def print_on_node_master(self, msg: str): + """ + Print message only from local rank 0. Local rank 0 refers to the 0th process running the current node. + + Args: + msg (str): message to print + """ + self._assert_local_rank_set() + if self.local_rank == 0: + print(msg) + + @contextmanager + def priority_execution(self, executor_rank: int = 0, process_group: ProcessGroup = None): + """ + This context manager is used to allow one process to execute while blocking all + other processes in the same process group. This is often useful when downloading is required + as we only want to download in one process to prevent file corruption. + + Example: + >>> from colossalai.cluster import DistCoordinator + >>> dist_coordinator = DistCoordinator() + >>> with dist_coordinator.priority_execution(): + >>> dataset = CIFAR10(root='./data', download=True) + + Args: + executor_rank (int): the process rank to execute without blocking, all other processes will be blocked + process_group (ProcessGroup, optional): process group to use for the executor rank check. Defaults to None, which refers to the default process group. + """ + rank = dist.get_rank(group=process_group) + should_block = rank != executor_rank + + if should_block: + self.block_all(process_group) + + yield + + if not should_block: + self.block_all(process_group) + + def destroy(self, process_group: ProcessGroup = None): + """ + Destroy the distributed process group. + + Args: + process_group (ProcessGroup, optional): process group to destroy. Defaults to None, which refers to the default process group. + """ + dist.destroy_process_group(process_group) + + def block_all(self, process_group: ProcessGroup = None): + """ + Block all processes in the process group. + + Args: + process_group (ProcessGroup, optional): process group to block. Defaults to None, which refers to the default process group. + """ + dist.barrier(group=process_group) + + def on_master_only(self, process_group: ProcessGroup = None): + """ + A function wrapper that only executes the wrapped function on the master process (rank 0). + + Example: + >>> from colossalai.cluster import DistCoordinator + >>> dist_coordinator = DistCoordinator() + >>> + >>> @dist_coordinator.on_master_only() + >>> def print_on_master(msg): + >>> print(msg) + """ + is_master = self.is_master(process_group) + + # define an inner function + def decorator(func): + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if is_master: + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/colossalai/cluster/process_group_manager.py b/colossalai/cluster/process_group_manager.py new file mode 100644 index 000000000000..e52661846f3e --- /dev/null +++ b/colossalai/cluster/process_group_manager.py @@ -0,0 +1,75 @@ +from typing import List + +import torch.distributed as dist +from torch.distributed import ProcessGroup + + +class ProcessGroupManager: + """ + ProcessGroupManager is used to manage the process groups in the cluster. + + There are some terms used in this class: + - pg: the short name for process group + - pg_name: the name of the process group + - pg_size: the world size of the process group + - rank: the rank of the current process in the process group + - world_size: the total number of processes in the process group + """ + + def __init__(self): + self.pg_store = dict() + + def create_process_group(self, name: str, ranks: List[int], backend: str = 'nccl') -> ProcessGroup: + """ + Get a process group by name. If the process group does not exist, it will be created. + + Args: + name (str): name of the process group + ranks (List[int]): ranks of the process group + backend (str, optional): backend of the process group. Defaults to 'nccl'. + + Returns: + ProcessGroup: the process group + """ + if name not in self.pg_store: + pg = dist.new_group(ranks=ranks, backend=backend) + self.pg_store[name] = pg + return pg + else: + raise ValueError(f'Process group {name} already exists.') + + def get(self, name: str) -> ProcessGroup: + """ + Get a process group by name. + + Args: + name (str): name of the process group + + Returns: + ProcessGroup: the process group + """ + if name in self.pg_store: + return self.pg_store[name] + else: + raise ValueError(f'Process group {name} does not exist.') + + def destroy(self, name: str) -> None: + """ + Destroy a process group by name. + + Args: + name (str): name of the process group + """ + if name in self.pg_store: + dist.destroy_process_group(self.pg_store[name]) + del self.pg_store[name] + else: + raise ValueError(f'Process group {name} does not exist.') + + def destroy_all(self) -> None: + """ + Destroy all process groups. + """ + for name in self.pg_store: + dist.destroy_process_group(self.pg_store[name]) + self.pg_store.clear() diff --git a/colossalai/communication/p2p.py b/colossalai/communication/p2p.py index 6dd4d0d6608d..1f20fca4f74d 100644 --- a/colossalai/communication/p2p.py +++ b/colossalai/communication/p2p.py @@ -103,10 +103,10 @@ def _communicate(object_send_next: Union[torch.Tensor, List[torch.Tensor]] = Non previous rank. recv_next (bool): boolean for whether tensor should be received from next rank. - recv_prev_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): shape of the tensor to be received from the previous stage, defualts to None. - recv_next_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): shape of the tensor to be received from the next stage, defualts to None. - prev_rank (int): the rank of the previous pipeline stage, defualts to None, - next_rank (int): the rank of the next pipeline stage, defualts to None, + recv_prev_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): shape of the tensor to be received from the previous stage, defaults to None. + recv_next_shape (Union[:class:`torch.Size`, List[:class:`torch.Size`]]): shape of the tensor to be received from the next stage, defaults to None. + prev_rank (int): the rank of the previous pipeline stage, defaults to None, + next_rank (int): the rank of the next pipeline stage, defaults to None, dtype (torch.dtype): data type of intermediate buffers, defaults to None scatter_gather_tensors (bool): whether to scatter and gather tensor between pipeline stages, defaults to False @@ -217,7 +217,7 @@ def recv_backward(output_grad_shape, next_rank (int, optional): The rank of the source of the tensor. Returns: - Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: The input gradient tensor or gradident tensor list. + Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: The input gradient tensor or gradient tensor list. """ if gpc.is_pipeline_last_stage(): output_tensor_grad = None diff --git a/colossalai/communication/p2p_v2.py b/colossalai/communication/p2p_v2.py index 4223f78d58cd..090311cb35f2 100644 --- a/colossalai/communication/p2p_v2.py +++ b/colossalai/communication/p2p_v2.py @@ -19,7 +19,7 @@ def init_process_group(): - """intialise process group by dist.new_group in the adjacent stages + """initialise process group by dist.new_group in the adjacent stages Args: None @@ -230,7 +230,7 @@ def recv_backward(next_rank: int = None) -> Any: next_rank (int, optional): The rank of the source of the tensor. Returns: - Any: The input gradient tensor or gradident tensor list. + Any: The input gradient tensor or gradient tensor list. """ if gpc.is_pipeline_last_stage(): output_tensor_grad = None diff --git a/colossalai/context/moe_context.py b/colossalai/context/moe_context.py index 1d7a883b1552..b41f4072a405 100644 --- a/colossalai/context/moe_context.py +++ b/colossalai/context/moe_context.py @@ -64,7 +64,7 @@ def setup(self, seed: int, use_kernel_optim: bool = True): from colossalai.core import global_context as gpc self.max_ep_size = gpc.config.get('max_ep_size', self.world_size) assert self.world_size % self.max_ep_size == 0, \ - "Maximum epxert parallel size must be a factor of the number of GPUs" + "Maximum expert parallel size must be a factor of the number of GPUs" self.min_dp_size = self.world_size // self.max_ep_size # Enabling kernel optimization may raise error in some cases diff --git a/colossalai/context/parallel_context.py b/colossalai/context/parallel_context.py index dd12dad6d347..003f0cdd91b6 100644 --- a/colossalai/context/parallel_context.py +++ b/colossalai/context/parallel_context.py @@ -10,15 +10,16 @@ import numpy as np import torch import torch.distributed as dist + from colossalai.constants import ALLOWED_MODES, INITIALIZER_MAPPING from colossalai.context.config import Config +from colossalai.context.singleton_meta import SingletonMeta from colossalai.global_variables import tensor_parallel_env as env from colossalai.logging import get_dist_logger from colossalai.registry import DIST_GROUP_INITIALIZER from .parallel_mode import ParallelMode from .random import add_seed, get_seeds, set_mode -from colossalai.context.singleton_meta import SingletonMeta class ParallelContext(metaclass=SingletonMeta): @@ -43,7 +44,7 @@ def __init__(self): # load config from file self._config = None - # default 3D parallel args, will be overwritten during process group intialization + # default 3D parallel args, will be overwritten during process group initialization self.world_size = 1 self.data_parallel_size = 1 self.pipeline_parallel_size = 1 @@ -263,7 +264,7 @@ def _add_world_size(self, parallel_mode: ParallelMode, world_size: int): """Adds world size for `parallel_mode`. Args: - parallel_mode (:class:`colossalai.context.ParallelMode`): The parallel mode correponding to the process group + parallel_mode (:class:`colossalai.context.ParallelMode`): The parallel mode corresponding to the process group world_size (int): The world size to be added Raises: diff --git a/colossalai/context/process_group_initializer/initializer_3d.py b/colossalai/context/process_group_initializer/initializer_3d.py index b752b8f45654..1ed8eec86efc 100644 --- a/colossalai/context/process_group_initializer/initializer_3d.py +++ b/colossalai/context/process_group_initializer/initializer_3d.py @@ -4,6 +4,7 @@ import math import torch.distributed as dist + from colossalai.global_variables import tensor_parallel_env as env from colossalai.registry import DIST_GROUP_INITIALIZER @@ -213,7 +214,8 @@ def init_dist_group(self): for h in range(self.num_group): for k in range(self.depth): ranks = [ - h * self.depth**3 + i + self.depth * (j + self.depth * k) for j in range(self.depth) + h * self.depth**3 + i + self.depth * (j + self.depth * k) + for j in range(self.depth) for i in range(self.depth) ] group = dist.new_group(ranks) @@ -266,7 +268,8 @@ def init_dist_group(self): for h in range(self.num_group): for j in range(self.depth): ranks = [ - h * self.depth**3 + i + self.depth * (j + self.depth * k) for k in range(self.depth) + h * self.depth**3 + i + self.depth * (j + self.depth * k) + for k in range(self.depth) for i in range(self.depth) ] group = dist.new_group(ranks) diff --git a/colossalai/context/process_group_initializer/initializer_data.py b/colossalai/context/process_group_initializer/initializer_data.py index 0b8b0d91fcb9..9715ebff7f00 100644 --- a/colossalai/context/process_group_initializer/initializer_data.py +++ b/colossalai/context/process_group_initializer/initializer_data.py @@ -4,8 +4,9 @@ from torch import distributed as dist from colossalai.registry import DIST_GROUP_INITIALIZER -from .process_group_initializer import ProcessGroupInitializer + from ..parallel_mode import ParallelMode +from .process_group_initializer import ProcessGroupInitializer @DIST_GROUP_INITIALIZER.register_module diff --git a/colossalai/context/process_group_initializer/initializer_sequence.py b/colossalai/context/process_group_initializer/initializer_sequence.py index eaacb14d2282..251a2940778a 100644 --- a/colossalai/context/process_group_initializer/initializer_sequence.py +++ b/colossalai/context/process_group_initializer/initializer_sequence.py @@ -91,11 +91,11 @@ def init_dist_group(self): parallel_setting = [] - local_rank, group_world_size, process_group, cpu_grop, ranks_in_group, mode = \ + local_rank, group_world_size, process_group, cpu_group, ranks_in_group, mode = \ self._sequence_initializer.init_dist_group() # change mode to sequence mode = ParallelMode.SEQUENCE - parallel_setting.append((local_rank, group_world_size, process_group, cpu_grop, ranks_in_group, mode)) + parallel_setting.append((local_rank, group_world_size, process_group, cpu_group, ranks_in_group, mode)) parallel_setting.append(self._sequence_dp_initializer.init_dist_group()) return parallel_setting diff --git a/colossalai/context/random/__init__.py b/colossalai/context/random/__init__.py index 422c3676c09d..d64b993257c1 100644 --- a/colossalai/context/random/__init__.py +++ b/colossalai/context/random/__init__.py @@ -1,5 +1,16 @@ -from ._helper import (seed, set_mode, with_seed, add_seed, get_seeds, get_states, get_current_mode, set_seed_states, - sync_states, moe_set_seed, reset_seeds) +from ._helper import ( + add_seed, + get_current_mode, + get_seeds, + get_states, + moe_set_seed, + reset_seeds, + seed, + set_mode, + set_seed_states, + sync_states, + with_seed, +) __all__ = [ 'seed', 'set_mode', 'with_seed', 'add_seed', 'get_seeds', 'get_states', 'get_current_mode', 'set_seed_states', diff --git a/colossalai/context/random/seed_manager.py b/colossalai/context/random/seed_manager.py index 3c84aaafc179..956f9001200d 100644 --- a/colossalai/context/random/seed_manager.py +++ b/colossalai/context/random/seed_manager.py @@ -59,23 +59,23 @@ def set_mode(self, parallel_mode: ParallelMode): self._current_mode = parallel_mode torch.cuda.set_rng_state(self._seed_states[parallel_mode]) - def add_seed(self, parallel_mode: ParallelMode, seed: int, overwrtie: bool = False): + def add_seed(self, parallel_mode: ParallelMode, seed: int, overwrite: bool = False): """Adds a seed to the seed manager for `parallel_mode`. Args: parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode. seed (int): The seed to be added. - overwrtie (bool, optional): Whether allows to overwrite the seed that has been set already + overwrite (bool, optional): Whether allows to overwrite the seed that has been set already Raises: AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of :class:`colossalai.context.ParallelMode` or the seed for `parallel_mode` has been added. """ assert isinstance(parallel_mode, ParallelMode), 'A valid ParallelMode must be provided' - if overwrtie is False: + if overwrite is False: assert parallel_mode not in self._seed_states, f'The seed for {parallel_mode} has been added' elif parallel_mode in self._seed_states: - print(f"Warnning: {parallel_mode} seed has been overwritten.", flush=True) + print(f"Warning: {parallel_mode} seed has been overwritten.", flush=True) current_state = torch.cuda.get_rng_state() torch.cuda.manual_seed(seed) diff --git a/colossalai/device/alpha_beta_profiler.py b/colossalai/device/alpha_beta_profiler.py index af2b10928c6f..f4e6cfffbcdf 100644 --- a/colossalai/device/alpha_beta_profiler.py +++ b/colossalai/device/alpha_beta_profiler.py @@ -197,7 +197,7 @@ def get_max_nbytes(process_group: Tuple[int], pg_handler: dist.ProcessGroup): dist.broadcast_object_list(broadcast_list, src=process_group[0]) alpha_beta_dict[process_group] = tuple(broadcast_list) - # add symmetry pair to the apha_beta_dict + # add symmetry pair to the alpha_beta_dict symmetry_ab_dict = {} for process_group, alpha_beta_pair in alpha_beta_dict.items(): symmetry_process_group = (process_group[1], process_group[0]) @@ -381,7 +381,7 @@ def _extract_alpha_beta(pg, pg_handler): first_latency, first_bandwidth = _extract_alpha_beta(first_axis, first_axis_process_group) second_latency, second_bandwidth = _extract_alpha_beta(second_axis, second_axis_process_group) mesh_alpha = [first_latency, second_latency] - # The beta values have been enlarged by 1e10 times temporarilly because the computation cost + # The beta values have been enlarged by 1e10 times temporarily because the computation cost # is still estimated in the unit of TFLOPs instead of time. We will remove this factor in future. mesh_beta = [1e10 / first_bandwidth, 1e10 / second_bandwidth] diff --git a/colossalai/device/device_mesh.py b/colossalai/device/device_mesh.py index 2a5f747fbc23..267c4529eb95 100644 --- a/colossalai/device/device_mesh.py +++ b/colossalai/device/device_mesh.py @@ -3,11 +3,19 @@ with some changes. """ import operator +from dataclasses import dataclass from functools import reduce -from typing import List, Tuple +from typing import Dict, List, Union import torch import torch.distributed as dist +from torch.distributed import ProcessGroup + + +@dataclass +class ProcessGroupContainer: + process_group: ProcessGroup + ranks: List[int] # modified from alpa LogicalDeviceMesh(https://github.com/alpa-projects/alpa/blob/main/alpa/shard_parallel/auto_sharding.py) @@ -27,9 +35,11 @@ class DeviceMesh: during initializing the DeviceMesh instance if the init_process_group set to True. Otherwise, users need to call create_process_groups_for_logical_mesh manually to init logical process group. (default: False) - need_flatten(bool, optional): initialize flatten_device_mesh during initializing the DeviceMesh instance if the need_flatten set to True. + device (str): the device for the process groups used by the DeviceMesh instance. (default: 'cuda') """ + _DIST_BACKEND = {"cuda": "nccl", "cpu": "gloo"} + def __init__(self, physical_mesh_id: torch.Tensor, mesh_shape: torch.Size = None, @@ -37,160 +47,442 @@ def __init__(self, mesh_alpha: List[float] = None, mesh_beta: List[float] = None, init_process_group: bool = False, - need_flatten: bool = True): - self.physical_mesh_id = physical_mesh_id + device: str = 'cuda'): + # ============================ + # Physical & Logical Mesh IDs + # ============================ + self._physical_mesh_id = physical_mesh_id + assert physical_mesh_id.dim() == 1, "physical_mesh_id should be a 1D tensor." + + # logical mesh ids can be obtained via two ways + # 1. provide physical mesh id and provide mesh shape + # 2. directly supply the logical mesh id + assert mesh_shape is None or logical_mesh_id is None, \ + "Only one of mesh_shape and logical_mesh_id can be specified." \ + "Logical mesh IDs are obtained from either mesh_shape + phyiscal_mesh_id or directly from the user-supplied logical_mesh_id" + if logical_mesh_id is None: - self.mesh_shape = mesh_shape - self._logical_mesh_id = self.physical_mesh_id.reshape(self.mesh_shape) + self._mesh_shape = mesh_shape + self._logical_mesh_id = self._physical_mesh_id.reshape(self._mesh_shape) else: self._logical_mesh_id = logical_mesh_id - self.mesh_shape = self._logical_mesh_id.shape + self._mesh_shape = self._logical_mesh_id.shape + + # ensure two things: + # 1. logical and physical mesh IDs should contain the same elements + # 2. there is no duplicate IDs in each mesh, e.g. [2, 2] is not allowed + assert torch.equal(torch.unique(self._physical_mesh_id), torch.unique(self.logical_mesh_id)), \ + "physical and logical mesh IDs should contain the same elements, please check if you have consistent physical_mesh_id and logical_mesh_id." + assert torch.unique(self._physical_mesh_id).numel() == self._physical_mesh_id.numel(), \ + "Found duplicate IDs in the phyiscal_mesh_id and this is not allowed, please check your physical_mesh_id again." + assert torch.unique(self.logical_mesh_id).numel() == self.logical_mesh_id.numel(), \ + "Found duplicate IDs in the logical_mesh_id and this is not allowed, please check your logical_mesh_id again." - # map global rank into logical rank - self.convert_map = {} - self._global_rank_to_logical_rank_map(self._logical_mesh_id, []) + # =============================================== # coefficient for alpha-beta communication model + # alpha is latency and beta is bandwidth + # =============================================== + # if the values are not provided, we assume they are 1 for simplicity if mesh_alpha is None: - mesh_alpha = [1] * len(self.mesh_shape) + mesh_alpha = [1] * len(self._mesh_shape) if mesh_beta is None: - mesh_beta = [1] * len(self.mesh_shape) + mesh_beta = [1] * len(self._mesh_shape) + self.mesh_alpha = tuple(mesh_alpha) self.mesh_beta = tuple(mesh_beta) - self.init_process_group = init_process_group - self.need_flatten = need_flatten - if self.init_process_group: - self.process_groups_dict = self.create_process_groups_for_logical_mesh() - if self.need_flatten and self._logical_mesh_id.dim() > 1: - self.flatten_device_mesh = self.flatten() - # Create a new member `flatten_device_meshes` to distinguish from original flatten methods (Because I'm not sure if there are functions that rely on the self.flatten()) - # self.flatten_device_meshes = FlattenDeviceMesh(self.physical_mesh_id, self.mesh_shape, self.mesh_alpha, - # self.mesh_beta) + + # ensure the alpha and beta have the same shape + assert len(self.mesh_alpha) == len(self.mesh_beta), \ + "mesh_alpha and mesh_beta should have the same length, please check your mesh_alpha and mesh_beta again." + + # ========================= + # Device for Process Group + # ========================= + self._device = device + self._dist_backend = self._DIST_BACKEND[device] + + # ========================= + # Process Group Management + # ========================= + # the _global_to_local_rank_mapping is structured as follows + # { + # : [ , , , ...] + # } + self._global_to_local_rank_mapping = dict() + self._init_global_to_logical_rank_mapping(mapping=self._global_to_local_rank_mapping, + tensor=self.logical_mesh_id) + + # create process group + self._process_group_dict = {} + self._ranks_in_the_process_group = {} + self._global_rank_of_current_process = None + self._is_initialized = False + + # attribute used to inidicate whether this objectd + # is created using DeviceMesh.from_process_group + # this attribute can be used to do some check in methods + # such get_process_group as no global rank information + # is known if created with from_process_group + self._is_init_from_process_group = False + + # initialize process group if specified + self._init_ranks_in_the_same_group() + self._init_process_group = init_process_group + if init_process_group: + self.init_logical_process_group() @property - def shape(self): - return self.mesh_shape + def shape(self) -> torch.Size: + """ + Return the shape of the logical mesh. + """ + return self._mesh_shape @property - def num_devices(self): - return reduce(operator.mul, self.physical_mesh_id.shape, 1) + def num_devices(self) -> int: + """ + Return the number of devices contained in the device mesh. + """ + return reduce(operator.mul, self._physical_mesh_id.shape, 1) @property - def logical_mesh_id(self): + def logical_mesh_id(self) -> torch.Tensor: + """ + Return the logical mesh id. + """ return self._logical_mesh_id - def __deepcopy__(self, memo): + @property + def is_initialized(self) -> bool: + """ + Return whether the process group is initialized. + """ + return self._is_initialized + + @staticmethod + def from_process_group(process_group: Union[ProcessGroup, List[ProcessGroup]]) -> "DeviceMesh": + """ + Create a DeviceMesh instance from the current process group. Please note that the DeviceMesh object created with this method + will not have information about the physical mesh id, and thus will not be able to query for other ranks and perform alpha-beta communication. + + Args: + process_group (Union[ProcessGroup, List[ProcessGroup]]): the process group or a list of process groups for the device mesh. + If the input is a ProcessGroup object, a 1D DeviceMesh object will be created. If the input is a list of ProcessGroup objects, + the ProcessGroup at the ith index will correspond to the process group in the ith axis of the device mesh. + + Returns: + DeviceMesh: the device mesh instance. + """ + + def _get_device_by_backend(process_group): + """ + Get the device type given a process group's backend. + """ + backend = dist.get_backend(process_group) + for _device, _backend in DeviceMesh._DIST_BACKEND.items(): + if _backend == backend: + return _device + return None + + if isinstance(process_group, ProcessGroup): + process_group = [process_group] + + # get mesh shape + mesh_shape = [dist.get_world_size(pg) for pg in process_group] + + # get device + device_list = [_get_device_by_backend(pg) for pg in process_group] + + # make sure all devices are the same + assert all([device == device_list[0] for device in device_list]), \ + "All devices should be the same, please check your input process groups are created with the same distributed backend." + + # create a fake physical mesh id + # as we only get the process group associated with the current process, + # we cannot get the global ranks for all processes in the mesh + # therefore, we only use this fake physical mesh id to create the device mesh + # and will remove this fake physical mesh id later + fake_physical_mesh_id = torch.arange(reduce(operator.mul, mesh_shape, 1)) + + # create the device mesh + device_mesh = DeviceMesh(physical_mesh_id=fake_physical_mesh_id, mesh_shape=mesh_shape, device=device_list[0]) + + # hack the device attribute + device_mesh._physical_mesh_id = None + device_mesh._logical_mesh_id = None + device_mesh._global_rank_of_current_process = dist.get_rank() + device_mesh._is_initialized = False + device_mesh._process_group_dict = { + device_mesh._global_rank_of_current_process: {axis: pg for axis, pg in enumerate(process_group)} + } + + return device_mesh + + def get_process_group(self, axis: int, global_rank: int = None) -> ProcessGroup: + """ + Return the process group on the specified axis. + + Args: + axis (int): the axis of the process group. + global_rank (int, optional): the global rank of the process group. If not specified, the current process is used. (default: None) + """ + if global_rank is None: + global_rank = self._global_rank_of_current_process + elif self._is_init_from_process_group: + raise RuntimeError( + "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known." + ) + return self._process_group_dict[global_rank][axis] + + def get_process_group_for_all_axes(self, global_rank: int = None) -> Dict[int, ProcessGroup]: + """ + Return the process groups for all axes. + + Args: + global_rank (int, optional): the global rank of the process + """ + if global_rank is None: + global_rank = self._global_rank_of_current_process + elif self._is_init_from_process_group: + raise RuntimeError( + "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known." + ) + return self._process_group_dict[global_rank] + + def get_ranks_in_process_group(self, axis: int, global_rank: int = None) -> List[int]: + """ + Return the ranks in the process group on the specified axis. + + Args: + axis (int): the axis of the process group. + global_rank (int, optional): the global rank of the process + """ + if global_rank is None: + global_rank = self._global_rank_of_current_process + elif self._is_init_from_process_group: + raise RuntimeError( + "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known." + ) + return self._ranks_in_the_process_group[global_rank][axis] + + def __deepcopy__(self, memo) -> "DeviceMesh": cls = self.__class__ result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): - if k != 'process_groups_dict': + if k != '_process_group_dict': setattr(result, k, __import__("copy").deepcopy(v, memo)) else: + # process group cannot be copied + # thus, we share them directly setattr(result, k, v) - return result - def flatten(self): - """ - Flatten the logical mesh into an effective 1d logical mesh, + def _init_global_to_logical_rank_mapping(self, + mapping: Dict, + tensor: torch.Tensor, + index_list: List[int] = []) -> Dict[int, List[int]]: """ - flatten_mesh_shape_size = len(self.mesh_shape) - flatten_mesh_shape = [self.num_devices] - return DeviceMesh(self.physical_mesh_id, - tuple(flatten_mesh_shape), - mesh_alpha=[max(self.mesh_alpha)] * (flatten_mesh_shape_size - 1), - mesh_beta=[max(self.mesh_beta)] * (flatten_mesh_shape_size - 1), - init_process_group=self.init_process_group, - need_flatten=False) + Build a global rank to local rank mapping for each process group in different axis in the logical device mesh. - def _global_rank_to_logical_rank_map(self, tensor, index_list): - ''' - This method is a helper function to build convert_map recursively. - ''' + Args: + mapping (Dict): a dictionary that maps the global rank to the local rank in the logical device mesh. + tensor (torch.Tensor): the tensor that contains the logical mesh ids. + index_list (List[int]) + + Returns: + mapping (Dict): a dictionary that maps the global rank to the local rank in the logical device mesh. + The value is a list of integers and each integer represents the local rank in the indexed axis. + """ for index, inner_tensor in enumerate(tensor): + # index means the local rank in the current axis + # inner_tensor refers to the processes with the same local rank + if inner_tensor.numel() == 1: - self.convert_map[int(inner_tensor)] = index_list + [index] + # if the inner_tensor only has one element, it means that + # it already reaches the last axis + # we append its local_rank in the last axis to the index_list + # and assign to the mapping + # the value of the mapping is the the local rank at the indexed axis of the device mesh + mapping[int(inner_tensor)] = index_list + [index] else: - self._global_rank_to_logical_rank_map(inner_tensor, index_list + [index]) + # we recursively go into the function until we reach the last axis + # meanwhile, we should add the local rank in the current axis in the index_list + self._init_global_to_logical_rank_mapping(mapping, inner_tensor, index_list + [index]) - def create_process_groups_for_logical_mesh(self): + def init_logical_process_group(self): ''' This method is used to initialize the logical process groups which will be used in communications among logical device mesh. Note: if init_process_group set to False, you have to call this method manually. Otherwise, the communication related function, such as ShapeConsistencyManager.apply will raise errors. ''' - process_groups_dict = {} - check_duplicate_list = [] - global_rank_flatten_list = self.physical_mesh_id.view(-1).tolist() + # sanity check + assert dist.is_initialized, "The torch.distributed should be initialized before calling init_logical_process_group" + assert not self._is_initialized, "The logical process group has been initialized, do not call init_logical_process_group twice" + + # update the global rank of the current process + self._global_rank_of_current_process = dist.get_rank() + duplicate_check_list = [] + + # flatten the global ranks to 1D list + global_rank_flatten_list = self._physical_mesh_id.view(-1).tolist() + for global_rank in global_rank_flatten_list: - process_groups = self.global_rank_to_process_groups_with_global_rank(global_rank) - for axis, process_group in process_groups.items(): - if axis not in process_groups_dict: - process_groups_dict[axis] = [] - if process_group not in check_duplicate_list: - check_duplicate_list.append(process_group) - process_group_handler = dist.new_group(process_group) - process_groups_dict[axis].append((process_group, process_group_handler)) + # find the other ranks which are in the same process group as global_rank + ranks_in_same_group_by_axis = self._collate_global_ranks_in_same_process_group(global_rank) - return process_groups_dict + for axis, ranks_in_same_group in ranks_in_same_group_by_axis.items(): + # skip duplicated process group creation + if ranks_in_same_group in duplicate_check_list: + continue - def global_rank_to_logical_rank(self, rank): - return self.convert_map[rank] + # create the process group + pg_handler = dist.new_group(ranks=ranks_in_same_group, backend=self._dist_backend) - def global_rank_to_process_groups_with_logical_rank(self, rank): - ''' - Give a global rank and return all logical process groups of this rank. - for example: - physical_mesh_id = torch.arange(0, 16).reshape(2, 8) - mesh_shape = (4, 4) - # [[0, 1, 2, 3], - # [4, 5, 6, 7], - # [8, 9, 10,11], - # [12,13,14,15]] - device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - print(device_mesh.global_rank_to_process_groups_with_logical_rank(0)) - output: - # key is axis name - # value is a list of logical ranks in same axis with rank 0 - {0: [[0, 0], [1, 0], [2, 0], [3, 0]], 1: [[0, 0], [0, 1], [0, 2], [0, 3]]} - ''' - process_groups = {} - for d in range(self.logical_mesh_id.dim()): - for replacer in range(self.logical_mesh_id.shape[d]): - if d not in process_groups: - process_groups[d] = [] - process_group_member = self.convert_map[rank].copy() - process_group_member[d] = replacer - process_groups[d].append(process_group_member) - return process_groups - - def global_rank_to_process_groups_with_global_rank(self, rank): + # keep this process group in the process_groups_dict + for rank in ranks_in_same_group: + if rank not in self._process_group_dict: + self._process_group_dict[rank] = dict() + self._process_group_dict[rank][axis] = pg_handler + + # update the init flag + # we only allow init for once + self._is_initialized = True + + def _init_ranks_in_the_same_group(self): + """ + This method is used to initialize the ranks_in_the_same_group dictionary. + """ + # flatten the global ranks to 1D list + global_rank_flatten_list = self._physical_mesh_id.view(-1).tolist() + + for global_rank in global_rank_flatten_list: + # find the other ranks which are in the same process group as global_rank + ranks_in_same_group_by_axis = self._collate_global_ranks_in_same_process_group(global_rank) + + for axis, ranks_in_same_group in ranks_in_same_group_by_axis.items(): + # create dict for each rank + if global_rank not in self._process_group_dict: + self._ranks_in_the_process_group[global_rank] = dict() + + # keep this process group in the process_groups_dict + self._ranks_in_the_process_group[global_rank][axis] = ranks_in_same_group + + def global_rank_to_local_rank(self, rank: int, axis: int = None) -> Union[List[int], int]: + """ + Return the local rank of the given global rank in the logical device mesh. + + Args: + rank (int): the global rank in the logical device mesh. + axis (int): the axis of the logical device mesh. + """ + if self._is_init_from_process_group: + raise RuntimeError( + "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known." + ) + + local_ranks = self._global_to_local_rank_mapping[rank] + if axis: + return local_ranks[axis] + else: + return local_ranks + + def _collate_global_ranks_in_same_process_group(self, global_rank): ''' - Give a global rank and return all process groups of this rank. - for example: - physical_mesh_id = torch.arange(0, 16).reshape(2, 8) - mesh_shape = (4, 4) - # [[0, 1, 2, 3], - # [4, 5, 6, 7], - # [8, 9, 10,11], - # [12,13,14,15]] - device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - print(device_mesh.global_rank_to_process_groups_with_global_rank(0)) - output: - # key is axis name - # value is a list of global ranks in same axis with rank 0 - {0: [0, 4, 8, 12], 1: [0, 1, 2, 3]} + Give a global rank and return all global ranks involved in its associated process group in each axis. + + Example: + + ```python + sphysical_mesh_id = torch.arange(0, 16) + mesh_shape = (4, 4) + + # logical mesh will look like + # [[0, 1, 2, 3], + # [4, 5, 6, 7], + # [8, 9, 10,11], + # [12,13,14,15]] + + device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) + print(device_mesh.collate_global_ranks_in_same_process_group(0)) + + # key is axis name + # value is a list of global ranks in same axis with rank 0 + # output will look like + # { + 0: [0, 4, 8, 12], + 1: [0, 1, 2, 3] + # } ''' - logical_process_groups = self.global_rank_to_process_groups_with_logical_rank(rank) - process_groups = {} - for dim, logical_ranks in logical_process_groups.items(): - process_groups[dim] = [] - for logical_rank in logical_ranks: - for g_rank, l_rank in self.convert_map.items(): - if l_rank == logical_rank: - process_groups[dim].append(g_rank) - return process_groups + # We have init the global rank to local rank by calling _init_global_to_logical_rank_mapping + # for self._global_to_local_rank_mapping + # the key is the global rank + # the value is the list of local ranks corresponding to the global rank with respect of different axes + # we can see the list of local ranks as the process coordinates for simplicity + # the key and value are all unique, therefore, + # we can also to use the coordinates to find the global rank + + # ========================================================================= + # Step 1 + # find all the process_coordinates for processes in the same process group + # as the given global rank + # ========================================================================= + + # each + processes_in_the_same_process_group = {} + + for dim in range(self.logical_mesh_id.dim()): + # iterate over the dimension size so that we can include all processes + # in the same process group in the given axis + # the _local_rank refers to the local rank of the current process + for _local_rank in range(self.logical_mesh_id.shape[dim]): + + # if this dimension is not initailized yet, + # initialize it with an empty array + if dim not in processes_in_the_same_process_group: + processes_in_the_same_process_group[dim] = [] + + # get the local rank corresponding to the global rank + process_coordinates = self._global_to_local_rank_mapping[global_rank].copy() + + # replace the local rank in the given dimension with the + # lcoal rank of the current process iterated + process_coordinates[dim] = _local_rank + processes_in_the_same_process_group[dim].append(process_coordinates) + + # ================================================================= + # Step 2 + # Use local rank combination to find its corresponding global rank + # ================================================================= + # the key of the dict is the axis + # the value is the list of global ranks which are in the same process group as the given global rank + global_pg_ranks = {} + for dim, coordinates_of_all_processes in processes_in_the_same_process_group.items(): + global_pg_ranks[dim] = [] + for process_coordinates in coordinates_of_all_processes: + # find the global rank by local rank combination + for _global_rank, _process_coordinates in self._global_to_local_rank_mapping.items(): + if process_coordinates == _process_coordinates: + global_pg_ranks[dim].append(_global_rank) + return global_pg_ranks + + def flatten(self): + """ + Flatten the logical mesh into an effective 1d logical mesh, + """ + if self._is_init_from_process_group: + raise RuntimeError( + "The logical device mesh is create with DeviceMesh.from_process_group, this method is not supported for this creation method as no global rank information is known." + ) + + flatten_mesh_shape_size = len(self._mesh_shape) + flatten_mesh_shape = [self.num_devices] + return DeviceMesh(self._physical_mesh_id, + tuple(flatten_mesh_shape), + mesh_alpha=[max(self.mesh_alpha)] * (flatten_mesh_shape_size - 1), + mesh_beta=[max(self.mesh_beta)] * (flatten_mesh_shape_size - 1), + init_process_group=self._init_process_group) def all_gather_cost(self, num_bytes, mesh_dim): num_devices = self.logical_mesh_id.shape[mesh_dim] @@ -212,38 +504,3 @@ def all_to_all_cost(self, num_bytes, mesh_dim): penalty_factor = num_devices / 2.0 return (self.mesh_alpha[mesh_dim] + self.mesh_beta[mesh_dim] * (num_devices - 1) / num_devices / num_devices * num_bytes * penalty_factor + 0.001) - - -class FlattenDeviceMesh(DeviceMesh): - - def __init__(self, physical_mesh_id, mesh_shape, mesh_alpha=None, mesh_beta=None): - super().__init__(physical_mesh_id, - mesh_shape, - mesh_alpha, - mesh_beta, - init_process_group=False, - need_flatten=False) - # Different from flatten(), mesh_shape leaves unchanged, mesh_alpha and mesh_beta are scalars - self.mesh_alpha = max(self.mesh_alpha) - self.mesh_beta = min(self.mesh_beta) - # Different from original process_groups_dict, rank_list is not stored - self.process_number_dict = self.create_process_numbers_for_logical_mesh() - - def create_process_numbers_for_logical_mesh(self): - ''' - Build 1d DeviceMesh in column-major(0) and row-major(1) - for example: - mesh_shape = (2,4) - # [[0, 1, 2, 3], - # [4, 5, 6, 7]] - # return {0: [0, 4, 1, 5, 2, 6, 3, 7], 1: [0, 1, 2, 3, 4, 5, 6, 7]} - ''' - num_devices = reduce(operator.mul, self.mesh_shape, 1) - process_numbers_dict = {} - process_numbers_dict[0] = torch.arange(num_devices).reshape(self.mesh_shape).transpose(1, 0).flatten().tolist() - process_numbers_dict[1] = torch.arange(num_devices).reshape(self.mesh_shape).flatten().tolist() - return process_numbers_dict - - def mix_gather_cost(self, num_bytes): - num_devices = reduce(operator.mul, self.mesh_shape, 1) - return (self.mesh_alpha + self.mesh_beta * (num_devices - 1) / num_devices * num_bytes + 0.1) diff --git a/colossalai/engine/_base_engine.py b/colossalai/engine/_base_engine.py index 59d8e1058652..db27ad0e8abe 100644 --- a/colossalai/engine/_base_engine.py +++ b/colossalai/engine/_base_engine.py @@ -10,9 +10,9 @@ from colossalai.engine.gradient_handler import BaseGradientHandler from colossalai.engine.schedule import BaseSchedule, InterleavedPipelineSchedule, NonPipelineSchedule, PipelineSchedule -from colossalai.gemini.ophooks import BaseOpHook, register_ophooks_recursively from colossalai.logging import get_dist_logger - +from colossalai.zero.legacy.gemini import BaseOpHook, register_ophooks_recursively +from colossalai.nn.optimizer import ColossalaiOptimizer class Engine: """Basic engine class for training and evaluation. It runs a specific process method diff --git a/colossalai/engine/gradient_accumulation/__init__.py b/colossalai/engine/gradient_accumulation/__init__.py index 4585b9a2529c..4cb6f4ad7384 100644 --- a/colossalai/engine/gradient_accumulation/__init__.py +++ b/colossalai/engine/gradient_accumulation/__init__.py @@ -1,10 +1,17 @@ +from typing import Iterable, List + import torch.nn as nn -from typing import List -from colossalai.engine import BaseGradientHandler -from typing import Iterable from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler -from ._gradient_accumulation import GradAccumDataloader, GradAccumOptimizer, GradAccumLrSchedulerByStep, GradAccumGradientHandler + +from colossalai.engine import BaseGradientHandler + +from ._gradient_accumulation import ( + GradAccumDataloader, + GradAccumGradientHandler, + GradAccumLrSchedulerByStep, + GradAccumOptimizer, +) __all__ = [ 'accumulate_gradient', 'GradAccumDataloader', 'GradAccumOptimizer', 'GradAccumLrSchedulerByStep', diff --git a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py index 89c28c3be87a..cf66be1cd821 100644 --- a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py +++ b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py @@ -1,21 +1,22 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from typing import Union +from typing import Any, Iterable, Tuple, Union + import torch.nn as nn from torch import Tensor -from typing import Iterable, Any, Tuple -from colossalai.nn.optimizer import ColossalaiOptimizer from torch.nn.parallel.distributed import DistributedDataParallel from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler from torch.utils.data import DataLoader -from colossalai.utils import conditional_context + from colossalai.engine import BaseGradientHandler +from colossalai.nn.optimizer import ColossalaiOptimizer +from colossalai.utils import conditional_context class GradAccumOptimizer(ColossalaiOptimizer): - """A wrapper for the optimizer to enable gradient accumulation by skipping the steps + """A wrapper for the optimizer to enable gradient accumulation by skipping the steps before accumulation size is reached. Args: @@ -161,7 +162,7 @@ def __next__(self) -> Union[Tensor, Tuple[Tensor]]: class GradAccumLrSchedulerByStep(_LRScheduler): - """A wrapper for the LR scheduler to enable gradient accumulation by skipping the steps + """A wrapper for the LR scheduler to enable gradient accumulation by skipping the steps before accumulation size is reached. Args: diff --git a/colossalai/engine/gradient_handler/__init__.py b/colossalai/engine/gradient_handler/__init__.py index 6177da69ba5b..2dea768bad7e 100644 --- a/colossalai/engine/gradient_handler/__init__.py +++ b/colossalai/engine/gradient_handler/__init__.py @@ -1,10 +1,9 @@ from ._base_gradient_handler import BaseGradientHandler from ._data_parallel_gradient_handler import DataParallelGradientHandler -from ._zero_gradient_handler import ZeROGradientHandler -from ._sequence_parallel_gradient_handler import SequenceParallelGradientHandler -from ._pipeline_parallel_gradient_handler import PipelineSharedModuleGradientHandler from ._moe_gradient_handler import MoeGradientHandler +from ._pipeline_parallel_gradient_handler import PipelineSharedModuleGradientHandler from ._sequence_parallel_gradient_handler import SequenceParallelGradientHandler +from ._zero_gradient_handler import ZeROGradientHandler __all__ = [ 'BaseGradientHandler', 'DataParallelGradientHandler', 'ZeROGradientHandler', 'PipelineSharedModuleGradientHandler', diff --git a/colossalai/engine/gradient_handler/_base_gradient_handler.py b/colossalai/engine/gradient_handler/_base_gradient_handler.py index c212359867d1..7d96dd8a88a6 100644 --- a/colossalai/engine/gradient_handler/_base_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_base_gradient_handler.py @@ -5,7 +5,7 @@ class BaseGradientHandler(ABC): - """A basic helper class to handle all-reduce operations of gradients across different parallel groups + """A basic helper class to handle all-reduce operations of gradients across different parallel groups before optimization. Args: diff --git a/colossalai/engine/gradient_handler/_data_parallel_gradient_handler.py b/colossalai/engine/gradient_handler/_data_parallel_gradient_handler.py index d113fc516459..5cc7169c5a9f 100644 --- a/colossalai/engine/gradient_handler/_data_parallel_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_data_parallel_gradient_handler.py @@ -1,16 +1,17 @@ from colossalai.core import global_context as gpc from colossalai.registry import GRADIENT_HANDLER -from ._base_gradient_handler import BaseGradientHandler + from ...context.parallel_mode import ParallelMode +from ._base_gradient_handler import BaseGradientHandler from .utils import bucket_allreduce @GRADIENT_HANDLER.register_module class DataParallelGradientHandler(BaseGradientHandler): """A helper class to handle all-reduce operations in a data parallel group. - A all-reduce collective communication will be operated in + A all-reduce collective communication will be operated in :func:`handle_gradient` among a data parallel group. - For better performance, it bucketizes the gradients of all parameters that are + For better performance, it bucketizes the gradients of all parameters that are the same type to improve the efficiency of communication. Args: diff --git a/colossalai/engine/gradient_handler/_moe_gradient_handler.py b/colossalai/engine/gradient_handler/_moe_gradient_handler.py index 02cea5e67a12..b499345d4e18 100644 --- a/colossalai/engine/gradient_handler/_moe_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_moe_gradient_handler.py @@ -1,45 +1,46 @@ -from colossalai.core import global_context as gpc -from colossalai.registry import GRADIENT_HANDLER -from colossalai.utils.moe import get_moe_epsize_param_dict -from ._base_gradient_handler import BaseGradientHandler -from ...context.parallel_mode import ParallelMode -from .utils import bucket_allreduce -from colossalai.context.moe_context import MOE_CONTEXT - - -@GRADIENT_HANDLER.register_module -class MoeGradientHandler(BaseGradientHandler): - """A helper class to handle all-reduce operations in a data parallel group and - moe model parallel. A all-reduce collective communication will be operated in - :func:`handle_gradient` among a data parallel group. - For better performance, it bucketizes the gradients of all parameters that are - the same type to improve the efficiency of communication. - - Args: - model (Module): Model where the gradients accumulate. - optimizer (Optimizer): Optimizer for updating the parameters. - """ - - def __init__(self, model, optimizer=None): - super().__init__(model, optimizer) - - def handle_gradient(self): - """A method running an all-reduce operation in a data parallel group. - Then running an all-reduce operation for all parameters in experts - across moe model parallel group - """ - global_data = gpc.data_parallel_size - - if global_data > 1: - epsize_param_dict = get_moe_epsize_param_dict(self._model) - - # epsize is 1, indicating the params are replicated among processes in data parallelism - # use the ParallelMode.DATA to get data parallel group - # reduce gradients for all parameters in data parallelism - if 1 in epsize_param_dict: - bucket_allreduce(param_list=epsize_param_dict[1], group=gpc.get_group(ParallelMode.DATA)) - - for ep_size in epsize_param_dict: - if ep_size != 1 and ep_size != MOE_CONTEXT.world_size: - bucket_allreduce(param_list=epsize_param_dict[ep_size], - group=MOE_CONTEXT.parallel_info_dict[ep_size].dp_group) +from colossalai.context.moe_context import MOE_CONTEXT +from colossalai.core import global_context as gpc +from colossalai.registry import GRADIENT_HANDLER +from colossalai.utils.moe import get_moe_epsize_param_dict + +from ...context.parallel_mode import ParallelMode +from ._base_gradient_handler import BaseGradientHandler +from .utils import bucket_allreduce + + +@GRADIENT_HANDLER.register_module +class MoeGradientHandler(BaseGradientHandler): + """A helper class to handle all-reduce operations in a data parallel group and + moe model parallel. A all-reduce collective communication will be operated in + :func:`handle_gradient` among a data parallel group. + For better performance, it bucketizes the gradients of all parameters that are + the same type to improve the efficiency of communication. + + Args: + model (Module): Model where the gradients accumulate. + optimizer (Optimizer): Optimizer for updating the parameters. + """ + + def __init__(self, model, optimizer=None): + super().__init__(model, optimizer) + + def handle_gradient(self): + """A method running an all-reduce operation in a data parallel group. + Then running an all-reduce operation for all parameters in experts + across moe model parallel group + """ + global_data = gpc.data_parallel_size + + if global_data > 1: + epsize_param_dict = get_moe_epsize_param_dict(self._model) + + # epsize is 1, indicating the params are replicated among processes in data parallelism + # use the ParallelMode.DATA to get data parallel group + # reduce gradients for all parameters in data parallelism + if 1 in epsize_param_dict: + bucket_allreduce(param_list=epsize_param_dict[1], group=gpc.get_group(ParallelMode.DATA)) + + for ep_size in epsize_param_dict: + if ep_size != 1 and ep_size != MOE_CONTEXT.world_size: + bucket_allreduce(param_list=epsize_param_dict[ep_size], + group=MOE_CONTEXT.parallel_info_dict[ep_size].dp_group) diff --git a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py index 83f5c00cf2af..5b49a9c0360d 100644 --- a/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_pipeline_parallel_gradient_handler.py @@ -4,9 +4,10 @@ import torch import torch.distributed as dist +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors + from colossalai.core import global_context as gpc from colossalai.registry import GRADIENT_HANDLER -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors from ._base_gradient_handler import BaseGradientHandler @@ -14,9 +15,9 @@ @GRADIENT_HANDLER.register_module class PipelineSharedModuleGradientHandler(BaseGradientHandler): """A helper class to handle all-reduce operations in sub parallel groups. - A all-reduce collective communication will be operated in + A all-reduce collective communication will be operated in :func:`handle_gradient` among all sub pipeline parallel groups. - For better performance, it bucketizes the gradients of all parameters that are + For better performance, it bucketizes the gradients of all parameters that are the same type to improve the efficiency of communication. Args: diff --git a/colossalai/engine/gradient_handler/_sequence_parallel_gradient_handler.py b/colossalai/engine/gradient_handler/_sequence_parallel_gradient_handler.py index 53a8ea935a42..ea4f0fbb1c71 100644 --- a/colossalai/engine/gradient_handler/_sequence_parallel_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_sequence_parallel_gradient_handler.py @@ -1,16 +1,17 @@ from colossalai.core import global_context as gpc from colossalai.registry import GRADIENT_HANDLER -from ._base_gradient_handler import BaseGradientHandler + from ...context.parallel_mode import ParallelMode +from ._base_gradient_handler import BaseGradientHandler from .utils import bucket_allreduce @GRADIENT_HANDLER.register_module class SequenceParallelGradientHandler(BaseGradientHandler): """A helper class to handle all-reduce operations in a data parallel group. - A all-reduce collective communication will be operated in + A all-reduce collective communication will be operated in :func:`handle_gradient` among a data parallel group. - For better performance, it bucketizes the gradients of all parameters that are + For better performance, it bucketizes the gradients of all parameters that are the same type to improve the efficiency of communication. Args: diff --git a/colossalai/engine/gradient_handler/_zero_gradient_handler.py b/colossalai/engine/gradient_handler/_zero_gradient_handler.py index f85303e75184..19fd1e97f86f 100644 --- a/colossalai/engine/gradient_handler/_zero_gradient_handler.py +++ b/colossalai/engine/gradient_handler/_zero_gradient_handler.py @@ -1,4 +1,5 @@ from colossalai.registry import GRADIENT_HANDLER + from ._base_gradient_handler import BaseGradientHandler diff --git a/colossalai/engine/schedule/__init__.py b/colossalai/engine/schedule/__init__.py index 54170286e99b..0f2c039d7057 100644 --- a/colossalai/engine/schedule/__init__.py +++ b/colossalai/engine/schedule/__init__.py @@ -1,5 +1,5 @@ from ._base_schedule import BaseSchedule -from ._pipeline_schedule import PipelineSchedule, InterleavedPipelineSchedule, get_tensor_shape from ._non_pipeline_schedule import NonPipelineSchedule +from ._pipeline_schedule import InterleavedPipelineSchedule, PipelineSchedule, get_tensor_shape __all__ = ['BaseSchedule', 'NonPipelineSchedule', 'PipelineSchedule', 'InterleavedPipelineSchedule', 'get_tensor_shape'] diff --git a/colossalai/engine/schedule/_base_schedule.py b/colossalai/engine/schedule/_base_schedule.py index ba797bad9778..a2d50041127a 100644 --- a/colossalai/engine/schedule/_base_schedule.py +++ b/colossalai/engine/schedule/_base_schedule.py @@ -2,10 +2,10 @@ # -*- encoding: utf-8 -*- from abc import ABC, abstractmethod +from typing import Callable, Iterable import torch -from typing import Iterable, Callable from colossalai.logging import get_dist_logger from colossalai.utils import get_current_device diff --git a/colossalai/engine/schedule/_non_pipeline_schedule.py b/colossalai/engine/schedule/_non_pipeline_schedule.py index c62bfb7d7375..b9239d928a7b 100644 --- a/colossalai/engine/schedule/_non_pipeline_schedule.py +++ b/colossalai/engine/schedule/_non_pipeline_schedule.py @@ -1,13 +1,14 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from typing import Iterable +import inspect +from typing import Callable, Iterable import torch -import inspect -from ._base_schedule import BaseSchedule + from colossalai.utils import conditional_context -from typing import Callable + +from ._base_schedule import BaseSchedule class NonPipelineSchedule(BaseSchedule): diff --git a/colossalai/engine/schedule/_pipeline_schedule.py b/colossalai/engine/schedule/_pipeline_schedule.py index 712ae8242409..9fc301a26559 100644 --- a/colossalai/engine/schedule/_pipeline_schedule.py +++ b/colossalai/engine/schedule/_pipeline_schedule.py @@ -152,12 +152,12 @@ def _get_data_slice(self, data, offset): raise TypeError(f"Expected data to be of type torch.Tensor, list, tuple, or dict, but got {type(data)}") def load_micro_batch(self): - mciro_batch_data = self._get_data_slice(self.batch_data, self.microbatch_offset) + micro_batch_data = self._get_data_slice(self.batch_data, self.microbatch_offset) self.microbatch_offset += self.microbatch_size - return self._move_to_device(mciro_batch_data) + return self._move_to_device(micro_batch_data) def pre_processing(self, engine): - from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 + from colossalai.zero.legacy import ShardedModelV2 # TODO: remove this after testing new zero with pipeline parallelism model = engine.model diff --git a/colossalai/engine/schedule/_pipeline_schedule_v2.py b/colossalai/engine/schedule/_pipeline_schedule_v2.py index 50a87aafad02..89e45c7aacec 100644 --- a/colossalai/engine/schedule/_pipeline_schedule_v2.py +++ b/colossalai/engine/schedule/_pipeline_schedule_v2.py @@ -1,11 +1,12 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from typing import Tuple, Iterable +from typing import Iterable, Tuple -from colossalai import engine -import colossalai.communication.p2p_v2 as comm import torch.cuda + +import colossalai.communication.p2p_v2 as comm +from colossalai import engine from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.utils.cuda import get_current_device @@ -35,7 +36,7 @@ def pack_return_tensors(return_tensors): class PipelineScheduleV2(PipelineSchedule): """Derived class of PipelineSchedule, the only difference is that forward_backward_step is reconstructed with p2p_v2 - + Args: num_microbatches (int): The number of microbatches. data_process_func (Callable, optional): @@ -43,9 +44,9 @@ class PipelineScheduleV2(PipelineSchedule): tensor_shape (torch.Size, optional): Specified shape in pipeline communication. scatter_gather_tensors (bool, optional): If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization. - + Example: - + # this shows an example of customized data_process_func def data_process_func(stage_output, dataloader_output): output1, output2 = stage_output @@ -83,7 +84,7 @@ def forward_backward_step(self, 'The argument \'return_loss\' has to be True when \'forward_only\' is False, but got False.' self.load_batch(data_iter) - # num_warmup_microbatches is the step when not all the processers are working + # num_warmup_microbatches is the step when not all the processes are working num_warmup_microbatches = \ (gpc.get_world_size(ParallelMode.PIPELINE) - gpc.get_local_rank(ParallelMode.PIPELINE) - 1) diff --git a/colossalai/fx/_compatibility.py b/colossalai/fx/_compatibility.py index 126403270301..0444a4816273 100644 --- a/colossalai/fx/_compatibility.py +++ b/colossalai/fx/_compatibility.py @@ -2,11 +2,19 @@ import torch -try: - from . import _meta_registrations - META_COMPATIBILITY = True -except: +TORCH_MAJOR = int(torch.__version__.split('.')[0]) +TORCH_MINOR = int(torch.__version__.split('.')[1]) + +if TORCH_MAJOR == 1 and TORCH_MINOR < 12: META_COMPATIBILITY = False +elif TORCH_MAJOR == 1 and TORCH_MINOR == 12: + from . import _meta_regist_12 + META_COMPATIBILITY = True +elif TORCH_MAJOR == 1 and TORCH_MINOR == 13: + from . import _meta_regist_13 + META_COMPATIBILITY = True +elif TORCH_MAJOR == 2: + META_COMPATIBILITY = True def compatibility(is_backward_compatible: bool = False) -> Callable: diff --git a/colossalai/fx/_meta_registrations.py b/colossalai/fx/_meta_regist_12.py similarity index 99% rename from colossalai/fx/_meta_registrations.py rename to colossalai/fx/_meta_regist_12.py index 153214447223..52e8d63ae543 100644 --- a/colossalai/fx/_meta_registrations.py +++ b/colossalai/fx/_meta_regist_12.py @@ -386,7 +386,7 @@ def meta_local_scalar_dense(self: torch.Tensor): @register_meta(aten.where.self) def meta_where_self(condition: torch.Tensor, self: torch.Tensor, other: torch.Tensor): result_type = torch.result_type(self, other) - return torch.empty_like(self, dtype=result_type) + return torch.empty_like(condition + self + other, dtype=result_type) @register_meta(aten.index.Tensor) diff --git a/colossalai/fx/_meta_regist_13.py b/colossalai/fx/_meta_regist_13.py new file mode 100644 index 000000000000..6caa87c449ab --- /dev/null +++ b/colossalai/fx/_meta_regist_13.py @@ -0,0 +1,57 @@ +import torch +from torch._meta_registrations import register_meta +from torch._prims_common import check + +aten = torch.ops.aten + + +# since we fix the torch version to 1.13.1, we have to add unimplemented meta ops +# all these functions are from here https://github.com/pytorch/pytorch/blob/master/torch/_meta_registrations.py +@register_meta([aten.convolution_backward.default]) +def meta_convolution_backward( + grad_output_, + input_, + weight_, + bias_sizes_opt, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output_mask, +): + # High level logic taken from slow_conv3d_backward_cpu which should + # be representative of all convolution_backward impls + backend_grad_input = None + backend_grad_weight = None + backend_grad_bias = None + + if output_mask[0]: + backend_grad_input = grad_output_.new_empty(input_.size()) + if output_mask[1]: + backend_grad_weight = grad_output_.new_empty(weight_.size()) + if output_mask[2]: + backend_grad_bias = grad_output_.new_empty(bias_sizes_opt) + + return (backend_grad_input, backend_grad_weight, backend_grad_bias) + + +@register_meta(aten._adaptive_avg_pool2d_backward.default) +def meta__adaptive_avg_pool2d_backward(grad_out, self): + ndim = grad_out.ndim + for i in range(1, ndim): + check( + grad_out.size(i) > 0, + lambda: f"adaptive_avg_pool2d_backward(): Expected grad_output to have non-zero \ + size for non-batch dimensions, {grad_out.shape} with dimension {i} being empty", + ) + check( + ndim == 3 or ndim == 4, + lambda: f"adaptive_avg_pool2d_backward(): Expected 3D or 4D tensor, but got {self.shape}", + ) + check( + self.dtype == grad_out.dtype, + lambda: f"expected dtype {self.dtype} for `grad_output` but got dtype {grad_out.dtype}", + ) + return self.new_empty(self.shape) diff --git a/colossalai/fx/codegen/activation_checkpoint_codegen.py b/colossalai/fx/codegen/activation_checkpoint_codegen.py index 492ebf918a9c..33b164800262 100644 --- a/colossalai/fx/codegen/activation_checkpoint_codegen.py +++ b/colossalai/fx/codegen/activation_checkpoint_codegen.py @@ -305,7 +305,7 @@ def emit_ckpt_func(body, delete_unused_value_func, level=0, in_ckpt=False): - """Emit ckpt fuction in nested way + """Emit ckpt function in nested way Args: body: forward code, in recursive calls, this part will be checkpoint functions code @@ -523,7 +523,7 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func, # append code text to body for idx, node in enumerate(node_list): # if this is the first node of the ckpt region - # append the ckpt function defition + # append the ckpt function definition if idx in start_idx: label = start_idx.index(idx) ckpt_fn_def = _gen_ckpt_fn_def(label, input_vars[label]) diff --git a/colossalai/fx/passes/adding_split_node_pass.py b/colossalai/fx/passes/adding_split_node_pass.py index 2c7b842b530c..245ba5d776da 100644 --- a/colossalai/fx/passes/adding_split_node_pass.py +++ b/colossalai/fx/passes/adding_split_node_pass.py @@ -206,7 +206,7 @@ def avgcompute_split_pass(gm: torch.fx.GraphModule, pp_size: int): def avgnode_split_pass(gm: torch.fx.GraphModule, pp_size: int): """ - In avgnode_split_pass, simpliy split graph by node number. + In avgnode_split_pass, simply split graph by node number. """ mod_graph = gm.graph avg_num_node = len(mod_graph.nodes) // pp_size diff --git a/colossalai/fx/passes/experimental/adding_shape_consistency_pass.py b/colossalai/fx/passes/experimental/adding_shape_consistency_pass.py index f28d65e2668a..4571bd93a790 100644 --- a/colossalai/fx/passes/experimental/adding_shape_consistency_pass.py +++ b/colossalai/fx/passes/experimental/adding_shape_consistency_pass.py @@ -16,7 +16,7 @@ def apply(*args, **kwargs): return shape_consistency_manager.apply(*args, **kwargs) -def solution_annotatation_pass(gm: torch.fx.GraphModule, solution: List[int], device_mesh): +def solution_annotation_pass(gm: torch.fx.GraphModule, solution: List[int], device_mesh): mod_graph = gm.graph nodes = tuple(mod_graph.nodes) diff --git a/colossalai/fx/passes/meta_info_prop.py b/colossalai/fx/passes/meta_info_prop.py index 2b4a8749cfd7..ab203dfd7440 100644 --- a/colossalai/fx/passes/meta_info_prop.py +++ b/colossalai/fx/passes/meta_info_prop.py @@ -31,7 +31,7 @@ class TensorMetadata(NamedTuple): numel: int is_tensor: bool # TODO: we can add a list of sharding spec here, and record the sharding - # behaviour by appending sharding spec into list. + # behavior by appending sharding spec into list. def _extract_tensor_metadata(result: torch.Tensor) -> TensorMetadata: diff --git a/colossalai/fx/passes/passes_for_gpt2_test.py b/colossalai/fx/passes/passes_for_gpt2_test.py index f98fcd686ea4..efdd34a01fe0 100644 --- a/colossalai/fx/passes/passes_for_gpt2_test.py +++ b/colossalai/fx/passes/passes_for_gpt2_test.py @@ -1,14 +1,15 @@ +import inspect +from typing import Any, Callable, Dict, List, Optional + import torch -from torch.fx.graph_module import GraphModule -from typing import Callable, List, Dict, Any, Optional -from torch.fx._compatibility import compatibility from packaging import version +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule +from torch.fx.node import Node + +from colossalai.fx.passes.adding_split_node_pass import balanced_split_pass, pipe_split from colossalai.fx.passes.meta_info_prop import TensorMetadata -import inspect -from typing import List from colossalai.fx.passes.split_module import Partition -from colossalai.fx.passes.adding_split_node_pass import pipe_split, balanced_split_pass -from torch.fx.node import Node def customized_split_pass_for_gpt2(gm: torch.fx.GraphModule, pp_size: int, partition_list: List[int]): @@ -229,7 +230,7 @@ def record_cross_partition_use(def_node: torch.fx.node.Node, use_partition.partitions_dependent_on.setdefault(def_partition_name) node_process_list = list(m.graph.nodes) - # split nodes into parititons + # split nodes into partitions while node_process_list: node = node_process_list.pop(0) orig_nodes[node.name] = node @@ -276,7 +277,7 @@ def record_cross_partition_use(def_node: torch.fx.node.Node, if len(sorted_partitions) != len(partitions): raise RuntimeError("cycle exists between partitions!") - # add placeholders to parititons + # add placeholders to partitions for partition_name in sorted_partitions: partition = partitions[partition_name] for input in partition.inputs: diff --git a/colossalai/fx/passes/split_module.py b/colossalai/fx/passes/split_module.py index bc257edc8c89..61ed037ab7a1 100644 --- a/colossalai/fx/passes/split_module.py +++ b/colossalai/fx/passes/split_module.py @@ -1,9 +1,10 @@ +import inspect +from typing import Any, Callable, Dict, List, Optional + import torch -from torch.fx.graph_module import GraphModule -from typing import Callable, List, Dict, Any, Optional -from torch.fx._compatibility import compatibility from packaging import version -import inspect +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule @compatibility(is_backward_compatible=True) @@ -28,8 +29,8 @@ def __repr__(self) -> str: f" nodes: {self.node_names},\n" \ f" inputs: {self.inputs},\n" \ f" outputs: {self.outputs},\n" \ - f" partitions depenent on: {self.partitions_dependent_on},\n" \ - f" parition dependents: {self.partition_dependents}" + f" partitions dependent on: {self.partitions_dependent_on},\n" \ + f" partition dependents: {self.partition_dependents}" # Creates subgraphs out of main graph @@ -38,7 +39,7 @@ def split_module( m: GraphModule, root_m: torch.nn.Module, split_callback: Callable[[torch.fx.node.Node], int], - merge_output = False, + merge_output=False, ): """ Adapted from https://github.com/pytorch/pytorch/blob/master/torch/fx/passes/split_module.py @@ -132,10 +133,8 @@ def record_cross_partition_use(def_node: torch.fx.node.Node, use_partition.inputs.setdefault(def_node.name) if def_partition_name is not None: use_partition.partitions_dependent_on.setdefault(def_partition_name) - - def record_output( - def_node: torch.fx.node.Node, use_node: Optional[torch.fx.node.Node] - ): # noqa: B950 + + def record_output(def_node: torch.fx.node.Node, use_node: Optional[torch.fx.node.Node]): # noqa: B950 def_partition_name = getattr(def_node, "_fx_partition", None) use_partition_name = getattr(use_node, "_fx_partition", None) if def_partition_name != use_partition_name: @@ -156,7 +155,7 @@ def record_output( use_partition = partitions[use_partition_name] use_partition.outputs.setdefault(def_node.name) - # split nodes into parititons + # split nodes into partitions for node in m.graph.nodes: orig_nodes[node.name] = node @@ -199,7 +198,7 @@ def record_output( if len(sorted_partitions) != len(partitions): raise RuntimeError("cycle exists between partitions!") - # add placeholders to parititons + # add placeholders to partitions for partition_name in sorted_partitions: partition = partitions[partition_name] for input in partition.inputs: @@ -291,7 +290,7 @@ def record_output( for partition_name in sorted_partitions: partition = partitions[partition_name] - + new_gm = torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph) return new_gm diff --git a/colossalai/fx/profiler/experimental/profiler_module/embedding.py b/colossalai/fx/profiler/experimental/profiler_module/embedding.py index dca6f9453af3..a1ade5d3ad93 100644 --- a/colossalai/fx/profiler/experimental/profiler_module/embedding.py +++ b/colossalai/fx/profiler/experimental/profiler_module/embedding.py @@ -1,5 +1,7 @@ from typing import Tuple + import torch + from ..registry import meta_profiler_module @@ -8,4 +10,4 @@ def torch_nn_embedding(self: torch.nn.Embedding, input: torch.Tensor) -> Tuple[i # nn.Embedding is a dictionary lookup, so technically it has 0 FLOPs. (https://discuss.pytorch.org/t/correct-way-to-calculate-flops-in-model/67198/6) flops = 0 macs = 0 - return flops, macs \ No newline at end of file + return flops, macs diff --git a/colossalai/fx/profiler/opcount.py b/colossalai/fx/profiler/opcount.py index 407a6bed5200..ba090a2ec51b 100644 --- a/colossalai/fx/profiler/opcount.py +++ b/colossalai/fx/profiler/opcount.py @@ -223,7 +223,8 @@ def zero_flop_jit(*args): return 0 -if version.parse(torch.__version__) >= version.parse('1.12.0'): +if version.parse(torch.__version__) >= version.parse('1.12.0') and version.parse( + torch.__version__) < version.parse('2.0.0'): flop_mapping = { # gemm, gemv and dot aten.mm.default: matmul_flop_jit, diff --git a/colossalai/fx/proxy.py b/colossalai/fx/proxy.py index 06272c48f852..7317072c6298 100644 --- a/colossalai/fx/proxy.py +++ b/colossalai/fx/proxy.py @@ -1,7 +1,9 @@ import operator +from typing import Any, List, Union + import torch -from torch.fx.proxy import Proxy, Attribute -from typing import List, Union, Any +from torch.fx.proxy import Attribute, Proxy + from colossalai.fx.tracer.meta_patch import meta_patched_function __all__ = ['ColoProxy'] diff --git a/colossalai/fx/tracer/_tracer_utils.py b/colossalai/fx/tracer/_tracer_utils.py index 0ec49a90a133..e160497a7444 100644 --- a/colossalai/fx/tracer/_tracer_utils.py +++ b/colossalai/fx/tracer/_tracer_utils.py @@ -1,6 +1,8 @@ -from typing import List, Union, Any -from ..proxy import ColoProxy, ColoAttribute +from typing import Any, List, Union + import torch + +from ..proxy import ColoAttribute, ColoProxy from .meta_patch import meta_patched_function, meta_patched_module __all__ = ['is_element_in_list', 'extract_meta'] diff --git a/colossalai/fx/tracer/bias_addition_patch/patched_bias_addition_module/bias_addition_module.py b/colossalai/fx/tracer/bias_addition_patch/patched_bias_addition_module/bias_addition_module.py index 85f1553e304c..591485fdb1ca 100644 --- a/colossalai/fx/tracer/bias_addition_patch/patched_bias_addition_module/bias_addition_module.py +++ b/colossalai/fx/tracer/bias_addition_patch/patched_bias_addition_module/bias_addition_module.py @@ -51,7 +51,7 @@ def extract_kwargs_from_mod(self): For example: The kwargs for conv2d module is {} because the attributes like 'padding' or 'groups' are - considered during module initilizing. However, we need to consider those attributes as kwargs + considered during module initializing. However, we need to consider those attributes as kwargs in F.conv2d. """ pass diff --git a/colossalai/fx/tracer/experimental.py b/colossalai/fx/tracer/experimental.py index 88b65b6188fa..22a67d1ceccc 100644 --- a/colossalai/fx/tracer/experimental.py +++ b/colossalai/fx/tracer/experimental.py @@ -295,7 +295,7 @@ class PatchedCheckpointFunction(torch.autograd.Function): @staticmethod def forward(ctx, run_function, preserve_rng_state, *args): - # signal that the current tracing occurs within activaton checkpoint part + # signal that the current tracing occurs within activation checkpoint part self.inside_torch_checkpoint_func = True out = run_function(*args) self.inside_torch_checkpoint_func = False diff --git a/colossalai/fx/tracer/tracer.py b/colossalai/fx/tracer/tracer.py index 1ae31f958975..28965a1b8e74 100644 --- a/colossalai/fx/tracer/tracer.py +++ b/colossalai/fx/tracer/tracer.py @@ -92,7 +92,7 @@ def create_proxy(self, kind, target, args, kwargs, name=None, type_expr=None, pr return proxy # if graph is traced for auto parallelism module, some extra node will be added during - # graph construction to deal with the compatability between bias addition and all reduce. + # graph construction to deal with the compatibility between bias addition and all reduce. # if no extra manipulation is applied, we just pass the origin arguments to create_proxy function # to create node on computation graph @@ -208,7 +208,7 @@ def _configure_tracer_type(self, tracer_type: TracerType): self.proxy_cls = ColoProxy self.tracer_type = TracerType.META else: - raise ValueError(f"Unrecognised tracer type {tracer_type}") + raise ValueError(f"Unrecognized tracer type {tracer_type}") def _meta_data_computing(self, kind, target, args, kwargs): @@ -445,7 +445,7 @@ class PatchedCheckpointFunction(torch.autograd.Function): @staticmethod def forward(ctx, run_function, preserve_rng_state, *args): - # signal that the current tracing occurs within activaton checkpoint part + # signal that the current tracing occurs within activation checkpoint part self.inside_torch_checkpoint_func = True out = run_function(*args) self.inside_torch_checkpoint_func = False diff --git a/colossalai/gemini/__init__.py b/colossalai/gemini/__init__.py deleted file mode 100644 index 7a5a44ebb1ef..000000000000 --- a/colossalai/gemini/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration -from .gemini_mgr import GeminiManager -from .stateful_tensor_mgr import StatefulTensorMgr -from .tensor_placement_policy import TensorPlacementPolicyFactory - -__all__ = [ - 'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager', - 'search_chunk_configuration' -] diff --git a/colossalai/global_variables.py b/colossalai/global_variables.py index e3575ea12ad0..61b31965e2e6 100644 --- a/colossalai/global_variables.py +++ b/colossalai/global_variables.py @@ -1,56 +1,56 @@ -from typing import Optional - - -class TensorParallelEnv(object): - _instance = None - - def __new__(cls, *args, **kwargs): - if cls._instance is None: - cls._instance = object.__new__(cls, *args, **kwargs) - return cls._instance - - def __init__(self, *args, **kwargs): - self.load(*args, **kwargs) - - def load(self, - mode: Optional[str] = None, - vocab_parallel: bool = False, - parallel_input_1d: bool = False, - summa_dim: int = None, - tesseract_dim: int = None, - tesseract_dep: int = None, - depth_3d: int = None, - input_group_3d=None, - weight_group_3d=None, - output_group_3d=None, - input_x_weight_group_3d=None, - output_x_weight_group_3d=None): - self.mode = mode - self.vocab_parallel = vocab_parallel - self.parallel_input_1d = parallel_input_1d - self.summa_dim = summa_dim - self.tesseract_dim = tesseract_dim - self.tesseract_dep = tesseract_dep - self.depth_3d = depth_3d - self.input_group_3d = input_group_3d - self.weight_group_3d = weight_group_3d - self.output_group_3d = output_group_3d - self.input_x_weight_group_3d = input_x_weight_group_3d - self.output_x_weight_group_3d = output_x_weight_group_3d - - def save(self): - return dict(mode=self.mode, - vocab_parallel=self.vocab_parallel, - parallel_input_1d=self.parallel_input_1d, - summa_dim=self.summa_dim, - tesseract_dim=self.tesseract_dim, - tesseract_dep=self.tesseract_dep, - depth_3d=self.depth_3d, - input_group_3d=self.input_group_3d, - weight_group_3d=self.weight_group_3d, - output_group_3d=self.output_group_3d, - input_x_weight_group_3d=self.input_x_weight_group_3d, - output_x_weight_group_3d=self.output_x_weight_group_3d) - - -tensor_parallel_env = TensorParallelEnv() +from typing import Optional + + +class TensorParallelEnv(object): + _instance = None + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = object.__new__(cls, *args, **kwargs) + return cls._instance + + def __init__(self, *args, **kwargs): + self.load(*args, **kwargs) + + def load(self, + mode: Optional[str] = None, + vocab_parallel: bool = False, + parallel_input_1d: bool = False, + summa_dim: int = None, + tesseract_dim: int = None, + tesseract_dep: int = None, + depth_3d: int = None, + input_group_3d=None, + weight_group_3d=None, + output_group_3d=None, + input_x_weight_group_3d=None, + output_x_weight_group_3d=None): + self.mode = mode + self.vocab_parallel = vocab_parallel + self.parallel_input_1d = parallel_input_1d + self.summa_dim = summa_dim + self.tesseract_dim = tesseract_dim + self.tesseract_dep = tesseract_dep + self.depth_3d = depth_3d + self.input_group_3d = input_group_3d + self.weight_group_3d = weight_group_3d + self.output_group_3d = output_group_3d + self.input_x_weight_group_3d = input_x_weight_group_3d + self.output_x_weight_group_3d = output_x_weight_group_3d + + def save(self): + return dict(mode=self.mode, + vocab_parallel=self.vocab_parallel, + parallel_input_1d=self.parallel_input_1d, + summa_dim=self.summa_dim, + tesseract_dim=self.tesseract_dim, + tesseract_dep=self.tesseract_dep, + depth_3d=self.depth_3d, + input_group_3d=self.input_group_3d, + weight_group_3d=self.weight_group_3d, + output_group_3d=self.output_group_3d, + input_x_weight_group_3d=self.input_x_weight_group_3d, + output_x_weight_group_3d=self.output_x_weight_group_3d) + + +tensor_parallel_env = TensorParallelEnv() diff --git a/colossalai/initialize.py b/colossalai/initialize.py index f3719dcb47b3..dc0df0517508 100644 --- a/colossalai/initialize.py +++ b/colossalai/initialize.py @@ -29,13 +29,12 @@ PipelineSchedule, get_tensor_shape, ) -from colossalai.gemini.ophooks import BaseOpHook from colossalai.logging import get_dist_logger from colossalai.nn.optimizer.colossalai_optimizer import ColossalaiOptimizer from colossalai.utils import get_current_device, is_using_ddp, is_using_pp, is_using_sequence, sync_model_param from colossalai.utils.moe import sync_moe_model_param -from colossalai.zero import convert_to_zero_v2 -from colossalai.zero.sharded_optim.sharded_optim_v2 import ShardedOptimizerV2 +from colossalai.zero.legacy import ShardedOptimizerV2, convert_to_zero_v2 +from colossalai.zero.legacy.gemini.ophooks import BaseOpHook def get_default_parser(): @@ -239,7 +238,7 @@ def initialize(model: nn.Module, loaded into gpc.config. Args: - model (:class:`torch.nn.Module` or Callbale): Your model instance or a function to build the model. + model (:class:`torch.nn.Module` or Callable): Your model instance or a function to build the model. optimizer (:class:`torch.optim.optimizer.Optimizer` or :class:`Type[torch.optim.optimizer]`): Your optimizer instance. criterion (:class:`torch.nn.modules.loss._Loss`, optional): Your criterion instance. diff --git a/colossalai/interface/__init__.py b/colossalai/interface/__init__.py new file mode 100644 index 000000000000..8c658e375146 --- /dev/null +++ b/colossalai/interface/__init__.py @@ -0,0 +1,4 @@ +from .model import ModelWrapper +from .optimizer import OptimizerWrapper + +__all__ = ['OptimizerWrapper', 'ModelWrapper'] diff --git a/colossalai/interface/model.py b/colossalai/interface/model.py new file mode 100644 index 000000000000..a067d7671ce7 --- /dev/null +++ b/colossalai/interface/model.py @@ -0,0 +1,25 @@ +import torch.nn as nn + + +class ModelWrapper(nn.Module): + """ + A wrapper class to define the common interface used by booster. + + Args: + module (nn.Module): The model to be wrapped. + """ + + def __init__(self, module: nn.Module) -> None: + super().__init__() + self.module = module + + def unwrap(self): + """ + Unwrap the model to return the original model for checkpoint saving/loading. + """ + if isinstance(self.module, ModelWrapper): + return self.module.unwrap() + return self.module + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) diff --git a/colossalai/booster/interface/optimizer.py b/colossalai/interface/optimizer.py similarity index 96% rename from colossalai/booster/interface/optimizer.py rename to colossalai/interface/optimizer.py index dd9acab17584..0eaf2e1ef8ba 100644 --- a/colossalai/booster/interface/optimizer.py +++ b/colossalai/interface/optimizer.py @@ -119,3 +119,9 @@ def unscale_grad(self): """ raise NotImplementedError( "The method unscale_grad is only available for optimizers with mixed precision training") + + def unwrap(self): + """ + Unwrap the optimizer for checkpoint saving/loading. + """ + return self.optim diff --git a/colossalai/kernel/cuda_native/csrc/type_shim.h b/colossalai/kernel/cuda_native/csrc/type_shim.h index 2f180a7783ec..03ccc02635fa 100644 --- a/colossalai/kernel/cuda_native/csrc/type_shim.h +++ b/colossalai/kernel/cuda_native/csrc/type_shim.h @@ -171,6 +171,21 @@ using g_scalar_t_##LEVEL = at::Half; \ using p_scalar_t_##LEVEL = at::Half; \ __VA_ARGS__; \ + } else if (GTYPE == at::ScalarType::Float && \ + PTYPE == at::ScalarType::BFloat16) { \ + using g_scalar_t_##LEVEL = float; \ + using p_scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + } else if (GTYPE == at::ScalarType::BFloat16 && \ + PTYPE == at::ScalarType::Float) { \ + using g_scalar_t_##LEVEL = at::BFloat16; \ + using p_scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + } else if (GTYPE == at::ScalarType::BFloat16 && \ + PTYPE == at::ScalarType::BFloat16) { \ + using g_scalar_t_##LEVEL = at::BFloat16; \ + using p_scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ } else { \ AT_ERROR(#NAME, "not implemented for '", toString(GTYPE), toString(PTYPE), \ "'"); \ diff --git a/colossalai/kernel/cuda_native/flash_attention.py b/colossalai/kernel/cuda_native/flash_attention.py index 907fa640d826..3db7374509a0 100644 --- a/colossalai/kernel/cuda_native/flash_attention.py +++ b/colossalai/kernel/cuda_native/flash_attention.py @@ -1,12 +1,6 @@ """ -The triton-based flash attention implementation is copied from the OpenAI/triton repository - -You can find the repository in Triton https://github.com/openai/triton -You can find the source file in https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py - -Reference: -1. Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf -2. Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf +A general attention module using the flash attention kernels from xformers: +https://github.com/facebookresearch/xformers/tree/main/xformers/ops/fmha """ import math @@ -15,6 +9,159 @@ import torch +try: + from xformers.ops.fmha import memory_efficient_attention + HAS_MEM_EFF_ATTN = True +except ImportError: + HAS_MEM_EFF_ATTN = False + print('please install xformers from https://github.com/facebookresearch/xformers') + +if HAS_MEM_EFF_ATTN: + + from typing import Optional + + from einops import rearrange + from xformers.ops.fmha import MemoryEfficientAttentionCutlassOp + from xformers.ops.fmha.attn_bias import BlockDiagonalMask, LowerTriangularMask, LowerTriangularMaskWithTensorBias + + from .scaled_softmax import AttnMaskType + + allow_alibi = True + for op in MemoryEfficientAttentionCutlassOp: + allow_alibi = allow_alibi & (LowerTriangularMaskWithTensorBias in op.SUPPORTED_ATTN_BIAS_TYPES) + + class Unpad(torch.autograd.Function): + """ + Adapted from + https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py + """ + + @staticmethod + def forward(ctx, tensor: torch.Tensor, indices: torch.Tensor): + ctx.save_for_backward(indices) + # [b, s, ...] + assert tensor.ndim >= 3 + ctx.bsz = tensor.shape[0] + out = rearrange(tensor, 'b s ... -> (b s) ...') + ctx.shape = out.shape + # [1, ntokens, ...] + return out[indices].unsqueeze(0) + + @staticmethod + def backward(ctx, grad_output): + indices, = ctx.saved_tensors + # [b*s, ...] + grad = torch.zeros(ctx.shape, dtype=grad_output.dtype, device=grad_output.device) + grad[indices] = grad_output.squeeze(0) + grad = rearrange(grad, '(b s) ... -> b s ...', b=ctx.bsz) + # [b, s, ...] + return grad, None + + class Repad(torch.autograd.Function): + """ + Adapted from + https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py + """ + + @staticmethod + def forward(ctx, tensor: torch.Tensor, indices: torch.Tensor, batch_size: int, seq_len: int): + ctx.save_for_backward(indices) + # [ntokens, ...] + tensor = tensor.squeeze(0) + out = torch.zeros((batch_size * seq_len, *tensor.shape[1:]), dtype=tensor.dtype, device=tensor.device) + # [b*s, ...] + out[indices] = tensor + # [b, s, ...] + out = rearrange(out, '(b s) ... -> b s ...', b=batch_size) + return out + + @staticmethod + def backward(ctx, grad_output): + indices, = ctx.saved_tensors + # [b*s, ...] + grad_output = rearrange(grad_output, 'b s ... -> (b s) ...') + grad = grad_output[indices] + # [1, ntokens, ...] + return grad.unsqueeze(0), None, None, None + + class ColoAttention(torch.nn.Module): + + def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0): + super().__init__() + assert embed_dim % num_heads == 0, \ + f"the embed dim ({embed_dim}) is not divisible by the number of attention heads ({num_heads})." + self.scale = 1 / math.sqrt(embed_dim // num_heads) + self.dropout = dropout + + @staticmethod + def get_seq_info_from_mask(attn_mask: torch.Tensor): + indices = torch.nonzero(attn_mask.flatten(), as_tuple=False).flatten() + seqlens = attn_mask.sum(dim=-1, dtype=torch.int32).flatten().tolist() + return indices, seqlens + + @staticmethod + def unpad(tensor: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + return Unpad.apply(tensor, indices) + + @staticmethod + def repad(tensor: torch.Tensor, indices: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor: + return Repad.apply(tensor, indices, batch_size, seq_len) + + def forward(self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + attn_mask_type: Optional[AttnMaskType] = None, + bias: Optional[torch.Tensor] = None): + batch_size, tgt_len, src_len = query.shape[0], query.shape[1], key.shape[1] + attn_bias = None + if attn_mask_type == AttnMaskType.padding: # bert style + assert attn_mask is not None, \ + f"attention mask {attn_mask} is not valid for attention mask type {attn_mask_type}." + assert attn_mask.dim() == 2, \ + "attention mask is supposed to have shape (batch_size, seq_len), " + \ + f"but got {attn_mask.dim()} dimensions." + if tgt_len == src_len: + q_indices, q_seqlen = self.get_seq_info_from_mask(attn_mask) + kv_seqlen = None + if batch_size > 1: + query, key, value = self.unpad(torch.stack([query, key, value], dim=2), q_indices).unbind(dim=2) + else: + q_indices = torch.arange(batch_size * tgt_len, dtype=torch.int32, device=query.device) + q_seqlen = torch.LongTensor([tgt_len] * batch_size, device=query.device) + kv_indices, kv_seqlen = self.get_seq_info_from_mask(attn_mask) + if batch_size > 1: + query = rearrange(query, "b s ... -> c (b s) ...", c=1) + key, value = self.unpad(torch.stack([query, key, value], dim=2), kv_indices).unbind(dim=2) + attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen) + elif attn_mask_type == AttnMaskType.causal: # gpt style + attn_bias = LowerTriangularMask() + + if bias is not None: # alibi / relative position embedding + assert allow_alibi, "flash attention with bias is not supported in this system." + assert attn_mask_type == AttnMaskType.causal, \ + "attention with bias is only supported for causal attention so far." + attn_bias = attn_bias.add_bias(bias) + + out = memory_efficient_attention(query, key, value, attn_bias=attn_bias, p=self.dropout, scale=self.scale) + + if attn_mask_type == AttnMaskType.padding and batch_size > 1: + out = self.repad(out, q_indices, batch_size, tgt_len) + + out = rearrange(out, 'b s h d -> b s (h d)') + return out + + +########################################################################## +# the flash attention functions below that are copied +# from the OpenAI/triton repository will be deprecated +# You can find the repository in Triton https://github.com/openai/triton +# You can find the source file in https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py +# Reference: +# 1. Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf +# 2. Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf + def triton_cuda_check(): cuda_home = os.getenv("CUDA_HOME", default="/usr/local/cuda") @@ -52,13 +199,6 @@ def triton_cuda_check(): HAS_FLASH_ATTN = False print('please install flash_attn from https://github.com/HazyResearch/flash-attention') -try: - from xformers.ops.fmha import memory_efficient_attention - HAS_MEM_EFF_ATTN = True -except ImportError: - HAS_MEM_EFF_ATTN = False - print('please install xformers from https://github.com/facebookresearch/xformers') - if HAS_TRITON: # the following functions are adapted from the OpenAI Triton tutorial # https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py @@ -422,25 +562,6 @@ def triton_flash_attention(q, k, v, sm_scale): if HAS_FLASH_ATTN: - from einops import rearrange - - class MaskedFlashAttention(torch.nn.Module): - - def __init__(self, num_attention_heads: int, attention_head_size: int, attention_dropout: float) -> None: - super().__init__() - self.num_attention_heads = num_attention_heads - self.attention_head_size = attention_head_size - self.attention_func = FlashAttention(softmax_scale=math.sqrt(attention_head_size), - attention_dropout=attention_dropout) - - def forward(self, query_key_value: torch.Tensor, attention_mask: torch.Tensor, causal=False): - if attention_mask.dtype is not torch.bool: - attention_mask = attention_mask.bool() - qkv = rearrange(query_key_value, 'b s (three h d) -> b s three h d', three=3, h=self.num_attention_heads) - context, _ = self.attention_func(qkv, key_padding_mask=attention_mask, causal=causal) - context = rearrange(context, 'b s h d -> b s (h d)') - return context - def flash_attention_qkv(qkv, sm_scale, batch_size, seq_len, dropout_p=0., causal=False): """ Arguments: @@ -511,20 +632,4 @@ def flash_attention_q_k_v(q, k, v, sm_scale, batch_size, q_seqlen, kv_seqlen, dr causal) -if HAS_MEM_EFF_ATTN: - - from einops import rearrange - from xformers.ops.fmha import LowerTriangularMask - - class MemoryEfficientAttention(torch.nn.Module): - - def __init__(self, hidden_size: int, num_attention_heads: int, attention_dropout: float = 0.0): - super().__init__() - attention_head_size = hidden_size // num_attention_heads - self.scale = 1 / attention_head_size**0.5 - self.dropout = attention_dropout - - def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor): - context = memory_efficient_attention(query, key, value, attention_mask, self.dropout, self.scale) - context = rearrange(context, 'b s h d -> b s (h d)') - return context +########################################################################## diff --git a/colossalai/kernel/cuda_native/multihead_attention.py b/colossalai/kernel/cuda_native/multihead_attention.py index 7df53731edc5..69246f2f3854 100644 --- a/colossalai/kernel/cuda_native/multihead_attention.py +++ b/colossalai/kernel/cuda_native/multihead_attention.py @@ -43,7 +43,7 @@ class Config: attn_prob_dropout_ratio: float # attention score dropout ratio hidden_dropout_ratio: float # dropout ration before residual norm_first: bool # norm_first - fp16: bool # fp16 presion + fp16: bool # fp16 precision class MultiHeadAttention1DFunc(Function): @@ -111,7 +111,7 @@ class MultiHeadAttention(nn.Module): Arguments: hidden_size: Total dimension of hidden_size. nhead: Number of parallel attention heads. - batch_size: Batch Size for one foward + batch_size: Batch Size for one forward max_seq_len: Max length of input sequence dropout: Dropout probability norm_first: perform LayerNorms before attention diff --git a/colossalai/kernel/jit/bias_dropout_add.py b/colossalai/kernel/jit/bias_dropout_add.py index 3687dde79a08..32965c1ebd69 100644 --- a/colossalai/kernel/jit/bias_dropout_add.py +++ b/colossalai/kernel/jit/bias_dropout_add.py @@ -1,4 +1,5 @@ import torch +from torch import Tensor def bias_dropout_add(x, bias, residual, prob, training): diff --git a/colossalai/kernel/jit/option.py b/colossalai/kernel/jit/option.py index aa41f57678fc..e20c08b051ed 100644 --- a/colossalai/kernel/jit/option.py +++ b/colossalai/kernel/jit/option.py @@ -43,7 +43,7 @@ def warmup_jit_fusion(batch_size: int, seq_length: int = 512, vocab_size: int = 32768, dtype: torch.dtype = torch.float32): - """ Compilie JIT functions before the main training steps """ + """ Compile JIT functions before the main training steps """ embed = Embedding(vocab_size, hidden_size).to(get_current_device()) linear_1 = Linear(hidden_size, hidden_size * 4, skip_bias_add=True).to(get_current_device()) diff --git a/colossalai/kernel/triton/ops.py b/colossalai/kernel/triton/ops.py new file mode 100644 index 000000000000..5e8d4ba3ec99 --- /dev/null +++ b/colossalai/kernel/triton/ops.py @@ -0,0 +1,209 @@ +import torch +from torch import nn + +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + print("please install triton from https://github.com/openai/triton") + +if HAS_TRITON: + from .qkv_matmul_kernel import qkv_gemm_4d_kernel + from .softmax_kernel import softmax_kernel + + def self_attention_forward_without_fusion(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, input_mask: torch.Tensor, scale: float): + r""" A function to do QKV Attention calculation by calling GEMM and softmax triton kernels + Args: + q (torch.Tensor): Q embedding in attention layer, shape should be (batch, seq_len, num_heads, head_size) + k (torch.Tensor): K embedding in attention layer, shape should be (batch, seq_len, num_heads, head_size) + v (torch.Tensor): V embedding in attention layer, shape should be (batch, seq_len, num_heads, head_size) + input_mask (torch.Tensor): mask for softmax layer, shape should be (batch, num_heads, seq_lem, seq_len) + scale: the float scale value which is used to multiply with Q*K^T before doing softmax + + Return: + output (Torch.Tensor): The output shape is (batch, seq_len, num_heads, head_size) + """ + assert len(q.shape) == 4, "the shape of q val must be 4" + batches, M, H, K = q.shape + assert q.shape == k.shape, "the shape of q and the shape of k must be equal" + assert q.shape == v.shape, "the shape of q and the shape of v must be equal" + assert q.shape[-1] == k.shape[-1], "the last dimension of q and k must be equal" + + N = k.shape[1] + + # head_size * num_of_head + d_model = q.shape[-1] * q.shape[-2] + + score_output = torch.empty( + (batches, H, M, N), device=q.device, dtype=q.dtype) + + grid = lambda meta: ( + batches, + H, + triton.cdiv(M, meta["BLOCK_SIZE_M"]) * + triton.cdiv(N, meta["BLOCK_SIZE_N"]), + ) + + qkv_gemm_4d_kernel[grid]( + q, k, score_output, + M, N, K, + q.stride(0), q.stride(2), q.stride(1), q.stride(3), + k.stride(0), k.stride(2), k.stride(3), k.stride(1), + score_output.stride(0), score_output.stride(1), score_output.stride(2), score_output.stride(3), + scale=scale, + # currently manually setting, later on we can use auto-tune config to match best setting + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=32, + BLOCK_SIZE_K=32, + GROUP_SIZE_M=8, + ) + + softmax_output = torch.empty( + score_output.shape, device=score_output.device, dtype=score_output.dtype) + score_output_shape = score_output.shape + + score_output = score_output.view(-1, score_output.shape[-1]) + n_rows, n_cols = score_output.shape + + if n_rows <= 350000: + + block_size = max(triton.next_power_of_2(n_cols), 2) + num_warps = 4 + if block_size >= 4096: + num_warps = 16 + elif block_size >= 2048: + num_warps = 8 + else: + num_warps = 4 + + softmax_kernel[(n_rows, )]( + softmax_output, + score_output, + score_output.stride(0), + n_cols, + mask_ptr = input_mask, + num_warps=num_warps, + BLOCK_SIZE=block_size, + ) + + else: + #TODO: change softmax kernel functions to make it suitable for large size dimension + softmax_output = torch.nn.functional.softmax(score_output, dim=-1) + softmax_output = softmax_output.view(*score_output_shape) + + batches, H, M, K = softmax_output.shape + N = v.shape[-1] + + output = torch.empty( + (batches, M, H, N), device=softmax_output.device, dtype=softmax_output.dtype) + + grid = lambda meta: ( + batches, + H, + triton.cdiv(M, meta["BLOCK_SIZE_M"]) * + triton.cdiv(N, meta["BLOCK_SIZE_N"]), + ) + + qkv_gemm_4d_kernel[grid]( + softmax_output, v, output, + M, N, K, + softmax_output.stride(0), + softmax_output.stride(1), + softmax_output.stride(2), + softmax_output.stride(3), + v.stride(0), + v.stride(2), + v.stride(1), + v.stride(3), + output.stride(0), + output.stride(2), + output.stride(1), + output.stride(3), + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=64, + BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, + scale=-1, + ) + return output.view(batches, -1, d_model) + + + def self_attention_compute_using_triton(qkv, + input_mask, + layer_past, + alibi, + scale, + head_size, + triangular=False, + use_flash=False): + + assert qkv.is_contiguous() + assert alibi is None, "current triton self-attention does not support alibi" + batches = qkv.shape[0] + d_model = qkv.shape[-1] // 3 + num_of_heads = d_model // head_size + + q = qkv[:, :, :d_model] + k = qkv[:, :, d_model:d_model * 2] + v = qkv[:, :, d_model * 2:] + q = q.view(batches, -1, num_of_heads, head_size) + k = k.view(batches, -1, num_of_heads, head_size) + v = v.view(batches, -1, num_of_heads, head_size) + + data_output_triton = self_attention_forward_without_fusion( + q, k, v, input_mask, scale) + + return data_output_triton + + + def softmax(input: torch.Tensor, mask: torch.Tensor = None, dim=-1) -> torch.Tensor: + if mask is not None: + assert input[-1] == mask[-1], "the last dimentions should be the same for input and mask" + assert dim == -1 or dim == len(input.shape)-1, "currently softmax layer only support last dimention" + + hidden_dim = input.shape[-1] + output = torch.empty_like(input) + input = input.view(-1, hidden_dim) + if mask is not None: + mask = mask.view(-1, hidden_dim) + assert input.shape[0] == mask.shape[0], "the fist dimention of mask and input should be the same" + + num_rows, num_cols = input.shape + block_size = max(triton.next_power_of_2(num_cols), 2) + num_warps = 16 + if block_size >= 4096: + num_warps = 16 + elif block_size >= 2048: + num_warps = 8 + else: + num_warps = 4 + + if num_rows <= 350000: + grid = (num_rows,) + softmax_kernel[grid](output, input, input.stride(0), num_cols, mask, BLOCK_SIZE = block_size, num_warps=num_warps) + else: + grid = lambda meta: () + + grid = lambda meta: ( + triton.cdiv(num_rows, meta["BLOCK_M"]), + ) + + BLOCK_M = 32 + if block_size >= 4096: + BLOCK_M = 4 + elif block_size >= 2048: + BLOCK_M = 8 + + softmax_kernel_2[grid](output_ptr = output, + input_ptr = input, + row_stride = input.stride(0), + n_rows = num_rows, + n_cols = num_cols, + mask_ptr = mask, + # currently manually setting up size + BLOCK_M = 32, + BLOCK_SIZE = block_size) + + return output \ No newline at end of file diff --git a/colossalai/kernel/triton/qkv_matmul_kernel.py b/colossalai/kernel/triton/qkv_matmul_kernel.py new file mode 100644 index 000000000000..62fc6bba0360 --- /dev/null +++ b/colossalai/kernel/triton/qkv_matmul_kernel.py @@ -0,0 +1,109 @@ +import torch +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + print("please install triton from https://github.com/openai/triton") + + +if HAS_TRITON: + ''' + this kernel function is modified from https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html + ''' + @triton.jit + def qkv_gemm_4d_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_ab, + stride_ah, + stride_am, + stride_ak, + stride_bb, + stride_bh, + stride_bk, + stride_bn, + stride_cb, + stride_ch, + stride_cm, + stride_cn, + scale, + # Meta-parameters + BLOCK_SIZE_M : tl.constexpr = 64, + BLOCK_SIZE_N : tl.constexpr = 32, + BLOCK_SIZE_K : tl.constexpr = 32, + GROUP_SIZE_M : tl.constexpr = 8, + ): + r""" A kernel function which is used to do batch-matmul for Q*K^T or score_matrix * V for attention layer, + where score_matrix is softmax(Q*V^T/sqrt(hidden_size)) + Args: + a_ptr(torch.Tensor): pointer to input tensor array (bs, M, h, K) or (bs, h, M, K) + b_ptr(torch.Tensor): pointer to input tensor array (bs, N, h, K) or (bs, h, N, K) + c_ptr(torch.Tensor): pointer to output tensor array (bs, M, h, N) or (bs, h, M, N) + stride_ab(tl.constexpr): stride for bs-dimention for tensor array A + stride_ah(tl.constexpr): stride for h-dimention for tensor array A + stride_am(tl.constexpr): stride for m-dimention for tensor array A + stride_ak(tl.constexpr): stride for k-dimention for tensor array A + stride_bb(tl.constexpr): stride for bs-dimention for tensor array B + stride_bh(tl.constexpr): stride for h-dimention for tensor array B + stride_bk(tl.constexpr): stride for k-dimention for tensor array B + stride_bn(tl.constexpr): stride for n-dimention for tensor array B + stride_cb(tl.constexpr): stride for bs-dimention for tensor array output + stride_ch(tl.constexpr): stride for h-dimention for tensor array output + stride_cm(tl.constexpr): stride for m-dimention for tensor array output + stride_cn(tl.constexpr): stride for n-dimention for tensor array output + BLOCK_SIZE_M : tiling size for M-dimension of tensor Array a + BLOCK_SIZE_N : tiling size for N-dimension of tensor Array b + BLOCK_SIZE_K : tiling size for K-dimension of a and b + GROUP_SIZE_M : group size for reducing cache miss, more details: + """ + + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + batch = tl.program_id(axis = 0) + head = tl.program_id(axis = 1) + pid = tl.program_id(axis = 2) + + # the following is from tutorial: https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = (a_ptr + batch * stride_ab + head * stride_ah + + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)) + b_ptrs = (b_ptr + batch * stride_bb + head * stride_bh + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, K, BLOCK_SIZE_K): + a_mask = (offs_am[:, None] < M) & (offs_k[None, :] + k < K) + b_mask = (offs_k[:, None] + k < K) & (offs_bn[None, :] < N) + a = tl.load(a_ptrs, mask=a_mask, other=0.) + b = tl.load(b_ptrs, mask=b_mask, other=0.) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + accumulator = accumulator.to(c_ptr.dtype.element_ty) + if scale > 0: + accumulator = accumulator * scale.to(c_ptr.dtype.element_ty) + + + offs_accumu_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_accumu_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = (c_ptr + batch * stride_cb + head * stride_ch + stride_cm * offs_accumu_m[:, None] + + stride_cn * offs_accumu_n[None, :]) + accumulator_mask = (offs_accumu_m[:, None] < M) & (offs_accumu_n[None, :] < N) + tl.store(c_ptrs, accumulator, mask=accumulator_mask) diff --git a/colossalai/kernel/triton/softmax_kernel.py b/colossalai/kernel/triton/softmax_kernel.py new file mode 100644 index 000000000000..c215890badff --- /dev/null +++ b/colossalai/kernel/triton/softmax_kernel.py @@ -0,0 +1,44 @@ +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + print("please install triton from https://github.com/openai/triton") + +if HAS_TRITON: + ''' + softmax kernel is modified based on + https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/02-fused-softmax.py + ''' + @triton.jit + def softmax_kernel(output_ptr, input_ptr, row_stride, n_cols, mask_ptr, BLOCK_SIZE: tl.constexpr): + r""" the kernel function for implementing softmax operator + Args: + output_ptr: the output after finishing softmax operation, (N, hidden_dim) + input_ptr: the tensor of input, shape should be (N, hidden_dim) + n_cols(tl.constexpr): the number of cols of input + BLOCK_SIZE(tl.constexpr): the block_size of your hidden_dim dimension, typically BLOCK_SIZE >= hidden_dim + """ + row_idx = tl.program_id(0) + row_start_ptr = input_ptr + row_idx * row_stride + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32) + row_minus_max = row - tl.max(row, axis=0) + + if mask_ptr is not None: + # load mask into SRAM + mask_ptrs = (mask_ptr + (row_indx * row_stride)) + col_offsets + mask = tl.load(mask_ptrs, mask=col_offsets < n_cols, other=0).to(tl.float32) + + # update + row_minus_max = row_minus_max + mask + + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + output_row_start_ptr = output_ptr + row_idx * row_stride + output_ptrs = output_row_start_ptr + col_offsets + # Write back output to DRAM + tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols) \ No newline at end of file diff --git a/colossalai/lazy/__init__.py b/colossalai/lazy/__init__.py new file mode 100644 index 000000000000..4387107bf773 --- /dev/null +++ b/colossalai/lazy/__init__.py @@ -0,0 +1,6 @@ +from .lazy_init import LazyInitContext, LazyTensor + +__all__ = [ + 'LazyInitContext', + 'LazyTensor', +] diff --git a/colossalai/utils/model/experimental.py b/colossalai/lazy/lazy_init.py similarity index 60% rename from colossalai/utils/model/experimental.py rename to colossalai/lazy/lazy_init.py index 00cb532d9c1d..1f5345015bf2 100644 --- a/colossalai/utils/model/experimental.py +++ b/colossalai/lazy/lazy_init.py @@ -1,17 +1,23 @@ -from typing import Callable, List, Optional, Union +from contextlib import contextmanager +from types import MethodType +from typing import Callable, Dict, Optional, Union import torch +import torch.distributed as dist import torch.nn as nn from torch import Tensor from torch.utils._pytree import tree_map -from colossalai.fx.profiler.tensor import MetaTensor +from colossalai._analyzer._subclasses import MetaTensor +from colossalai.device.device_mesh import DeviceMesh +from colossalai.tensor.d_tensor import distribute_tensor +from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec # reference: https://pytorch.org/cppdocs/notes/tensor_creation.html _NORMAL_FACTORY = [ "arange", - "empty", "full", + "empty", "linspace", "logspace", "ones", @@ -30,6 +36,11 @@ _EARLY_MATERIALIZED_OPS = ['__getitem__', 'split'] +# If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset) +# without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block. +# These ops cannot be unwrapped using .data +_CHANGE_META_OPS = ['_cudnn_rnn_flatten_weight', 'requires_grad_', '__get__', '__set__', 'numel', 'size', 'dim'] + _LEGACY_TENSOR_CONSTRUCTOR = { 'FloatTensor': torch.float, 'DoubleTensor': torch.double, @@ -43,18 +54,23 @@ 'BoolTensor': torch.bool, } +_EMPTY_DATA = torch.empty(0) + class _MyTensor(Tensor): """This class is only for correctness verification. """ _pre_op_fn: Callable[['LazyTensor'], None] = lambda *args: None + default_device: Optional[torch.device] = None + def __new__(cls, func, *args, concrete_data=None, **kwargs) -> '_MyTensor': cls._pre_op_fn() if concrete_data is not None: # uniform api as LazyTensor data = concrete_data else: + kwargs['device'] = cls.default_device data = func(*args, **kwargs) return Tensor._make_subclass(cls, data, require_grad=data.requires_grad) @@ -64,6 +80,35 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): return super().__torch_function__(func, types, args, kwargs) +def _data_tolist(tensor: torch.Tensor) -> list: + """tolist() method is not allowed for a subclass of tensor. Tensor.data returns a Tensor. + """ + return tensor.data.tolist() + + +def _convert_cls(tensor: 'LazyTensor', target: torch.Tensor) -> torch.Tensor: + """Convert a lazy tensor's class to target's class, with target's data. + + The reason why we change the class of a lazy tensor in-place is that this can easily handle shared modules/parameters, which is common in huggingface models. + If we create a new tensor and update the module by ``setattr(module, name, param)``, the shared parameters will not be updated. And we have to track all shared parameters and update them manually. + + Args: + tensor (LazyTensor): the LazyTensor to be converted + target (torch.Tensor): target tensor + + Returns: + torch.Tensor: the converted tensor + """ + cls_to_become = nn.Parameter if isinstance(tensor, nn.Parameter) else torch.Tensor + tensor.__class__ = cls_to_become + tensor.data = target + tensor.requires_grad = target.requires_grad + # subclass of torch.Tensor does not have tolist() method + # overwrite this method after materialization or distribution + tensor.tolist = MethodType(_data_tolist, tensor) + return tensor + + class LazyTensor(torch.Tensor): """A naive implementation of LazyTensor (https://arxiv.org/pdf/2102.13267.pdf). @@ -101,6 +146,8 @@ class LazyTensor(torch.Tensor): _meta_data: Optional[MetaTensor] = None # shape, dtype, device _pre_op_fn: Callable[['LazyTensor'], None] = lambda *args: None + default_device: Optional[torch.device] = None + @staticmethod def __new__(cls, func, *args, meta_data=None, concrete_data=None, **kwargs): if concrete_data is not None: @@ -110,34 +157,43 @@ def __new__(cls, func, *args, meta_data=None, concrete_data=None, **kwargs): if meta_data is None: device = kwargs.get('device', 'cpu') elem = func(*args, **{**kwargs, 'device': 'meta'}) - meta_data = MetaTensor(elem, fake_device=device) + meta_data = MetaTensor(elem, device=device) elem = meta_data._tensor - r = torch.Tensor._make_wrapper_subclass(cls, - elem.size(), - strides=elem.stride(), - storage_offset=elem.storage_offset(), - dtype=elem.dtype, - layout=elem.layout, - device=elem.device, - requires_grad=elem.requires_grad) + # As a meta tensor cannot be modified __class__ to torch.Tensor, we should use an empty real tensor here + r = torch.Tensor._make_subclass(cls, _EMPTY_DATA, require_grad=elem.requires_grad) r._meta_data = meta_data return r def __init__(self, func, *args, meta_data=None, concrete_data=None, **kwargs): + if func.__name__ in _NORMAL_FACTORY: + kwargs = {**kwargs, 'device': LazyTensor.default_device} self._factory_method = (func, args, kwargs) # (func, args, kwargs) self._op_buffer = [] # (func, args, kwargs, replace) self._materialized_data: Optional[torch.Tensor] = concrete_data # materialized data def materialize(self) -> torch.Tensor: - """Materialize the ``LazyTensor`` to ``torch.Tensor``. + """Materialize the ``LazyTensor`` to ``torch.Tensor`` by modifying __class__ (inplace). Returns: - torch.Tensor: The materialized tensor. + torch.Tensor: The materialized tensor (self). """ target = self._materialize_data() - if isinstance(self, nn.Parameter): - target = nn.Parameter(target, requires_grad=self.requires_grad) - return target + self.clean() + return _convert_cls(self, target) + + def distribute(self, device_mesh: DeviceMesh, sharding_spec: ShardingSpec) -> torch.Tensor: + """Distribute the ``LazyTensor`` to ``torch.Tensor`` by modifying __class__ (inplace), according to the layout. + + Args: + layout (Layout): Distribution layout. + + Returns: + torch.Tensor: The distributed tensor (self). + """ + target = self._materialize_data() + self.clean() + local_tensor = distribute_tensor(target, device_mesh, sharding_spec) + return _convert_cls(self, local_tensor) def clean(self) -> None: """Clean all stored operations, meta data and materialized data, which prevents memory leaking. This should be called after all tensors are materialized. @@ -158,16 +214,11 @@ def _materialize_data(self) -> torch.Tensor: if self._materialized_data is None: # apply factory method func, args, kwargs = self._factory_method - # apply cached sequence self._pre_op_fn() - try: - init_val = func(*tree_map(self._replace_with_materialized, args), - **tree_map(self._replace_with_materialized, kwargs)) - except TypeError as e: - print(f'init fn: {func.__name__}') - raise e + init_val = func(*tree_map(self._replace_with_materialized, args), + **tree_map(self._replace_with_materialized, kwargs)) self._materialized_data = self._rerun_ops(init_val) return self._materialized_data @@ -214,7 +265,9 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): tree_map(cls._replace_with_materialized, args) tree_map(cls._replace_with_materialized, kwargs) is_inplace: bool = (func.__name__.endswith('_') and not (func.__name__.endswith('__')) - or func.__name__ == "__setitem__") + or func.__name__ in ('__setitem__', '__set__')) + + is_change_meta_op: bool = func.__name__ in _CHANGE_META_OPS if isinstance(func, torch._C.ScriptMethod): # FIXME(ver217): torch script functions are not verified @@ -239,10 +292,10 @@ def unwrap(x): if isinstance(x, LazyTensor): if x._materialized_data is not None: # for early materialized tensor, use its materialized data directly - return x._materialized_data.data + return x._materialized_data if is_change_meta_op else x._materialized_data.data t = x if is_inplace else x.clone() t._op_buffer.append((func, args, kwargs)) - meta = x._meta_data.data + meta = x._meta_data if is_change_meta_op else x._meta_data.data meta_to_lazy[meta] = t return meta return x @@ -255,6 +308,7 @@ def wrap(y, i=None): else: # out of place op, create new lazy tensor fn = lambda *a, **kw: func(*a, **kw) if i is None else func(*a, **kw)[i] + fn.__name__ = func.__name__ lazy_y = LazyTensor(fn, *args, meta_data=y, **kwargs) return lazy_y elif type(y) is Tensor: @@ -275,7 +329,9 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None): def clone(self) -> "LazyTensor": def factory_fn(): - return self.materialize().clone() + # if self is materialized, return self + new_tensor = self.materialize() if type(self) is LazyTensor else self + return new_tensor.clone() target = LazyTensor(factory_fn, meta_data=self._meta_data) @@ -284,19 +340,69 @@ def factory_fn(): def detach(self) -> Tensor: return self + def __deepcopy__(self, memo): + if not self.is_leaf: + raise RuntimeError("Only Tensors created explicitly by the user " + "(graph leaves) support the deepcopy protocol at the moment") + if id(self) in memo: + return memo[id(self)] + + def factory_fn(): + # if self is materialized, return self + new_tensor = self.materialize() if type(self) is LazyTensor else self + copied = new_tensor.detach().clone() + if new_tensor.requires_grad: + copied.requires_grad_() + return copied + + if self._materialized_data is not None: + # self is early materialized + copied = self._materialized_data.detach().clone() + if self.requires_grad: + copied.requires_grad_() + target = LazyTensor(lambda: None, concrete_data=copied) + else: + target = LazyTensor(factory_fn, meta_data=self._meta_data) + + memo[id(self)] = target + return target + @property def data(self): return self @data.setter def data(self, other: 'LazyTensor'): + """This is sightly different from oringinal `data` setter. + + E.g.: + >>> a = torch.randn(3, 3) # a is a Tensor + >>> b = torch.rand(2, 2) + >>> a.data = b + >>> b.add_(1) # this will affect a + >>> x = torch.randn(3, 3) # x is a LazyTensor + >>> y = torch.rand(2, 2) # y is a LazyTensor + >>> x.data = y + >>> y.add_(1) # this will not affect x + + """ if other is self: return - # TODO(ver217): to avoid infinity recursion, do early materialization - self._materialized_data = other._materialize_data() + + self._op_buffer.append(other._factory_method) + + def replace(x): + if x is other: + return self + return x + + for func, args, kwargs in other._op_buffer: + self._op_buffer.append((func, tree_map(replace, args), tree_map(replace, kwargs))) def tolist(self) -> list: - t = self.materialize() + # Though self.__class__ is modified to torch.Tensor, in C++ side, it is still a subclass of torch.Tensor + # And subclass of torch.Tensor does not have tolist() method + t = self._materialize_data() return t.tolist() def __hash__(self): @@ -333,14 +439,21 @@ class LazyInitContext: """ _replaced: bool = False - def __init__(self, tensor_cls: Union[_MyTensor, LazyTensor] = LazyTensor): + def __init__(self, + tensor_cls: Union[_MyTensor, LazyTensor] = LazyTensor, + default_device: Optional[Union[torch.device, str, int]] = None): + assert tensor_cls is LazyTensor or tensor_cls is _MyTensor self.overrides = {} self.tensor_cls = tensor_cls + self.old_default_device = LazyTensor.default_device + self.default_device = default_device def __enter__(self): if LazyInitContext._replaced: raise RuntimeError(f'LazyInitContext is not reentrant') LazyInitContext._replaced = True + self.old_default_device = self.tensor_cls.default_device + self.tensor_cls.default_device = self.default_device def wrap_factory_method(target): # factory functions (eg. torch.empty()) @@ -416,76 +529,93 @@ def wrapper(*args, **kwargs): setattr(torch, name, wrapper) def __exit__(self, exc_type, exc_val, exc_tb): + self.tensor_cls.default_device = self.old_default_device LazyInitContext._replaced = False for name, (wrapper, orig) in self.overrides.items(): setattr(torch, name, orig) @staticmethod - def materialize(module: torch.nn.Module, verbose: bool = False): - """Initialize all ``nn.Parameter`` from ``LazyTensor``. + def materialize(module: nn.Module, verbose: bool = False) -> nn.Module: + """Initialize all ``nn.Parameter`` from ``LazyTensor``. This function will modify the module in-place. Args: - module (torch.nn.Module): Target ``nn.Module`` + module (nn.Module): Target ``nn.Module`` verbose (bool): Whether to print lazy initialization rate. Defaults to False. """ - if verbose: - param_cnt = 0 - param_lazy_cnt = 0 - buf_cnt = 0 - buf_lazy_cnt = 0 - non_lazy_numel = 0 - - # do post cleaning to handle shared parameter - visited_lazy_tensors: List[LazyTensor] = [] - # handle shared module - visited_modules = set() - - @torch.no_grad() - def init_recursively(module: nn.Module): - nonlocal param_cnt, param_lazy_cnt, buf_cnt, buf_lazy_cnt, non_lazy_numel - # recursively initialize the module - for mod in module.children(): - if id(mod) not in visited_modules: - visited_modules.add(id(mod)) - init_recursively(mod) - - # initialize tensors directly attached to the current module - for name, param in module.named_parameters(recurse=False): - if verbose: - param_cnt += 1 - if getattr(param, '_materialized_data', False) is None: - # if no _materialized_data attr, the tensor is not lazy - param_lazy_cnt += 1 - else: - non_lazy_numel += param.numel() - if hasattr(param, 'materialize'): - # TODO(ver217): apex layers cannot be captured - visited_lazy_tensors.append(param) - setattr(module, name, param.materialize()) - - for name, buf in module.named_buffers(recurse=False): - if verbose: - buf_cnt += 1 - if getattr(buf, "_materialized_data", False) is None: - # if no _materialized_data attr, the tensor is not lazy - buf_lazy_cnt += 1 - else: - non_lazy_numel += buf.numel() - if hasattr(buf, 'materialize'): - # TODO(ver217): apex layers cannot be captured - visited_lazy_tensors.append(buf) - setattr(module, name, buf.materialize()) - init_recursively(module) + def apply_fn(name: str, p: LazyTensor): + p.materialize() - for t in visited_lazy_tensors: - t.clean() + return _apply_to_lazy_module(module, apply_fn, verbose) + @staticmethod + def distribute(module: nn.Module, + device_mesh: DeviceMesh, + sharding_spec_dict: Dict[str, ShardingSpec], + verbose: bool = False) -> nn.Module: + """Distribute all ``nn.Parameter`` from ``LazyTensor``. This function will modify the module in-place. + + Args: + module (nn.Module): Target ``nn.Module`` + layout_dict (dict): Dict of layout for each parameter/buffer. The key is the parameter/buffer name, and the value is the layout. + verbose (bool, optional): Whether to print lazy initialization rate. Defaults to False. + """ + + def apply_fn(name: str, p: LazyTensor): + p.distribute(device_mesh, sharding_spec_dict[name]) + + return _apply_to_lazy_module(module, apply_fn, verbose) + + +def _apply_to_lazy_module(module: nn.Module, + apply_fn: Callable[[str, torch.Tensor], None], + verbose: bool = False) -> nn.Module: + if verbose: + # verbose info + param_cnt = 0 + param_lazy_cnt = 0 + buf_cnt = 0 + buf_lazy_cnt = 0 + total_numel = 0 + non_lazy_numel = 0 + + for name, p in module.named_parameters(): + if verbose: + param_cnt += 1 + total_numel += p.numel() + if getattr(p, '_materialized_data', False) is None: + # if no _materialized_data attr, the tensor is not lazy + param_lazy_cnt += 1 + else: + non_lazy_numel += p.numel() + if isinstance(p, LazyTensor): + apply_fn(name, p) + + for name, buf in module.named_buffers(): if verbose: - print(f'Param lazy rate: {param_lazy_cnt}/{param_cnt}') - print(f'Buffer lazy rate: {buf_lazy_cnt}/{buf_cnt}') - print(f'Non-lazy numel: {non_lazy_numel} ({non_lazy_numel/1024**2:.3f} M)') - return module + buf_cnt += 1 + total_numel += buf.numel() + if getattr(buf, "_materialized_data", False) is None: + # if no _materialized_data attr, the tensor is not lazy + buf_lazy_cnt += 1 + else: + non_lazy_numel += buf.numel() + if isinstance(buf, LazyTensor): + apply_fn(name, buf) + + if verbose: + non_lazy_numel_ratio = non_lazy_numel / total_numel * 100 if non_lazy_numel != 0 else 0 + _print_rank_0(f'Param lazy rate: {param_lazy_cnt}/{param_cnt}') + _print_rank_0(f'Buffer lazy rate: {buf_lazy_cnt}/{buf_cnt}') + _print_rank_0( + f'Non lazy numel: {non_lazy_numel} ({non_lazy_numel/1024**2:.3f} M), ratio: {non_lazy_numel_ratio}%') + + return module + + +def _print_rank_0(*args, **kwargs): + if not dist.is_initialized() or dist.get_rank() == 0: + print(*args, **kwargs) def _is_int_tuple(args) -> bool: diff --git a/colossalai/nn/_ops/_utils.py b/colossalai/nn/_ops/_utils.py index 56bb5f465184..24877bbb552f 100644 --- a/colossalai/nn/_ops/_utils.py +++ b/colossalai/nn/_ops/_utils.py @@ -1,12 +1,11 @@ -import torch -from typing import Union, Optional, List -from colossalai.tensor import ColoTensor +from typing import List, Optional, Union + import torch import torch.distributed as dist -from colossalai.global_variables import tensor_parallel_env as env +from colossalai.global_variables import tensor_parallel_env as env from colossalai.nn.layer.utils import divide -from colossalai.tensor import ProcessGroup, ColoTensorSpec +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup GeneralTensor = Union[ColoTensor, torch.Tensor] Number = Union[int, float] @@ -135,7 +134,7 @@ def backward(ctx, grad_output): class _SplitForwardGatherBackward(torch.autograd.Function): """ Split the input and keep only the corresponding chuck to the rank. - + Args: input_: input matrix. process_group: parallel mode. diff --git a/colossalai/nn/_ops/addmm.py b/colossalai/nn/_ops/addmm.py index fe2eb0c999a1..660b48a71d57 100644 --- a/colossalai/nn/_ops/addmm.py +++ b/colossalai/nn/_ops/addmm.py @@ -1,9 +1,9 @@ import torch + +from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ReplicaSpec, ShardSpec, distspec from colossalai.tensor.op_wrapper import colo_op_impl -from colossalai.tensor import ComputePattern, ComputePattern, ComputeSpec, ColoTensor -from colossalai.tensor import distspec, ColoTensorSpec, ShardSpec, ReplicaSpec -from ._utils import GeneralTensor, Number, convert_to_colo_tensor -from ._utils import reduce_input, reduce_grad + +from ._utils import GeneralTensor, Number, convert_to_colo_tensor, reduce_grad, reduce_input def colo_addmm_1Drow(input_tensor: ColoTensor, mat1: ColoTensor, mat2: ColoTensor, beta: Number, @@ -69,9 +69,13 @@ def colo_addmm(input_tensor: GeneralTensor, if not mat2.has_compute_spec(): # No Model Parallel Applied assert mat2.is_replicate(), 'Invalid mat2 spec for native addmm op' assert input_tensor.is_replicate(), 'Invalid input spec for native addmm op' - ret_tensor = ColoTensor.from_torch_tensor( - tensor=torch.addmm(input_tensor, mat1, mat2, beta=beta, alpha=alpha, **kargs), - spec=ColoTensorSpec(mat2.get_process_group())) + ret_tensor = ColoTensor.from_torch_tensor(tensor=torch.addmm(input_tensor, + mat1, + mat2, + beta=beta, + alpha=alpha, + **kargs), + spec=ColoTensorSpec(mat2.get_process_group())) elif mat2.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied if mat2.is_shard_1drow() and input_tensor.is_replicate(): mode = 'row' diff --git a/colossalai/nn/_ops/embedding_bag.py b/colossalai/nn/_ops/embedding_bag.py index 0e8aa8fecb01..0026f579b6dc 100644 --- a/colossalai/nn/_ops/embedding_bag.py +++ b/colossalai/nn/_ops/embedding_bag.py @@ -88,7 +88,7 @@ def colo_embedding_bag(input_tensor: GeneralTensor, assert isinstance(weight, ColoTensor) input_tensor = convert_to_colo_tensor(input_tensor, weight.get_process_group()) - # Handle differen parallel actions. + # Handle different parallel actions. if not weight.has_compute_spec(): # No Model Parallel Applied assert weight.is_replicate(), 'Invalid weight spec for native embedding op' diff --git a/colossalai/nn/layer/base_layer.py b/colossalai/nn/layer/base_layer.py index c85f53cc44c3..4a06bdcb7629 100644 --- a/colossalai/nn/layer/base_layer.py +++ b/colossalai/nn/layer/base_layer.py @@ -1,14 +1,16 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- +from contextlib import contextmanager + import torch.nn as nn from colossalai.context import ParallelMode from colossalai.core import global_context as gpc -from contextlib import contextmanager class ParallelLayer(nn.Module): + global_state_dict: bool = True def __init__(self): diff --git a/colossalai/nn/layer/moe/__init__.py b/colossalai/nn/layer/moe/__init__.py index 2a51344c31a4..05333fe965f1 100644 --- a/colossalai/nn/layer/moe/__init__.py +++ b/colossalai/nn/layer/moe/__init__.py @@ -1,9 +1,10 @@ -from .experts import Experts, FFNExperts, TPExperts -from .layers import MoeLayer, MoeModule -from .routers import MoeRouter, Top1Router, Top2Router -from .utils import NormalNoiseGenerator, UniformNoiseGenerator, build_ffn_experts - -__all__ = [ - 'Experts', 'FFNExperts', 'TPExperts', 'Top1Router', 'Top2Router', 'MoeLayer', 'NormalNoiseGenerator', - 'UniformNoiseGenerator', 'build_ffn_experts', 'MoeModule', 'MoeRouter' -] +from .checkpoint import load_moe_model, save_moe_model +from .experts import Experts, FFNExperts, TPExperts +from .layers import MoeLayer, MoeModule +from .routers import MoeRouter, Top1Router, Top2Router +from .utils import NormalNoiseGenerator, UniformNoiseGenerator, build_ffn_experts + +__all__ = [ + 'Experts', 'FFNExperts', 'TPExperts', 'Top1Router', 'Top2Router', 'MoeLayer', 'NormalNoiseGenerator', + 'UniformNoiseGenerator', 'build_ffn_experts', 'MoeModule', 'MoeRouter', 'save_moe_model', 'load_moe_model' +] diff --git a/colossalai/nn/layer/moe/checkpoint.py b/colossalai/nn/layer/moe/checkpoint.py new file mode 100644 index 000000000000..efda1f22252d --- /dev/null +++ b/colossalai/nn/layer/moe/checkpoint.py @@ -0,0 +1,40 @@ +import torch +import torch.distributed as dist +import torch.nn as nn + +from .experts import MoeExperts + + +def save_moe_model(model: nn.Module, save_path: str): + state_dict = model.state_dict() + if dist.get_rank() == 0: + torch.save(state_dict, save_path) + dist.barrier() + + +def load_moe_model(model: nn.Module, load_path: str): + state_dict = torch.load(load_path) + + for prefix, module in model.named_modules(): + if prefix.endswith('.moe_layer.experts'): + # this module should be an Experts instance + assert isinstance(module, MoeExperts) + + ep_rank = dist.get_rank(module.dist_info.ep_group) + num_local = module.num_local_experts + for i in range(num_local): + expert_id = ep_rank * num_local + i + for name, _ in module.experts[i].named_parameters(): + cur_key = f'{prefix}.experts.{i}.{name}' + param_key = f'{prefix}.experts.{expert_id}.{name}' + load_param = state_dict[param_key] + state_dict[cur_key] = load_param + + for name, _ in module.experts[0].named_parameters(): + pop_pre = f'{prefix}.experts.' + pop_suf = f'.{name}' + for i in range(num_local, module.num_total_experts): + pop_key = f'{pop_pre}{i}{pop_suf}' + state_dict.pop(pop_key) + + model.load_state_dict(state_dict) diff --git a/colossalai/nn/layer/moe/experts.py b/colossalai/nn/layer/moe/experts.py index 055afded9a20..56b11f4d9e08 100644 --- a/colossalai/nn/layer/moe/experts.py +++ b/colossalai/nn/layer/moe/experts.py @@ -1,172 +1,203 @@ -import math - -import torch -import torch.nn as nn -from colossalai.context import ParallelMode, seed -from colossalai.utils import get_current_device -from colossalai.context.moe_context import MOE_CONTEXT -from colossalai.zero.init_ctx import no_shard_zero_decrator -from typing import Type - - -class MoeExperts(nn.Module): - """Basic class for experts in MoE. It stores what kind of communication expersts use - to exchange tokens, how many experts in a single GPU and parallel information such as - expert parallel size, data parallel size and their distributed communication groups. - """ - - def __init__(self, comm_name: str, num_experts: int): - super().__init__() - assert comm_name in {"all_to_all", "all_gather"}, \ - "This kind of communication has not been implemented yet.\n Please use Experts build function." - self.comm_name = comm_name - # Get the configuration of experts' deployment and parallel information from moe contex - self.num_local_experts, self.dist_info = MOE_CONTEXT.get_info(num_experts) - - -@no_shard_zero_decrator(is_replicated=False) -class Experts(MoeExperts): - """A wrapper class to create experts. It will create E experts across the - moe model parallel group, where E is the number of experts. Every expert - is a instence of the class, 'expert' in initialization parameters. - - Args: - expert_cls (:class:`torch.nn.Module`): The class of all experts - num_experts (int): The number of experts - expert_args: Args used to initialize experts, the args could be found in corresponding expert class - """ - - def __init__(self, expert_cls: Type[nn.Module], num_experts: int, **expert_args): - super().__init__("all_to_all", num_experts) - - # Use seed to make every expert different from others - with seed(ParallelMode.TENSOR): - self.experts = nn.ModuleList([expert_cls(**expert_args) for _ in range(self.num_local_experts)]) - - # Attach parallel information for all parameters in Experts - for exp in self.experts: - for param in exp.parameters(): - param.__setattr__('moe_info', self.dist_info) - - def forward(self, inputs: torch.Tensor): - # Split inputs for each expert - expert_input = torch.chunk(inputs, self.num_local_experts, dim=1) - expert_output = [] - - # Get outputs from each expert - for i in range(self.num_local_experts): - expert_output.append(self.experts[i](expert_input[i])) - - # Concatenate all outputs together - output = torch.cat(expert_output, dim=1).contiguous() - return output - - -class FFNExperts(MoeExperts): - """Use torch.bmm to speed up for multiple experts. - """ - - def __init__(self, num_experts: int, d_model: int, d_ff: int, activation=None, drop_rate: float = 0): - super().__init__("all_to_all", num_experts) - - self.w1 = nn.Parameter(torch.empty(self.num_local_experts, d_model, d_ff, device=get_current_device())) - self.b1 = nn.Parameter(torch.empty(self.num_local_experts, 1, d_ff, device=get_current_device())) - - self.w2 = nn.Parameter(torch.empty(self.num_local_experts, d_ff, d_model, device=get_current_device())) - self.b2 = nn.Parameter(torch.empty(self.num_local_experts, 1, d_model, device=get_current_device())) - - s1 = math.sqrt(0.1 / d_model) - s2 = math.sqrt(0.1 / d_ff) - - with seed(ParallelMode.TENSOR): - nn.init.trunc_normal_(self.w1, std=s1) - nn.init.trunc_normal_(self.b1, std=s1) - nn.init.trunc_normal_(self.w2, std=s2) - nn.init.trunc_normal_(self.b2, std=s2) - - self.act = nn.GELU() if activation is None else activation - self.drop = nn.Dropout(p=drop_rate) - - for param in self.parameters(): - param.__setattr__('moe_info', self.dist_info) - - def forward(self, inputs): # inputs [g, el, c, h] - - el = inputs.size(1) - h = inputs.size(-1) - - inputs = inputs.transpose(0, 1) - inshape = inputs.shape - inputs = inputs.reshape(el, -1, h) - - out_ff = torch.baddbmm(self.b1, inputs, self.w1) - out_act = self.act(out_ff) - with seed(ParallelMode.TENSOR): - out_inter = self.drop(out_act) - - out_model = torch.baddbmm(self.b2, out_inter, self.w2) - with seed(ParallelMode.TENSOR): - outputs = self.drop(out_model) # outputs [el, gc, h] - - outputs = outputs.reshape(inshape) - outputs = outputs.transpose(0, 1).contiguous() - return outputs - - -class TPExperts(MoeExperts): - """Use tensor parallelism to split each expert evenly, which can deploy experts in - case that the number of experts can't be divied by maximum expert parallel size or - maximum expert parallel size can't be divied by the number of experts. - """ - - def __init__(self, num_experts: int, d_model: int, d_ff: int, activation=None, drop_rate: float = 0): - super().__init__("all_gather", MOE_CONTEXT.max_ep_size) - - assert d_ff % MOE_CONTEXT.max_ep_size == 0, \ - "d_ff should be divied by maximum expert parallel size" - - p_ff = d_ff // MOE_CONTEXT.max_ep_size - - self.w1 = nn.Parameter(torch.empty(num_experts, d_model, p_ff, device=get_current_device())) - self.b1 = nn.Parameter(torch.empty(num_experts, 1, p_ff, device=get_current_device())) - - self.w2 = nn.Parameter(torch.empty(num_experts, p_ff, d_model, device=get_current_device())) - self.b2 = nn.Parameter(torch.empty(num_experts, 1, d_model, device=get_current_device())) - - s1 = math.sqrt(0.1 / d_model) - s2 = math.sqrt(0.1 / d_ff) - - with seed(ParallelMode.TENSOR): - nn.init.trunc_normal_(self.w1, std=s1) - nn.init.trunc_normal_(self.b1, std=s1) - nn.init.trunc_normal_(self.w2, std=s2) - - nn.init.trunc_normal_(self.b2, std=s2) - - self.act = nn.GELU() if activation is None else activation - self.drop = nn.Dropout(p=drop_rate) - - self.w1.__setattr__('moe_info', self.dist_info) - self.w2.__setattr__('moe_info', self.dist_info) - self.b1.__setattr__('moe_info', self.dist_info) - - def forward(self, inputs): # inputs [g, e, c, h] - - e = inputs.size(1) - h = inputs.size(-1) - - inputs = inputs.transpose(0, 1) - inshape = inputs.shape - inputs = inputs.reshape(e, -1, h) - - out_ff = torch.baddbmm(self.b1, inputs, self.w1) - out_act = self.act(out_ff) - with seed(ParallelMode.TENSOR): - out_inter = self.drop(out_act) - - out_model = torch.baddbmm(self.b2, out_inter, self.w2) - outputs = self.drop(out_model) # outputs [e, gc, h] - - outputs = outputs.reshape(inshape) - outputs = outputs.transpose(0, 1).contiguous() - return outputs # outputs [g, e, c, h] +import math +from copy import deepcopy +from typing import Type + +import torch +import torch.distributed as dist +import torch.nn as nn + +from colossalai.context import ParallelMode, seed +from colossalai.context.moe_context import MOE_CONTEXT +from colossalai.utils import get_current_device +from colossalai.zero.legacy.init_ctx import no_shard_zero_decrator + + +class MoeExperts(nn.Module): + """Basic class for experts in MoE. It stores what kind of communication experts use + to exchange tokens, how many experts in a single GPU and parallel information such as + expert parallel size, data parallel size and their distributed communication groups. + """ + + def __init__(self, comm_name: str, num_experts: int): + super().__init__() + assert comm_name in {"all_to_all", "all_gather"}, \ + "This kind of communication has not been implemented yet.\n Please use Experts build function." + self.comm_name = comm_name + self.num_total_experts = num_experts + # Get the configuration of experts' deployment and parallel information from moe context + self.num_local_experts, self.dist_info = MOE_CONTEXT.get_info(num_experts) + + +@no_shard_zero_decrator(is_replicated=False) +class Experts(MoeExperts): + """A wrapper class to create experts. It will create E experts across the + moe model parallel group, where E is the number of experts. Every expert + is a instance of the class, 'expert' in initialization parameters. + + Args: + expert_cls (:class:`torch.nn.Module`): The class of all experts + num_experts (int): The number of experts + expert_args: Args used to initialize experts, the args could be found in corresponding expert class + """ + + def __init__(self, expert_cls: Type[nn.Module], num_experts: int, **expert_args): + super().__init__("all_to_all", num_experts) + + # Use seed to make every expert different from others + with seed(ParallelMode.TENSOR): + self.experts = nn.ModuleList([expert_cls(**expert_args) for _ in range(self.num_local_experts)]) + + # Attach parallel information for all parameters in Experts + for exp in self.experts: + for param in exp.parameters(): + param.__setattr__('moe_info', self.dist_info) + + def forward(self, inputs: torch.Tensor): + # Split inputs for each expert + expert_input = torch.chunk(inputs, self.num_local_experts, dim=1) + expert_output = [] + + # Get outputs from each expert + for i in range(self.num_local_experts): + expert_output.append(self.experts[i](expert_input[i])) + + # Concatenate all outputs together + output = torch.cat(expert_output, dim=1).contiguous() + return output + + def state_dict(self, destination=None, prefix='', keep_vars=False): + assert keep_vars == False, "Only support keep_vars=False now" + dp_rank = dist.get_rank(self.dist_info.dp_group) + ep_rank = dist.get_rank(self.dist_info.ep_group) + submodule_dict = dict() + example_submodule = None + for name, subm in self.experts.named_modules(): + if subm is self.experts: + continue + module_number = self.num_local_experts * ep_rank + int(name) + submodule_dict[module_number] = subm + example_submodule = subm + + if dp_rank == 0: + local_prefix = prefix + 'experts.' + buffer_module = deepcopy(example_submodule) + for i in range(self.num_total_experts): + source_rank = i // self.num_local_experts + current_prefix = local_prefix + str(i) + '.' + comm_module = submodule_dict.get(i, buffer_module) + for name, param in comm_module.named_parameters(): + dist.broadcast(param.data, src=source_rank, group=self.dist_info.ep_group) + if ep_rank == 0: + destination[current_prefix + name] = param.data.cpu() + + dist.barrier() + + +class FFNExperts(MoeExperts): + """Use torch.bmm to speed up for multiple experts. + """ + + def __init__(self, num_experts: int, d_model: int, d_ff: int, activation=None, drop_rate: float = 0): + super().__init__("all_to_all", num_experts) + + self.w1 = nn.Parameter(torch.empty(self.num_local_experts, d_model, d_ff, device=get_current_device())) + self.b1 = nn.Parameter(torch.empty(self.num_local_experts, 1, d_ff, device=get_current_device())) + + self.w2 = nn.Parameter(torch.empty(self.num_local_experts, d_ff, d_model, device=get_current_device())) + self.b2 = nn.Parameter(torch.empty(self.num_local_experts, 1, d_model, device=get_current_device())) + + s1 = math.sqrt(0.1 / d_model) + s2 = math.sqrt(0.1 / d_ff) + + with seed(ParallelMode.TENSOR): + nn.init.trunc_normal_(self.w1, std=s1) + nn.init.trunc_normal_(self.b1, std=s1) + nn.init.trunc_normal_(self.w2, std=s2) + nn.init.trunc_normal_(self.b2, std=s2) + + self.act = nn.GELU() if activation is None else activation + self.drop = nn.Dropout(p=drop_rate) + + for param in self.parameters(): + param.__setattr__('moe_info', self.dist_info) + + def forward(self, inputs): # inputs [g, el, c, h] + + el = inputs.size(1) + h = inputs.size(-1) + + inputs = inputs.transpose(0, 1) + inshape = inputs.shape + inputs = inputs.reshape(el, -1, h) + + out_ff = torch.baddbmm(self.b1, inputs, self.w1) + out_act = self.act(out_ff) + with seed(ParallelMode.TENSOR): + out_inter = self.drop(out_act) + + out_model = torch.baddbmm(self.b2, out_inter, self.w2) + with seed(ParallelMode.TENSOR): + outputs = self.drop(out_model) # outputs [el, gc, h] + + outputs = outputs.reshape(inshape) + outputs = outputs.transpose(0, 1).contiguous() + return outputs + + +class TPExperts(MoeExperts): + """Use tensor parallelism to split each expert evenly, which can deploy experts in + case that the number of experts can't be divide by maximum expert parallel size or + maximum expert parallel size can't be divide by the number of experts. + """ + + def __init__(self, num_experts: int, d_model: int, d_ff: int, activation=None, drop_rate: float = 0): + super().__init__("all_gather", MOE_CONTEXT.max_ep_size) + + assert d_ff % MOE_CONTEXT.max_ep_size == 0, \ + "d_ff should be divide by maximum expert parallel size" + + p_ff = d_ff // MOE_CONTEXT.max_ep_size + + self.w1 = nn.Parameter(torch.empty(num_experts, d_model, p_ff, device=get_current_device())) + self.b1 = nn.Parameter(torch.empty(num_experts, 1, p_ff, device=get_current_device())) + + self.w2 = nn.Parameter(torch.empty(num_experts, p_ff, d_model, device=get_current_device())) + self.b2 = nn.Parameter(torch.empty(num_experts, 1, d_model, device=get_current_device())) + + s1 = math.sqrt(0.1 / d_model) + s2 = math.sqrt(0.1 / d_ff) + + with seed(ParallelMode.TENSOR): + nn.init.trunc_normal_(self.w1, std=s1) + nn.init.trunc_normal_(self.b1, std=s1) + nn.init.trunc_normal_(self.w2, std=s2) + + nn.init.trunc_normal_(self.b2, std=s2) + + self.act = nn.GELU() if activation is None else activation + self.drop = nn.Dropout(p=drop_rate) + + self.w1.__setattr__('moe_info', self.dist_info) + self.w2.__setattr__('moe_info', self.dist_info) + self.b1.__setattr__('moe_info', self.dist_info) + + def forward(self, inputs): # inputs [g, e, c, h] + + e = inputs.size(1) + h = inputs.size(-1) + + inputs = inputs.transpose(0, 1) + inshape = inputs.shape + inputs = inputs.reshape(e, -1, h) + + out_ff = torch.baddbmm(self.b1, inputs, self.w1) + out_act = self.act(out_ff) + with seed(ParallelMode.TENSOR): + out_inter = self.drop(out_act) + + out_model = torch.baddbmm(self.b2, out_inter, self.w2) + outputs = self.drop(out_model) # outputs [e, gc, h] + + outputs = outputs.reshape(inshape) + outputs = outputs.transpose(0, 1).contiguous() + return outputs # outputs [g, e, c, h] diff --git a/colossalai/nn/layer/moe/layers.py b/colossalai/nn/layer/moe/layers.py index 259f53f1adf5..03f55d91f3a8 100644 --- a/colossalai/nn/layer/moe/layers.py +++ b/colossalai/nn/layer/moe/layers.py @@ -1,203 +1,210 @@ -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F -from colossalai.context.moe_context import MOE_CONTEXT -from colossalai.utils import get_current_device -from colossalai.nn.layer.moe._operation import COL_MOE_KERNEL_FLAG, AllToAll, AllGather, \ - ReduceScatter, MoeDispatch, MoeCombine -from colossalai.nn.layer.moe.experts import MoeExperts, Experts -from colossalai.nn.layer.moe.utils import UniformNoiseGenerator, NormalNoiseGenerator -from colossalai.nn.layer.moe.routers import MoeRouter, Top1Router, Top2Router -from colossalai.zero.init_ctx import no_shard_zero_context, no_shard_zero_decrator -from typing import Optional, Type, Tuple - - -@no_shard_zero_decrator(is_replicated=True) -class MoeLayer(nn.Module): - """A MoE layer, that puts its input tensor to its gate and uses the output logits - to router all tokens, is mainly used to exchange all tokens for every expert across - the moe tensor group by all to all comunication. Then it will get the output of all - experts and exchange the output. At last returns the output of the moe system. - - Args: - dim_model (int): Dimension of model. - num_experts (int): The number of experts. - router (MoeRouter): Instance of router used in routing. - experts (MoeExperts): Instance of experts generated by Expert. - """ - - def __init__(self, dim_model: int, num_experts: int, router: MoeRouter, experts: MoeExperts): - super().__init__() - self.d_model = dim_model - self.num_experts = num_experts - self.gate_weight = torch.nn.Parameter(torch.empty(num_experts, dim_model)) - self.router: MoeRouter = router - self.experts: MoeExperts = experts - self.use_kernel = True if COL_MOE_KERNEL_FLAG and MOE_CONTEXT.use_kernel_optim else False - self.ep_group = experts.dist_info.ep_group - self.ep_size = experts.dist_info.ep_size - self.num_local_experts = experts.num_local_experts - - nn.init.trunc_normal_(self.gate_weight, std=math.sqrt(0.1 / dim_model)) - - def a2a_process(self, dispatch_data: torch.Tensor): - expert_input = AllToAll.apply(dispatch_data, self.ep_group) - input_shape = expert_input.shape - expert_input = expert_input.reshape(self.ep_size, self.num_local_experts, -1, self.d_model) - expert_output = self.experts(expert_input) - expert_output = expert_output.reshape(input_shape) - expert_output = AllToAll.apply(expert_output, self.ep_group) - return expert_output - - def tp_process(self, dispatch_data: torch.Tensor): - expert_in = AllGather.apply(dispatch_data, self.ep_group) - expert_out = self.experts(expert_in) - expert_out = ReduceScatter.apply(expert_out, self.ep_group) - return expert_out - - def forward(self, inputs: torch.Tensor) -> Tuple: - # reshape the input tokens - tokens = inputs.reshape(-1, self.d_model) - - # the data type of the inputs in the gating should be fp32 - fp32_input = tokens.to(torch.float) - fp32_weight = self.gate_weight.to(torch.float) - gate_output = F.linear(fp32_input, fp32_weight) - - # the result from the router - route_result_list = self.router(inputs=gate_output, use_kernel=self.use_kernel, ep_group=self.ep_group) - - if self.use_kernel: - dispatch_data = MoeDispatch.apply(tokens, *route_result_list[1:]) - dispatch_data = dispatch_data.reshape(self.num_experts, -1, self.d_model) - else: - sec_mask_f = route_result_list[1].type_as(inputs) - dispatch_data = torch.matmul(sec_mask_f.permute(1, 2, 0), tokens) - - # dispatch_data [e, c, h] - if self.experts.comm_name == "all_to_all": - expert_output = self.a2a_process(dispatch_data) - elif self.experts.comm_name == "all_gather": - expert_output = self.tp_process(dispatch_data) - else: - raise NotImplementedError("This kind of communication has not been implemented yet.\n Please use Experts " - "build function.") - # expert_output [e, c, h] - if self.use_kernel: - expert_output = expert_output.reshape(-1, self.d_model) - ans = MoeCombine.apply(expert_output, *route_result_list) - else: - combine_weights = route_result_list[0].type_as(inputs) - combine_weights = combine_weights.view(combine_weights.shape[0], -1) - expert_output = expert_output.view(-1, expert_output.shape[-1]) - ans = torch.matmul(combine_weights, expert_output) - - ans = ans.reshape(inputs.shape) - l_aux = self.router.pop_routing_loss() - return ans, l_aux - - -class MoeModule(nn.Module): - """A class for users to create MoE modules in their models. - - Args: - dim_model (int): Hidden dimension of training model - num_experts (int): The number experts - top_k (int, optional): The number of experts for dispatchment of each token - capacity_factor_train (float, optional): Capacity factor in routing during training - capacity_factor_eval (float, optional): Capacity factor in routing during evaluation - min_capacity (int, optional): The minimum number of the capacity of each expert - noisy_policy (str, optional): The policy of noisy function. Now we have 'Jitter' and 'Gaussian'. - 'Jitter' can be found in `Switch Transformer paper`_. - 'Gaussian' can be found in `ViT-MoE paper`_. - drop_tks (bool, optional): Whether drops tokens in evaluation - use_residual (bool, optional): Makes this MoE layer a Residual MoE. - More information can be found in `Microsoft paper`_. - residual_instance (nn.Module, optional): The instance of residual module in Resiual MoE - expert_instance (MoeExperts, optional): The instance of experts module in MoeLayer - expert_cls (Type[nn.Module], optional): The class of each expert when no instance is given - expert_args (optional): The args of expert when no instance is given - - .. _Switch Transformer paper: - https://arxiv.org/abs/2101.03961 - .. _ViT-MoE paper: - https://arxiv.org/abs/2106.05974 - .. _Microsoft paper: - https://arxiv.org/abs/2201.05596 - """ - - def __init__(self, - dim_model: int, - num_experts: int, - top_k: int = 1, - capacity_factor_train: float = 1.25, - capacity_factor_eval: float = 2.0, - min_capacity: int = 4, - noisy_policy: Optional[str] = None, - drop_tks: bool = True, - use_residual: bool = False, - residual_instance: Optional[nn.Module] = None, - expert_instance: Optional[MoeExperts] = None, - expert_cls: Optional[Type[nn.Module]] = None, - **expert_args): - super().__init__() - - noisy_func = None - if noisy_policy is not None: - if noisy_policy == 'Jitter': - noisy_func = UniformNoiseGenerator() - elif noisy_policy == 'Gaussian': - noisy_func = NormalNoiseGenerator(num_experts) - else: - raise NotImplementedError("Unsupported input noisy policy") - - if top_k == 1: - moe_router_cls = Top1Router - elif top_k == 2: - moe_router_cls = Top2Router - else: - raise NotImplementedError("top_k > 2 is not supported yet") - - self.moe_router = moe_router_cls(capacity_factor_train=capacity_factor_train, - capacity_factor_eval=capacity_factor_eval, - min_capacity=min_capacity, - noisy_func=noisy_func, - drop_tks=drop_tks) - self.use_residual = use_residual - if use_residual: - if residual_instance is not None: - self.residual_module = residual_instance - else: - assert expert_cls is not None, \ - "Expert class can't be None when residual instance is not given" - self.residual_module = expert_cls(**expert_args) - - with no_shard_zero_context(): - self.residual_combine = nn.Linear(dim_model, 2, device=get_current_device()) - - if expert_instance is not None: - self.experts = expert_instance - else: - assert expert_cls is not None, \ - "Expert class can't be None when experts instance is not given" - self.experts = Experts(expert_cls, num_experts, **expert_args) - - self.moe_layer = MoeLayer(dim_model=dim_model, - num_experts=num_experts, - router=self.moe_router, - experts=self.experts) - - def forward(self, inputs: torch.Tensor): - moe_output, l_aux = self.moe_layer(inputs) - - if self.use_residual: - residual_output = self.residual_module(inputs) - combine_coef = self.residual_combine(inputs) - combine_coef = F.softmax(combine_coef, dim=-1) - output = moe_output * combine_coef[..., 0:1] + residual_output * combine_coef[..., 1:] - else: - output = moe_output - - return output, l_aux +import math +from typing import Optional, Tuple, Type + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from colossalai.context.moe_context import MOE_CONTEXT +from colossalai.nn.layer.moe._operation import ( + COL_MOE_KERNEL_FLAG, + AllGather, + AllToAll, + MoeCombine, + MoeDispatch, + ReduceScatter, +) +from colossalai.nn.layer.moe.experts import Experts, MoeExperts +from colossalai.nn.layer.moe.routers import MoeRouter, Top1Router, Top2Router +from colossalai.nn.layer.moe.utils import NormalNoiseGenerator, UniformNoiseGenerator +from colossalai.utils import get_current_device +from colossalai.zero.legacy.init_ctx import no_shard_zero_context, no_shard_zero_decrator + + +@no_shard_zero_decrator(is_replicated=True) +class MoeLayer(nn.Module): + """A MoE layer, that puts its input tensor to its gate and uses the output logits + to router all tokens, is mainly used to exchange all tokens for every expert across + the moe tensor group by all to all communication. Then it will get the output of all + experts and exchange the output. At last returns the output of the moe system. + + Args: + dim_model (int): Dimension of model. + num_experts (int): The number of experts. + router (MoeRouter): Instance of router used in routing. + experts (MoeExperts): Instance of experts generated by Expert. + """ + + def __init__(self, dim_model: int, num_experts: int, router: MoeRouter, experts: MoeExperts): + super().__init__() + self.d_model = dim_model + self.num_experts = num_experts + self.gate_weight = torch.nn.Parameter(torch.empty(num_experts, dim_model)) + self.router: MoeRouter = router + self.experts: MoeExperts = experts + self.use_kernel = True if COL_MOE_KERNEL_FLAG and MOE_CONTEXT.use_kernel_optim else False + self.ep_group = experts.dist_info.ep_group + self.ep_size = experts.dist_info.ep_size + self.num_local_experts = experts.num_local_experts + + nn.init.trunc_normal_(self.gate_weight, std=math.sqrt(0.1 / dim_model)) + + def a2a_process(self, dispatch_data: torch.Tensor): + expert_input = AllToAll.apply(dispatch_data, self.ep_group) + input_shape = expert_input.shape + expert_input = expert_input.reshape(self.ep_size, self.num_local_experts, -1, self.d_model) + expert_output = self.experts(expert_input) + expert_output = expert_output.reshape(input_shape) + expert_output = AllToAll.apply(expert_output, self.ep_group) + return expert_output + + def tp_process(self, dispatch_data: torch.Tensor): + expert_in = AllGather.apply(dispatch_data, self.ep_group) + expert_out = self.experts(expert_in) + expert_out = ReduceScatter.apply(expert_out, self.ep_group) + return expert_out + + def forward(self, inputs: torch.Tensor) -> Tuple: + # reshape the input tokens + tokens = inputs.reshape(-1, self.d_model) + + # the data type of the inputs in the gating should be fp32 + fp32_input = tokens.to(torch.float) + fp32_weight = self.gate_weight.to(torch.float) + gate_output = F.linear(fp32_input, fp32_weight) + + # the result from the router + route_result_list = self.router(inputs=gate_output, use_kernel=self.use_kernel, ep_group=self.ep_group) + + if self.use_kernel: + dispatch_data = MoeDispatch.apply(tokens, *route_result_list[1:]) + dispatch_data = dispatch_data.reshape(self.num_experts, -1, self.d_model) + else: + sec_mask_f = route_result_list[1].type_as(inputs) + dispatch_data = torch.matmul(sec_mask_f.permute(1, 2, 0), tokens) + + # dispatch_data [e, c, h] + if self.experts.comm_name == "all_to_all": + expert_output = self.a2a_process(dispatch_data) + elif self.experts.comm_name == "all_gather": + expert_output = self.tp_process(dispatch_data) + else: + raise NotImplementedError("This kind of communication has not been implemented yet.\n Please use Experts " + "build function.") + # expert_output [e, c, h] + if self.use_kernel: + expert_output = expert_output.reshape(-1, self.d_model) + ans = MoeCombine.apply(expert_output, *route_result_list) + else: + combine_weights = route_result_list[0].type_as(inputs) + combine_weights = combine_weights.view(combine_weights.shape[0], -1) + expert_output = expert_output.view(-1, expert_output.shape[-1]) + ans = torch.matmul(combine_weights, expert_output) + + ans = ans.reshape(inputs.shape) + l_aux = self.router.pop_routing_loss() + return ans, l_aux + + +class MoeModule(nn.Module): + """A class for users to create MoE modules in their models. + + Args: + dim_model (int): Hidden dimension of training model + num_experts (int): The number experts + top_k (int, optional): The number of experts for dispatchment of each token + capacity_factor_train (float, optional): Capacity factor in routing during training + capacity_factor_eval (float, optional): Capacity factor in routing during evaluation + min_capacity (int, optional): The minimum number of the capacity of each expert + noisy_policy (str, optional): The policy of noisy function. Now we have 'Jitter' and 'Gaussian'. + 'Jitter' can be found in `Switch Transformer paper`_. + 'Gaussian' can be found in `ViT-MoE paper`_. + drop_tks (bool, optional): Whether drops tokens in evaluation + use_residual (bool, optional): Makes this MoE layer a Residual MoE. + More information can be found in `Microsoft paper`_. + residual_instance (nn.Module, optional): The instance of residual module in Residual MoE + expert_instance (MoeExperts, optional): The instance of experts module in MoeLayer + expert_cls (Type[nn.Module], optional): The class of each expert when no instance is given + expert_args (optional): The args of expert when no instance is given + + .. _Switch Transformer paper: + https://arxiv.org/abs/2101.03961 + .. _ViT-MoE paper: + https://arxiv.org/abs/2106.05974 + .. _Microsoft paper: + https://arxiv.org/abs/2201.05596 + """ + + def __init__(self, + dim_model: int, + num_experts: int, + top_k: int = 1, + capacity_factor_train: float = 1.25, + capacity_factor_eval: float = 2.0, + min_capacity: int = 4, + noisy_policy: Optional[str] = None, + drop_tks: bool = True, + use_residual: bool = False, + residual_instance: Optional[nn.Module] = None, + expert_instance: Optional[MoeExperts] = None, + expert_cls: Optional[Type[nn.Module]] = None, + **expert_args): + super().__init__() + + noisy_func = None + if noisy_policy is not None: + if noisy_policy == 'Jitter': + noisy_func = UniformNoiseGenerator() + elif noisy_policy == 'Gaussian': + noisy_func = NormalNoiseGenerator(num_experts) + else: + raise NotImplementedError("Unsupported input noisy policy") + + if top_k == 1: + moe_router_cls = Top1Router + elif top_k == 2: + moe_router_cls = Top2Router + else: + raise NotImplementedError("top_k > 2 is not supported yet") + + self.moe_router = moe_router_cls(capacity_factor_train=capacity_factor_train, + capacity_factor_eval=capacity_factor_eval, + min_capacity=min_capacity, + noisy_func=noisy_func, + drop_tks=drop_tks) + self.use_residual = use_residual + if use_residual: + if residual_instance is not None: + self.residual_module = residual_instance + else: + assert expert_cls is not None, \ + "Expert class can't be None when residual instance is not given" + self.residual_module = expert_cls(**expert_args) + + with no_shard_zero_context(): + self.residual_combine = nn.Linear(dim_model, 2, device=get_current_device()) + + if expert_instance is not None: + my_experts = expert_instance + else: + assert expert_cls is not None, \ + "Expert class can't be None when experts instance is not given" + my_experts = Experts(expert_cls, num_experts, **expert_args) + + self.moe_layer = MoeLayer(dim_model=dim_model, + num_experts=num_experts, + router=self.moe_router, + experts=my_experts) + + def forward(self, inputs: torch.Tensor): + moe_output, l_aux = self.moe_layer(inputs) + + if self.use_residual: + residual_output = self.residual_module(inputs) + combine_coef = self.residual_combine(inputs) + combine_coef = F.softmax(combine_coef, dim=-1) + output = moe_output * combine_coef[..., 0:1] + residual_output * combine_coef[..., 1:] + else: + output = moe_output + + return output, l_aux diff --git a/colossalai/nn/layer/moe/routers.py b/colossalai/nn/layer/moe/routers.py index c522c655a511..c5b8390bf047 100644 --- a/colossalai/nn/layer/moe/routers.py +++ b/colossalai/nn/layer/moe/routers.py @@ -60,7 +60,7 @@ def pop_routing_loss(self) -> torch.Tensor: class Top1Router(MoeRouter): """Top1 router that returns the dispatch mask [s, e, c] and combine weight [s, e, c] - for routing usage. More deailted function can be found in the paper about Switch Transformer + for routing usage. More detailed function can be found in the paper about Switch Transformer of Google. Args: capacity_factor_train (float, optional): Capacity factor in routing of training. @@ -143,7 +143,7 @@ def forward(self, inputs: torch.Tensor, use_kernel: bool = False, ep_group: Opti class Top2Router(MoeRouter): """Top2 router that returns the dispatch mask [s, e, c] and combine weight [s, e, c] - for routing usage. More deailted function can be found in the paper about ViT-MoE. + for routing usage. More detailed function can be found in the paper about ViT-MoE. Args: capacity_factor_train (float, optional): Capacity factor in routing of training. capacity_factor_eval (float, optional): Capacity factor in routing of evaluation. diff --git a/colossalai/nn/layer/moe/utils.py b/colossalai/nn/layer/moe/utils.py index 9362347414e0..4ca8bd703386 100644 --- a/colossalai/nn/layer/moe/utils.py +++ b/colossalai/nn/layer/moe/utils.py @@ -12,7 +12,7 @@ def half(self, memory_format=None): class NormalNoiseGenerator: - """Generates a random noisy mask for logtis tensor. + """Generates a random noisy mask for logits tensor. All noise is generated from a normal distribution :math:`(0, 1 / E^2)`, where `E = the number of experts`. @@ -32,7 +32,7 @@ def __call__(self, inputs: torch.Tensor): class UniformNoiseGenerator: - """Generates a random noisy mask for logtis tensor. + """Generates a random noisy mask for logits tensor. copied from mesh tensorflow: Multiply values by a random number between :math:`1-epsilon` and :math:`1+epsilon`. Makes models more resilient to rounding errors introduced by bfloat16. diff --git a/colossalai/nn/layer/parallel_1d/_operation.py b/colossalai/nn/layer/parallel_1d/_operation.py index 394334558275..300baf9c12ba 100644 --- a/colossalai/nn/layer/parallel_1d/_operation.py +++ b/colossalai/nn/layer/parallel_1d/_operation.py @@ -1,5 +1,6 @@ import torch import torch.distributed as dist + from colossalai.core import global_context as gpc try: diff --git a/colossalai/nn/layer/parallel_1d/layers.py b/colossalai/nn/layer/parallel_1d/layers.py index e96abd87ed10..406173a18c60 100644 --- a/colossalai/nn/layer/parallel_1d/layers.py +++ b/colossalai/nn/layer/parallel_1d/layers.py @@ -439,7 +439,7 @@ class Linear1D_Col(ParallelLayer): to all GPUs, otherwise, every GPU will have its output which is :math:`Y_i = XA_i`, defaults to False skip_bias_add (bool, optional): If set to ``True``, it will skip bias add for linear layer, - which is preserved for kernel fusion, defaults to Fals + which is preserved for kernel fusion, defaults to False weight_initializer (:class:`typing.Callable`, optional): The initializer of weight, defaults to kaiming uniform initializer. bias_initializer (:class:`typing.Callable`, optional): @@ -578,7 +578,7 @@ class Linear1D_Row(ParallelLayer): dtype (:class:`torch.dtype`, optional): The dtype of parameters, defaults to None. parallel_input (bool, optional): If set to ``True``, it's assumed that the input is split, defaults to False. skip_bias_add (bool, optional): If set to ``True``, it will skip bias add for linear layer, - which is preserved for kernel fusion, defaults to Fals + which is preserved for kernel fusion, defaults to False weight_initializer (:class:`typing.Callable`, optional): The initializer of weight, defaults to kaiming uniform initializer. bias_initializer (:class:`typing.Callable`, optional): @@ -994,11 +994,11 @@ class PatchEmbedding1D(ColossalaiModule): :type dtype: torch.dtype, optional :param flatten: whether to flatten output tensor, defaults to True :type flatten: bool, optional - :param weight_initializer: The intializer of weight, defaults to kaiming uniform initializer + :param weight_initializer: The initializer of weight, defaults to kaiming uniform initializer :type weight_initializer: typing.Callable, optional - :param bias_initializer: The intializer of bias, defaults to xavier uniform initializer + :param bias_initializer: The initializer of bias, defaults to xavier uniform initializer :type bias_initializer: typing.Callable, optional - :param position_embed_initializer: The intializer of position embedding, defaults to zero + :param position_embed_initializer: The initializer of position embedding, defaults to zero :type position_embed_initializer: typing.Callable, optional """ diff --git a/colossalai/nn/layer/parallel_sequence/layers.py b/colossalai/nn/layer/parallel_sequence/layers.py index d9486217bbc9..0887f8389dbe 100644 --- a/colossalai/nn/layer/parallel_sequence/layers.py +++ b/colossalai/nn/layer/parallel_sequence/layers.py @@ -195,7 +195,7 @@ class _Linear(nn.Module): keep_master_weight_for_test: This was added for testing and should be set to False. It returns the master weights used for initialization. - skip_bias_add: This was added to enable performance optimations where bias + skip_bias_add: This was added to enable performance optimizations where bias can be fused with other elementwise operations. we skip adding bias but instead return it. """ diff --git a/colossalai/nn/loss/loss_1d.py b/colossalai/nn/loss/loss_1d.py index 2fabd954f8fb..dd548c1d3dd4 100644 --- a/colossalai/nn/loss/loss_1d.py +++ b/colossalai/nn/loss/loss_1d.py @@ -21,7 +21,7 @@ def forward(ctx, vocab_parallel_logits, targets, process_group): # Subtract the maximum value. vocab_parallel_logits.sub_(logits_max.unsqueeze(dim=-1)) - # Get the partition's vocab indecies + # Get the partition's vocab indices partition_vocab_size = vocab_parallel_logits.size()[-1] rank = dist.get_rank(process_group) vocab_start_index = partition_vocab_size * rank @@ -61,10 +61,10 @@ def forward(ctx, vocab_parallel_logits, targets, process_group): @custom_bwd def backward(ctx, grad_output): - # Retreive tensors from the forward path. + # Retrieve tensors from the forward path. softmax, target_mask, masked_target_1d = ctx.saved_tensors - # All the inputs have softmax as thier gradient. + # All the inputs have softmax as their gradient. grad_input = softmax # For simplicity, work with the 2D gradient. partition_vocab_size = softmax.size()[-1] diff --git a/colossalai/nn/loss/loss_2d.py b/colossalai/nn/loss/loss_2d.py index cb12e723c323..7da8b2d697fa 100644 --- a/colossalai/nn/loss/loss_2d.py +++ b/colossalai/nn/loss/loss_2d.py @@ -106,7 +106,7 @@ def forward(ctx, logits, targets): @staticmethod @custom_bwd def backward(ctx, output_grad): - # Retreive tensors from the forward path. + # Retrieve tensors from the forward path. softmax, target_mask, masked_target = ctx.saved_tensors # All the inputs have softmax as their gradient. diff --git a/colossalai/nn/loss/loss_2p5d.py b/colossalai/nn/loss/loss_2p5d.py index f8e3324fc5ff..63dc4f33ad32 100644 --- a/colossalai/nn/loss/loss_2p5d.py +++ b/colossalai/nn/loss/loss_2p5d.py @@ -100,7 +100,7 @@ def forward(ctx, logits, targets): @staticmethod @custom_bwd def backward(ctx, output_grad): - # Retreive tensors from the forward path. + # Retrieve tensors from the forward path. softmax, target_mask, masked_target = ctx.saved_tensors # All the inputs have softmax as their gradient. diff --git a/colossalai/nn/loss/loss_3d.py b/colossalai/nn/loss/loss_3d.py index e76439191fdb..f27d57ad6c99 100644 --- a/colossalai/nn/loss/loss_3d.py +++ b/colossalai/nn/loss/loss_3d.py @@ -99,10 +99,10 @@ def forward(ctx, logits, targets, output_parallel_mode): @staticmethod @custom_bwd def backward(ctx, output_grad): - # Retreive tensors from the forward path. + # Retrieve tensors from the forward path. softmax, target_mask, masked_target = ctx.saved_tensors - # All the inputs have softmax as thier gradient. + # All the inputs have softmax as their gradient. input_grad = softmax # For simplicity, work with the 2D gradient. partition_vocab_size = softmax.size()[-1] diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py index 54036973e1e3..3a6d37103398 100644 --- a/colossalai/nn/optimizer/cpu_adam.py +++ b/colossalai/nn/optimizer/cpu_adam.py @@ -13,7 +13,7 @@ class CPUAdam(NVMeOptimizer): """Implements Adam algorithm. - Supports parameters updating on both GPU and CPU, depanding on the device of paramters. + Supports parameters updating on both GPU and CPU, depending on the device of parameters. But the parameters and gradients should on the same device: * Parameters on CPU and gradients on CPU is allowed. * Parameters on GPU and gradients on GPU is allowed. @@ -21,7 +21,7 @@ class CPUAdam(NVMeOptimizer): `CPUAdam` requires CUDA extensions which can be built during installation or runtime. - This version of CPU Adam accelates parameters updating on CPU with SIMD. + This version of CPU Adam accelerates parameters updating on CPU with SIMD. Support of AVX2 or AVX512 is required. The GPU part is implemented in an naive way. @@ -93,8 +93,7 @@ def torch_adam_update(self, bias_correction1, bias_correction2, use_adamw=False): - # FIXME(ver217): remove the below line when replace torch adam with fused adam - grad = grad.float() + grad = grad.to(data.dtype) if weight_decay != 0: if use_adamw: @@ -133,10 +132,12 @@ def step(self, closure=None, div_scale: float = -1): if len(state) == 0: state['step'] = 0 + # FIXME(ver217): CPU adam kernel only supports fp32 states now + assert p.dtype is torch.float, "CPUAdam only support fp32 parameters" # gradient momentums - state['exp_avg'] = torch.zeros_like(p, dtype=torch.float, device=target_device) + state['exp_avg'] = torch.zeros_like(p, device=target_device) # gradient variances - state['exp_avg_sq'] = torch.zeros_like(p, dtype=torch.float, device=target_device) + state['exp_avg_sq'] = torch.zeros_like(p, device=target_device) self._post_state_init(p) state['step'] += 1 @@ -147,9 +148,17 @@ def step(self, closure=None, div_scale: float = -1): assert state['exp_avg'].device.type == 'cpu', "exp_avg should stay on cpu" assert state['exp_avg_sq'].device.type == 'cpu', "exp_avg should stay on cpu" self._pre_update(p, 'exp_avg', 'exp_avg_sq') - self.cpu_adam_op.step(state['step'], group['lr'], beta1, beta2, group['eps'], group['weight_decay'], - group['bias_correction'], p.data, p.grad.data, state['exp_avg'], - state['exp_avg_sq'], div_scale) + if p.grad.dtype is torch.bfloat16: + # cpu adam kernel does not support bf16 now + bias_correction1 = 1 - beta1**state['step'] + bias_correction2 = 1 - beta2**state['step'] + self.torch_adam_update(p.data, p.grad.data, state['exp_avg'], state['exp_avg_sq'], group['lr'], + beta1, beta2, group['eps'], group['weight_decay'], bias_correction1, + bias_correction2, self.adamw_mode) + else: + self.cpu_adam_op.step(state['step'], group['lr'], beta1, beta2, group['eps'], + group['weight_decay'], group['bias_correction'], p.data, p.grad.data, + state['exp_avg'], state['exp_avg_sq'], div_scale) self._post_update(p, 'exp_avg', 'exp_avg_sq') elif target_device.type == 'cuda': assert div_scale == -1, "div_scale should remain default" diff --git a/colossalai/nn/optimizer/fused_adam.py b/colossalai/nn/optimizer/fused_adam.py index 987af8a968b7..82a6250f1fd1 100644 --- a/colossalai/nn/optimizer/fused_adam.py +++ b/colossalai/nn/optimizer/fused_adam.py @@ -134,8 +134,8 @@ def step(self, closure=None, grads=None, output_params=None, scale=None, grad_no # Exponential moving average of squared gradient values state['exp_avg_sq'] = torch.zeros_like(p) - if p.dtype not in [torch.float16, torch.float32]: - raise RuntimeError('FusedAdam only support fp16 and fp32.') + if p.dtype not in [torch.float16, torch.float32, torch.bfloat16]: + raise RuntimeError('FusedAdam only support fp16, fp32 and bf16.') g_l.append(p.grad.data) p_l.append(p.data) diff --git a/colossalai/nn/optimizer/gemini_optimizer.py b/colossalai/nn/optimizer/gemini_optimizer.py deleted file mode 100644 index 31d161612600..000000000000 --- a/colossalai/nn/optimizer/gemini_optimizer.py +++ /dev/null @@ -1,15 +0,0 @@ -from typing import Any - -import torch - -from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer - -__all__ = ['GeminiAdamOptimizer'] - - -class GeminiAdamOptimizer(ZeroOptimizer): - - def __init__(self, model: torch.nn.Module, **defaults: Any) -> None: - optimizer = HybridAdam(model.parameters(), **defaults) - super().__init__(optimizer, model, **defaults) diff --git a/colossalai/nn/optimizer/hybrid_adam.py b/colossalai/nn/optimizer/hybrid_adam.py index 1d0fb92de499..84903ac36832 100644 --- a/colossalai/nn/optimizer/hybrid_adam.py +++ b/colossalai/nn/optimizer/hybrid_adam.py @@ -1,31 +1,32 @@ from typing import Any, Optional import torch +from torch.optim import Adam -from colossalai.kernel.op_builder import CPUAdamBuilder, FusedOptimBuilder +from colossalai.kernel.op_builder import FusedOptimBuilder from colossalai.registry import OPTIMIZERS from colossalai.utils import multi_tensor_applier -from .nvme_optimizer import NVMeOptimizer +from .cpu_adam import CPUAdam @OPTIMIZERS.register_module -class HybridAdam(NVMeOptimizer): +class HybridAdam(CPUAdam): """Implements Adam algorithm. - Supports parameters updating on both GPU and CPU, depanding on the device of paramters. + Supports parameters updating on both GPU and CPU, depending on the device of parameters. But the parameters and gradients should on the same device: * Parameters on CPU and gradients on CPU is allowed. * Parameters on GPU and gradients on GPU is allowed. * Parameters on GPU and gradients on CPU is **not** allowed. - `HybriadAdam` requires CUDA extensions which can be built during installation or runtime. + `HybridAdam` requires CUDA extensions which can be built during installation or runtime. This version of Hybrid Adam is an hybrid of CPUAdam and FusedAdam. * For parameters updating on CPU, it uses CPUAdam. * For parameters updating on GPU, it uses FusedAdam. - * Hybird precision calculation of fp16 and fp32 is supported, eg fp32 parameters and fp16 gradients. + * Hybrid precision calculation of fp16 and fp32 is supported, eg fp32 parameters and fp16 gradients. :class:`colossalai.nn.optimizer.HybridAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``, or ``torch.optim.Adam`` with ``adamw_mode=False`` @@ -74,15 +75,9 @@ def __init__(self, nvme_offload_dir: Optional[str] = None, **defaults: Any): - default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction) - super(HybridAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir) - self.adamw_mode = adamw_mode - - # build during runtime if not found - cpu_optim = CPUAdamBuilder().load() + super().__init__(model_params, lr, bias_correction, betas, eps, weight_decay, adamw_mode, nvme_offload_fraction, + nvme_offload_dir) fused_optim = FusedOptimBuilder().load() - self.cpu_adam_op = cpu_optim.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode) - self.gpu_adam_op = fused_optim.multi_tensor_adam self._dummy_overflow_buf = torch.cuda.IntTensor([0]) @@ -108,10 +103,12 @@ def step(self, closure=None, div_scale: float = -1): if len(state) == 0: state['step'] = 0 + # FIXME(ver217): CPU adam kernel only supports fp32 states now + assert p.dtype is torch.float, "HybridAdam only support fp32 parameters" # gradient momentums - state['exp_avg'] = torch.zeros_like(p, dtype=torch.float, device=target_device) + state['exp_avg'] = torch.zeros_like(p, device=target_device) # gradient variances - state['exp_avg_sq'] = torch.zeros_like(p, dtype=torch.float, device=target_device) + state['exp_avg_sq'] = torch.zeros_like(p, device=target_device) self._post_state_init(p) state['step'] += 1 @@ -122,16 +119,24 @@ def step(self, closure=None, div_scale: float = -1): assert state['exp_avg'].device.type == 'cpu', "exp_avg should stay on cpu" assert state['exp_avg_sq'].device.type == 'cpu', "exp_avg should stay on cpu" self._pre_update(p, 'exp_avg', 'exp_avg_sq') - self.cpu_adam_op.step(state['step'], group['lr'], beta1, beta2, group['eps'], group['weight_decay'], - group['bias_correction'], p.data, p.grad.data, state['exp_avg'], - state['exp_avg_sq'], div_scale) + if p.grad.dtype is torch.bfloat16: + # cpu adam kernel does not support bf16 now + bias_correction1 = 1 - beta1**state['step'] + bias_correction2 = 1 - beta2**state['step'] + self.torch_adam_update(p.data, p.grad.data, state['exp_avg'], state['exp_avg_sq'], group['lr'], + beta1, beta2, group['eps'], group['weight_decay'], bias_correction1, + bias_correction2, self.adamw_mode) + else: + self.cpu_adam_op.step(state['step'], group['lr'], beta1, beta2, group['eps'], + group['weight_decay'], group['bias_correction'], p.data, p.grad.data, + state['exp_avg'], state['exp_avg_sq'], div_scale) self._post_update(p, 'exp_avg', 'exp_avg_sq') elif target_device.type == 'cuda': assert state['exp_avg'].device.type == 'cuda', "exp_avg should stay on cuda" assert state['exp_avg_sq'].device.type == 'cuda', "exp_avg should stay on cuda" - # record the state by gruop and update at once + # record the state by group and update at once g_l.append(p.grad.data) p_l.append(p.data) m_l.append(state['exp_avg']) diff --git a/colossalai/nn/optimizer/lamb.py b/colossalai/nn/optimizer/lamb.py index 7ac2109572a4..399ad39b6658 100644 --- a/colossalai/nn/optimizer/lamb.py +++ b/colossalai/nn/optimizer/lamb.py @@ -59,7 +59,7 @@ def step(self, closure=None): continue grad = p.grad.data if grad.is_sparse: - raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instad.') + raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instead.') state = self.state[p] diff --git a/colossalai/nn/optimizer/nvme_optimizer.py b/colossalai/nn/optimizer/nvme_optimizer.py index cbb435a90f61..fb3a4d87be60 100644 --- a/colossalai/nn/optimizer/nvme_optimizer.py +++ b/colossalai/nn/optimizer/nvme_optimizer.py @@ -1,9 +1,10 @@ -import torch +import math import os import tempfile -import math +from typing import Callable, Dict, List, Optional + +import torch from torch.nn.parameter import Parameter -from typing import Optional, List, Dict, Callable class NVMeOptimizer(torch.optim.Optimizer): @@ -42,8 +43,9 @@ def __init__(self, self.offloader = None self.is_on_nvme: Dict[Parameter, bool] = {} self.offloaded_numel: int = 0 - self.total_numel: int = self._get_numel() - self.can_offload_numel = math.floor(self.total_numel * self.nvme_offload_fraction) + # As param may be not materialized here, these attributes are initialized when the first step + self.total_numel: Optional[int] = None + self.can_offload_numel: Optional[int] = None self.prefetch_params: List[Parameter] = [] self.param_to_prefetch_idx: Dict[Parameter, int] = {} @@ -77,6 +79,9 @@ def _setup_prefetch_params(self) -> List[Parameter]: self.prefetch_params.append(p) def _pre_step(self, *state_keys: str) -> None: + if self.total_numel is None: + self.total_numel = self._get_numel() + self.can_offload_numel = math.floor(self.total_numel * self.nvme_offload_fraction) self._setup_prefetch_params() if self.offloader is None or len(self.prefetch_params) == 0: return diff --git a/colossalai/nn/optimizer/zero_optimizer.py b/colossalai/nn/optimizer/zero_optimizer.py deleted file mode 100644 index 422ebb7a3944..000000000000 --- a/colossalai/nn/optimizer/zero_optimizer.py +++ /dev/null @@ -1,318 +0,0 @@ -# this code is inspired by the DeepSpeed library and implemented with our own design from scratch -import math -import warnings -from enum import Enum -from typing import Any, Dict, Set, Tuple - -import torch -import torch.distributed as dist -from torch.nn import Parameter -from torch.optim import Optimizer - -from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler -from colossalai.gemini.chunk import Chunk, ChunkManager -from colossalai.logging import get_dist_logger -from colossalai.nn.optimizer import ColossalaiOptimizer, CPUAdam, FusedAdam, HybridAdam -from colossalai.nn.parallel.data_parallel import ZeroDDP -from colossalai.utils import disposable, get_current_device, is_ddp_ignored - -_AVAIL_OPTIM_LIST = {FusedAdam, CPUAdam, HybridAdam} - - -class OptimState(Enum): - SCALED = 0 - UNSCALED = 1 - - -class ZeroOptimizer(ColossalaiOptimizer): - """A wrapper for optimizer. ``ZeroDDP`` and ``ZeroOptimizer`` implement Zero Redundancy Optimizer (ZeRO state-3). - - Note: - You must use ``ZeroDDP`` with ``ZeroOptimizer``. - - Note: - Make sure you set ``placement_policy`` of ``GeminiManager`` to `"auto"`, - if you set ``gpu_margin_mem_ratio > 0``. - - Args: - optim (Optimizer): An Optimizer instance. - module (ZeroDDP): A ``ZeroDDP`` instance. - gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward) - which will be used when using hybrid CPU optimizer. - This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto". - Defaults to 0.0. - initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32. - min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. - growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2. - backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5. - growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000. - hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2. - max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32. - """ - - def __init__(self, - optim: Optimizer, - module: ZeroDDP, - gpu_margin_mem_ratio: float = 0.0, - initial_scale: float = 2**32, - min_scale: float = 1, - growth_factor: float = 2, - backoff_factor: float = 0.5, - growth_interval: int = 1000, - hysteresis: int = 2, - max_scale: float = 2**32, - clipping_norm: float = 0.0, - norm_type: float = 2.0, - **defaults: Any): - super().__init__(optim) - assert isinstance(module, ZeroDDP) - assert type(optim) in _AVAIL_OPTIM_LIST, "You should use an optimizer in the available list:\n" \ - f"{_AVAIL_OPTIM_LIST}" - self.module = module - self.gemini_manager = module.gemini_manager - self.chunk_manager: ChunkManager = self.gemini_manager.chunk_manager - self.optim_state = OptimState.UNSCALED - self.param_to_range: Dict[Parameter, Tuple[int, int]] = dict() - self.param_to_chunk32: Dict[Parameter, Chunk] = dict() - self.chunk16_set: Set[Chunk] = set() - self.clipping_flag = clipping_norm > 0.0 - self.max_norm = clipping_norm - - if self.clipping_flag: - assert norm_type == 2.0, "ZeroOptimizer only supports L2 norm now" - - ddp_param_list = [] - for name, param in module.named_parameters(): - if is_ddp_ignored(param): - if param.requires_grad: - warnings.warn(f"Parameter `{name}` is ignored by DDP but requires gradient! " - "You should handle its optimizer update by yourself!") - else: - ddp_param_list.append(param) - - for p, fp32_p in zip(ddp_param_list, module.fp32_params): - chunk_16 = self.chunk_manager.get_chunk(p) - if chunk_16 not in self.chunk16_set: - chunk_16.l2_norm_flag = self.clipping_flag - self.chunk16_set.add(chunk_16) - - self.__init__optimizer() - - # Grad scaler - self.grad_scaler = DynamicGradScaler(initial_scale=initial_scale, - min_scale=min_scale, - growth_factor=growth_factor, - backoff_factor=backoff_factor, - growth_interval=growth_interval, - hysteresis=hysteresis, - max_scale=max_scale) - self._found_overflow: torch.Tensor = torch.zeros(1, dtype=torch.int64, device=get_current_device()) - self._logger = get_dist_logger() - - self.gpu_margin_mem_ratio: float = float(gpu_margin_mem_ratio) - assert 0.0 <= self.gpu_margin_mem_ratio <= 1.0, f'gpu_margin_mem_ratio must >=0.0 and <=1.0' - # Only move fp32 shards from CPU to GPU when user allows and inner optimizer is valid - # Inner optimizer must support optimizing hybrid (CPU and CUDA) tensors, - # and it must set `num_fp32_shards_per_param` correctly - self._should_move_fp32_params_h2d: bool = self.gemini_manager.is_cuda_margin_mem_avail and self.gpu_margin_mem_ratio > 0.0 and getattr( - optim, 'num_fp32_shards_per_param', 0) >= 2 - if self.gpu_margin_mem_ratio > 0.0 and not self.gemini_manager.is_cuda_margin_mem_avail: - self._logger.warning(f'gpu_margin_mem_ratio is meaningless when placement_policy is not "auto"', ranks=[0]) - - self._register_states = disposable(self._register_states_) - - def _set_grad_ptr(self): - for group in self.param_groups: - for fake_param in group['params']: - chunk32 = self.param_to_chunk32[fake_param] - begin, end = self.param_to_range[fake_param] - chunk16 = chunk32.paired_chunk - - fake_param.data = chunk16.payload[begin:end] - fake_param.grad = fake_param.data - fake_param.data = chunk32.payload[begin:end] - - def _update_fp16_params(self): - none_tensor = torch.empty([0]) - for group in self.param_groups: - for fake_param in group['params']: - assert fake_param.grad is None - fake_param.data = none_tensor.to(fake_param.device) - - for chunk16 in self.chunk16_set: - chunk16.optim_update() - - def _check_overflow(self): - # clear previous overflow record - self._found_overflow.fill_(self.module.overflow_counter) - - # all-reduce across global group - dist.all_reduce(self._found_overflow) - - return self._found_overflow.item() > 0 - - def _clear_global_norm(self) -> None: - for c16 in self.chunk16_set: - c16.l2_norm = None - - def _calc_global_norm(self) -> float: - norm_sqr: float = 0.0 - group_to_norm = dict() - for c16 in self.chunk16_set: - assert c16.l2_norm is not None - - if c16.is_gathered: - norm_sqr += c16.l2_norm - else: - # this chunk is sharded, use communication to collect total norm - if c16.torch_pg not in group_to_norm: - group_to_norm[c16.torch_pg] = 0.0 - group_to_norm[c16.torch_pg] += c16.l2_norm - - c16.l2_norm = None # clear l2 norm - - comm_buffer = torch.zeros(1, dtype=torch.float, device=get_current_device()) - for group, part_norm in group_to_norm.items(): - comm_buffer.fill_(part_norm) - dist.all_reduce(comm_buffer, group=group) - norm_sqr += comm_buffer.item() - - global_norm = math.sqrt(norm_sqr) - return global_norm - - def _get_combined_scale(self): - loss_scale = 1 - - if self.optim_state == OptimState.SCALED: - loss_scale = self.loss_scale - self.optim_state = OptimState.UNSCALED - - combined_scale = loss_scale - if self.clipping_flag: - total_norm = self._calc_global_norm() - clip = ((total_norm / loss_scale) + 1e-6) / self.max_norm - if clip > 1: - combined_scale = clip * loss_scale - - if combined_scale == 1: - return -1 - else: - return combined_scale - - @property - def loss_scale(self): - return self.grad_scaler.scale.item() - - def zero_grad(self, *args, **kwargs): - self.module.overflow_counter = 0 - return self.optim.zero_grad(set_to_none=True) - - def step(self, *args, **kwargs): - self._maybe_move_fp32_params() - self._set_grad_ptr() - - found_inf = self._check_overflow() - if found_inf: - self.optim_state = OptimState.UNSCALED # no need to unscale grad - self.grad_scaler.update(found_inf) # update gradient scaler - self._logger.info(f'Found overflow. Skip step') - self._clear_global_norm() # clear recorded norm - self.zero_grad() # reset all gradients - self._update_fp16_params() - return - - # get combined scale. combined scale = loss scale * clipping norm - # so that gradient = gradient / combined scale - combined_scale = self._get_combined_scale() - self.grad_scaler.update(found_inf) - - ret = self.optim.step(div_scale=combined_scale, *args, **kwargs) - self._register_states() - self.zero_grad() - self._update_fp16_params() - return ret - - def clip_grad_norm(self, model: torch.nn.Module, max_norm: float, norm_type: float = 2.0): - raise NotImplementedError - - def backward(self, loss: torch.Tensor): - loss = self.loss_scale * loss - self.optim_state = OptimState.SCALED - self.module.backward(loss) - - def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tensor): - # This function is called except the last stage of pipeline parallel - # It receives the scaled grad from the previous rank - # No need to scale the grad again - # Need to unscale when optimizing - self.optim_state = OptimState.SCALED - self.module.backward_by_grad(tensor, grad) - - def _maybe_move_fp32_params(self): - if self._should_move_fp32_params_h2d: - self._should_move_fp32_params_h2d = False - available_cuda_margin_mem = self.gemini_manager.cuda_margin_mem * self.gpu_margin_mem_ratio - fp32_params_available_cuda_margin_mem = available_cuda_margin_mem / self.optim.num_fp32_shards_per_param - fp32_params_used_cuda_margin_mem = 0 - - for group in self.param_groups: - for fake_param in group['params']: - chunk32 = self.param_to_chunk32[fake_param] - chunk16 = chunk32.paired_chunk - - if chunk32.device_type == 'cuda': - continue - - if fp32_params_used_cuda_margin_mem + chunk32.payload_mem < fp32_params_available_cuda_margin_mem: - self.chunk_manager.move_chunk(chunk32, get_current_device()) - # stores grad now - self.chunk_manager.move_chunk(chunk16, get_current_device()) - self.module.set_chunk_grad_device(chunk16, get_current_device()) - fp32_params_used_cuda_margin_mem += chunk32.payload_mem - - for group in self.param_groups: - for fake_param in group['params']: - chunk32 = self.param_to_chunk32[fake_param] - if chunk32.device_type == 'cuda': - state = self.optim.state[fake_param] - for k, v in state.items(): - if isinstance(v, torch.Tensor): - state[k] = v.to(get_current_device()) - - def _register_states_(self): - for group in self.optim.param_groups: - for p in group['params']: - state = self.optim.state[p] - for val in state.values(): - if isinstance(val, torch.Tensor): - self.chunk_manager.add_extern_static_tensor(val) - - def __init__optimizer(self): - - def get_range_pair(local_chunk: Chunk, local_param: Parameter): - param_info = local_chunk.tensors_info[local_param] - if local_chunk.keep_gathered: - return param_info.offset, param_info.end - begin = max(0, param_info.offset - local_chunk.shard_begin) - end = min(local_chunk.shard_size, param_info.end - local_chunk.shard_begin) - return begin, end - - for group in self.optim.param_groups: - fake_params_list = list() - - for param in group['params']: - if is_ddp_ignored(param): - continue - chunk16 = self.chunk_manager.get_chunk(param) - range_pair = get_range_pair(chunk16, param) - if range_pair[0] >= range_pair[1]: - continue - - grad_device = self.module.grads_device[param] - fake_param = torch.nn.Parameter(torch.empty([0], device=grad_device)) - self.param_to_chunk32[fake_param] = chunk16.paired_chunk - self.param_to_range[fake_param] = range_pair - - fake_params_list.append(fake_param) - - group['params'] = fake_params_list diff --git a/colossalai/nn/parallel/__init__.py b/colossalai/nn/parallel/__init__.py index 2afc8f18c36f..17e010f478c9 100644 --- a/colossalai/nn/parallel/__init__.py +++ b/colossalai/nn/parallel/__init__.py @@ -1,5 +1,5 @@ -from .data_parallel import ColoDDP, ZeroDDP -from .gemini_parallel import GeminiDDP -from .zero_wrapper import zero_model_wrapper, zero_optim_wrapper +from .data_parallel import ColoDDP -__all__ = ['ColoDDP', 'ZeroDDP', 'GeminiDDP', 'zero_model_wrapper', 'zero_optim_wrapper'] +__all__ = [ + 'ColoDDP', +] diff --git a/colossalai/nn/parallel/data_parallel.py b/colossalai/nn/parallel/data_parallel.py index a9d001bd0a9c..f839d6b28444 100644 --- a/colossalai/nn/parallel/data_parallel.py +++ b/colossalai/nn/parallel/data_parallel.py @@ -1,31 +1,14 @@ -import itertools from collections import OrderedDict from functools import partial -from typing import Dict, Iterable, List, Optional, Set +from typing import Iterable, Optional, Set import torch import torch.distributed as dist -import torch.nn as nn -from colossalai.gemini.chunk import Chunk, ChunkManager, TensorState -from colossalai.gemini.gemini_mgr import GeminiManager -from colossalai.gemini.memory_tracer import OrderedParamGenerator -from colossalai.logging import get_dist_logger -from colossalai.nn.parallel.utils import get_temp_total_chunk_on_cuda from colossalai.tensor import ProcessGroup as ColoProcessGroup -from colossalai.tensor import ReplicaSpec -from colossalai.tensor.colo_parameter import ColoParameter, ColoTensor, ColoTensorSpec -from colossalai.tensor.param_op_hook import ColoParamOpHookManager -from colossalai.utils import get_current_device, is_ddp_ignored -from colossalai.zero.utils.gemini_hook import GeminiZeROHook +from colossalai.utils import is_ddp_ignored from .reducer import Reducer -from .utils import get_static_torch_model - -try: - from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, _IncompatibleKeys -except ImportError: - _EXTRA_STATE_KEY_SUFFIX = '_extra_state' def free_storage(data: torch.Tensor) -> None: @@ -189,507 +172,3 @@ def state_dict(self, destination=None, prefix='', keep_vars=False): def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict: bool = True): return self.module.load_state_dict(state_dict, strict) - - -class ZeroDDP(ColoDDP): - """ZeRO DDP for ColoTensor. - Warning: Nested ZeroDDP is not supported now. - It is designed to be used with ChunkManager and GeminiManager. - For more details, see the API reference of ``ChunkManager`` and ``GeminiManager``. - - Args: - module (torch.nn.Module): Module to apply ZeRO-DP. - gemini_manager (GeminiManager): Manages the chunk manager and heterogeneous momery space. - For more details, see the API reference of ``GeminiManager``. - pin_memory (bool): Chunks on CPU Memory use pin-memory. - force_outputs_fp32 (bool): If set to True, outputs will be fp32. Otherwise, outputs will be fp16. - Defaults to False. - strict_ddp_mode (bool): If set to True, there is no tensor sharding, each tensor is replicated. - Defaults to False. Users can set it to True, when they clearly know that they only need DDP. - """ - - def __init__(self, - module: torch.nn.Module, - gemini_manager: GeminiManager, - pin_memory: bool = False, - force_outputs_fp32: bool = False, - strict_ddp_mode: bool = False) -> None: - super().__init__(module, process_group=ColoProcessGroup()) - self.gemini_manager = gemini_manager - self.chunk_manager: ChunkManager = gemini_manager.chunk_manager - self.force_outputs_fp32 = force_outputs_fp32 - self.param_op_hook = GeminiZeROHook(gemini_manager) - self.fp32_params: List[ColoTensor] = list() - self.fp16_params: List[ColoParameter] = list() - self.overflow_counter = 0 - self.grads_device: Dict[torch.Tensor, torch.device] = dict() - self.param2name: Dict[nn.Parameter, str] = dict() - self.name2param: Dict[str, nn.Parameter] = dict() - - self._cast_buffers() - self._logger = get_dist_logger() - - if self.gemini_manager._premade_memstats_: - # build chunk in param runtime visited order. - param_order = self.gemini_manager.memstats()._param_runtime_order - else: - # build chunk in param initialized order. - # Note: in this way, it can not get filter unused params during runtime. - param_order = OrderedParamGenerator() - for p in module.parameters(): - param_order.append(p) - - self._init_chunks(param_order=param_order, - strict_ddp_mode=strict_ddp_mode, - cpu_offload=self.gemini_manager.policy_name != 'cuda', - pin_memory=pin_memory) - - for name, param in module.named_parameters(): - self.param2name[param] = name - for m_name, m_var in module.named_modules(): - for p_name, p_var in m_var.named_parameters(recurse=False): - param_name = m_name + '.' + p_name if m_name else p_name - self.name2param[param_name] = p_var - - def _post_forward(self): - """This function is only triggered for inference. - """ - access_list = list(self.chunk_manager.accessed_chunks) - # we need to scatter all accessed chunks and move them to their original places - for chunk in access_list: - if chunk.keep_gathered: - self.chunk_manager.fake_release_chunk(chunk) - else: - assert chunk.can_release - self.chunk_manager.release_chunk(chunk) - first_param = next(iter(chunk.tensors_info)) - self.chunk_manager.move_chunk(chunk, self.grads_device[first_param]) - assert self.chunk_manager.accessed_mem == 0 - # reset all recorded attributes - self.gemini_manager.reset_attributes() - - def forward(self, *args, **kwargs): - # check whether we are in a inference mode - grad_flag = torch.is_grad_enabled() - if not grad_flag: - assert not self.gemini_manager.need_warmup or not self.gemini_manager.is_warmup( - ), "You should run a completed iteration as your warmup iter" - - args, kwargs = _cast_float(args, torch.half), _cast_float(kwargs, torch.half) - self.module.zero_grad(set_to_none=True) - self.gemini_manager.pre_iter(*args) - with ColoParamOpHookManager.use_hooks(self.param_op_hook): - outputs = self.module(*args, **kwargs) - # scatter chunks in the inference mode - if not grad_flag: - self._post_forward() - - if self.force_outputs_fp32: - return _cast_float(outputs, torch.float) - return outputs - - def _setup_grads_ptr(self): - for p in self.module.parameters(): - if is_ddp_ignored(p): - continue - p.grad = None - - def _pre_backward(self): - # set a visit label for all parameters - # the label is used to check whether the parameter is correctly reduced - for param in self.param2name: - if not is_ddp_ignored(param): - setattr(param, "_gemini_reduced", False) - - def _post_backward(self): - if self.chunk_manager.accessed_mem != 0: - error_params = ["Reduction failed at followed parameters:"] - for param in self.param2name: - if not is_ddp_ignored(param) and not getattr(param, "_gemini_reduced"): - error_params.append(self.param2name[param]) - error_str = "\n\t".join(error_params) - raise RuntimeError("ZERO DDP error: the synchronization of gradients doesn't exit properly.", - "The most possible reason is that the model is not compatible with ZeroDDP.\n", - f"{error_str}") - self._setup_grads_ptr() - self._logger.debug( - f'comp cuda demand time: {self.gemini_manager._comp_cuda_demand_time}, layout time: {self.gemini_manager._layout_time}, evict time: {self.gemini_manager._evict_time}, CPU->CUDA vol: {self.gemini_manager._h2d_volume}B, CUDA->CPU vol: {self.gemini_manager._d2h_volume}' - ) - self.gemini_manager.post_iter() - - def backward(self, loss: torch.Tensor): - self._pre_backward() - with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook): - loss.backward() - self._post_backward() - - def backward_by_grad(self, tensor, grad): - with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook): - torch.autograd.backward(tensor, grad) - self._post_backward() - - def grad_handle(self, p, grad): - empty_grad = torch.empty_like(grad) - free_storage(empty_grad) - with torch._C.DisableTorchFunction(): - chunk = self.chunk_manager.get_chunk(p) - if chunk.tensors_info[p].state != TensorState.HOLD_AFTER_BWD: - raise RuntimeError(f"Parameter `{self.param2name[p]}` failed at the gradient reduction. " - "Some unsupported torch function is operated upon this parameter.") - self.chunk_manager.trans_tensor_state(p, TensorState.READY_FOR_REDUCE) - chunk.copy_tensor_to_chunk_slice(p, grad) - reduced = self.chunk_manager.reduce_chunk(chunk) - if reduced: - if chunk.is_gathered: - chunk.cuda_global_chunk.div_(chunk.pg_size) - else: - chunk.cuda_shard.div_(chunk.pg_size) - # check overflow elements - self.overflow_counter += chunk.has_inf_or_nan - # record l2 norm for gradient clipping - if chunk.l2_norm_flag: - chunk.set_l2_norm() - self.chunk_manager.move_chunk(chunk, self.grads_device[p], force_copy=True) - return empty_grad - - def zero_grad(self, set_to_none: bool = False) -> None: - self.module.zero_grad(set_to_none=True) - - def set_chunk_grad_device(self, chunk: Chunk, device: torch.device) -> None: - for tensor in chunk.get_tensors(): - self.grads_device[tensor] = device - - def state_dict(self, destination=None, prefix='', keep_vars=False, only_rank_0: bool = True): - """Returns a dictionary containing a whole state of the module. - - Both parameters and persistent buffers (e.g. running averages) are included. - Keys are corresponding parameter and buffer names. - Parameters and buffers set to ``None`` are not included. - - Warning: The non strict state dict would ignore the parameters if the tensors of the parameters - are shared with other parameters which have been included in the dictionary. - When you need to load the state dict, you should set the argument `strict` to False. - - Returns: - dict: - a dictionary containing a whole state of the module - """ - if destination is None: - destination = OrderedDict() - destination._metadata = OrderedDict() - destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version) - self._save_to_state_dict(destination, prefix, keep_vars, only_rank_0) - - for hook in self._state_dict_hooks.values(): - hook_result = hook(self, destination, prefix, local_metadata) - if hook_result is not None: - destination = hook_result - return destination - - def _get_param_to_save_data(self, param_list: List[torch.nn.Parameter], only_rank_0: bool) -> Dict: - """ - get param content from chunks. - - Args: - param_list (_type_): a list of torch.nn.Parameters - only_rank_0 (_type_): _description_ - - Returns: - Dict: a dict whose key is param name and value is param with correct payload - """ - # save parameters - param_to_save_data = dict() - chunk_list = self.chunk_manager.get_chunks(param_list) - for chunk in chunk_list: - temp_chunk = get_temp_total_chunk_on_cuda(chunk) - - for tensor, tensor_info in chunk.tensors_info.items(): - record_tensor = torch.empty([0]) - record_flag = (not only_rank_0) | (dist.get_rank(chunk.torch_pg) == 0) - if record_flag: - record_tensor = temp_chunk[tensor_info.offset:tensor_info.end].view(tensor.shape).cpu() - - assert tensor not in param_to_save_data - param_to_save_data[tensor] = record_tensor - - del temp_chunk - return param_to_save_data - - def _save_to_state_dict(self, destination, prefix, keep_vars, only_rank_0=True): - r"""Saves module state to `destination` dictionary, containing a state - of the module, but not its descendants. This is called on every - submodule in :meth:`~torch.nn.Module.state_dict`. - - In rare cases, subclasses can achieve class-specific behavior by - overriding this method with custom logic. - - Args: - destination (dict): a dict where state will be stored - prefix (str): the prefix for parameters and buffers used in this - module - """ - assert keep_vars is False, "`state_dict` with parameter, `keep_vars=True`, is not supported now." - - # get copies of fp32 parameters in CPU - param_to_save_data = self._get_param_to_save_data(self.fp32_params, only_rank_0) - # get the mapping between copies and fp16 parameters - p_mapping = dict() - for p, fp32_p in zip(self.fp16_params, self.fp32_params): - name = self.param2name[p] - assert fp32_p in param_to_save_data, "Parameter '{}' is neglected in the chunk list".format(name) - record_parameter = param_to_save_data[fp32_p] - p_mapping[p] = record_parameter - for name, param in self.name2param.items(): - if param is not None: - if is_ddp_ignored(param): - # deal with ddp ignored parameters - destination[prefix + name] = param if keep_vars else param.detach() - else: - destination[prefix + name] = p_mapping[param] - del p_mapping - del param_to_save_data - - # save all buffers - for name, buf in self.named_buffers(): - if buf is not None and name not in self._non_persistent_buffers_set: - destination[prefix + name] = buf if keep_vars else buf.detach() - # save extra states - extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX - if getattr(self.__class__, "get_extra_state", - torch.nn.Module.get_extra_state) is not torch.nn.Module.get_extra_state: - destination[extra_state_key] = self.get_extra_state() - - def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict: bool = True): - r"""Copies parameters and buffers from :attr:`state_dict` into - this module and its descendants. If :attr:`strict` is ``True``, then - the keys of :attr:`state_dict` must exactly match the keys returned - by this module's :meth:`~torch.nn.Module.state_dict` function. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - strict (bool, optional): whether to strictly enforce that the keys - in :attr:`state_dict` match the keys returned by this module's - :meth:`~torch.nn.Module.state_dict` function. Default: ``True`` - - Returns: - ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: - * **missing_keys** is a list of str containing the missing keys - * **unexpected_keys** is a list of str containing the unexpected keys - - Note: - If a parameter or buffer is registered as ``None`` and its corresponding key - exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a - ``RuntimeError``. - """ - missing_keys: List[str] = [] - unexpected_keys: List[str] = [] - error_msgs: List[str] = [] - - # copy state_dict so _load_from_state_dict can modify it - metadata = getattr(state_dict, '_metadata', None) - state_dict = state_dict.copy() - if metadata is not None: - # mypy isn't aware that "_metadata" exists in state_dict - state_dict._metadata = metadata # type: ignore[attr-defined] - - prefix = '' - local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) - self._load_from_state_dict(state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs) - - if strict: - if len(unexpected_keys) > 0: - error_msgs.insert( - 0, 'Unexpected key(s) in state_dict: {}. '.format(', '.join( - '"{}"'.format(k) for k in unexpected_keys))) - if len(missing_keys) > 0: - error_msgs.insert( - 0, 'Missing key(s) in state_dict: {}. '.format(', '.join('"{}"'.format(k) for k in missing_keys))) - - if len(error_msgs) > 0: - raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( - self.__class__.__name__, "\n\t".join(error_msgs))) - return _IncompatibleKeys(missing_keys, unexpected_keys) - - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, - error_msgs): - r"""Copies parameters and buffers from :attr:`state_dict` into only - this module, but not its descendants. This is called on every submodule - in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this - module in input :attr:`state_dict` is provided as :attr:`local_metadata`. - For state dicts without metadata, :attr:`local_metadata` is empty. - Subclasses can achieve class-specific backward compatible loading using - the version number at `local_metadata.get("version", None)`. - - .. note:: - :attr:`state_dict` is not the same object as the input - :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So - it can be modified. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - prefix (str): the prefix for parameters and buffers used in this - module - local_metadata (dict): a dict containing the metadata for this module. - See - strict (bool): whether to strictly enforce that the keys in - :attr:`state_dict` with :attr:`prefix` match the names of - parameters and buffers in this module - missing_keys (list of str): if ``strict=True``, add missing keys to - this list - unexpected_keys (list of str): if ``strict=True``, add unexpected - keys to this list - error_msgs (list of str): error messages should be added to this - list, and will be reported together in - :meth:`~torch.nn.Module.load_state_dict` - """ - for hook in self._load_state_dict_pre_hooks.values(): - hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) - - persistent_buffers = {k: v for k, v in self.named_buffers() if k not in self._non_persistent_buffers_set} - local_name_params = itertools.chain(self.named_parameters(), persistent_buffers.items()) - local_state = {k: v for k, v in local_name_params if v is not None} - - def load(param_name, dest_tensor, copy_func): - state_key = prefix + param_name - if state_key in state_dict: - input_param = state_dict[state_key] - # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+ - if len(dest_tensor.shape) == 0 and len(input_param.shape) == 1: - input_param = input_param[0] - if input_param.shape != dest_tensor.shape: - # local shape should match the one in checkpoint - error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, ' - 'the shape in current model is {}.'.format(state_key, input_param.shape, - dest_tensor.shape)) - return - try: - with torch.no_grad(): - copy_func(input_param) - except Exception as ex: - error_msgs.append('While copying the parameter named "{}", ' - 'whose dimensions in the model are {} and ' - 'whose dimensions in the checkpoint are {}, ' - 'an exception occurred : {}.'.format(state_key, dest_tensor.size(), - input_param.size(), ex.args)) - elif strict: - missing_keys.append(state_key) - - def load_fp32_parameter(chunk_slice, data): - chunk_slice.copy_(data.flatten()) - - for name, param in self.named_parameters(): - if is_ddp_ignored(param): - # deal with ddp ignored parameters - load(name, param, param.copy_) - - fp32_to_name = dict() - for p, fp32_p in zip(self.fp16_params, self.fp32_params): - if p is not None: - name = self.param2name[p] - fp32_to_name[fp32_p] = name - - chunk_list = self.chunk_manager.get_chunks(self.fp32_params) - for chunk in chunk_list: - temp_chunk = get_temp_total_chunk_on_cuda(chunk) - - for tensor, tensor_info in chunk.tensors_info.items(): - parameter_name = fp32_to_name[tensor] - parameter_slice = temp_chunk[tensor_info.offset:tensor_info.end] - load(parameter_name, tensor, partial(load_fp32_parameter, parameter_slice)) - - if chunk.is_gathered: - chunk.cuda_global_chunk.copy_(temp_chunk) - elif chunk.cuda_shard is not None: - chunk.cuda_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end]) - else: - chunk.cpu_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end]) - - del temp_chunk - - for chunk_32 in chunk_list: - chunk_16 = chunk_32.paired_chunk - assert chunk_16 is not None - chunk_16.optim_update() - - for name, buf in persistent_buffers.items(): - if buf is not None: - load(name, buf, buf.copy_) - - extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX - if getattr(self.__class__, "set_extra_state", - torch.nn.Module.set_extra_state) is not torch.nn.Module.set_extra_state: - if extra_state_key in state_dict: - self.set_extra_state(state_dict[extra_state_key]) - elif strict: - missing_keys.append(extra_state_key) - elif strict and (extra_state_key in state_dict): - unexpected_keys.append(extra_state_key) - - if strict: - for key in state_dict.keys(): - if key.startswith(prefix) and key != extra_state_key: - input_name = key[len(prefix):] - if input_name not in local_state: - unexpected_keys.append(key) - - def _init_chunks(self, param_order, strict_ddp_mode: bool, cpu_offload: bool, pin_memory: bool): - ddp_pg = ColoProcessGroup() - for p in param_order.generate(): - assert isinstance(p, ColoParameter) - - # gather sharded parameters in the strict ddp mode - if strict_ddp_mode: - if not p.is_replicate(): - p.set_dist_spec(ReplicaSpec()) - p.set_process_group(pg=ddp_pg) - - # ignore the parameters with no gradient - if not p.requires_grad: - self.set_params_to_ignore([p]) - - # move ignored parameters to CUDA - if is_ddp_ignored(p): - p.data = p.data.to(device=get_current_device(), dtype=torch.float16) - continue - - # create a fp32 parameter - fp32_data = p.data.float() - fp32_p = ColoTensor(fp32_data, spec=ColoTensorSpec(p.process_group)) - # create a fp16 parameter - p.data = p.data.half() - - # register the fp16 parameter and fp32 parameter in the chunk manager - dp_world_size = p.process_group.dp_world_size() - self.chunk_manager.register_tensor(tensor=p, - group_type='fp16_param', - config_key=dp_world_size, - cpu_offload=cpu_offload, - pin_memory=pin_memory) - self.chunk_manager.register_tensor(tensor=fp32_p, - group_type='fp32_param', - config_key=dp_world_size, - cpu_offload=cpu_offload, - pin_memory=pin_memory) - - self.fp16_params.append(p) - self.fp32_params.append(fp32_p) - self.grads_device[p] = self.gemini_manager.default_device - - self.chunk_manager.close_all_groups() - - for p, fp32_p in zip(self.fp16_params, self.fp32_params): - chunk_16 = self.chunk_manager.get_chunk(p) - chunk_32 = self.chunk_manager.get_chunk(fp32_p) - chunk_32.init_pair(chunk_16) - - # keep gathered chunks are in CUDA - if chunk_16.keep_gathered: - self.grads_device[p] = get_current_device() - - def _cast_buffers(self): - for buffer in self.module.buffers(): - buffer.data = buffer.cuda() - if torch.is_floating_point(buffer): - buffer.data = buffer.half() diff --git a/colossalai/nn/parallel/gemini_parallel.py b/colossalai/nn/parallel/gemini_parallel.py deleted file mode 100644 index 2c6e15d91736..000000000000 --- a/colossalai/nn/parallel/gemini_parallel.py +++ /dev/null @@ -1,63 +0,0 @@ -from typing import Optional - -import torch - -from colossalai.gemini.chunk import init_chunk_manager -from colossalai.gemini.gemini_mgr import GeminiManager -from colossalai.gemini.memory_tracer import MemStats - -from .data_parallel import ZeroDDP - - -class GeminiDDP(ZeroDDP): - - def __init__(self, - module: torch.nn.Module, - device: torch.device, - placement_policy: str = "cpu", - pin_memory: bool = False, - force_outputs_fp32: bool = False, - strict_ddp_mode: bool = False, - search_range_mb: int = 32, - hidden_dim: Optional[int] = None, - min_chunk_size_mb: float = 32, - memstats: Optional[MemStats] = None) -> None: - """ - A torch.Module warpper using ZeRO-DP and Genimi. - ZeRO is for parallel. Gemini is for memory management. - WARNING: The class will modify the module inline! - - Example: - model is initialized under the context of ColoInitContext - >>> model = GeminiDDP(model, torch.cuda.current_device(), "cuda") - >>> logits = model(x) - >>> loss = criterion(logits, labels) - >>> model.backward(loss) - - Args: - module (torch.nn.Module): the model to be wrapped. - device (torch.device): device to place the model. - placement_policy (str, optional): "cpu", "cuda", "auto". Defaults to "cpu". - pin_memory (bool, optional): use pin memory on CPU. Defaults to False. - force_outputs_fp32 (bool, optional): force outputs are fp32. Defaults to False. - search_range_mb (int, optional): chunk size searching range in MegaByte. Defaults to 32. - hidden_dim (int, optional): the hidden dimension of DNN. - Users can provide this argument to speed up searching. - If users do not know this argument before training, it is ok. We will use a default value 1024. - min_chunk_size_mb (float, optional): the minimum chunk size in MegaByte. - If the aggregate size of parameters is still samller than the minimum chunk size, - all parameters will be compacted into one small chunk. - memstats (MemStats, optional) the memory statistics collector by a runtime memory tracer. - """ - # some ugly hotfix for the compatibility with Lightning - if search_range_mb is None: - search_range_mb = 32 - - chunk_manager = init_chunk_manager(model=module, - init_device=device, - hidden_dim=hidden_dim, - search_range_mb=search_range_mb, - min_chunk_size_mb=min_chunk_size_mb, - strict_ddp_flag=strict_ddp_mode) - gemini_manager = GeminiManager(placement_policy, chunk_manager, memstats) - super().__init__(module, gemini_manager, pin_memory, force_outputs_fp32, strict_ddp_mode) diff --git a/colossalai/nn/parallel/layers/cache_embedding/cache_mgr.py b/colossalai/nn/parallel/layers/cache_embedding/cache_mgr.py index da043df368ae..a6159856dcce 100644 --- a/colossalai/nn/parallel/layers/cache_embedding/cache_mgr.py +++ b/colossalai/nn/parallel/layers/cache_embedding/cache_mgr.py @@ -20,8 +20,8 @@ def _wait_for_data(t, stream: Optional[torch.cuda.streams.Stream]) -> None: return torch.cuda.current_stream().wait_stream(stream) # As mentioned in https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html, - # PyTorch uses the "caching allocator" for memroy allocation for tensors. When a tensor is - # freed, its memory is likely to be reused by newly constructed tenosrs. By default, + # PyTorch uses the "caching allocator" for memory allocation for tensors. When a tensor is + # freed, its memory is likely to be reused by newly constructed tensors. By default, # this allocator traces whether a tensor is still in use by only the CUDA stream where it # was created. When a tensor is used by additional CUDA streams, we need to call record_stream # to tell the allocator about all these streams. Otherwise, the allocator might free the @@ -294,7 +294,7 @@ def print_comm_stats(self): print( f"CPU->CUDA BWD {self._cpu_to_cuda_numel * self.elem_size_in_byte / 1e6 / elapsed} MB/s {self._cpu_to_cuda_numel / 1e6} M elem" ) - print(f'cpu_to_cuda_elpase {elapsed} sec') + print(f'cpu_to_cuda_elapse {elapsed} sec') for k, v in self._elapsed_dict.items(): print(f'{k}: {v}') diff --git a/colossalai/nn/parallel/layers/cache_embedding/cached_embedding.py b/colossalai/nn/parallel/layers/cache_embedding/cached_embedding.py index a0c45d8e80c0..a74cb8d94bab 100644 --- a/colossalai/nn/parallel/layers/cache_embedding/cached_embedding.py +++ b/colossalai/nn/parallel/layers/cache_embedding/cached_embedding.py @@ -12,23 +12,23 @@ class CachedEmbeddingBag(BaseEmbeddingBag): Cached Embedding. Apply a GPU-based software cache approaches to dynamically manage the embedding table in the CPU and GPU memory space. It can leverage the id's frequency statistics of the target dataset, by passing a frequency list to param `ids_freq_mapping`. - You can also apply a navie LFU cache eviction strategy by setting `evict_strategy` as EvictionStrategy.LFU. + You can also apply a naive LFU cache eviction strategy by setting `evict_strategy` as EvictionStrategy.LFU. Args: num_embeddings (int): size of the dictionary of embeddings embedding_dim (int): the size of each embedding vector padding_idx (int, optional): If specified, the entries at padding_idx do not contribute to the gradient; therefore, the embedding vector at padding_idx is not updated during training, i.e. it remains as a fixed “pad”. For a newly constructed EmbeddingBag, the embedding vector at padding_idx will default to all zeros, but can be updated to another value to be used as the padding vector. Note that the embedding vector at padding_idx is excluded from the reduction. max_norm (float, optional): If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm - norm_type (str, optional): The p of the p-norm to compute for the max_norm option. Defaults to 2.. + norm_type (str, optional): The p of the p-norm to compute for the max_norm option. Defaults to 2. scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of the words in the mini-batch. Default False. Note: this option is not supported when mode="max". Defaults to False. sparse (bool, optional): if True, gradient w.r.t. weight matrix will be a sparse tensor. See Notes for more details regarding sparse gradients. Note: this option is not supported when mode="max".. Defaults to False. - _weight (torch.Tensor, optional): an embedding weight tensor. Concate multiple tables in a embedding bag as a single one. Defaults to None. + _weight (torch.Tensor, optional): an embedding weight tensor. Concatenate multiple tables in a embedding bag as a single one. Defaults to None. mode (str, optional): "sum", "mean" or "max". Specifies the way to reduce the bag. "sum" computes the weighted sum, taking per_sample_weights into consideration. "mean" computes the average of the values in the bag, "max" computes the max value over each bag. Default: "mean". Defaults to 'mean'. include_last_offset (bool, optional): if True, offsets has one additional element, where the last element is equivalent to the size of indices. This matches the CSR format.. Defaults to False. dtype (torch.dtype, optional): data type of the cpu weight initialization. Defaults to None meaning float32. device (torch.device, optional): device type to the cpu weight. Defaults to None meaning cpu. cache_ratio (float, float): cache ratio of the #cuda_weight_row / #cpu_weight_row - ids_freq_mapping (Union[List, torch.Tensor], optional): the frequency of each embedding vector occures in dataset. Defaults to None. + ids_freq_mapping (Union[List, torch.Tensor], optional): the frequency of each embedding vector occurs in dataset. Defaults to None. warmup_ratio (float, optional): the ratio of cuda cache is warmuped with. Defaults to 0.7. buffer_size (int, optional): the max number of vectors in transmitter buffer. If set to 0, the buffer is not used. Defaults to 0. pin_weight (bool, optional): pin the cpu weight. Defaults to False. @@ -145,7 +145,7 @@ def num_write_back_history(self): def swap_in_bandwidth(self): if self.cache_weight_mgr._cpu_to_cuda_numel > 0: return self.cache_weight_mgr._cpu_to_cuda_numel * self.cache_weight_mgr.elem_size_in_byte / 1e6 / \ - self.cache_weight_mgr._cpu_to_cuda_elpase + self.cache_weight_mgr._cpu_to_cuda_elapse else: return 0 diff --git a/colossalai/nn/parallel/layers/cache_embedding/copyer.py b/colossalai/nn/parallel/layers/cache_embedding/copyer.py index b586be1dc6d9..aa1f794482f9 100644 --- a/colossalai/nn/parallel/layers/cache_embedding/copyer.py +++ b/colossalai/nn/parallel/layers/cache_embedding/copyer.py @@ -17,7 +17,7 @@ def __init__(self, size: int) -> None: def index_copy(self, dim: int, src_index: LongTensor, tgt_index: LongTensor, src: torch.Tensor, tgt: torch.Tensor): """copy src tensor[src_index] -(index_select)-> tmp -(index_copy_)-> tgt tensor [tgt_index] - The valid rows in the src tensor are continous, while rows in tgt tensor is scattered. + The valid rows in the src tensor are continuous, while rows in tgt tensor is scattered. Args: dim (int): dimension along which to index diff --git a/colossalai/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py b/colossalai/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py index cb4647028d47..80a54b4fadd4 100644 --- a/colossalai/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py +++ b/colossalai/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py @@ -114,7 +114,7 @@ def forward(self, indices: torch.Tensor, offsets: torch.Tensor = None, per_sampl # get result of shape = (batch_size, (len(assigned_table_list)*embedding_dim)) local_output = torch.cat(local_output_list, 1) - # then concatenate those local_output on the second demension. + # then concatenate those local_output on the second dimension. # use all_to_all remains = batch_size % self.world_size scatter_strides = [batch_size // self.world_size + int(i < remains) for i in range(self.world_size)] diff --git a/colossalai/pipeline/pipelinable.py b/colossalai/pipeline/pipelinable.py index 9731530a6e15..79913987b7cc 100644 --- a/colossalai/pipeline/pipelinable.py +++ b/colossalai/pipeline/pipelinable.py @@ -83,7 +83,7 @@ def _post_init_method(self, module: torch.nn.Module, *args, **kwargs): for k, v in kwargs.items(): if isinstance(v, torch.nn.Module): v = self._layer_spec_dict[id(v)] - # (lyl)TODO: analyse ColoTensor as well + # (lyl)TODO: analyze ColoTensor as well modified_kwargs[k] = v # keep track of the module children @@ -117,7 +117,7 @@ def _post_init_method(self, module: torch.nn.Module, *args, **kwargs): def to_layer_list(self, exec_seq=None): """ Create a layer spec list and func list with execution sequence given by user. - If exec_seq is None, we will take the module initizing order as execution order. + If exec_seq is None, we will take the module initializing order as execution order. """ self._exec_seq = exec_seq @@ -177,7 +177,7 @@ def to_layer_list(self, exec_seq=None): def partition(self, num_chunks, pipeline_size, rank): """ - Partitioned model will be built respect to partion policy. + Partitioned model will be built respect to partition policy. The real module instance will be built in this method. """ if isinstance(self._policy, str): @@ -193,7 +193,7 @@ def partition(self, num_chunks, pipeline_size, rank): self.customized_parts = customized_partition(self._exec_seq) assert len(self.customized_parts) == gpc.get_world_size( ParallelMode.PIPELINE - ), f'World size is {gpc.get_world_size(ParallelMode.PIPELINE)}, but the number of partions is {len(self.customized_parts)}' + ), f'World size is {gpc.get_world_size(ParallelMode.PIPELINE)}, but the number of partitions is {len(self.customized_parts)}' parts = self.customized_parts[rank] else: raise ValueError("A string partition policy should be one of ['uniform', 'balanced', 'customized'].") diff --git a/colossalai/pipeline/rpc/_pipeline_base.py b/colossalai/pipeline/rpc/_pipeline_base.py index 2d7e25c82e7b..9e549df58214 100644 --- a/colossalai/pipeline/rpc/_pipeline_base.py +++ b/colossalai/pipeline/rpc/_pipeline_base.py @@ -123,7 +123,7 @@ def __init__(self, self.device = device self._initialize_outstanding_range() - # variable and const for context managment + # variable and const for context management self.outstanding = 0 self.forward_times = 0 self.backward_times = 0 @@ -226,7 +226,7 @@ def sync_global_worker_rrefs(self, pp_rank_to_worker_rref: Dict[int, PyRRef]) -> self.pp_rank_to_worker_rref = pp_rank_to_worker_rref # for some schedule need the other worker's info to initialise partition (like Chimera) - # construction of partition is executed after the registion of pp_rank_to_worker_rref + # construction of partition is executed after the registration of pp_rank_to_worker_rref self._initialize_partition() # res_use works for lifecycle counter, @@ -418,7 +418,7 @@ def subscribe_producer(self, microbatch_id: int, forward_only: bool): # On current PP middleware design for DAG, get_output_by_key used by _subscribe_producer # can only be executed once for every producer-consumer stage pair, which is necessary # to count the lifecycle of work_item. So, keeping the _subscribe_producer in the same - # lock of work_item queue operation gurantees the consistency of lifecycle counter. + # lock of work_item queue operation guarantees the consistency of lifecycle counter. work_item_from_producer = self._subscribe_producer(microbatch_id, forward_only) self.work_list[key] = work_item_from_producer self.work_list_condition_lock.notify_all() @@ -460,7 +460,7 @@ def subscribe_consumer(self, microbatch_id: int): # On current PP middleware design for DAG, get_output_by_key used by subscribe_consumer # can only be executed once for every producer-consumer stage pair, which is necessary # to count the lifecycle of work_item. So, keeping the subscribe_consumer in the same - # lock of work_item queue operation gurantees the consistency of lifecycle counter. + # lock of work_item queue operation guarantees the consistency of lifecycle counter. work_item_from_consumer = self._subscribe_consumer(microbatch_id) self.work_list[key] = work_item_from_consumer self.work_list_condition_lock.notify_all() @@ -508,7 +508,7 @@ def _get_producer_consumer(self) -> None: assert self.producer_stage_ids is None, f"all the producers of rank {rank} has been subscribed" assert self.consumer_stage_ids is None, f"all the consumers of rank {rank} has been subscribed" - # should be aranged in order, the order of the input of current forward + # should be arranged in order, the order of the input of current forward self.producer_stage_ids = self.get_producer_stage_ids() self.consumer_stage_ids = self.get_consumer_stage_ids() diff --git a/colossalai/pipeline/rpc/_pipeline_schedule.py b/colossalai/pipeline/rpc/_pipeline_schedule.py index 0d572231d378..6eda8f3b34b7 100644 --- a/colossalai/pipeline/rpc/_pipeline_schedule.py +++ b/colossalai/pipeline/rpc/_pipeline_schedule.py @@ -123,7 +123,7 @@ def _get_producer_consumer(self) -> None: assert self.producer_stage_ids is None, f"all the producers of rank {rank} has been subscribed" assert self.consumer_stage_ids is None, f"all the consumers of rank {rank} has been subscribed" - # should be aranged in order, the order of the input of current forward + # should be arranged in order, the order of the input of current forward self.producer_stage_ids = [] self.consumer_stage_ids = [] @@ -174,7 +174,7 @@ def _initialize_partition(self): else: # if it is down pipeline, create partition by origin method co_up_pp_worker_rref = self.pp_rank_to_worker_rref[pp_rank - stage_num] - # get the coresponding model state dict and wait for its init + # get the corresponding model state dict and wait for its init state_dict = co_up_pp_worker_rref.rpc_sync().get_partition_state_dict() super()._initialize_partition() self.module_partition.load_state_dict(state_dict) @@ -228,7 +228,7 @@ def _hook_before_step(self): stage_num = self.actual_stage_num co_pp_rank = (pp_rank + stage_num) % (2 * stage_num) - # if currrent pp_rank is not the first to do step + # if current pp_rank is not the first to do step # wait its previous pp_rank finish step grads = self.get_parameter_gradients() diff --git a/colossalai/pipeline/utils.py b/colossalai/pipeline/utils.py index df7226644a7a..ac8a3ad7d1db 100644 --- a/colossalai/pipeline/utils.py +++ b/colossalai/pipeline/utils.py @@ -113,7 +113,7 @@ def _binary_search(weights, num): def partition_uniform(num_items, pipeline_parallel_size, num_chunks): assert num_items % num_chunks == 0, \ - "Layer length should be divided by the number of chunks, otherwise parameter method is recomended" + "Layer length should be divided by the number of chunks, otherwise parameter method is recommended" logger = get_dist_logger() parts = [[] for _ in range(pipeline_parallel_size)] diff --git a/colossalai/shardformer/README.md b/colossalai/shardformer/README.md new file mode 100644 index 000000000000..bf4215c52980 --- /dev/null +++ b/colossalai/shardformer/README.md @@ -0,0 +1,387 @@ +# ⚡️ ShardFormer + +## 📚 Table of Contents + +- [⚡️ ShardFormer](#️-shardformer) + - [📚 Table of Contents](#-table-of-contents) + - [🔗 Introduction](#-introduction) + - [🔨 Usage](#-usage) + - [Quick Start](#quick-start) + - [Write your own policy](#write-your-own-policy) + - [🗺 Roadmap](#-roadmap) + - [💡 API Design](#-api-design) + - [Distributed Modules](#distributed-modules) + - [Shard Config](#shard-config) + - [Policy](#policy) + - [Model Sharder](#model-sharder) + - [User-facing API](#user-facing-api) + - [⌨️ Development Notes](#️-development-notes) + - [Add New Policy to Shardformer](#add-new-policy-to-shardformer) + - [Write Your Unit Testing](#write-your-unit-testing) + - [📊 Benchmarking](#-benchmarking) + - [System Performance](#system-performance) + - [Convergence](#convergence) + +## 🔗 Introduction + +**Shardformer** is a module that automatically parallelizes the mainstream models in libraries such as HuggingFace and TIMM. This module aims to make parallelization hassle-free for users who are not from the system background. + +## 🔨 Usage + +### Quick Start + +The sample API usage is given below: + +```python +from colossalai.shardformer import ShardConfig, Shard +from transformers import BertForMaskedLM + +# launch colossalai +colossalai.launch_from_torch() + +# create model +config = BertConfig.from_pretrained('bert-base-uncased') +model = BertForMaskedLM.from_pretrained('bert-base-uncased', config=config) + +# create huggingface model as normal +shard_config = ShardConfig() +shard_former = ShardFormer(shard_config=shard_config) +sharded_model = shard_former.optimize(model).to('cuda') + +# do everything like normal +... +``` + +### Write your own policy + +If you have a custom model, you can also use Shardformer to parallelize it by writing your own sharding policy. More information about the sharding policy can be found in [API Design](#-api-design). + +```python +from colossalai.shardformer import Policy + +class MyPolicy(Policy): + # implement your own policy + ... + +# init model and shard former +... + +# use customized policy to shard model +my_policy = MyPolicy() +shard_former.optimize(model, my_policy) + + + +``` + +## 🗺 Roadmap + +We will follow this roadmap to develop Shardformer: + +- [x] API Design +- [x] API Implementation +- [x] Unit Testing +- [ ] Policy Implementation + - [ ] Hugging Face + - [ ] NLP + - [x] BERT + - [x] T5 + - [x] LlaMa + - [x] GPT2 + - [x] OPT + - [x] BLOOM + - [ ] GLM + - [ ] RoBERTa + - [ ] ALBERT + - [ ] ERNIE + - [ ] GPT Neo + - [ ] GPT-J + - [ ] CV + - [x] ViT + - [ ] BEiT + - [ ] SwinTransformer + - [ ] SwinTransformer V2 + - [ ] Audio + - [ ] Whisper + - [ ] Multi-modal + - [ ] To be added + +## 💡 API Design + +We will discuss the major components of `ShardFormer` below to help you better understand how things work. +This section serves as the design doc for Shardformer and the function signature might differ from the actual implementation. +Please refer to the code for more details. + +

+ +
+

+ +### Distributed Modules + +`ShardFormer` replaces the original PyTorch module with a distributed module. +The distributed module keeps the same attributes as the original module but replaces the original parameters with distributed parameters and defines a new `forward` function to execute distributed computation. +Each distributed module implements its `from_native_module` static method to convert the PyTorch module to its corresponding distributed module. + +````python +class ParallelModule(torch.nn.Module): + + @abstractmethod + def from_native_module(module: torch.nn.Module, process_group: Union[ProcessGroup, Tuple[ProcessGroup]]) -> ParallelModule + """ + Convert a native module to a parallelized + + Examples: + + ```python + # replace module + my_linear = Linear1D_Col.from_native_module(my_linear, process_group) + ``` + """ +```` + +### Shard Config + +`ShardConfig` is a simple data class to tell `ShardFormer` how sharding will be performed. + +```python +@dataclass +class ShardConfig: + tensor_parallel_process_group: ProcessGroup = None + enable_fused_normalization: bool = False + ... + + # Some possible future config fields + tensor_parallel_mode: Choice['1d', '2d', '2.5d', '3d'] # support different tensor parallel mode + inference_only: bool # only inject inference-suitable sharding policy + use_flash_attention: bool # whether to use flash attention to speed up attention +``` + +### Policy + +The `Policy` class describes how to handle the model sharding. +It is merely a description, the actual sharding will be performed by `ModelSharder`. +We abstract the policy into four stages: + +1. Preprocessing: call `Policy.preprocess` to do some prior work before sharding, for example, resizing the embedding +2. Providing `ModulePolicyDescription`: call `Policy.module_policy` to get a bunch of `ModulePolicyDescription` to tell `ModelSharder` how the submodules's attributes, child parameters, and deeper submodules will be substituted. +3. Postprocessing: call `Policy.postprocess` to perform some postprocessing work, for example, binding the embedding and classifier head weights of the BERT model. + +```python +@dataclass +class ModulePolicyDescription: + r""" + Describe how the attributes and parameters will be transformed in a policy. + + Args: + attribute_replacement (Dict[str, Any]): key is the attribute name, value is the attribute value after sharding + param_replacement (List[Callable]): a list of functions to perform in-place param replacement. The function must receive only one arguments: module. + sub_module_replacement (List[SubModuleReplacementDescription]): each element in the list is a ParamReplacementDescription + object which specifies the module to be replaced and the target module used to replacement. + method_replace (Dict[str, Callable]): key is the method name, value is the method for replacement + """ + attribute_replacement: Dict[str, Any] = None + param_replacement: List[Callable] = None + sub_module_replacement: List[SubModuleReplacementDescription] = None + method_replacement: Dict[str, Callable] = None + +@dataclass +class SubModuleReplacementDescription: + r""" + Describe how a submodule will be replaced + + Args: + suffix (str): used to get the submodule object + target_module (ParallelModule): specifies the module class used to replace to submodule + kwargs (Dict[str, Any]): the dictionary used to pass extra arguments to the `ParallelModule.from_native_module` method. + ignore_if_not_exist (bool): if the submodule does not exist, ignore it or raise an exception + """ + suffix: str + target_module: ParallelModule + kwargs: Dict[str, Any] = None + ignore_if_not_exist: bool = False + + +class Policy(ABC): + + def __init__(self) + self.model = None + + def set_model(self, model: nn.Module) -> None: + """ + Set model as an attribute of the Policy object so that we can access the model's attributes. + """ + self.model = model + + @abstractmethod + def preprocess(self) -> nn.Module: + """ + Perform some preprocessing on the model, such as resizing the embedding size + """ + ... + + @abstractmethod + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + """ + Return the dict for the modify policy, the key is the original layer class and the value is the + argument for the modify layer + """ + ... + + @abstractmethods + def postprocess(self) -> nn.Module: + """ + Perform some postprocessing on the model, such as binding the embedding with the weight of the classifier head + """ + ... +``` + +### Model Sharder + +`ModelSharder` is the class in charge of sharding the model based on the given policy. + +```python +class ModelSharder: + + def __init__(self, model: torch.nn.Module, shard_config: ShardConfig, Policy: ShardPolicy = None): + #TODO: input is a cls or a obj + ... + + def shard(self) -> None: + """ + Shard model with parallelism with the help of pre-processing, replace_model_class, replace_module, and post-processing. + """ + ... + + def replace_module(self) -> None: + """ + Replace the layer according to the policy. Call Policy.module_policy() to get the module. Call _replace_module recursively. + """ + ... +``` + +### User-facing API + +We only expose a limited number of APIs to the user to keep their user experience simple and clean. + +```python +class ShardFormer: + """ + Parallelize model based on the given config and policy + + Example: + + shard_former = ShardFormer(shard_config=shard_config) + shard_former.init_distributed() + model = shard_former.optimize(model, policy=policy) + dataloader = shard_former.shard_dataset(dataset) + + """ + + def __init__(self, shard_config: ShardConfig): + """ + Do two things: + 1. Create a colossalai.cluster.process_group_manager to manage process groups for dp, tp and pp + 2. serve as a store for shard config + """ + self.shard_config = shard_config + self.pg_manager = None + + def init_distributed(self) -> colossalai.cluster.ProcessGroupManager: + """ + Initialize the distributed process group according to the + """ + pg_manager = ... + self.pg_manager = pg_manager + return pg_manager + + def shard_model(self, model: torch.nn.Module，policy: Policy) -> torch.nn.Module: + """ + Shard model for TP and PP + """ + ... + + def shard_dataset(self, dataset: Dataset) -> Dataloader: + """ + Shard dataset for DP + """ + ... +``` + +## ⌨️ Development Notes + +### Add New Policy to Shardformer + +This section serves as the guideline for writing new policies and register them into `shardformer`. + +- Step 1. Write your own model policy + +You can create a new file in the `colossalai/shardformer/policies` folder and name the file with the model name. You can implement your policy in this file. You should not import the any model zoo library at the header section of the file because we do not want to import the library when we do not use the policy. Libraries such as `transformers` should be imported only in the function body when needed. + +Please follow the following protocols when writing your policy: + +- You have to make a clear decision what you want to replace exactly in the original PyTorch module + - Use `ModulePolicyDescription.attribute_replacement` to replace the module attributes + - Use `ModulePolicyDescription.param_replacement` to replace the module parameters + - Use `ModulePolicyDescription.sub_module_replacement` to replace the submodules completely. The target module should implement the `from_native_module` for the replacement. + - Use `ModulePolicyDescription.method_replacement` to replace the module methods. **These replacement methods should be put in the `shardformer/modeling/.py`**. +- You can implement the `ParallelModule` for primitive modules in the `shardformer/layer/.py` file. Primitive modules refer to modules which are not composed of other modules. For example, the `torch.nn.Linear` module is a primitive module while modules such as `BertEncoder` module in the `transformers` library is a composite module. Primitive modules do not nested inner `nn.Module` members. For composite modules, you should consider using `ModulePolicyDescription` to implement your replacement. +- `ParallelModule` is meant to be used in two ways: `ParallelModule.from_native_module` to convert native PyTorch module to the `ParallelModule` and `ParallelModule(...)` to instantiate the module directly just like a normal PyTorch module. `ParallelModule` should be only implemented for modules whose weights are sharded. If you want to make your module compatible with the `ModulePolicyDescription.sub_module_replacement` and there is no weight sharding in your module, you can just implement the `from_native_module` method without inheriting the `ParallelModule` like `colossalai/shardformer/layer/normalization.py`. +- **Do not import any file in the `colossalai/shardformer/policies` and `colossalai/shardformer/modeling` to avoid unwanted import error**. For example, a file in these folders accidentally imports `transformers` library at the top of the file, then the user will have to install `transformers` library even if they do not use this file. Any file in the `modeling` folder should be only imported by the policy file. A policy implementation should be only imported dynamically via the autopolicy or manually via the `ShardFormer` module. +- Try to keep your import statement on third-party libraries such as `transformers` within the function body instead of the header section of the file. This is because we do not want to import the library when we do not use the policy. + +- Step 2. Register your policy to the autopolicy + +Next, you need to register your policy in the `colossalai/shardformer/policies/autopolicy.py` file. + +For example, if we register the policy for the BERT model, we just add a key-value in the `_POLICY_LIST` dictionary. The key if the `qualname` of the model object (you can get it by model.\_\_class\_\_.\_\_qualname\_\_). The value is a `PolicyLocation` object, which contains the file name and the class name of the policy. We do not import the policy directly because the policy file may contain libraries (such as `transformers`) which we do not want to import when we do not use the policy. + +```python +_POLICY_LIST = { + # BERT + "transformers.models.bert.modeling_bert.BertModel": + PolicyLocation(file_name="bert", class_name="BertModelPolicy"), +} +``` + +### Write Your Unit Testing + +This section serves as the guideline for testing the `shardformer` module. + +- Step 1. Add your model to the model zoo in the test kits. + +Add your model to the `tests/kit/model_zoo` file. This allows you to define test-related components for this model. You can take `tests/kit/model_zoo/transformers/llama.py` as an example for reference. + +- Step 2. Write your unit testing for the model + +Next, implement your unit test in the `tests/test_shardformer` folder. Please refer to other similar tests for style consistency. + +- Step 3. Execute your test + +When you run tests locally, you should run tests for both your newly-added test file and the whole `shardformer` module tests. + +```bash +# test for your own test file +pytest tests/test_shardformer/test_model/.py + +# test for the whole shardformer module +pytest tests/test_shardformer +``` + +## 📊 Benchmarking + +### System Performance + +To be added. + +### Convergence + +To validate that training the model using shardformers does not impact its convergence. We [fine-tuned the BERT model](./examples/shardformer_benchmark.py) using both shardformer and non-shardformer approaches. We compared the accuracy, loss, F1 score of the training results. + +| accuracy | f1 | loss | GPU number | model shard | +| :------: | :-----: | :-----: | :--------: | :---------: | +| 0.82594 | 0.87441 | 0.09913 | 4 | True | +| 0.81884 | 0.87299 | 0.10120 | 2 | True | +| 0.81855 | 0.87124 | 0.10357 | 1 | False | + +Overall, the results demonstrate that using shardformers during model training does not affect the convergence. diff --git a/colossalai/shardformer/__init__.py b/colossalai/shardformer/__init__.py new file mode 100644 index 000000000000..77c2af8d18f7 --- /dev/null +++ b/colossalai/shardformer/__init__.py @@ -0,0 +1 @@ +from .shard import ShardConfig, ShardFormer diff --git a/colossalai/shardformer/_utils.py b/colossalai/shardformer/_utils.py new file mode 100644 index 000000000000..4ad877e72357 --- /dev/null +++ b/colossalai/shardformer/_utils.py @@ -0,0 +1,80 @@ +import re + + +def get_obj_list_element(obj, a): + r""" + Get the element of the list in the object + """ + re_pattern = r'\[\d+\]' + prog = re.compile(re_pattern) + result = prog.search(a) + if result: + matched_brackets = result.group() + matched_index = matched_brackets.replace('[', '') + matched_index = matched_index.replace(']', '') + a_ = a.replace(matched_brackets, '') + container_obj = getattr(obj, a_) + obj = container_obj[int(matched_index)] + else: + obj = getattr(obj, a) + return obj + + +def hasattr_(obj, attr: str): + r""" + Check whether the object has the multi sublevel attr + + Args: + obj (object): The object to check + attr (str): The multi level attr to check + """ + attrs = attr.split('.') + for a in attrs: + try: + obj = get_obj_list_element(obj, a) + except AttributeError: + return False + return True + + +def setattr_(obj, attr: str, value, ignore: bool = False): + r""" + Set the object's multi sublevel attr to value, if ignore, ignore when it doesn't exist + + Args: + obj (object): The object to set + attr (str): The multi level attr to set + value (Any): The value to set + ignore (bool): Whether to ignore when the attr doesn't exist + """ + + attrs = attr.split('.') + for a in attrs[:-1]: + try: + obj = get_obj_list_element(obj, a) + except AttributeError: + if ignore: + return + raise AttributeError(f"Object {obj.__class__.__name__} has no attribute {attr}") + setattr(obj, attrs[-1], value) + + +def getattr_(obj, attr: str, ignore: bool = False): + r""" + Get the object's multi sublevel attr + + Args: + obj (object): The object to set + attr (str): The multi level attr to set + ignore (bool): Whether to ignore when the attr doesn't exist + """ + + attrs = attr.split('.') + for a in attrs: + try: + obj = get_obj_list_element(obj, a) + except AttributeError: + if ignore: + return None + raise AttributeError(f"Object {obj.__class__.__name__} has no attribute {attr}") + return obj diff --git a/colossalai/shardformer/examples/data.py b/colossalai/shardformer/examples/data.py new file mode 100644 index 000000000000..6296d4be4eb0 --- /dev/null +++ b/colossalai/shardformer/examples/data.py @@ -0,0 +1,146 @@ +import datasets +from torch.utils.data import DataLoader +from transformers import AutoTokenizer, PreTrainedTokenizer + +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase + + +class GLUEDataBuilder: + + task_text_field_map = { + "cola": ["sentence"], + "sst2": ["sentence"], + "mrpc": ["sentence1", "sentence2"], + "qqp": ["question1", "question2"], + "stsb": ["sentence1", "sentence2"], + "mnli": ["premise", "hypothesis"], + "qnli": ["question", "sentence"], + "rte": ["sentence1", "sentence2"], + "wnli": ["sentence1", "sentence2"], + "ax": ["premise", "hypothesis"], + } + + glue_task_num_labels = { + "cola": 2, + "sst2": 2, + "mrpc": 2, + "qqp": 2, + "stsb": 1, + "mnli": 3, + "qnli": 2, + "rte": 2, + "wnli": 2, + "ax": 3, + } + + loader_columns = [ + "datasets_idx", + "input_ids", + "token_type_ids", + "attention_mask", + "start_positions", + "end_positions", + "labels", + ] + + def __init__( + self, + model_name_or_path: str, + plugin: DPPluginBase = None, + task_name: str = "mrpc", + max_seq_length: int = 128, + train_batch_size: int = 32, + eval_batch_size: int = 32, + **kwargs, + ): + super().__init__() + self.model_name_or_path = model_name_or_path + self.task_name = task_name + self.max_seq_length = max_seq_length + self.train_batch_size = train_batch_size + self.eval_batch_size = eval_batch_size + self.plugin = plugin + + self.text_fields = self.task_text_field_map[task_name] + self.num_labels = self.glue_task_num_labels[task_name] + self.tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + self.setup() + + def setup(self): + self.dataset = datasets.load_dataset("glue", self.task_name) + + for split in self.dataset.keys(): + self.dataset[split] = self.dataset[split].map( + self.convert_to_features, + batched=True, + remove_columns=["label"], + ) + self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns] + self.dataset[split].set_format(type="torch", columns=self.columns) + + self.eval_splits = [x for x in self.dataset.keys() if "validation" in x] + + def prepare_data(self): + datasets.load_dataset("glue", self.task_name) + AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + + def train_dataloader(self): + if self.plugin == None: + return self.native_prepare_dataloader(self.dataset["train"], + batch_size=self.train_batch_size, + shuffle=True, + drop_last=True) + return self.plugin.prepare_dataloader(self.dataset["train"], + batch_size=self.train_batch_size, + shuffle=True, + drop_last=True) + + def val_dataloader(self): + if self.plugin == None: + return self.native_prepare_dataloader(self.dataset["validation"], batch_size=self.eval_batch_size) + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["validation"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def test_dataloader(self): + if self.plugin == None: + return self.native_prepare_dataloader(self.dataset["test"], batch_size=self.train_batch_size) + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["test"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def convert_to_features(self, example_batch): + + # Either encode single sentence or sentence pairs + if len(self.text_fields) > 1: + texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]])) + else: + texts_or_text_pairs = example_batch[self.text_fields[0]] + + # Tokenize the text/text pairs + features = self.tokenizer.batch_encode_plus(texts_or_text_pairs, + max_length=self.max_seq_length, + padding='max_length', + truncation=True) + + # Rename label to labels to make it easier to pass to model forward + features["labels"] = example_batch["label"] + + return features + + def native_prepare_dataloader(self, dataset, batch_size, shuffle=False, drop_last=False, pin_memory=False): + + return DataLoader(dataset, + batch_size=batch_size, + sampler=None, + shuffle=shuffle, + drop_last=drop_last, + pin_memory=pin_memory) diff --git a/colossalai/shardformer/examples/shardformer_benchmark.py b/colossalai/shardformer/examples/shardformer_benchmark.py new file mode 100644 index 000000000000..de82305b2547 --- /dev/null +++ b/colossalai/shardformer/examples/shardformer_benchmark.py @@ -0,0 +1,154 @@ +import argparse +import math +from typing import Any, List, Union + +import evaluate +import torch +import torch.distributed as dist +from data import GLUEDataBuilder +from torch import nn +from torch.optim import Adam, AdamW, Optimizer +from torch.utils._pytree import tree_map +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import BertConfig, BertForSequenceClassification, get_linear_schedule_with_warmup + +import colossalai +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam +from colossalai.shardformer import ShardConfig, ShardFormer + + +def to_device(x: Any, device: torch.device) -> Any: + + def _to(t: Any): + if isinstance(t, torch.Tensor): + return t.to(device) + return t + + return tree_map(_to, x) + + +def train(args): + colossalai.launch_from_torch(config={}, seed=42) + coordinator = DistCoordinator() + + # prepare for data and dataset + data_builder = GLUEDataBuilder(model_name_or_path=args.pretrain, + task_name=args.task, + train_batch_size=args.batch_size, + eval_batch_size=args.batch_size) + train_dataloader = data_builder.train_dataloader() + test_dataloader = data_builder.test_dataloader() + + if args.model == "bert": + cfg = BertConfig.from_pretrained(args.pretrain, num_labels=data_builder.num_labels) + model = BertForSequenceClassification.from_pretrained(args.pretrain, config=cfg) + + model.to(torch.cuda.current_device()) + + # if multiple GPUs, shard the model + if dist.get_world_size() > 1: + shard_config = ShardConfig(enable_fused_normalization=args.fused_layernorm) + shard_former = ShardFormer(shard_config=shard_config) + model = shard_former.optimize(model) + + optim = Adam(model.parameters(), lr=args.lr) + num_update_steps_per_epoch = len(train_dataloader) // args.accumulation_steps + max_steps = math.ceil(args.max_epochs * num_update_steps_per_epoch) + lr_scheduler = get_linear_schedule_with_warmup( + optim, + num_warmup_steps=math.ceil(max_steps * args.warmup_fraction), + num_training_steps=max_steps, + ) + fit(model, optim, lr_scheduler, train_dataloader, args.max_epochs, args.accumulation_steps, args.batch_size, + coordinator) + results = evaluate_model(model, test_dataloader, data_builder.num_labels, args.task, data_builder.eval_splits, + coordinator) + if coordinator.is_master(): + print(results) + if args.target_f1 is not None and 'f1' in results: + assert results['f1'] >= args.target_f1, f'f1 score {results["f1"]} is lower than target {args.target_f1}' + + +def fit(model: nn.Module, optimizer: Optimizer, scheduler, train_dataloader, max_epochs, accumulation_steps, batch_size, + coordinator): + step_bar = tqdm(range(len(train_dataloader) // accumulation_steps * max_epochs), + desc=f'steps', + disable=not coordinator.is_master()) + total_loss = 0 + for epoch in range(max_epochs): + model.train() + for batch_id, batch in enumerate(train_dataloader): + batch = to_device(batch, torch.cuda.current_device()) + outputs = model(**batch) + loss = outputs.loss + loss = loss / accumulation_steps + loss.backward() + total_loss += loss.item() + if (batch_id + 1) % accumulation_steps == 0: + optimizer.step() + scheduler.step() + optimizer.zero_grad() + step_bar.set_postfix({ + 'epoch': epoch, + 'loss': total_loss / batch_size, + 'lr': scheduler.get_last_lr()[0] + }) + total_loss = 0 + step_bar.update() + + +# evaluate +@torch.no_grad() +def evaluate_model(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoader]], num_labels: int, + task_name: str, eval_splits: List[str], coordinator: DistCoordinator): + metric = evaluate.load("glue", task_name, process_id=coordinator.rank, num_process=coordinator.world_size) + model.eval() + + def evaluate_subset(dataloader: DataLoader): + accum_loss = torch.zeros(1, device=torch.cuda.current_device()) + for batch in dataloader: + batch = to_device(batch, torch.cuda.current_device()) + outputs = model(**batch) + val_loss, logits = outputs[:2] + accum_loss.add_(val_loss) + + if num_labels > 1: + preds = torch.argmax(logits, axis=1) + elif num_labels == 1: + preds = logits.squeeze() + + labels = batch["labels"] + metric.add_batch(predictions=preds, references=labels) + + results = metric.compute() + if coordinator.is_master(): + results['loss'] = accum_loss.item() / (len(dataloader) * dataloader.batch_size) + return results + + if isinstance(test_dataloader, DataLoader): + return evaluate_subset(test_dataloader) + else: + assert len(test_dataloader) == len(eval_splits) + final_results = {} + for split, sub_loader in zip(eval_splits, test_dataloader): + results = evaluate_subset(sub_loader) + final_results.update({f'{k}_{split}': v for k, v in results.items()}) + return final_results + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('-t', '--task', default='mrpc', help="GLUE task to run") + parser.add_argument('--model', type=str, default="bert") + parser.add_argument('--pretrain', type=str, default="bert-base-uncased") + parser.add_argument('--max_epochs', type=int, default=1) + parser.add_argument('--batch_size', type=int, default=4) + parser.add_argument('--lr', type=float, default=2.4e-5) + parser.add_argument('--fused_layernorm', type=bool, default=False) + parser.add_argument('--accumulation_steps', type=int, default=8) + parser.add_argument('--warmup_fraction', type=float, default=0.03) + parser.add_argument('--target_f1', type=float, default=None) + args = parser.parse_args() + train(args) diff --git a/colossalai/shardformer/examples/shardformer_benchmark.sh b/colossalai/shardformer/examples/shardformer_benchmark.sh new file mode 100644 index 000000000000..f42b19a32d35 --- /dev/null +++ b/colossalai/shardformer/examples/shardformer_benchmark.sh @@ -0,0 +1,9 @@ +torchrun --standalone --nproc_per_node=4 shardformer_benchmark.py \ + --model "bert" \ + --pretrain "bert-base-uncased" \ + --max_epochs 1 \ + --batch_size 2 \ + --lr 2.4e-5 \ + --fused_layernorm False \ + --accumulation_steps 8 \ + --warmup_fraction 0.03 diff --git a/colossalai/shardformer/layer/__init__.py b/colossalai/shardformer/layer/__init__.py new file mode 100644 index 000000000000..7fad4948dfd0 --- /dev/null +++ b/colossalai/shardformer/layer/__init__.py @@ -0,0 +1,12 @@ +from .dropout import DropoutForParallelInput, DropoutForReplicatedInput +from .embedding import Embedding1D, VocabParallelEmbedding1D +from .linear import Linear1D_Col, Linear1D_Row +from .loss import cross_entropy_1d +from .normalization import FusedLayerNorm, FusedRMSNorm +from .qkv_fused_linear import GPT2FusedLinearConv1D_Col, GPT2FusedLinearConv1D_Row + +__all__ = [ + "Embedding1D", "VocabParallelEmbedding1D", "Linear1D_Col", "Linear1D_Row", 'GPT2FusedLinearConv1D_Col', + 'GPT2FusedLinearConv1D_Row', 'DropoutForParallelInput', 'DropoutForReplicatedInput', "cross_entropy_1d", + 'FusedLayerNorm', 'FusedRMSNorm' +] diff --git a/colossalai/shardformer/layer/_operation.py b/colossalai/shardformer/layer/_operation.py new file mode 100644 index 000000000000..7e97bee01b33 --- /dev/null +++ b/colossalai/shardformer/layer/_operation.py @@ -0,0 +1,290 @@ +import torch +import torch.distributed as dist +import torch.nn.functional as F + +try: + import fused_mix_prec_layer_norm_cuda +except: + fused_mix_prec_layer_norm_cuda = None + + +class FusedLayerNormAffineFunction1D(torch.autograd.Function): + r"""Layernorm + + Args: + input: input matrix. + weight: weight matrix. + bias: bias matrix. + normalized_shape: input shape from an expected input of size. + :math:`[* \times \text{normalized_shape}[0] \times \text{normalized_shape}[1] \times \ldots \times \text{normalized_shape}[-1]]` + If a single integer is used, it is treated as a singleton list, and this module will + normalize over the last dimension which is expected to be of that specific size. + eps: a value added to the denominator for numerical stability + """ + + @staticmethod + def forward(ctx, input, weight, bias, normalized_shape, eps): + ctx.normalized_shape = normalized_shape + ctx.eps = eps + input_ = input.contiguous() + weight_ = weight.contiguous() + bias_ = bias.contiguous() + output, mean, invvar = fused_mix_prec_layer_norm_cuda.forward_affine(input_, ctx.normalized_shape, weight_, + bias_, ctx.eps) + ctx.save_for_backward(input_, weight_, bias_, mean, invvar) + return output + + @staticmethod + def backward(ctx, grad_output): + input_, weight_, bias_, mean, invvar = ctx.saved_tensors + grad_input = grad_weight = grad_bias = None + grad_input, grad_weight, grad_bias \ + = fused_mix_prec_layer_norm_cuda.backward_affine( + grad_output.contiguous(), mean, invvar, + input_, ctx.normalized_shape, + weight_, bias_, ctx.eps) + + return grad_input, grad_weight, grad_bias, None, None + + +class MatmulWithAsyncCommunication(torch.autograd.Function): + """ + Linear layer execution with asynchronous communication in backprop. + """ + + @staticmethod + def forward(ctx, input_, weight, bias, process_group, async_grad_allreduce): + ctx.save_for_backward(input_, weight) + ctx.use_bias = bias is not None + ctx.process_group = process_group + ctx.async_grad_allreduce = async_grad_allreduce + + output = torch.matmul(input_, weight) + + if bias is not None: + output = output + bias + return output + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + use_bias = ctx.use_bias + + total_input = input + grad_input = grad_output.matmul(weight.T) + grad_output = grad_output.contiguous() + # Convert the tensor shapes to 2D for execution compatibility + if len(grad_output.shape) > 2: + grad_output = grad_output.view(-1, grad_output.shape[-1]) + total_input = total_input.view(-1, total_input.shape[-1]) + + if ctx.async_grad_allreduce: + # Asynchronous all-reduce + handle = dist.all_reduce(grad_input, group=ctx.process_group, async_op=True) + # Delay the start of weight gradient computation shortly (3us) to have + # all-reduce scheduled first and have GPU resources allocated + _ = torch.empty(1, device=grad_output.device) + 1 + + grad_weight = total_input.t().matmul(grad_output) + grad_bias = grad_output.sum(dim=0) if use_bias else None + + if ctx.async_grad_allreduce: + handle.wait() + + return grad_input, grad_weight, grad_bias, None, None, None + + +class LinearWithAsyncCommunication(torch.autograd.Function): + """ + Linear layer execution with asynchronous communication in backprop. + """ + + @staticmethod + def forward(ctx, input_, weight, bias, process_group, async_grad_allreduce): + ctx.save_for_backward(input_, weight) + ctx.use_bias = bias is not None + ctx.process_group = process_group + ctx.async_grad_allreduce = async_grad_allreduce + + if bias is not None: + output = F.linear(input_, weight, bias) + else: + output = F.linear(input_, weight) + return output + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + use_bias = ctx.use_bias + + total_input = input + grad_input = grad_output.matmul(weight) + grad_output = grad_output.contiguous() + # Convert the tensor shapes to 2D for execution compatibility + if len(grad_output.shape) > 2: + grad_output = grad_output.view(-1, grad_output.shape[-1]) + total_input = total_input.view(-1, total_input.shape[-1]) + + if ctx.async_grad_allreduce: + # Asynchronous all-reduce + handle = dist.all_reduce(grad_input, group=ctx.process_group, async_op=True) + # Delay the start of weight gradient computation shortly (3us) to have + # all-reduce scheduled first and have GPU resources allocated + _ = torch.empty(1, device=grad_output.device) + 1 + + grad_weight = grad_output.t().matmul(total_input) + grad_bias = grad_output.sum(dim=0) if use_bias else None + + if ctx.async_grad_allreduce: + handle.wait() + + return grad_input, grad_weight, grad_bias, None, None, None + + +class _SplitForwardGatherBackward(torch.autograd.Function): + """ + Split the input and keep only the corresponding chuck to the rank. + + Args: + input_ (`torch.Tensor`): input matrix. + dim (int): the dimension to perform split and gather + process_group (`torch.distributed.ProcessGroup`): the process group used for collective communication + + """ + + @staticmethod + def forward(ctx, input_, dim, process_group): + ctx.process_group = process_group + ctx.dim = dim + return _split(input_, dim, process_group) + + @staticmethod + def backward(ctx, grad_output): + return _gather(grad_output, ctx.dim, ctx.process_group), None, None + + +class _ReduceForward(torch.autograd.Function): + """ + All-reduce the input from the model parallel region. + + Args: + input_: input matrix. + parallel_mode: parallel mode. + """ + + @staticmethod + def forward(ctx, input_, process_group): + return _reduce(input_, process_group) + + @staticmethod + def backward(ctx, grad_output): + return grad_output, None + + +class _ReduceBackward(torch.autograd.Function): + """ + All-reduce the input from the model parallel region. + + Args: + input_: input matrix. + parallel_mode: parallel mode. + """ + + @staticmethod + def forward(ctx, input_, process_group): + ctx.process_group = process_group + return input_ + + @staticmethod + def backward(ctx, grad_output): + return _reduce(grad_output, ctx.process_group), None + + +def _reduce(input_, process_group): + # skip if only one rank involved + if dist.get_world_size(process_group) == 1: + return input_ + else: + dist.all_reduce(input_, group=process_group) + return input_ + + +def _split(input_, dim=-1, process_group=None): + # skip if only one rank involved + world_size = dist.get_world_size(process_group) + if world_size == 1: + return input_ + + # Split along last dimension. + dim_size = input_.size(dim) + assert dim_size % world_size == 0, \ + f'The dimension to split ({dim_size}) is not a multiple of world size ({world_size}), ' \ + f'cannot split tensor evenly' + + tensor_list = torch.split(input_, dim_size // world_size, dim=dim) + rank = dist.get_rank(process_group) + output = tensor_list[rank].contiguous() + + return output + + +def _gather(input_, dim=-1, process_group=None): + # skip if only one rank involved + world_size = dist.get_world_size(process_group) + if world_size == 1: + return input_ + + # all gather + rank = dist.get_rank(process_group) + tensor_list = [torch.empty_like(input_) for _ in range(world_size)] + tensor_list[rank] = input_ + torch.distributed.all_gather(tensor_list, input_, group=process_group) + + # concat + output = torch.cat(tensor_list, dim=dim).contiguous() + + return output + + +class _GatherForwardSplitBackward(torch.autograd.Function): + """Gather the input from model parallel region and concatenate. + + Args: + input_: input matrix. + parallel_mode: parallel mode. + dim: dimension + """ + + @staticmethod + def forward(ctx, input_, dim, process_group): + ctx.process_group = process_group + ctx.dim = dim + return _gather(input_, dim, process_group) + + @staticmethod + def backward(ctx, grad_output): + return _split(grad_output, ctx.dim, ctx.process_group), None, None + + +def matmul_with_async_comm(input_, weight, bias, process_group, async_grad_allreduce): + return MatmulWithAsyncCommunication.apply(input_, weight, bias, process_group, async_grad_allreduce) + + +def linear_with_async_comm(input_, weight, bias, process_group, async_grad_allreduce): + return LinearWithAsyncCommunication.apply(input_, weight, bias, process_group, async_grad_allreduce) + + +def gather_forward_split_backward(input_, dim, process_group): + return _GatherForwardSplitBackward.apply(input_, dim, process_group) + + +def split_forward_gather_backward(input_, dim, process_group): + return _SplitForwardGatherBackward.apply(input_, dim, process_group) + + +def reduce_forward(input_, process_group): + return _ReduceForward.apply(input_, process_group) + + +def reduce_backward(input_, process_group): + return _ReduceBackward.apply(input_, process_group) diff --git a/colossalai/shardformer/layer/dropout.py b/colossalai/shardformer/layer/dropout.py new file mode 100644 index 000000000000..2625fe97889a --- /dev/null +++ b/colossalai/shardformer/layer/dropout.py @@ -0,0 +1,83 @@ +from typing import List, Union + +import torch +import torch.nn as nn +from torch.distributed import ProcessGroup + +from .parallel_module import ParallelModule +from .utils import create_randomizer_with_offset + +__all__ = ['DropoutForParallelInput', 'DropoutForReplicatedInput'] + + +class DropoutForParallelInput(ParallelModule, nn.Dropout): + """ + The Dropout Layer will apply dropout mask to the input tensor. The dropout mask is generated with + randomness on different ranks of the given process group. This can avoid the same dropout mask is generated + and applied on the same position of different ranks, leading to poor convergence performance. + + Args: + p (float): probability of an element to be zeroed. Defaults to 0.5. + inplace (bool): If set to True, will do this operation in-place. Defaults to False. + process_group (ProcessGroup): the process group to be used for generating randomness. Defaults to None. + """ + + def __init__(self, p: float = 0.5, inplace: bool = False, process_group: ProcessGroup = None): + # init with nn.Dropout + super(nn.Dropout, self).__init__(p=p, inplace=inplace) + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=process_group) + + @staticmethod + def from_native_module(module: nn.Dropout, + process_group: Union[ProcessGroup, List[ProcessGroup]] = None) -> "DropoutForParallelInput": + """ + Create a DropoutForParallelInput layer from a native dropout layer. + """ + p = module.p + inplace = module.inplace + return DropoutForParallelInput(p=p, inplace=inplace, process_group=process_group) + + def forward(self, input): + with self.randomizer.fork_rng(): + input = super().forward(input) + return input + + +class DropoutForReplicatedInput(ParallelModule, nn.Dropout): + """ + The Dropout Layer will apply dropout mask to the input tensor. The dropout mask is generated with + randomness on different ranks of the given process group. This can avoid the same dropout mask is generated + and applied on the same position of different ranks, leading to poor convergence performance. + + Args: + p (float): probability of an element to be zeroed. Defaults to 0.5. + inplace (bool): If set to True, will do this operation in-place. Defaults to False. + process_group (ProcessGroup): the process group to be used for generating randomness. Defaults to None. + """ + + def __init__(self, p: float = 0.5, inplace: bool = False, process_group: ProcessGroup = None): + # init with nn.Dropout + super(nn.Dropout, self).__init__(p=p, inplace=inplace) + + # offset the seed with randomizer index only + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=process_group, offset_by_rank=False) + + @staticmethod + def from_native_module( + module: nn.Dropout, + process_group: Union[ProcessGroup, List[ProcessGroup]] = None) -> "DropoutForReplicatedInput": + """ + Create a Dropout1D layer from a native dropout layer. + """ + p = module.p + inplace = module.inplace + return DropoutForReplicatedInput(p=p, inplace=inplace, process_group=process_group) + + def forward(self, input): + with self.randomizer.fork_rng(): + input = super().forward(input) + return input diff --git a/colossalai/shardformer/layer/embedding.py b/colossalai/shardformer/layer/embedding.py new file mode 100644 index 000000000000..db39a457b7fd --- /dev/null +++ b/colossalai/shardformer/layer/embedding.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +from typing import Callable, List, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.distributed import ProcessGroup +from torch.nn.parameter import Parameter + +from colossalai.nn import init as init +from colossalai.nn.layer.utils import divide +from colossalai.tensor.d_tensor.api import shard_colwise, shard_rowwise, sharded_tensor_to_param + +from ._operation import gather_forward_split_backward, reduce_forward +from .parallel_module import ParallelModule +from .utils import create_randomizer_with_offset + +__all__ = ['Embedding1D', 'VocabParallelEmbedding1D'] + + +class Embedding1D(ParallelModule): + r"""Embedding for 1D parallelism. + + Args: + num_embeddings (int): number of embeddings. + embedding_dim (int): dimension of embedding. + padding_idx (int, optional): If specified, the entries at padding_idx do not contribute to the gradient; + therefore, the embedding vector at padding_idx is not updated during training, + i.e. it remains as a fixed “pad”, defaults to None. + dtype (:class:`torch.dtype`, optional): The dtype of parameters, defaults to None. + weight_initializer (:class:`typing.Callable`, optional): + he initializer of weight, defaults to normal initializer. + + The ``args`` and ``kwargs`` used in :class:`torch.nn.functional.embedding` should contain: + :: + + max_norm (float, optional): If given, each embedding vector with norm larger than max_norm is + renormalized to have norm max_norm. Note: this will modify weight in-place. + norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default 2. + scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse + of frequency of the words in the mini-batch. Default False. + sparse (bool, optional): If True, gradient w.r.t. weight will be a sparse tensor. Default False. + + More details about ``args`` and ``kwargs`` could be found in + `Embedding `_. + + More details about ``initializer`` please refer to + `init `_ + """ + + def __init__(self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int = None, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + gather_output: bool = True, + weight_initializer: Callable = init.normal_(), + *args, + **kwargs): + super().__init__() + + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self.process_group = process_group + + self.padding_idx = padding_idx + self.embed_args = args + self.embed_kwargs = kwargs + self.gather_output = gather_output + + # Parameters. + factory_kwargs = {'device': device, 'dtype': dtype} + weight = torch.empty((num_embeddings, self.embedding_dim), **factory_kwargs) + sharded_weight = shard_colwise(weight, process_group) + self.weight = sharded_tensor_to_param(sharded_weight) + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + + with self.randomizer.fork_rng(enable_cpu=True): + self.reset_parameters(weight_initializer) + + @staticmethod + def from_native_module(module: nn.Embedding, + process_group: Union[ProcessGroup, List[ProcessGroup]] = None, + *args, + **kwargs) -> "Embedding1D": + r""" + Build a 1D parallelized Embedding from a native nn.Embedding module. + """ + # get the attributes + num_embedding = module.num_embeddings + embedding_dim = module.embedding_dim + padding_idx = module.padding_idx + max_norm = module.max_norm + norm_type = module.norm_type + scale_grad_by_freq = module.scale_grad_by_freq + sparse = module.sparse + dtype = module.weight.dtype + device = module.weight.device + + # sparse is not support yet + if sparse: + raise NotImplementedError("The Embedding1D module does not support sparse embedding yet.") + + embedding = Embedding1D(num_embeddings=num_embedding, + embedding_dim=embedding_dim, + padding_idx=padding_idx, + process_group=process_group, + dtype=dtype, + device=device, + max_norm=max_norm, + norm_type=norm_type, + scale_grad_by_freq=scale_grad_by_freq, + sparse=sparse, + *args, + **kwargs) + + # copy the weight + with torch.no_grad(): + sharded_weight = shard_colwise(module.weight.data, process_group) + embedding.weight.copy_(sharded_weight) + + return embedding + + def reset_parameters(self, weight_initializer) -> None: + fan_in, fan_out = self.num_embeddings, self.embedding_dim + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + self._fill_padding_idx_with_zero() + + def _fill_padding_idx_with_zero(self) -> None: + if self.padding_idx is not None: + with torch.no_grad(): + self.weight[self.padding_idx].fill_(0) + + def forward(self, input_: Tensor) -> Tensor: + output_parallel = F.embedding(input_, self.weight, self.padding_idx, *self.embed_args, **self.embed_kwargs) + + if self.gather_output: + output = gather_forward_split_backward(output_parallel, dim=-1, process_group=self.process_group) + return output + else: + return output_parallel + + +class VocabParallelEmbedding1D(ParallelModule): + r"""Embedding parallelized in the vocabulary dimension. + + Args: + num_embeddings (int): number of embeddings. + embedding_dim (int): dimension of embedding. + padding_idx (int, optional): If specified, the entries at padding_idx do not contribute to the gradient; + therefore, the embedding vector at padding_idx is not updated during training, + i.e. it remains as a fixed “pad”, defaults to None. + dtype (:class:`torch.dtype`, optional): The dtype of parameters, defaults to None. + weight_initializer (:class:`typing.Callable`, optional): + he initializer of weight, defaults to normal initializer. + + The ``args`` and ``kwargs`` used in :class:``torch.nn.functional.embedding`` should contain: + :: + + max_norm (float, optional): If given, each embedding vector with norm larger than max_norm is + renormalized to have norm max_norm. Note: this will modify weight in-place. + norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default 2. + scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse + of frequency of the words in the mini-batch. Default False. + sparse (bool, optional): If True, gradient w.r.t. weight will be a sparse tensor. Default False. + + More details about ``args`` and ``kwargs`` could be found in + `Embedding `_. + + More details about initializer please refer to + `init `_. + """ + + def __init__(self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int = None, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + weight_initializer: Callable = init.normal_(), + *args, + **kwargs): + super().__init__() + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self.padding_idx = padding_idx + self.embed_args = args + self.embed_kwargs = kwargs + self.process_group = process_group + + tensor_parallel_size = dist.get_world_size(group=process_group) + tensor_parallel_rank = dist.get_rank(group=process_group) + + self.num_embeddings_per_partition = divide(num_embeddings, tensor_parallel_size) + self.num_embeddings = self.num_embeddings_per_partition + self.vocab_start_index = tensor_parallel_rank * self.num_embeddings_per_partition + self.vocab_end_index = self.vocab_start_index + self.num_embeddings_per_partition + + # parameter + factory_kwargs = {'device': device, 'dtype': dtype} + weight = torch.empty((num_embeddings, self.embedding_dim), **factory_kwargs) + sharded_weight = shard_rowwise(weight, process_group) + self.weight = sharded_tensor_to_param(sharded_weight) + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + self.reset_parameters(weight_initializer) + + @staticmethod + def from_native_module(module: nn.Embedding, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, + **kwargs) -> ParallelModule: + r""" + Convert a native pytorch embedding module to a parallel module. + """ + # get the origin attributes + num_embeddings = module.num_embeddings + embedding_dim = module.embedding_dim + padding_idx = module.padding_idx + device = module.weight.device + + # ensure only one process group is used + if isinstance(process_group, (list, tuple)): + assert len(process_group) == 1, \ + f'Expected only one process group, got {len(process_group)}.' + process_group = process_group[0] + + # create the parallel module + vocab_embedding_1d = VocabParallelEmbedding1D(num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=padding_idx, + device=device, + process_group=process_group, + *args, + **kwargs) + with torch.no_grad(): + # shard and slice the weight along the vocabulary(num_embeddings) dimension + # the shape of the weight is (num_embeddings, embedding_dim) + shard_weight = shard_rowwise(module.weight.data, process_group) + vocab_embedding_1d.weight.data.copy_(shard_weight) + + return vocab_embedding_1d + + def reset_parameters(self, weight_initializer) -> None: + with self.randomizer.fork_rng(enable_cpu=True): + fan_in, fan_out = self.num_embeddings, self.embedding_dim + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + self._fill_padding_idx_with_zero() + + def _fill_padding_idx_with_zero(self) -> None: + if self.padding_idx is not None and \ + self.padding_idx >= self.vocab_start_index and self.padding_idx < self.vocab_end_index: + with torch.no_grad(): + self.weight[self.padding_idx - self.vocab_start_index].fill_(0) + + def forward(self, input_: Tensor) -> Tensor: + # Build the mask. + input_mask = (input_ < self.vocab_start_index) | (input_ >= self.vocab_end_index) + # Mask the input. + masked_input = input_.clone() - self.vocab_start_index + masked_input[input_mask] = 0 + + output_parallel = F.embedding(masked_input, self.weight, self.padding_idx, *self.embed_args, + **self.embed_kwargs) + + # Mask the output embedding. + output_parallel[input_mask, :] = 0. + # Reduce across all the model parallel GPUs. + output = reduce_forward(output_parallel, self.process_group) + return output diff --git a/colossalai/shardformer/layer/linear.py b/colossalai/shardformer/layer/linear.py new file mode 100644 index 000000000000..26ba5883c64f --- /dev/null +++ b/colossalai/shardformer/layer/linear.py @@ -0,0 +1,336 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import math +from typing import Callable, List, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.distributed import ProcessGroup +from torch.nn.parameter import Parameter + +from colossalai.nn import init as init +from colossalai.nn.layer.utils import divide +from colossalai.tensor.d_tensor import shard_colwise, shard_rowwise, sharded_tensor_to_param + +from ._operation import ( + gather_forward_split_backward, + linear_with_async_comm, + reduce_forward, + split_forward_gather_backward, +) +from .parallel_module import ParallelModule +from .utils import create_randomizer_with_offset + +__all__ = ['Linear1D_Col', 'Linear1D_Row'] + + +class Linear1D_Col(ParallelModule): + r"""Linear layer with column parallelism. + + The linear layer is defined as :math:`Y = XA + b`. A is parallelized along + its second dimension as :math:`A = [A_1, ..., A_p]`. + + Args: + in_features (int): size of each input sample. + out_features (int): size of each output sample. + bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``. + dtype (`torch.dtype`): The dtype of parameters, defaults to None. + device (`torch.device`): The device of parameters, defaults to None. + process_group (`torch.distributed.ProcessGroup`): The process group to be used for weight sharding and communication, defaults to None. + gather_output (bool, optional): If true, call all-gather on output and make Y available + to all GPUs, otherwise, every GPU will have its output + which is :math:`Y_i = XA_i`, defaults to False + skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer, + which is preserved for kernel fusion, defaults to False + weight_initializer (`typing.Callable`): + The initializer of weight, defaults to kaiming uniform initializer. + bias_initializer (`typing.Callable`): + The initializer of bias, defaults to xavier uniform initializer. + + More details about ``initializer`` please refer to + `init `_. + """ + + def __init__(self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + gather_output: bool = False, + skip_bias_add: bool = False, + weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)), + bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1)): + super().__init__() + + # Keep input parameters + self.in_features = in_features + self.out_features = out_features + self.gather_output = gather_output + self.skip_bias_add = skip_bias_add + self.device = device + self.process_group = process_group + + if skip_bias_add and not bias: + raise ValueError('cannot skip bias addition if bias is None') + + # Parameters. + factory_kwargs = {'device': device, 'dtype': dtype} + + weight = torch.empty(self.out_features, self.in_features, **factory_kwargs) + sharded_weight = shard_rowwise(weight, self.process_group) + self.weight = sharded_tensor_to_param(sharded_weight) + + if bias: + bias = torch.empty(self.out_features, **factory_kwargs) + sharded_bias = shard_colwise(bias, self.process_group) + self.bias = sharded_tensor_to_param(sharded_bias) + else: + self.bias = None + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + + # init weights + self.reset_parameters(weight_initializer, bias_initializer) + + @staticmethod + def from_native_module(module: nn.Linear, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, + **kwargs) -> ParallelModule: + r""" + Convert a native PyTorch linear layer to a parallelized linear layer. + """ + # get the attributes + in_features = module.in_features + out_features = module.out_features + bias = module.bias is not None + device = module.weight.device + + # ensure only one process group is passed + if isinstance(process_group, (list, tuple)): + assert len(process_group) == 1, \ + f'Expected only one process group, got {len(process_group)}.' + process_group = process_group[0] + + linear_1d = Linear1D_Col(in_features=in_features, + out_features=out_features, + bias=bias, + device=device, + process_group=process_group, + *args, + **kwargs) + + with torch.no_grad(): + # the weigh to the linear layer is a transpose + # thus shard on row is equal to shard on column + sharded_weight = shard_rowwise(module.weight.data, process_group) + linear_1d.weight.data.copy_(sharded_weight) + if bias: + sharded_bias = shard_colwise(module.bias.data, process_group) + linear_1d.bias.copy_(sharded_bias) + return linear_1d + + def reset_parameters(self, weight_initializer, bias_initializer) -> None: + with self.randomizer.fork_rng(enable_cpu=True): + fan_in, fan_out = self.in_features, self.out_features + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + if self.bias is not None: + bias_initializer(self.bias, fan_in=fan_in) + + def forward(self, input_: Tensor) -> Tuple[Tensor, Tensor]: + assert input_.shape[-1] == self.weight.shape[-1], \ + 'Invalid shapes in Linear1D_Col forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1]) + + # Set up backprop all-reduce. + input_parallel = input_ + + # Matrix multiply. + bias = self.bias if not self.skip_bias_add else None + output_parallel = linear_with_async_comm(input_parallel, self.weight, bias, self.process_group, True) + + if self.gather_output: + # All-gather across the partitions. + output = gather_forward_split_backward(output_parallel, dim=-1, process_group=self.process_group) + else: + output = output_parallel + + if self.skip_bias_add: + return output, self.bias + else: + return output + + +class Linear1D_Row(ParallelModule): + r""" Linear layer with row parallelism + + Args: + in_features (int): size of each input sample. + out_features (int): size of each output sample. + bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``. + dtype (`torch.dtype`): The dtype of parameters, defaults to None. + parallel_input (bool): If set to ``True``, it's assumed that the input is split, defaults to False. + skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer, + which is preserved for kernel fusion, defaults to False + weight_initializer (:class:`typing.Callable`, optional): + The initializer of weight, defaults to kaiming uniform initializer. + bias_initializer (:class:`typing.Callable`, optional): + The initializer of bias, defaults to xavier uniform initializer. + + More details about ``initializer`` please refer to + `init `_. + """ + + def __init__(self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + parallel_input: bool = True, + skip_bias_add: bool = False, + weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)), + bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1), + stream_chunk_num: int = 1): + super().__init__() + + self.stream_chunk_num = stream_chunk_num + + # Keep input parameters + self.in_features = in_features + self.out_features = out_features + self.parallel_input = parallel_input + self.skip_bias_add = skip_bias_add + self.process_group = process_group + self.num_partitions = dist.get_world_size(self.process_group) + + if skip_bias_add and not bias: + raise ValueError('cannot skip bias addition if bias is None') + + # Parameters. + # Initialize weight. + factory_kwargs = {'device': device, 'dtype': dtype} + weight = torch.empty(self.out_features, self.in_features, **factory_kwargs) + sharded_weight = shard_colwise(weight, self.process_group) + self.weight = sharded_tensor_to_param(sharded_weight) + + if self.stream_chunk_num > 1: + # TODO() work for inference only + self.chunk_weight() + if bias: + self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs)) + else: + self.bias = None + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + + with self.randomizer.fork_rng(enable_cpu=True): + self.reset_parameters(weight_initializer, bias_initializer) + + @staticmethod + def from_native_module(module: nn.Linear, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, + **kwargs) -> ParallelModule: + r""" + Convert a native PyTorch linear layer to a parallelized linear layer. + """ + # get the attributes + in_features = module.in_features + out_features = module.out_features + bias = module.bias is not None + device = module.weight.device + + # ensure only one process group is passed + if isinstance(process_group, (list, tuple)): + assert len(process_group) == 1, \ + f'Expected only one process group, got {len(process_group)}.' + process_group = process_group[0] + + linear_1d = Linear1D_Row(in_features=in_features, + out_features=out_features, + bias=bias, + device=device, + process_group=process_group, + *args, + **kwargs) + + # TODO: copy the sharded weights + with torch.no_grad(): + # the weigh to the linear layer is a transpose + # thus shard on col is equal to shard on row + sharded_weight = shard_colwise(module.weight.data, process_group) + linear_1d.weight.data.copy_(sharded_weight) + + if bias: + linear_1d.bias.copy_(module.bias.data) + + return linear_1d + + def chunk_weight(self): + self.weight_list = torch.chunk(self.weight, self.stream_chunk_num, dim=0) + + @torch.no_grad() + def reset_parameters(self, weight_initializer, bias_initializer) -> None: + fan_in, fan_out = self.in_features, self.out_features + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + + if self.bias is not None: + bias_initializer(self.bias, fan_in=fan_in) + if self.process_group is None: + src_rank = 0 + else: + src_rank = dist.distributed_c10d._get_global_rank(self.process_group, 0) + + origin_device = self.bias.device + bias = self.bias.cuda() + dist.broadcast(bias, src=src_rank, group=self.process_group) + bias = bias.to(origin_device) + self.bias.copy_(bias) + + def forward(self, input_: Tensor) -> Tensor: + # Set up backprop all-reduce. + if self.parallel_input: + assert input_.shape[-1] == self.weight.shape[-1], \ + 'Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1]) + input_ = input_ + else: + assert divide(input_.shape[-1], self.num_partitions) == self.weight.shape[-1], \ + 'Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1] * self.num_partitions) + input_ = split_forward_gather_backward(input_, dim=-1, process_group=self.process_group) + + if self.stream_chunk_num > 1: + if self.training: + raise RuntimeError("use stream_chunk_num=1 in Linear1D_Row for training!") + with torch.no_grad(): + output_parallel_list = [None for i in range(self.stream_chunk_num)] + handle_list = [] + for i in range(self.stream_chunk_num): + output_parallel_list[i] = F.linear(input_, self.weight_list[i]) + handle = torch.distributed.all_reduce(output_parallel_list[i], + group=self.process_group, + async_op=True) + handle_list.append(handle) + # output_parallel_list[i] = reduce_input(output_parallel_list[i], ParallelMode.PARALLEL_1D) + for handle in handle_list: + handle.wait() + output = torch.cat(output_parallel_list, dim=-1) + else: + output_parallel = F.linear(input_, self.weight) + output = reduce_forward(output_parallel, self.process_group) + + if not self.skip_bias_add: + if self.bias is not None: + output = output + self.bias + return output + else: + return output, self.bias diff --git a/colossalai/shardformer/layer/loss.py b/colossalai/shardformer/layer/loss.py new file mode 100644 index 000000000000..7e3f6926b6d4 --- /dev/null +++ b/colossalai/shardformer/layer/loss.py @@ -0,0 +1,109 @@ +import torch +import torch.distributed as dist +from torch.autograd import Function +from torch.distributed import ProcessGroup + +__all__ = ['DistCrossEntropy', 'cross_entropy_1d'] + + +class DistCrossEntropy(Function): + r""" + Overwrite the forward and backward function to calculate the cross entropy loss before gather + + Args: + Function (:class:`torch.autograd.Function`): default + """ + + @staticmethod + def forward(ctx, vocab_logits: torch.Tensor, target: torch.Tensor, ignore_index: int, process_group: ProcessGroup): + r""" + Calculate the cross entropy loss before gather, the origin loss function is as follows: + loss = -log(exp(x[class])/sum(exp(x[i])) + and can be rewrite as: + loss = log(sum(exp(x[i])) - x[class] + + To avoid the `nan` of log(sum(exp(x[i]))), we minus the max of x[i] + + Args: + vocab_logits (:class:`torch.Tensor`): The logits of the vocabulary, shape is + [batch_size, seq_len, vocab_size] + labels (:class:`torch.Tensor`): The labels of the vocabulary, shape is + [batch_size, seq_len] + + Returns: + :class:`torch.Tensor`: The cross entropy loss + """ + # get the max + logits_max = torch.max(vocab_logits, dim=-1)[0] + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group) + + # minus the max to avoid the result of sum of exp is too large and the log is nan + vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1) + + # mask the target in the local device + partition_vocab_size = vocab_logits.size()[-1] + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + global_vocab_size = partition_vocab_size * world_size + + # [down, up) => false, other device and -100 => true + delta = (global_vocab_size + world_size - 1) // world_size + down_threshold = rank * delta + up_threshold = down_threshold + delta + mask = (target < down_threshold) | (target >= up_threshold) + masked_target = target.clone() - down_threshold + masked_target[mask] = 0 + + # reshape the logits and target + # reshape the vocab_logits to [bath_size * seq_len, vocab_size] + # reshape the labels to [bath_size * seq_len] + logits_2d = vocab_logits.view(-1, partition_vocab_size) + masked_target_1d = masked_target.view(-1) + + # extract the x[class] and set the x[other device] to zero + pred_logits_1d = logits_2d[torch.arange(start=0, end=logits_2d.shape[0], device=logits_2d.device), + masked_target_1d] + pred_logits_1d = pred_logits_1d.clone().contiguous() + pred_logits = pred_logits_1d.view_as(target) + pred_logits[mask] = 0.0 + + # allreduce the get all x(i,y) + dist.all_reduce(pred_logits, op=dist.ReduceOp.SUM, group=process_group) + exp_logits = vocab_logits + torch.exp(vocab_logits, out=exp_logits) + sum_exp_logits = torch.sum(exp_logits, dim=-1) + dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=process_group) + + # calculate the loss + # loss = log(sum(exp(x[i]))) - x[class] + loss = torch.where(target == ignore_index, 0.0, torch.log(sum_exp_logits) - pred_logits) + loss = torch.sum(loss).div_(torch.sum(loss != 0.0)) + + # calculate the softmax + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + ctx.save_for_backward(exp_logits, mask, masked_target_1d) + + return loss + + @staticmethod + def backward(ctx, grad_output): + # retrieve the saved tensors + exp_logits, mask, masked_target_1d = ctx.saved_tensors + + # use exp logits as the input grad + grad_logits = exp_logits + partion_vocab_size = grad_logits.shape[-1] + grad_logits_2d = grad_logits.view(-1, partion_vocab_size) + + update = 1.0 - mask.view(-1).float() + grad_logits_2d[torch.arange(0, grad_logits_2d.shape[0]), masked_target_1d] -= update + + grad_logits.mul_(grad_output.unsqueeze(dim=-1)) + return grad_logits, None, None + + +def cross_entropy_1d(vocab_logits: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + process_group: ProcessGroup = None) -> torch.Tensor: + return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group) diff --git a/colossalai/shardformer/layer/normalization.py b/colossalai/shardformer/layer/normalization.py new file mode 100644 index 000000000000..b27307154a76 --- /dev/null +++ b/colossalai/shardformer/layer/normalization.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn as nn + +__all__ = ['FusedLayerNorm', 'FusedRMSNorm'] + +FAST_LAYERNORM_SUPPORTED_SIZE = [ + 1024, 1536, 2048, 2304, 3072, 3840, 4096, 5120, 6144, 8192, 10240, 12288, 12800, 15360, 16384, 18432, 20480, 24576, + 25600, 30720, 32768, 40960, 49152, 65536 +] + + +class FusedLayerNorm(): + r""" + This is a wrapper around the apex fused layernorm implementation. It is meant to be used only with the from_native_module interface. + """ + + def __init__(self) -> None: + raise NotImplementedError( + 'FusedLayerNorm is not implemented as a physical class. ' + 'It is meant to be used only with the from_native_module interface to wrap the fused layernorm implementation provided by apex.' + ) + + @staticmethod + def from_native_module(module: nn.LayerNorm, *args, **kwargs) -> nn.Module: + r""" + Convert a native pytorch layer norm module to colossalai layer norm module + """ + # check if apex is installed + try: + import apex + except ImportError: + raise ImportError( + 'Please install apex from source (https://github.com/NVIDIA/apex) to use the fused layernorm kernel') + + # get the attributes of the module + normalized_shape = module.normalized_shape + eps = module.eps + elementwise_affine = module.elementwise_affine + dtype = module.weight.dtype + device = module.weight.device + + # pick the suitable layernorm implementation + use_fast_ln = normalized_shape in FAST_LAYERNORM_SUPPORTED_SIZE + + if use_fast_ln: + try: + from apex.contrib.layer_norm.layer_norm import FastLayerNorm as ApexFusedLayerNorm + except ImportError: + # fall back to the normal fused layernorm is not built + from apex.normalization import FusedLayerNorm as ApexFusedLayerNorm + else: + from apex.normalization import FusedLayerNorm as ApexFusedLayerNorm + + layernorm = ApexFusedLayerNorm(normalized_shape, eps=eps, + elementwise_affine=elementwise_affine).to(dtype).to(device) + + with torch.no_grad(): + # copy weight and bias + layernorm.weight.copy_(module.weight) + layernorm.bias.copy_(module.bias) + return layernorm + + +class FusedRMSNorm(): + """ + This is a wrapper around the apex fused rms norm implementation. It is meant to be used only with the from_native_module interface. + """ + + def __init__(self) -> None: + raise NotImplementedError( + 'FusedRMSNorm is not implemented as a physical class. ' + 'It is meant to be used only with the from_native_module interface to wrap the fused rms norm implementation provided by apex.' + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + try: + from apex.normalization import FusedRMSNorm as ApexFusedRMSNorm + except ImportError: + raise ImportError( + 'Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMS normalization kernel' + ) + + # to check if it is huggingface LlamaRMSNorm + if module.__class__.__name__ == "LlamaRMSNorm": + normalized_shape = module.weight.shape[0] + eps = module.variance_epsilon + elementwise_affine = True + else: + # get the attributes of the module + normalized_shape = module.normalized_shape + eps = module.eps + elementwise_affine = module.elementwise_affine + + rmsnorm = ApexFusedRMSNorm(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine) + + with torch.no_grad(): + # copy weight and bias + rmsnorm.weight.copy_(module.weight) + + return rmsnorm diff --git a/colossalai/shardformer/layer/parallel_module.py b/colossalai/shardformer/layer/parallel_module.py new file mode 100644 index 000000000000..bda147b121ab --- /dev/null +++ b/colossalai/shardformer/layer/parallel_module.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import itertools +from abc import ABC, abstractmethod +from typing import List, Union + +import torch +import torch.nn as nn +from torch.distributed import ProcessGroup +from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, Module + +from colossalai.tensor.d_tensor import ( + distribute_tensor, + distribute_tensor_with_customization, + get_device_mesh, + get_sharding_spec, + is_customized_distributed_tensor, + is_distributed_tensor, + sharded_tensor_to_param, + to_global, + to_global_for_customized_distributed_tensor, +) + +__all__ = ['ParallelModule'] + + +class ParallelModule(nn.Module, ABC): + + @abstractmethod + def from_native_module(module: nn.Module, + process_group: Union[ProcessGroup, List[ProcessGroup]] = None) -> "ParallelModule": + """ + Convert a native PyTorch module to a parallelized module. + + Args: + module (nn.Module): the module to be converted. + process_group (ProcessGroup or list[ProcessGroup]): the process group(s) to be used for communication. + If this is a list, the process group at the ith index of the list will correspond to the process group + in the ith axis of the device mesh. Defaults to None, which means the global process group. + """ + pass + + def _save_to_state_dict(self, destination, prefix, keep_vars): + r"""Saves module state to `destination` dictionary, containing a state + of the module, but not its descendants. This is called on every + submodule in :meth:`~torch.nn.Module.state_dict`. + + In rare cases, subclasses can achieve class-specific behavior by + overriding this method with custom logic. + + Args: + destination (dict): a dict where state will be stored + prefix (str): the prefix for parameters and buffers used in this + module + """ + for name, param in self._parameters.items(): + if param is not None: + param_ = param if keep_vars else param.detach() + if is_distributed_tensor(param_): + destination[prefix + name] = to_global(param_) + elif is_customized_distributed_tensor(param_): + destination[prefix + name] = to_global_for_customized_distributed_tensor(param_) + else: + destination[prefix + name] = param_ + + for name, buf in self._buffers.items(): + if buf is not None and name not in self._non_persistent_buffers_set: + destination[prefix + name] = buf if keep_vars else buf.detach() + extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX + if getattr(self.__class__, "get_extra_state", Module.get_extra_state) is not Module.get_extra_state: + destination[extra_state_key] = self.get_extra_state() + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, + error_msgs): + r"""Copies parameters and buffers from :attr:`state_dict` into only + this module, but not its descendants. This is called on every submodule + in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this + module in input :attr:`state_dict` is provided as :attr:`local_metadata`. + For state dicts without metadata, :attr:`local_metadata` is empty. + Subclasses can achieve class-specific backward compatible loading using + the version number at `local_metadata.get("version", None)`. + + .. note:: + :attr:`state_dict` is not the same object as the input + :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So + it can be modified. + + Args: + state_dict (dict): a dict containing parameters and + persistent buffers. + prefix (str): the prefix for parameters and buffers used in this + module + local_metadata (dict): a dict containing the metadata for this module. + See + strict (bool): whether to strictly enforce that the keys in + :attr:`state_dict` with :attr:`prefix` match the names of + parameters and buffers in this module + missing_keys (list of str): if ``strict=True``, add missing keys to + this list + unexpected_keys (list of str): if ``strict=True``, add unexpected + keys to this list + error_msgs (list of str): error messages should be added to this + list, and will be reported together in + :meth:`~torch.nn.Module.load_state_dict` + """ + for hook in self._load_state_dict_pre_hooks.values(): + hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + + persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set} + local_name_params = itertools.chain(self._parameters.items(), persistent_buffers.items()) + local_state = {k: v for k, v in local_name_params if v is not None} + + for name, param in local_state.items(): + key = prefix + name + + if key in state_dict: + input_param = state_dict[key] + if not torch.overrides.is_tensor_like(input_param): + error_msgs.append('While copying the parameter named "{}", ' + 'expected torch.Tensor or Tensor-like object from checkpoint but ' + 'received {}'.format(key, type(input_param))) + continue + + if is_distributed_tensor(param): + # shard the input param + device_mesh = get_device_mesh(param) + sharding_spec = get_sharding_spec(param) + sharded_tensor = distribute_tensor(input_param, device_mesh, sharding_spec) + input_param = sharded_tensor_to_param(sharded_tensor) + elif is_customized_distributed_tensor(param): + input_param = distribute_tensor_with_customization(input_param, param.shard_fn, param.gather_fn) + + # This is used to avoid copying uninitialized parameters into + # non-lazy modules, since they dont have the hook to do the checks + # in such case, it will error when accessing the .shape attribute. + is_param_lazy = torch.nn.parameter.is_lazy(param) + # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+ + if not is_param_lazy and len(param.shape) == 0 and len(input_param.shape) == 1: + input_param = input_param[0] + + if not is_param_lazy and input_param.shape != param.shape: + # local shape should match the one in checkpoint + error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, ' + 'the shape in current model is {}.'.format(key, input_param.shape, param.shape)) + continue + + try: + with torch.no_grad(): + param.copy_(input_param) + except Exception as ex: + error_msgs.append('While copying the parameter named "{}", ' + 'whose dimensions in the model are {} and ' + 'whose dimensions in the checkpoint are {}, ' + 'an exception occurred : {}.'.format(key, param.size(), input_param.size(), + ex.args)) + elif strict: + missing_keys.append(key) + + extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX + if getattr(self.__class__, "set_extra_state", Module.set_extra_state) is not Module.set_extra_state: + if extra_state_key in state_dict: + self.set_extra_state(state_dict[extra_state_key]) + elif strict: + missing_keys.append(extra_state_key) + elif strict and (extra_state_key in state_dict): + unexpected_keys.append(extra_state_key) + + if strict: + for key in state_dict.keys(): + if key.startswith(prefix) and key != extra_state_key: + input_name = key[len(prefix):] + input_name = input_name.split('.', 1)[0] # get the name of param/buffer/child + if input_name not in self._modules and input_name not in local_state: + unexpected_keys.append(key) diff --git a/colossalai/shardformer/layer/qkv_fused_linear.py b/colossalai/shardformer/layer/qkv_fused_linear.py new file mode 100644 index 000000000000..9d51670c65dd --- /dev/null +++ b/colossalai/shardformer/layer/qkv_fused_linear.py @@ -0,0 +1,473 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import math +from typing import Callable, List, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.distributed import ProcessGroup +from torch.nn.parameter import Parameter + +from colossalai.nn import init as init +from colossalai.nn.layer.utils import divide +from colossalai.tensor.d_tensor.api import ( + customized_distributed_tensor_to_param, + distribute_tensor_with_customization, + shard_rowwise, + sharded_tensor_to_param, +) + +from ._operation import ( + gather_forward_split_backward, + matmul_with_async_comm, + reduce_backward, + reduce_forward, + split_forward_gather_backward, +) +from .parallel_module import ParallelModule +from .utils import create_randomizer_with_offset + +__all__ = ['FusedLinear1D_Col', 'FusedLinear1D_Row'] + +# ==================================== +# For GPT Only +# ==================================== + + +def split_fused_qkv_in_gpt2_style(qkv: torch.Tensor, + n_fused: int, + process_group: ProcessGroup, + is_transposed: bool = False): + """ + The fused qkv tensor looks like [Q1, Q2, K1, K2, V1, V2], this function will split them into [Q1, K1, V1] and [Q2, K2, V2]. + + Args: + qkv (torch.Tensor): The fused qkv tensor. + n_fused (int): The number items fused together, defaults to 3 (query, key and value). + process_group (ProcessGroup): The process group for distributed communication. + is_transposed (bool): generally the tensor is the shape of (out_features, in_features). Set this to True if the tensor is in the shape (in_features, out_features). + """ + # get the number of slice for the fused qkv + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + order = torch.arange(world_size * n_fused) + + # split the fused qkv + # from + # [Q, K, V] + # to + # [Q1, Q2, K1, K2, V1, V2] + if is_transposed: + weight_chunks = torch.chunk(qkv, world_size * n_fused, dim=-1) + else: + weight_chunks = torch.chunk(qkv, world_size * n_fused, dim=0) + + # rearrange the slice into the final order + # from + # [Q1, Q2, K1, K2, V1, V2] + # to + # [Q1, K1, V1], [Q2, K2, V2] + weight_chunks_of_current_rank = [weight_chunks[i] for i in order[rank::world_size]] + + if is_transposed: + weight_of_current_rank = torch.cat(weight_chunks_of_current_rank, dim=-1) + else: + weight_of_current_rank = torch.cat(weight_chunks_of_current_rank, dim=0) + return weight_of_current_rank + + +def gather_fused_qkv_in_gpt2_style(qkv: torch.Tensor, + n_fused: int, + process_group: ProcessGroup, + is_transposed: bool = False): + """ + The splitted qkv tensor looks like [Q1, K1, V1] and [Q2, K2, V2], this function will gather them into [Q1, Q2, K1, K2, V1, V2]. + + Args: + qkv (torch.Tensor): The fused qkv tensor. + n_fused (int): The number items fused together, defaults to 3 (query, key and value). + process_group (ProcessGroup): The process group for distributed communication. + is_transposed (bool): generally the tensor is the shape of (out_features, in_features). Set this to True if the tensor is in the shape (in_features, out_features). + """ + world_size = dist.get_world_size(group=process_group) + + # gather the tensors + # from + # [Q1, K1, V1], [Q2, K2, V2] + # to + # [Q1, K1, V1, Q2, K2, V2] + origin_device = qkv.device + qkv = qkv.cuda() + gather_list = [torch.zeros_like(qkv) for _ in range(world_size)] + dist.all_gather(gather_list, qkv, group=process_group) + + if is_transposed: + gather_weight = torch.cat(gather_list, dim=-1) + else: + gather_weight = torch.cat(gather_list, dim=0) + gather_weight = gather_weight.to(origin_device) + qkv = qkv.to(origin_device) + + # rearrange the tensor slices + # from + # [Q1, K1, V1, Q2, K2, V2] + # to + # [Q1, Q2, K1, K2, V1, V2] + if is_transposed: + weight_chunks = torch.chunk(gather_weight, world_size * n_fused, dim=-1) + else: + weight_chunks = torch.chunk(gather_weight, world_size * n_fused, dim=0) + + reordered_chunk_list = [] + for i in range(n_fused): + reordered_chunk_list.extend(weight_chunks[i::n_fused]) + + if is_transposed: + reordered_gather_weight = torch.cat(reordered_chunk_list, dim=-1) + else: + reordered_gather_weight = torch.cat(reordered_chunk_list, dim=0) + return reordered_gather_weight + + +class GPT2FusedLinearConv1D_Col(ParallelModule): + r"""Linear layer with column parallelism. + + The linear layer is defined as :math:`Y = XA + b`. A is parallelized along + its second dimension as :math:`A = [A_1, ..., A_p]`. This layer is used to fit `Conv1D` layer (Fused QKV) in gpt2 of huggingface. + + Args: + in_features (int): size of each input sample. + out_features (int): size of each output sample. + bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``. + dtype (`torch.dtype`): The dtype of parameters, defaults to None. + device (`torch.device`): The device of parameters, defaults to None. + n_fused (int): The number items fused, defaults to 3 (QKV). + process_group (`torch.distributed.ProcessGroup`): The process group to be used for weight sharding and communication, defaults to None. + gather_output (bool, optional): If true, call all-gather on output and make Y available + to all GPUs, otherwise, every GPU will have its output + which is :math:`Y_i = XA_i`, defaults to False + skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer, + which is preserved for kernel fusion, defaults to False + weight_initializer (`typing.Callable`): + The initializer of weight, defaults to kaiming uniform initializer. + bias_initializer (`typing.Callable`): + The initializer of bias, defaults to xavier uniform initializer. + + More details about ``initializer`` please refer to + `init `_. + """ + + def __init__(self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + async_communication: bool = False, + gather_output: bool = False, + skip_bias_add: bool = False, + n_fused: int = 3, + weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)), + bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1)): + super().__init__() + + # Keep input parameters + self.in_features = in_features + self.out_features = out_features + self.gather_output = gather_output + self.skip_bias_add = skip_bias_add + self.device = device + self.n_fused = n_fused + self.process_group = process_group + self.async_communication = async_communication + + if skip_bias_add and not bias: + raise ValueError('cannot skip bias addition if bias is None') + + # Parameters. + # Initialize weight. + factory_kwargs = {'device': device, 'dtype': dtype} + weight = torch.empty(self.in_features, self.out_features, **factory_kwargs) + + def shard_fn(tensor): + return split_fused_qkv_in_gpt2_style(tensor, self.n_fused, self.process_group, True) + + def gather_fn(tensor): + return gather_fused_qkv_in_gpt2_style(tensor, 3, self.process_group, True) + + with torch.no_grad(): + sharded_weight = distribute_tensor_with_customization(weight, shard_fn, gather_fn) + self.weight = customized_distributed_tensor_to_param(sharded_weight) + + if bias: + bias = torch.empty(self.out_features, **factory_kwargs) + + with torch.no_grad(): + sharded_bias = distribute_tensor_with_customization(bias, shard_fn, gather_fn) + self.bias = customized_distributed_tensor_to_param(sharded_bias) + else: + self.bias = None + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + + # init weights + self.reset_parameters(weight_initializer, bias_initializer) + + @staticmethod + def from_native_module(module: nn.Module, process_group: Union[ProcessGroup, List[ProcessGroup]], n_fused: int, + *args, **kwargs) -> ParallelModule: + r""" + Convert a huggingface layer `Conv1D` in gpt2 to a parallelized linear layer. + + Args: + module (`nn.Linear`): The module to be converted. + process_group (`Union[ProcessGroup, List[ProcessGroup]]`): The process group to be used for weight sharding and communication. + n_fused (int): The number of layers to be fused. In GPT2, Q,K,V are fused in one weight. + """ + # get the attributes + in_features = module.weight.shape[0] + out_features = module.weight.shape[1] + bias = module.bias is not None + device = module.weight.device + + # ensure only one process group is passed + if isinstance(process_group, (list, tuple)): + assert len(process_group) == 1, \ + f'Expected only one process group, got {len(process_group)}.' + process_group = process_group[0] + + linear_1d = GPT2FusedLinearConv1D_Col(in_features=in_features, + out_features=out_features, + bias=bias, + device=device, + process_group=process_group, + *args, + **kwargs) + + # TODO: copy the sharded weights + with torch.no_grad(): + sharded_weight = split_fused_qkv_in_gpt2_style(module.weight.data, + n_fused=n_fused, + process_group=process_group, + is_transposed=True) + linear_1d.weight.data.copy_(sharded_weight.data) + + if bias: + sharded_bias = split_fused_qkv_in_gpt2_style(module.bias.data, + n_fused=n_fused, + process_group=process_group, + is_transposed=True) + linear_1d.bias.data.copy_(sharded_bias.data) + + return linear_1d + + def reset_parameters(self, weight_initializer, bias_initializer) -> None: + with self.randomizer.fork_rng(enable_cpu=True): + fan_in, fan_out = self.in_features, self.out_features + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + if self.bias is not None: + bias_initializer(self.bias, fan_in=fan_in) + + def forward(self, input_: Tensor) -> Tuple[Tensor, Tensor]: + assert input_.shape[-1] == self.weight.shape[0], \ + 'Invalid shapes in Linear1D_Col forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1]) + # Set up backprop all-reduce. + input_parallel = reduce_backward(input_, self.process_group) + # input_parallel = input_ + + # Matrix multiply. + bias = self.bias if not self.skip_bias_add else None + + output_parallel = matmul_with_async_comm(input_parallel, self.weight, bias, self.process_group, + self.async_communication) + + if self.gather_output: + # All-gather across the partitions. + output = gather_forward_split_backward(output_parallel, dim=-1, process_group=self.process_group) + else: + output = output_parallel + + if self.skip_bias_add: + return output, self.bias + else: + return output + + +class GPT2FusedLinearConv1D_Row(ParallelModule): + r""" Linear layer with row parallelism. + This layer is used to fit `Conv1D` layer (Fused QKV) in gpt2 of huggingface. + + Args: + in_features (int): size of each input sample. + out_features (int): size of each output sample. + bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``. + dtype (`torch.dtype`): The dtype of parameters, defaults to None. + parallel_input (bool): If set to ``True``, it's assumed that the input is split, defaults to False. + skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer, + which is preserved for kernel fusion, defaults to False + weight_initializer (:class:`typing.Callable`, optional): + The initializer of weight, defaults to kaiming uniform initializer. + bias_initializer (:class:`typing.Callable`, optional): + The initializer of bias, defaults to xavier uniform initializer. + + More details about ``initializer`` please refer to + `init `_. + """ + + def __init__(self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + process_group: ProcessGroup = None, + parallel_input: bool = True, + skip_bias_add: bool = False, + weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)), + bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1), + stream_chunk_num: int = 1): + super().__init__() + + self.stream_chunk_num = stream_chunk_num + + # Keep input parameters + self.in_features = in_features + self.out_features = out_features + self.parallel_input = parallel_input + self.skip_bias_add = skip_bias_add + self.process_group = process_group + self.num_partitions = dist.get_world_size(self.process_group) + + if skip_bias_add and not bias: + raise ValueError('cannot skip bias addition if bias is None') + + # Divide the weight matrix along the last dimension. + self.input_size_per_partition = divide(in_features, self.num_partitions) + + # Parameters. + # Initialize weight. + factory_kwargs = {'device': device, 'dtype': dtype} + weight = torch.empty(self.in_features, self.out_features, **factory_kwargs) + sharded_weight = shard_rowwise(weight, self.process_group) + self.weight = sharded_tensor_to_param(sharded_weight) + + if self.stream_chunk_num > 1: + # TODO() work for inference only + self.chunk_weight() + if bias: + self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs)) + else: + self.bias = None + + # offset the seed with randomizer index and rank + seed = torch.random.initial_seed() + self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group) + + # init weights + self.reset_parameters(weight_initializer, bias_initializer) + + @staticmethod + def from_native_module(module: nn.Linear, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, + **kwargs) -> ParallelModule: + r""" + Convert a native PyTorch linear layer to a parallelized linear layer. + """ + # get the attributes + in_features = module.weight.shape[0] + out_features = module.weight.shape[1] + bias = module.bias is not None + device = module.weight.device + + # ensure only one process group is passed + if isinstance(process_group, (list, tuple)): + assert len(process_group) == 1, \ + f'Expected only one process group, got {len(process_group)}.' + process_group = process_group[0] + + linear_1d = GPT2FusedLinearConv1D_Row(in_features=in_features, + out_features=out_features, + bias=bias, + device=device, + process_group=process_group, + *args, + **kwargs) + + # TODO: copy the sharded weights + with torch.no_grad(): + # the weigh to the linear layer is a transpose + # thus shard on col is equal to shard on row + sharded_weight = shard_rowwise(module.weight.data, process_group) + linear_1d.weight.data.copy_(sharded_weight.data) + + if bias: + linear_1d.bias.copy_(module.bias.data) + + return linear_1d + + def chunk_weight(self): + self.weight_list = torch.chunk(self.weight, self.stream_chunk_num, dim=0) + + def reset_parameters(self, weight_initializer, bias_initializer) -> None: + with self.randomizer.fork_rng(enable_cpu=True): + fan_in, fan_out = self.in_features, self.out_features + weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out) + + if self.bias is not None: + bias_initializer(self.bias, fan_in=fan_in) + if self.process_group is None: + src_rank = 0 + else: + src_rank = dist.distributed_c10d._get_global_rank(self.process_group, 0) + + origin_device = self.bias.device + self.bias = self.bias.cuda() + dist.broadcast(self.bias, src=src_rank, group=self.process_group) + self.bias = self.bias.to(origin_device) + + def forward(self, input_: Tensor) -> Tensor: + # Set up backprop all-reduce. + if self.parallel_input: + assert input_.shape[-1] == self.weight.shape[0], \ + 'Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1]) + input_ = input_ + else: + assert divide(input_.shape[-1], self.num_partitions) == self.weight.shape[0], \ + 'Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.'.format( + input_.shape, self.weight.shape, self.weight.shape[-1] * self.num_partitions) + input_ = split_forward_gather_backward(input_, dim=-1, process_group=self.process_group) + + if self.stream_chunk_num > 1: + if self.training: + raise RuntimeError("use stream_chunk_num=1 in Linear1D_Row for training!") + with torch.no_grad(): + output_parallel_list = [None for i in range(self.stream_chunk_num)] + handle_list = [] + for i in range(self.stream_chunk_num): + output_parallel_list[i] = torch.matmul(input_, self.weight_list[i]) + handle = torch.distributed.all_reduce(output_parallel_list[i], + group=self.process_group, + async_op=True) + handle_list.append(handle) + # output_parallel_list[i] = reduce_input(output_parallel_list[i], ParallelMode.PARALLEL_1D) + for handle in handle_list: + handle.wait() + output = torch.cat(output_parallel_list, dim=-1) + else: + output_parallel = torch.matmul(input_, self.weight) + output = reduce_forward(output_parallel, self.process_group) + + if not self.skip_bias_add: + if self.bias is not None: + output = output + self.bias + return output + else: + return output, self.bias diff --git a/colossalai/shardformer/layer/utils.py b/colossalai/shardformer/layer/utils.py new file mode 100644 index 000000000000..f2ac6563c46f --- /dev/null +++ b/colossalai/shardformer/layer/utils.py @@ -0,0 +1,202 @@ +from contextlib import contextmanager + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.distributed.distributed_c10d import _get_global_rank + + +class Randomizer: + """ + Randomizer enables the program to be executed under a different seed within the context. + + Example: + + ```python + randomizer = Randomizer(seed=1024) + + with randomizer.fork(): + # do something here with seed 1024 + do_something() + ``` + + Args: + seed (int): The random seed to set. + enable_cpu (bool): fork the CPU RNG state as well. + with_index (bool): whether to use the index of the randomizer. + """ + + _INDEX = 0 + + def __init__(self, seed: int): + # TODO: remove colossalai.context.random + + self.seed = seed + + # Handle CUDA rng state + # 1. get the current rng state + # 2. set the seed and store the rng state + # 3. recover the original rng state + cuda_original_rng_state = torch.cuda.get_rng_state() + torch.cuda.manual_seed(seed) + self.cuda_rng_state = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(cuda_original_rng_state) + + # to the same for cpu rng state + cpu_original_rng_state = torch.get_rng_state() + torch.manual_seed(seed) + self.cpu_rng_state = torch.get_rng_state() + torch.set_rng_state(cpu_original_rng_state) + + def _set_cuda_rng_state(self, rng_state): + torch.cuda.set_rng_state(rng_state) + + def _get_cuda_rng_state(self): + current_state = torch.cuda.get_rng_state() + return current_state + + def _set_cpu_rng_state(self, rng_state): + torch.set_rng_state(rng_state) + + def _get_cpu_rng_state(self): + current_state = torch.get_rng_state() + return current_state + + @contextmanager + def fork_rng(self, enable_cpu: bool = False): + """ + This is a context manager to change the dropout state and recover the original state. + + Usage: + :: + >>> with _seed_manager.dropout_mode(): + >>> input = super().forward(input) + """ + try: + current_cuda_rng_state = self._get_cuda_rng_state() + self._set_cuda_rng_state(self.cuda_rng_state) + + if enable_cpu: + current_cpu_rng_state = self._get_cpu_rng_state() + self._set_cpu_rng_state(self.cpu_rng_state) + yield + finally: + self.cuda_rng_state = self._get_cuda_rng_state() + self._set_cuda_rng_state(current_cuda_rng_state) + + if enable_cpu: + self.cpu_rng_state = self._get_cpu_rng_state() + self._set_cpu_rng_state(current_cpu_rng_state) + + @staticmethod + def index(): + """ + Return the index of the randomizer. The index is useful when the user wants + to introduce some randomness in the program. + + Note: + The index will increment by one each time this method is called. + + Example: + + ```python + # assume we need a randomizer to init the weight of different layers + # we can use the index of the randomizer to do so that + # each layer has its own randomizer with a different seed + base_seed = torch.random.initial_seed() + seed = base_seed + Randomizer.index() + randomizer = Randomizer(seed) + + with randomizer.fork(): + init_weights() + ``` + + """ + idx = Randomizer._INDEX + return idx + + @staticmethod + def increment_index(): + """ + Increment the index of the randomizer by one. + """ + Randomizer._INDEX += 1 + + @staticmethod + def is_randomizer_index_synchronized(process_group: ProcessGroup = None): + """ + Return whether the randomizer index is synchronized across processes. + """ + index = Randomizer.index() + if dist.is_initialized(): + # convert the index to tensor + index_tensor = torch.tensor(index, dtype=torch.int32).cuda() + + # all gather the index + gathered_index = [torch.zeros_like(index_tensor) for _ in range(dist.get_world_size(process_group))] + dist.all_gather(gathered_index, index_tensor, process_group) + + # make sure all the gathered index are the same + for i in range(1, dist.get_world_size(process_group)): + if gathered_index[i] != gathered_index[0]: + return False + + return True + + @staticmethod + def synchronize_index(process_group: ProcessGroup = None): + """ + All gather the index and pick the largest value. + """ + index = Randomizer.index() + + if dist.is_initialized(): + # convert the index to tensor + index_tensor = torch.tensor(index, dtype=torch.int32).cuda() + + # all gather the index + gathered_index = [torch.zeros_like(index_tensor) for _ in range(dist.get_world_size(process_group))] + dist.all_gather(gathered_index, index_tensor, process_group) + + # pick the largest index + for i in range(1, dist.get_world_size(process_group)): + if gathered_index[i] > index_tensor: + index_tensor = gathered_index[i] + + # set the index + Randomizer._INDEX = index_tensor.item() + + +def create_randomizer_with_offset(seed: int, + process_group: ProcessGroup = None, + offset_by_rank: bool = True, + offset_by_index: bool = True): + """ + Create a randomizer with an offset. The offset is equal to the rank of the process and the index of the randomizer. + + Args: + seed (int): The base random seed to set. + process_group (ProcessGroup): the process group to get the rank from. + offset_by_rank (bool): whether to offset by the rank of the process, i.e., the rank of the process will be added to the seed. Default: True. + offset_by_index (bool): whether to offset by the index of the randomizer, i.e., the index of the randomizer will be added to the seed. Default: True. + + Returns: + Randomizer: the randomizer with offset. + """ + base_seed = seed + + if offset_by_rank and dist.is_initialized(): + rank = dist.get_rank(process_group) + base_seed += rank + + if offset_by_index: + # check if the randomizer index is synchronized + is_synchronized = Randomizer.is_randomizer_index_synchronized(process_group) + assert is_synchronized, ("We detect that the randomizer index is not synchronized across processes." + "This is not allowed when we want to create a randomizer with offset by index." + "Please call Randomizer.synchronize_index() first.") + + base_seed += Randomizer.index() + Randomizer.increment_index() + + return Randomizer(seed=base_seed) diff --git a/colossalai/shardformer/modeling/__init__.py b/colossalai/shardformer/modeling/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/colossalai/shardformer/modeling/bloom.py b/colossalai/shardformer/modeling/bloom.py new file mode 100644 index 000000000000..a3d774ff2abb --- /dev/null +++ b/colossalai/shardformer/modeling/bloom.py @@ -0,0 +1,69 @@ +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + + +def build_bloom_alibi_tensor_fn(process_group: ProcessGroup) -> torch.Tensor: + + def build_bloom_alibi_tensor(self, attention_mask: torch.Tensor, num_heads: int, + dtype: torch.dtype) -> torch.Tensor: + """ + Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value + `softmax(l+a) = softmax(l)`. Based on + https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 + TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. + + Args: + Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) + attention_mask (`torch.Tensor`): + Token-wise attention mask, this should be of shape (batch_size, max_seq_len). + num_heads (`int`, *required*): + number of heads + dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): + dtype of the output tensor + """ + import math + + if dist.is_initialized(): + world_size = dist.get_world_size(process_group) + num_heads = num_heads * world_size + + batch_size, seq_length = attention_mask.shape + closest_power_of_2 = 2**math.floor(math.log2(num_heads)) + base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != num_heads: + extra_base = torch.tensor(2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))), + device=attention_mask.device, + dtype=torch.float32) + num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) + extra_powers = torch.arange(1, + 1 + 2 * num_remaining_heads, + 2, + device=attention_mask.device, + dtype=torch.int32) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + + # Note: alibi will added to the attention bias that will be applied to the query, key product of attention + # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length) + # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length) + # => the query_length dimension will then be broadcasted correctly + # This is more or less identical to T5's relative position bias: + # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527 + arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] + alibi = slopes[..., None] * arange_tensor + if dist.is_initialized(): + num_heads_per_rank = int(num_heads / dist.get_world_size(process_group)) + offset = dist.get_rank(process_group) * num_heads_per_rank + alibi = alibi.view(batch_size, num_heads, 1, seq_length) + alibi = alibi[:, offset:num_heads_per_rank + offset, :, :] + return alibi.reshape(batch_size * num_heads_per_rank, 1, seq_length).to(dtype) + else: + return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype) + + return build_bloom_alibi_tensor diff --git a/colossalai/shardformer/policies/__init__.py b/colossalai/shardformer/policies/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/colossalai/shardformer/policies/autopolicy.py b/colossalai/shardformer/policies/autopolicy.py new file mode 100644 index 000000000000..085e3150c697 --- /dev/null +++ b/colossalai/shardformer/policies/autopolicy.py @@ -0,0 +1,137 @@ +import importlib +from dataclasses import dataclass + +import torch.nn as nn + +from .basepolicy import Policy + +__all__ = ["PolicyLocation", "get_autopolicy", "import_policy"] + + +@dataclass +class PolicyLocation: + """ + PolicyLocation describes the location of a policy class. + + Args: + file_name (str): The file name of the policy under colossalai.shardformer.policies + class_name (str): The class name of the policy class + """ + file_name: str + class_name: str + + +# we don't want to import all policies here +# as each policy file imports its own model zoo library +# we will allow the user to only import the policy file needed +_POLICY_LIST = { + # BERT + "transformers.models.bert.modeling_bert.BertModel": + PolicyLocation(file_name="bert", class_name="BertModelPolicy"), + "transformers.models.bert.modeling_bert.BertForPreTraining": + PolicyLocation(file_name="bert", class_name="BertForPretrainingPolicy"), + "transformers.models.bert.modeling_bert.BertLMHeadModel": + PolicyLocation(file_name="bert", class_name="BertLMHeadModelPolicy"), + "transformers.models.bert.modeling_bert.BertForMaskedLM": + PolicyLocation(file_name="bert", class_name="BertForMaskedLMPolicy"), + "transformers.models.bert.modeling_bert.BertForSequenceClassification": + PolicyLocation(file_name="bert", class_name="BertForSequenceClassificationPolicy"), + "transformers.models.bert.modeling_bert.BertForTokenClassification": + PolicyLocation(file_name="bert", class_name="BertForTokenClassificationPolicy"), + "transformers.models.bert.modeling_bert.BertForNextSentencePrediction": + PolicyLocation(file_name="bert", class_name="BertForNextSentencePredictionPolicy"), + "transformers.models.bert.modeling_bert.BertForMultipleChoice": + PolicyLocation(file_name="bert", class_name="BertForMultipleChoicePolicy"), + + # LLaMA + "transformers.models.llama.modeling_llama.LlamaModel": + PolicyLocation(file_name="llama", class_name="LlamaPolicy"), + "transformers.models.llama.modeling_llama.LlamaForCausalLM": + PolicyLocation(file_name="llama", class_name="LlamaForCausalLMPolicy"), + "transformers.models.llama.modeling_llama.LlamaForSequenceClassification": + PolicyLocation(file_name="llama", class_name="LlamaForSequenceClassificationPolicy"), + + # T5 + "transformers.models.t5.modeling_t5.T5Model": + PolicyLocation(file_name="t5", class_name="T5ModelPolicy"), + "transformers.models.t5.modeling_t5.T5ForConditionalGeneration": + PolicyLocation(file_name="t5", class_name="T5ForConditionalGenerationPolicy"), + "transformers.models.t5.modeling_t5.T5EncoderModel": + PolicyLocation(file_name="t5", class_name="T5EncoderPolicy"), + + # GPT2 + "transformers.models.gpt2.modeling_gpt2.GPT2Model": + PolicyLocation(file_name="gpt2", class_name="GPT2ModelPolicy"), + "transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel": + PolicyLocation(file_name="gpt2", class_name="GPT2LMHeadModelPolicy"), + "transformers.models.gpt2.modeling_gpt2.GPT2DoubleHeadsModel": + PolicyLocation(file_name="gpt2", class_name="GPT2DoubleHeadsModelPolicy"), + "transformers.models.gpt2.modeling_gpt2.GPT2ForTokenClassification": + PolicyLocation(file_name="gpt2", class_name="GPT2ForTokenClassificationPolicy"), + "transformers.models.gpt2.modeling_gpt2.GPT2ForSequenceClassification": + PolicyLocation(file_name="gpt2", class_name="GPT2ForSequenceClassificationPolicy"), + + # OPT + "transformers.models.opt.modeling_opt.OPTModel": + PolicyLocation(file_name="opt", class_name="OPTModelPolicy"), + "transformers.models.opt.modeling_opt.OPTForCausalLM": + PolicyLocation(file_name="opt", class_name="OPTForCausalLMPolicy"), + "transformers.models.opt.modeling_opt.OPTForSequenceClassification": + PolicyLocation(file_name="opt", class_name="OPTForSequenceClassificationPolicy"), + "transformers.models.opt.modeling_opt.OPTForQuestionAnswering": + PolicyLocation(file_name="opt", class_name="OPTForQuestionAnsweringPolicy"), + + # Bloom + "transformers.models.bloom.modeling_bloom.BloomModel": + PolicyLocation(file_name="bloom", class_name="BloomModelPolicy"), + "transformers.models.bloom.modeling_bloom.BloomForCausalLM": + PolicyLocation(file_name="bloom", class_name="BloomForCausalLMPolicy"), + "transformers.models.bloom.modeling_bloom.BloomForSequenceClassification": + PolicyLocation(file_name="bloom", class_name="BloomForSequenceClassificationPolicy"), + "transformers.models.bloom.modeling_bloom.BloomForTokenClassification": + PolicyLocation(file_name="bloom", class_name="BloomForTokenClassificationPolicy"), + "transformers.models.bloom.modeling_bloom.BloomForQuestionAnswering": + PolicyLocation(file_name="bloom", class_name="BloomForQuestionAnsweringPolicy"), +} + + +def import_policy(policy_location: PolicyLocation) -> Policy: + """ + Dynamically import a Policy class based on the policy location. + """ + module_name = f"colossalai.shardformer.policies.{policy_location.file_name}" + module = importlib.import_module(module_name) + return getattr(module, policy_location.class_name) + + +def _fullname(obj): + """ + Return the full name of an object, including the module name. + """ + klass = obj.__class__ + module = klass.__module__ + if module == 'builtins': + return klass.__qualname__ # avoid outputs like 'builtins.str' + return module + '.' + klass.__qualname__ + + +def get_autopolicy(model: nn.Module) -> Policy: + r""" + Return the auto policy for the model + + Args: + model (:class:`nn.Module`): The model to get the auto policy + + Return: + :class:`Policy`: The auto policy for the model + """ + full_name = _fullname(model) + policy_location = _POLICY_LIST.get(full_name, None) + + if policy_location is None: + raise NotImplementedError( + f"Auto policy for {model.__class__.__qualname__} is not implemented\n. Supported models are {list(_POLICY_LIST.keys())}" + ) + else: + policy = import_policy(policy_location) + return policy() diff --git a/colossalai/shardformer/policies/basepolicy.py b/colossalai/shardformer/policies/basepolicy.py new file mode 100644 index 000000000000..2d347542fa7a --- /dev/null +++ b/colossalai/shardformer/policies/basepolicy.py @@ -0,0 +1,153 @@ +# part of code modified from https://github.com/tunib-ai/parallelformers + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Type, Union + +import torch.nn as nn + +from ..shard.shard_config import ShardConfig + +__all__ = ["ParallelModule", "SubModuleReplacementDescription", "ModulePolicyDescription", "Policy"] + + +class ParallelModule(): + + def __init__(self): + pass + + +@dataclass +class SubModuleReplacementDescription: + r""" + Describe how a submodule will be replaced + + Args: + suffix (str): used to get the submodule object + target_module (ParallelModule): specifies the module class used to replace to submodule + kwargs (Dict[str, Any]): the dictionary used to pass extra arguments to the `ParallelModule.from_native_module` method. + ignore_if_not_exist (bool): if the submodule does not exist, ignore it or raise an exception + """ + suffix: str + target_module: ParallelModule + kwargs: Dict[str, Any] = None + ignore_if_not_exist: bool = False + + +@dataclass +class ModulePolicyDescription: + r""" + Describe how the attributes and parameters will be transformed in a policy. + + Args: + attribute_replacement (Dict[str, Any]): key is the attribute name, value is the attribute value after sharding + param_replacement (List[Callable]): a list of functions to perform in-place param replacement. The function + must receive only one arguments: module. One example is + + ```python + def example_replace_weight(module: torch.nn.Module): + weight = module.weight + new_weight = shard_rowwise(weight, process_group) + module.weight = torch.nn.Parameter(new_weight) + ``` + sub_module_replacement (List[SubModuleReplacementDescription]): each element in the list is a ParamReplacementDescription + object which specifies the module to be replaced and the target module used to replacement. + method_replace (Dict[str, Callable]): key is the method name, value is the method for replacement + """ + attribute_replacement: Dict[str, Any] = None + param_replacement: List[Callable] = None + sub_module_replacement: List[SubModuleReplacementDescription] = None + method_replacement: Dict[str, Callable] = None + + +class Policy(ABC): + r""" + The base class for all the policies. For each different model, it should have a different policy class, + like BertPolicy for Bert Model or OPTPolicy for OPT model. + + Shardformer has provided many built-in sharding policies for the mainstream models. You can use the + built-in policies by setting `policy = None`, which is already the default argument for `Shardformer.optimize`. + If you want to define your own policy, you can inherit from this class and overwrite the methods you want to modify. + """ + + def __init__(self) -> None: + self.shard_config = None + self.model = None + self.shard_config = None + + def set_model(self, model: nn.Module) -> None: + r""" + Set model as an attribute of the Policy object so that we can access the model's attributes. + + Args: + model (:class:`nn.Module`): The model to be perform + """ + self.model = model + + def set_shard_config(self, shard_config: ShardConfig) -> None: + r""" + Set shard config as an attribute of the Policy object. + + Args: + shard_config (:class:`ShardConfig`): The shard config to be perform + """ + self.shard_config = shard_config + self.config_sanity_check() + + @abstractmethod + def config_sanity_check(self): + """ + Check if the shard config is valid for the model. Raise an exception if the config is invalid. + This method is made abstractmethod with no default implementation because we want to the policy writer + to take note of the feature supported by his/her model and policy. + """ + pass + + @abstractmethod + def preprocess(self) -> nn.Module: + r""" + Perform some preprocessing of the model, like reshaping the embedding layer. + """ + pass + + @abstractmethod + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + r""" + This method returns the module policy, which is a dictionary. The key is the module name or the module object, + and the value is the ModulePolicyDescription object. The ModulePolicyDescription object describes how the module + will be transformed. + """ + pass + + @abstractmethod + def postprocess(self) -> nn.Module: + r""" + Perform some postprocessing of the model, like binding the weight of embedding layer with + the classifier layer + """ + pass + + def append_or_create_submodule_replacement( + self, description: Union[SubModuleReplacementDescription, + List[SubModuleReplacementDescription]], policy: Dict[Union[str, nn.Module], + ModulePolicyDescription], + target_key: Union[str, nn.Module]) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + r""" + Append or create a new submodule replacement description to the policy for the given key. + + Args: + submodule_replace_desc (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended + policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated + target_key (Union[str, nn.Module]): the key of the policy to be updated + """ + # convert to list + if isinstance(description, SubModuleReplacementDescription): + description = [description] + + # append or create a new description + if target_key in policy: + policy[target_key].sub_module_replacement.extend(description) + else: + policy[target_key] = ModulePolicyDescription(sub_module_replacement=description) + + return policy diff --git a/colossalai/shardformer/policies/bert.py b/colossalai/shardformer/policies/bert.py new file mode 100644 index 000000000000..9c2736cc64d3 --- /dev/null +++ b/colossalai/shardformer/policies/bert.py @@ -0,0 +1,293 @@ +import torch.nn as nn + +import colossalai.shardformer.layer as col_nn + +from .._utils import getattr_, setattr_ +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + 'BertPolicy', 'BertModelPolicy', 'BertForPretrainingPolicy', 'BertLMHeadModelPolicy', 'BertForMaskedLMPolicy', + 'BertForNextSentencePredictionPolicy', 'BertForSequenceClassificationPolicy', 'BertForTokenClassificationPolicy', + 'BertForMultipleChoicePolicy' +] + + +class BertPolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + # TODO: + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + return self.model + + def module_policy(self): + from transformers.models.bert.modeling_bert import BertEmbeddings, BertLayer + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[BertLayer] = ModulePolicyDescription(attribute_replacement={ + "attention.self.all_head_size": + self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "crossattention.self.all_head_size": + self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "attention.self.num_attention_heads": + self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + "crossattention.self.num_attention_heads": + self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.self.query", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.self.key", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.self.value", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.self.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ]) + + policy[BertEmbeddings] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="word_embeddings", + target_module=col_nn.VocabParallelEmbedding1D, + ), + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ) + ]) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + # Handle bert layer + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="attention.output.LayerNorm", + target_module=col_nn.FusedLayerNorm, + ), + SubModuleReplacementDescription( + suffix="output.LayerNorm", + target_module=col_nn.FusedLayerNorm, + ) + ], + policy=policy, + target_key=BertLayer) + + # handle embedding layer + self.append_or_create_submodule_replacement( + description=[SubModuleReplacementDescription( + suffix="LayerNorm", + target_module=col_nn.FusedLayerNorm, + )], + policy=policy, + target_key=BertEmbeddings) + return policy + + def add_lm_head_policy(self, base_policy): + from transformers.models.bert.modeling_bert import BertLMPredictionHead + + # optimize for tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="decoder", target_module=col_nn.Linear1D_Col, kwargs={"gather_output": True}), + policy=base_policy, + target_key=BertLMPredictionHead) + + # optimize with fused normalization + if self.shard_config.enable_fused_normalization: + # Handle bert lm prediction head + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="transform.LayerNorm", + target_module=col_nn.FusedLayerNorm, + ), + policy=base_policy, + target_key=BertLMPredictionHead) + return base_policy + + def postprocess(self): + return self.model + + +# BertModel +class BertModelPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + +# BertForPreTraining +class BertForPretrainingPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + module_policy = super().module_policy() + module_policy = self.add_lm_head_policy(module_policy) + return module_policy + + def postprocess(self): + binding_map = {"bert.embeddings.word_embeddings.weight": "cls.predictions.decoder.weight"} + for k, v in binding_map.items(): + param = getattr_(self.model, k) + setattr_(self.model, v, param) + return self.model + + +# BertLMHeadModel +class BertLMHeadModelPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + module_policy = super().module_policy() + module_policy = self.add_lm_head_policy(module_policy) + return module_policy + + def postprocess(self): + binding_map = {"bert.embeddings.word_embeddings.weight": "cls.predictions.decoder.weight"} + for k, v in binding_map.items(): + param = getattr_(self.model, k) + setattr_(self.model, v, param) + return self.model + + +# BertForMaskedLM +class BertForMaskedLMPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + module_policy = super().module_policy() + module_policy = self.add_lm_head_policy(module_policy) + return module_policy + + def postprocess(self): + binding_map = {"bert.embeddings.word_embeddings.weight": "cls.predictions.decoder.weight"} + for k, v in binding_map.items(): + param = getattr_(self.model, k) + setattr_(self.model, v, param) + return self.model + + +# BertForSequenceClassification +class BertForSequenceClassificationPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForSequenceClassification + + module_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForSequenceClassification: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ]) + } + module_policy.update(addon_module) + return module_policy + + +# BertForTokenClassification +class BertForTokenClassificationPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForTokenClassification + + module_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForTokenClassification: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ]) + } + module_policy.update(addon_module) + return module_policy + + +# BertForNextSentencePrediction +class BertForNextSentencePredictionPolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + +# BertForMultipleChoice +class BertForMultipleChoicePolicy(BertPolicy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForMultipleChoice + + module_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForMultipleChoice: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ]) + } + module_policy.update(addon_module) + return module_policy diff --git a/colossalai/shardformer/policies/bloom.py b/colossalai/shardformer/policies/bloom.py new file mode 100644 index 000000000000..a0b5340f72bc --- /dev/null +++ b/colossalai/shardformer/policies/bloom.py @@ -0,0 +1,185 @@ +import torch.nn as nn + +import colossalai.shardformer.layer as col_nn + +from .._utils import getattr_, setattr_ +from ..modeling.bloom import build_bloom_alibi_tensor_fn +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + + +class BloomPolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + return self.model + + def module_policy(self): + from transformers.models.bloom.modeling_bloom import BloomBlock, BloomModel + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[BloomBlock] = ModulePolicyDescription(attribute_replacement={ + "self_attention.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attention.split_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attention.num_heads": self.model.config.n_head // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attention.query_key_value", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="self_attention.dense", + target_module=col_nn.Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="self_attention.attention_dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="mlp.dense_h_to_4h", + target_module=col_nn.Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="mlp.dense_4h_to_h", + target_module=col_nn.Linear1D_Row, + ), + ]) + + policy[BloomModel] = ModulePolicyDescription( + attribute_replacement={ + "num_heads": self.model.config.n_head // self.shard_config.tensor_parallel_size, + }, + method_replacement={ + "build_alibi_tensor": build_bloom_alibi_tensor_fn(self.shard_config.tensor_parallel_process_group) + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="word_embeddings", + target_module=col_nn.VocabParallelEmbedding1D, + ) + ]) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + # handle bloom model + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="ln_f", + target_module=col_nn.FusedLayerNorm, + ), + SubModuleReplacementDescription( + suffix="word_embeddings_layernorm", + target_module=col_nn.FusedLayerNorm, + ) + ], + policy=policy, + target_key=BloomModel) + + # handle bloom block + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=col_nn.FusedLayerNorm, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=col_nn.FusedLayerNorm, + ) + ], + policy=policy, + target_key=BloomBlock) + + return policy + + def postprocess(self): + return self.model + + +class BloomModelPolicy(BloomPolicy): + pass + + +class BloomForCausalLMPolicy(BloomPolicy): + + def module_policy(self): + from transformers.models.bloom.modeling_bloom import BloomForCausalLM + policy = super().module_policy() + + # handle tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="lm_head", target_module=col_nn.Linear1D_Col, kwargs=dict(gather_output=True)), + policy=policy, + target_key=BloomForCausalLM) + + return policy + + def postprocess(self): + binding_map = {"transformer.word_embeddings.weight": "lm_head.weight"} + + for k, v in binding_map.items(): + param = getattr_(self.model, k) + + if not isinstance(param, nn.Parameter): + param = nn.Parameter(param) + + # tie weights + setattr_(self.model, v, param) + return self.model + + +class BloomForSequenceClassificationPolicy(BloomPolicy): + + def module_policy(self): + from transformers.models.bloom.modeling_bloom import BloomForSequenceClassification + policy = super().module_policy() + + # handle tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="score", target_module=col_nn.Linear1D_Col, kwargs=dict(gather_output=True)), + policy=policy, + target_key=BloomForSequenceClassification) + + return policy + + +class BloomForTokenClassificationPolicy(BloomPolicy): + + def module_policy(self): + from transformers.models.bloom.modeling_bloom import BloomForTokenClassification + policy = super().module_policy() + + # handle tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription(suffix="classifier", + target_module=col_nn.Linear1D_Col, + kwargs=dict(gather_output=True)), + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ], + policy=policy, + target_key=BloomForTokenClassification) + + return policy + + +class BloomForQuestionAnsweringPolicy(BloomPolicy): + # No head sharding as the output features is only 2 + pass diff --git a/colossalai/shardformer/policies/gpt2.py b/colossalai/shardformer/policies/gpt2.py new file mode 100644 index 000000000000..549cdbf87a80 --- /dev/null +++ b/colossalai/shardformer/policies/gpt2.py @@ -0,0 +1,193 @@ +import torch.nn as nn + +import colossalai.shardformer.layer as col_nn + +from .._utils import getattr_, setattr_ +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + 'GPT2Policy', 'GPT2ModelPolicy', 'GPT2LMHeadModelPolicy', 'GPT2DoubleHeadsModelPolicy', + 'GPT2ForTokenClassificationPolicy', 'GPT2ForSequenceClassificationPolicy' +] + + +class GPT2Policy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + return self.model + + def module_policy(self): + from transformers.models.gpt2.modeling_gpt2 import GPT2Block, GPT2Model + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[GPT2Model] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="wte", + target_module=col_nn.VocabParallelEmbedding1D, + ), + ]) + policy[GPT2Block] = ModulePolicyDescription(attribute_replacement={ + "attn.embed_dim": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "attn.split_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attn.c_attn", + target_module=col_nn.GPT2FusedLinearConv1D_Col, + kwargs={ + "n_fused": 3, + }, + ), + SubModuleReplacementDescription( + suffix="attn.c_proj", + target_module=col_nn.GPT2FusedLinearConv1D_Row, + ), + SubModuleReplacementDescription( + suffix="mlp.c_fc", + target_module=col_nn.GPT2FusedLinearConv1D_Col, + kwargs={ + "n_fused": 1, + }, + ), + SubModuleReplacementDescription( + suffix="mlp.c_proj", + target_module=col_nn.GPT2FusedLinearConv1D_Row, + ), + SubModuleReplacementDescription( + suffix="attn.attn_dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attn.resid_dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="mlp.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + ]) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="ln_f", + target_module=col_nn.FusedLayerNorm, + ), + policy=policy, + target_key=GPT2Model) + + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="ln_1", + target_module=col_nn.FusedLayerNorm, + ), + SubModuleReplacementDescription( + suffix="ln_2", + target_module=col_nn.FusedLayerNorm, + ), + SubModuleReplacementDescription(suffix="ln_cross_attn", + target_module=col_nn.FusedLayerNorm, + ignore_if_not_exist=True) + ], + policy=policy, + target_key=GPT2Block) + return policy + + def postprocess(self): + return self.model + + +# GPT2Model +class GPT2ModelPolicy(GPT2Policy): + + def __init__(self) -> None: + super().__init__() + + +# GPT2LMHeadModel +class GPT2LMHeadModelPolicy(GPT2Policy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel + + module_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + GPT2LMHeadModel: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", target_module=col_nn.Linear1D_Col, kwargs={"gather_output": True}) + ]) + } + module_policy.update(addon_module) + return module_policy + + def postprocess(self): + binding_map = {"transformer.wte.weight": "lm_head.weight"} + for k, v in binding_map.items(): + param = getattr_(self.model, k) + setattr_(self.model, v, param) + return self.model + + +# GPT22DoubleHeadsModel +class GPT2DoubleHeadsModelPolicy(GPT2Policy): + + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + from transformers.models.gpt2.modeling_gpt2 import GPT2DoubleHeadsModel + + module_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + GPT2DoubleHeadsModel: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", target_module=col_nn.Linear1D_Col, kwargs={"gather_output": True}) + ]) + } + module_policy.update(addon_module) + return module_policy + + def postprocess(self): + binding_map = {"transformer.wte.weight": "lm_head.weight"} + for k, v in binding_map.items(): + param = getattr_(self.model, k) + setattr_(self.model, v, param) + return self.model + + +# GPT2ForTokenClassification +class GPT2ForTokenClassificationPolicy(GPT2Policy): + + def __init__(self) -> None: + super().__init__() + + +# GPT2ForSequenceClassification +class GPT2ForSequenceClassificationPolicy(GPT2Policy): + + def __init__(self) -> None: + super().__init__() diff --git a/colossalai/shardformer/policies/llama.py b/colossalai/shardformer/policies/llama.py new file mode 100644 index 000000000000..157785bdcf13 --- /dev/null +++ b/colossalai/shardformer/policies/llama.py @@ -0,0 +1,145 @@ +from typing import Dict, Union + +import torch.nn as nn + +from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col, Linear1D_Row, VocabParallelEmbedding1D + +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ['LlamaPolicy', 'LlamaForCausalLMPolicy', 'LlamaForSequenceClassificationPolicy'] + + +class LlamaPolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # Resize embedding + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaModel + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[LlamaDecoderLayer] = ModulePolicyDescription( + attribute_replacement={ + "self_attn.hidden_size": + self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": + self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + ) + ], + ) + + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=VocabParallelEmbedding1D, + ), + policy=policy, + target_key=LlamaModel) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=FusedRMSNorm, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=FusedRMSNorm, + ) + ], + policy=policy, + target_key=LlamaDecoderLayer) + + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="norm", + target_module=FusedRMSNorm, + ), + policy=policy, + target_key=LlamaModel) + + return policy + + def postprocess(self): + return self.model + + +class LlamaForCausalLMPolicy(LlamaPolicy): + + def module_policy(self): + from transformers import LlamaForCausalLM + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + # add a new item for casual lm + new_item = { + LlamaForCausalLM: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", target_module=Linear1D_Col, kwargs=dict(gather_output=True)) + ]) + } + policy.update(new_item) + return policy + + +class LlamaForSequenceClassificationPolicy(LlamaPolicy): + + def module_policy(self): + from transformers import LlamaForSequenceClassification + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + LlamaForSequenceClassification: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", target_module=Linear1D_Col, kwargs=dict(gather_output=True)) + ]) + } + policy.update(new_item) + return policy diff --git a/colossalai/shardformer/policies/opt.py b/colossalai/shardformer/policies/opt.py new file mode 100644 index 000000000000..b87db53f45f1 --- /dev/null +++ b/colossalai/shardformer/policies/opt.py @@ -0,0 +1,140 @@ +from colossalai.shardformer.layer import FusedLayerNorm, Linear1D_Col, Linear1D_Row, VocabParallelEmbedding1D + +from .._utils import getattr_, setattr_ +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + 'OPTPolicy', 'OPTModelPolicy', 'OPTForCausalLMPolicy', 'OPTForSequenceClassificationPolicy', + 'OPTForQuestionAnsweringPolicy' +] + + +class OPTPolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + return self.model + + def module_policy(self): + from transformers.models.opt.modeling_opt import OPTAttention, OPTDecoder, OPTDecoderLayer + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[OPTDecoder] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=VocabParallelEmbedding1D, + ) + ]) + policy[OPTDecoderLayer] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="fc1", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=Linear1D_Row, + ) + ]) + + policy[OPTAttention] = ModulePolicyDescription(attribute_replacement={ + "embed_dim": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="q_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="k_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="v_proj", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="out_proj", + target_module=Linear1D_Row, + ), + ]) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="final_layer_norm", target_module=FusedLayerNorm, ignore_if_not_exist=True), + policy=policy, + target_key=OPTDecoder) + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription(suffix="self_attn_layer_norm", + target_module=FusedLayerNorm, + ignore_if_not_exist=True), + SubModuleReplacementDescription(suffix="final_layer_norm", + target_module=FusedLayerNorm, + ignore_if_not_exist=True) + ], + policy=policy, + target_key=OPTDecoderLayer) + + return policy + + def postprocess(self): + return self.model + + +class OPTModelPolicy(OPTPolicy): + + def __init__(self) -> None: + super().__init__() + + +class OPTForCausalLMPolicy(OPTPolicy): + + def module_policy(self): + from transformers.models.opt.modeling_opt import OPTForCausalLM + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="lm_head", target_module=Linear1D_Col, kwargs=dict(gather_output=True)), + policy=policy, + target_key=OPTForCausalLM) + return policy + + def postprocess(self): + binding_map = { + 'model.decoder.embed_tokens': 'lm_head', + } + + for k, v in binding_map.items(): + src_mod = getattr_(self.model, k) + dst_mod = getattr_(self.model, v) + dst_mod.weight = src_mod.weight + + return self.model + + +class OPTForSequenceClassificationPolicy(OPTPolicy): + + def __init__(self) -> None: + super().__init__() + + +class OPTForQuestionAnsweringPolicy(OPTPolicy): + + def __init__(self) -> None: + super().__init__() diff --git a/colossalai/shardformer/policies/t5.py b/colossalai/shardformer/policies/t5.py new file mode 100644 index 000000000000..cde59ab77042 --- /dev/null +++ b/colossalai/shardformer/policies/t5.py @@ -0,0 +1,249 @@ +from colossalai.shardformer.layer import ( + DropoutForParallelInput, + Embedding1D, + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + VocabParallelEmbedding1D, +) +from colossalai.shardformer.policies.basepolicy import ModulePolicyDescription + +from .._utils import getattr_, setattr_ +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["T5ModelPolicy", "T5ForConditionalGenerationPolicy", "T5EncoderPolicy"] + + +class T5BasePolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + return self.model + + def module_policy(self): + from transformers.models.t5.modeling_t5 import ( + T5Attention, + T5DenseActDense, + T5DenseGatedActDense, + T5LayerCrossAttention, + T5LayerFF, + T5LayerSelfAttention, + T5Stack, + ) + + policy = {} + + if self.shard_config.enable_tensor_parallelism: + policy[T5Stack] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=Embedding1D, + ) + ]) + policy[T5LayerSelfAttention] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ), + ]) + policy[T5LayerCrossAttention] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ) + ]) + policy[T5Attention] = ModulePolicyDescription(attribute_replacement={ + "d_model": + self.model.config.d_model // self.shard_config.tensor_parallel_size, + "n_heads": + self.model.config.num_heads // self.shard_config.tensor_parallel_size, + "inner_dim": + self.model.config.num_heads * self.model.config.d_kv // self.shard_config.tensor_parallel_size + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="q", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="k", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="v", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="o", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="relative_attention_bias", + target_module=Embedding1D, + kwargs=dict(gather_output=False), + ignore_if_not_exist=True) + ]) + policy[T5LayerFF] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ), + ]) + policy[T5DenseGatedActDense] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="wi_0", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="wi_1", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="wo", target_module=Linear1D_Col, kwargs=dict(gather_output=True)), + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ) + ]) + policy[T5DenseActDense] = ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="wi", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="wo", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForParallelInput, + ) + ]) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="layer_norm", + target_module=FusedRMSNorm, + ), + policy=policy, + target_key=T5LayerFF) + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="layer_norm", + target_module=FusedRMSNorm, + ), + policy=policy, + target_key=T5LayerFF) + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="layer_norm", target_module=FusedRMSNorm), + policy=policy, + target_key=T5LayerSelfAttention) + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="layer_norm", target_module=FusedRMSNorm), + policy=policy, + target_key=T5LayerCrossAttention) + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="final_layer_norm", target_module=FusedRMSNorm), + policy=policy, + target_key=T5Stack) + return policy + + def postprocess(self): + binding_map = [["shared", "encoder.embed_tokens"], ["shared", "decoder.embed_tokens"]] + + for k, v in binding_map: + mod = getattr_(self.model, k) + setattr_(self.model, v, mod) + return self.model + + +class T5ModelPolicy(T5BasePolicy): + + def module_policy(self): + from transformers import T5Model + base_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="shared", + target_module=VocabParallelEmbedding1D, + ), + policy=base_policy, + target_key=T5Model) + return base_policy + + +class T5ForConditionalGenerationPolicy(T5BasePolicy): + + def module_policy(self): + from transformers import T5ForConditionalGeneration + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=[ + SubModuleReplacementDescription( + suffix="shared", + target_module=VocabParallelEmbedding1D, + ), + SubModuleReplacementDescription(suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict(gather_output=True)) + ], + policy=policy, + target_key=T5ForConditionalGeneration) + return policy + + def postprocess(self): + super().postprocess() + + binding_map = {"shared": "lm_head"} + + for k, v in binding_map.items(): + src_mod = getattr_(self.model, k) + dst_mod = getattr_(self.model, v) + dst_mod.weight = src_mod.weight + + return self.model + + +class T5EncoderPolicy(T5BasePolicy): + + def module_policy(self): + from transformers import T5EncoderModel + + base_policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement(description=SubModuleReplacementDescription( + suffix="shared", + target_module=VocabParallelEmbedding1D, + ), + policy=base_policy, + target_key=T5EncoderModel) + return base_policy + + def postprocess(self): + binding_map = [ + ["shared", "encoder.embed_tokens"], + ] + + for k, v in binding_map: + mod = getattr_(self.model, k) + setattr_(self.model, v, mod) + return self.model diff --git a/colossalai/shardformer/policies/vit.py b/colossalai/shardformer/policies/vit.py new file mode 100644 index 000000000000..eaebe2eee0ba --- /dev/null +++ b/colossalai/shardformer/policies/vit.py @@ -0,0 +1,110 @@ +from typing import Dict, Union + +import torch.nn as nn + +from colossalai.shardformer.layer import DropoutForReplicatedInput, FusedLayerNorm, Linear1D_Col, Linear1D_Row + +from .basepolicy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ['ViTPolicy'] + + +class ViTPolicy(Policy): + + def config_sanity_check(self): + pass + + def preprocess(self): + # Resize embedding + vocab_size = self.model.config.vocab_size + world_size = self.shard_config.tensor_parallel_size + + if vocab_size % world_size != 0: + new_vocab_size = vocab_size + world_size - vocab_size % world_size + self.model.resize_token_embeddings(new_vocab_size) + + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer + + base_policy = { + ViTEmbeddings: + ModulePolicyDescription(sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForReplicatedInput, + ) + ]), + ViTLayer: + ModulePolicyDescription(attribute_replacement={ + "attention.attention.num_attention_heads": + self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + "attention.attention.all_head_size": + self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.attention.query", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.attention.key", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.attention.value", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="attention.attention.dropout", + target_module=DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=Linear1D_Col, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=Linear1D_Row, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=DropoutForParallelInput, + ), + ]), + } + + # optimization configuration + if self.shard_config.enable_fused_normalization: + base_policy[ViTAttention].sub_module_replacement.extend([ + SubModuleReplacementDescription( + suffix="layernorm_before", + target_module=FusedLayerNorm, + ), + SubModuleReplacementDescription( + suffix="layernorm_after", + target_module=FusedLayerNorm, + ) + ]) + base_policy[ViTModel].sub_module_replacement.append( + SubModuleReplacementDescription( + suffix="layernorm", + target_module=FusedLayerNorm, + )) + + return base_policy + + def new_model_class(self): + return None + + def postprocess(self): + return self.model diff --git a/colossalai/shardformer/shard/__init__.py b/colossalai/shardformer/shard/__init__.py new file mode 100644 index 000000000000..7abdd45ec7c5 --- /dev/null +++ b/colossalai/shardformer/shard/__init__.py @@ -0,0 +1,5 @@ +from .shard_config import ShardConfig +from .sharder import ModelSharder +from .shardformer import ShardFormer + +__all__ = ['ShardConfig', 'ModelSharder', 'ShardFormer'] diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py new file mode 100644 index 000000000000..83c08d275df3 --- /dev/null +++ b/colossalai/shardformer/shard/shard_config.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass + +import torch.distributed as dist +from torch.distributed import ProcessGroup + +__all__ = ['ShardConfig'] + + +@dataclass +class ShardConfig: + r""" + The config for sharding the huggingface model + + Args: + tensor_parallel_process_group (int): The process group for tensor parallelism, defaults to None, which is the global process group. + enable_tensor_parallelism (bool): Whether to turn on tensor parallelism, default is True. + enable_fused_normalization (bool): Whether to use fused layernorm, default is False. + enable_all_optimization (bool): Whether to turn on all optimization, default is False. + """ + tensor_parallel_process_group: ProcessGroup = None + enable_tensor_parallelism: bool = True + enable_fused_normalization: bool = False + enable_all_optimization: bool = False + + # TODO: add support for tensor parallel + # pipeline_parallel_size: int + # data_parallel_size: int + # tensor_parallel_mode: Literal['1d', '2d', '2.5d', '3d'] + # inference_only: bool = True + # gather_output: bool = True + + @property + def tensor_parallel_size(self): + return self._tensor_parallel_size + + def __post_init__(self): + if not self.enable_tensor_parallelism: + self._tensor_parallel_size = 1 + else: + # get the parallel size + self._tensor_parallel_size = dist.get_world_size(self.tensor_parallel_process_group) + + # turn on all optimization if all_optimization is set to True + if self.enable_all_optimization: + self._turn_on_all_optimization() + + def _turn_on_all_optimization(self): + """ + Turn on all optimization. + """ + # you can add all the optimization flag here + self.enable_fused_normalization = True diff --git a/colossalai/shardformer/shard/sharder.py b/colossalai/shardformer/shard/sharder.py new file mode 100644 index 000000000000..201e0a08cbfe --- /dev/null +++ b/colossalai/shardformer/shard/sharder.py @@ -0,0 +1,174 @@ +from typing import Any, Callable, Dict, List, Union + +import torch.nn as nn + +from .._utils import getattr_, setattr_ +from ..policies.autopolicy import get_autopolicy +from ..policies.basepolicy import Policy, SubModuleReplacementDescription +from .shard_config import ShardConfig + +__all__ = ['ModelSharder', 'shard_model'] + + +class ModelSharder(object): + r""" + Shard the original huggingface model according to the policy + + Args: + policy (:class:`Policy`): The policy to shard the model + model (:class:`torch.Module`): The model to shard + shard_config: The setting of distributed model + """ + + def __init__(self, model: nn.Module, policy: Policy, shard_config: ShardConfig = None) -> None: + self.model = model + self.policy = get_autopolicy(self.model) if policy is None else policy + self.shard_config = shard_config + + def shard(self) -> None: + r""" + Shard the model according to the policy + """ + self.policy.set_model(self.model) + self.policy.set_shard_config(self.shard_config) + self._preprocess() + self._replace_module() + self._postprocess() + + def _preprocess(self) -> None: + self.model = self.policy.preprocess() + + def _postprocess(self) -> None: + self.model = self.policy.postprocess() + + def _replace_module(self,) -> None: + r""" + Replace the module according to the policy, and replace the module one by one + + Args: + model (:class:`torch.nn.Module`): The model to shard + """ + module_descriptions = self.policy.module_policy() + for layer_cls, module_description in module_descriptions.items(): + attr_replacement = module_description.attribute_replacement + param_replacement = module_description.param_replacement + sub_module_replacement = module_description.sub_module_replacement + method_replacement = module_description.method_replacement + self._recursive_replace_layer(self.model, layer_cls, attr_replacement, param_replacement, + method_replacement, sub_module_replacement) + + def _recursive_replace_layer( + self, + module: nn.Module, + origin_cls: Union[str, nn.Module], + attr_replacement: Dict[str, Any], + param_replacement: List[Callable], + method_replacement: Dict[str, Callable], + sub_module_replacement: List[Callable], + ) -> None: + r""" + Reverse the replace layer operation + + Args: + layer (torch.nn.Module): The object of layer to shard + origin_cls (Union[str, torch.nn.Module]): The origin layer class or a string of layer class name. + attr_replacement (Dict): The attribute dict to modify + param_replacement (List[Callable]): The function list to get parameter shard information in policy + sub_module_replacement (List[Callable]): The function list to get sub module shard information in policy + """ + if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or \ + (module.__class__ == origin_cls): + if attr_replacement is not None: + self._replace_attr(module, attr_replacement) + + if param_replacement is not None: + self._replace_param(module, param_replacement) + + if method_replacement is not None: + self._replace_method(module, method_replacement) + + if sub_module_replacement is not None: + self._replace_sub_module(module, sub_module_replacement) + + for name, child in module.named_children(): + self._recursive_replace_layer(child, origin_cls, attr_replacement, param_replacement, method_replacement, + sub_module_replacement) + + def _replace_attr( + self, + module: nn.Module, + attr_replacement: Dict[str, Any], + ) -> None: + r""" + Replace the attribute of the layer + + Args: + layer (:class:`torch.nn.Module`): The object of layer to shard + attr_replacement (Dict): The attribute dict to modify + """ + for k, v in attr_replacement.items(): + setattr_(module, k, v, ignore=True) + + def _replace_param( + self, + module: nn.Module, + param_replacement: List[Callable], + ) -> None: + r""" + Replace the parameter of the layer + + Args: + layer (:class:`torch.nn.Module`): The object of layer to shard + param_replacement (List[Callable]): The function list to get parameter shard information in policy + """ + for param_func in param_replacement: + param_func(module) + + def _replace_method(self, module: nn.Module, method_replacement: Dict[str, Callable]): + for method_name, new_method in method_replacement.items(): + # bind the new method to the module + setattr(module, method_name, new_method.__get__(module, module.__class__)) + + def _replace_sub_module( + self, + org_layer: nn.Module, + sub_module_replacement: List[SubModuleReplacementDescription], + ) -> None: + r""" + Shard one layer according to the policy, the layer should be the same class as the key in policy's argument_policy return dict + + Args: + org_layer (torch.nn.Module): The origin layer object to shard + sub_module_replacement (List[SubModuleReplacementDescription]): The sub module replacement description list + + """ + for description in sub_module_replacement: + suffix = description.suffix + target_module = description.target_module + kwargs = {} if description.kwargs is None else description.kwargs + + assert target_module is not None, 'target_module should not be None' + + # TODO: support different parallel mode + native_sub_module = getattr_(org_layer, suffix, ignore=True) + + assert not isinstance(native_sub_module, target_module), \ + f"The module with suffix {suffix} has been replaced, please check the policy" + + # if it is None and we are allowed to ignore this module + # just skip + if description.ignore_if_not_exist and native_sub_module is None: + continue + + try: + replace_layer = target_module.from_native_module(native_sub_module, + self.shard_config.tensor_parallel_process_group, + **kwargs) + except Exception as e: + raise RuntimeError( + f"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}" + f" with {target_module.__qualname__} with the exception: {e}. " + "Please check your model configuration or sharding policy, you can set up an issue for us to help you as well." + ) + + setattr_(org_layer, suffix, replace_layer) diff --git a/colossalai/shardformer/shard/shardformer.py b/colossalai/shardformer/shard/shardformer.py new file mode 100644 index 000000000000..3fce12463414 --- /dev/null +++ b/colossalai/shardformer/shard/shardformer.py @@ -0,0 +1,46 @@ +import torch.nn as nn + +from colossalai.cluster import DistCoordinator + +from ..policies.basepolicy import Policy +from .shard_config import ShardConfig +from .sharder import ModelSharder + + +class ShardFormer: + """ + Parallelize model based on the given config and policy + + Example: + + ```python + from colossalai.shardformer import ShardFormer, ShardConfig + from transformers import BertForMaskedLM + import colossalai + import torch + + colossalai.launch_from_torch(config={}) + + org_model = BertForMaskedLM.from_pretrained('bert-base-uncased') + shard_config = ShardConfig() + shard_former = ShardFormer(shard_config=shard_config) + model = shard_former.optimize(org_model) + ``` + """ + + def __init__(self, shard_config: ShardConfig): + self.coordinator = DistCoordinator() + self.shard_config = shard_config + + def optimize(self, model: nn.Module, policy: Policy = None): + r""" + This method will optimize the model based on the given policy. + + Args: + model (`torch.nn.Model`): the origin huggingface model + shard_config (`ShardConfig`): the config for distribute information + policy (`Policy`): the custom policy for sharding + """ + sharder = ModelSharder(model=model, shard_config=self.shard_config, policy=policy) + sharder.shard() + return model diff --git a/colossalai/tensor/colo_tensor.py b/colossalai/tensor/colo_tensor.py index bbed8847abbc..4d762076461d 100644 --- a/colossalai/tensor/colo_tensor.py +++ b/colossalai/tensor/colo_tensor.py @@ -138,6 +138,15 @@ def set_process_group(self, pg: ProcessGroup): def get_tp_world_size(self) -> int: return self.process_group.tp_world_size() + def get_dp_world_size(self) -> int: + """get_dp_world_size + get the dp world size of the tensor. + + Returns: + int: dp world size + """ + return self.process_group.dp_world_size() + def set_dist_spec(self, dist_spec: _DistSpec): """set_dist_spec set dist spec and change the payloads. @@ -175,7 +184,7 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): # we have to capture the `backward` function # and make sure that it does not in `torch._C.DisableTorchFunction()` context if func is torch.Tensor.backward: - assert len(args) == 1 # only has 1 paramter + assert len(args) == 1 # only has 1 parameter backward_tensor = torch.Tensor(args[0]) tensor_kwargs = {k: torch.Tensor(v) if torch.is_tensor(v) else v for k, v in kwargs.items()} return backward_tensor.backward(**tensor_kwargs) @@ -219,7 +228,7 @@ def redistribute(self, dist_spec: _DistSpec, pg: Optional[ProcessGroup] = None) 2. If the pg is not not None and not equal to the current process group. First, convert the tensor as replicated among the TP process group. Second, reset the process group to the new pg. - Third, conver the tensor (new replicated both among the tp process group) to the new dist_spec. + Third, convert the tensor (new replicated both among the tp process group) to the new dist_spec. Args: dist_spec (_DistSpec): the new dist spec. @@ -288,7 +297,7 @@ def size_local(self, *args) -> torch.Size: def size_global(self, *args) -> torch.Size: """size_global - override the torch buildin size() + override the torch building size() the shape passed in must be in a replicate placement. Returns: diff --git a/colossalai/tensor/comm_spec.py b/colossalai/tensor/comm_spec.py index 0d8de1062d42..204f81343199 100644 --- a/colossalai/tensor/comm_spec.py +++ b/colossalai/tensor/comm_spec.py @@ -16,69 +16,66 @@ def _all_gather(tensor, comm_spec): ''' Implement all gather operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - tensor_list = [ - torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) - for _ in range(comm_spec.device_mesh.mesh_shape[comm_spec.logical_process_axis]) - ] - # without this contiguous operation, the all gather may get some unexpected results. - tensor = tensor.contiguous() - dist.all_gather(tensor_list, tensor, group=process_group) - output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous() - return output + process_groups = comm_spec.device_mesh.get_process_group_for_all_axes() + process_group = process_groups[comm_spec.logical_process_axis] + + tensor_list = [ + torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) + for _ in range(comm_spec.device_mesh.shape[comm_spec.logical_process_axis]) + ] + # without this contiguous operation, the all gather may get some unexpected results. + tensor = tensor.contiguous() + dist.all_gather(tensor_list, tensor, group=process_group) + output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous() + return output def _split(tensor, comm_spec): ''' Implement shard operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, _ in process_groups_list: - if dist.get_rank() in rank_list: - dim = comm_spec.shard_dim - length = tensor.shape[comm_spec.shard_dim] // len(rank_list) - start = length * rank_list.index(dist.get_rank()) - output = torch.narrow(tensor, dim, start, length).contiguous() - return output + process_groups = comm_spec.device_mesh.get_process_group_for_all_axes() + process_group = process_groups[comm_spec.logical_process_axis] + + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // dist.get_world_size(process_group) + start = length * dist.get_rank(process_group) + output = torch.narrow(tensor, dim, start, length).contiguous() + return output def _all_to_all(tensor, comm_spec): ''' Implement all to all operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - new_shape = list(tensor.shape) - new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // len(rank_list) - new_shape = torch.Size(new_shape) - output_tensor_list = [ - torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list)) - ] - dim = comm_spec.shard_dim - length = tensor.shape[comm_spec.shard_dim] // len(rank_list) - input_tensor_list = [ - torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(len(rank_list)) - ] - group = process_group - dist.all_to_all(output_tensor_list, input_tensor_list, group) - output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous() - return output + process_groups = comm_spec.device_mesh.get_process_group_for_all_axes() + process_group = process_groups[comm_spec.logical_process_axis] + world_size = dist.get_world_size(process_group) + + new_shape = list(tensor.shape) + new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // world_size + new_shape = torch.Size(new_shape) + output_tensor_list = [torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(world_size)] + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // world_size + input_tensor_list = [torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(world_size)] + group = process_group + dist.all_to_all(output_tensor_list, input_tensor_list, group) + output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous() + return output def _all_reduce(tensor, comm_spec, async_op=False): ''' Implement all reduce operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group, async_op=async_op) - return tensor + process_groups = comm_spec.device_mesh.get_process_group_for_all_axes() + process_group = process_groups[comm_spec.logical_process_axis] + + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group, async_op=async_op) + return tensor def _mix_gather(tensor, comm_spec): @@ -128,7 +125,7 @@ def _mix_gather(tensor, comm_spec): process_group = "[0, 1, 2, 3, 4, 5, 6, 7]" tensor_list = [(0,0),(1,0),(2,0),(3,0),(4,0),(5,0),(6,0),(7,0)] ''' - total_slices = comm_spec.device_mesh.mesh_shape[0] + total_slices = comm_spec.device_mesh.shape[0] tensor_list = [torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) for _ in range(total_slices)] leading_group_dim = comm_spec.logical_process_axes[0] assert len(comm_spec.device_mesh.process_groups_dict) == 1 @@ -149,7 +146,7 @@ def _mix_gather(tensor, comm_spec): if comm_spec.logical_process_axes[0] == comm_spec.logical_process_axes[1]: output = torch.cat(tuple(tensor_list), comm_spec.gather_dim[0]).contiguous() else: - mesh_shape = comm_spec.device_meshes.mesh_shape + mesh_shape = comm_spec.device_meshes.shape cat_slice = [mesh_shape[comm_spec.logical_process_axes[0]], mesh_shape[comm_spec.logical_process_axes[1]]] tmp_tensor_shape = list(tensor.shape) tmp_tensor_shape[comm_spec.gather_dim[0]] *= cat_slice[0] @@ -181,9 +178,9 @@ def _mix_split(tensor, comm_spec): # [4, 5, 6, 7]] # return {0: [0, 4, 1, 5, 2, 6, 3, 7], 1: [0, 1, 2, 3, 4, 5, 6, 7]} ''' - mesh_shape = comm_spec.device_meshes.mesh_shape + mesh_shape = comm_spec.device_meshes.shape dim = comm_spec.gather_dim - total_slices = comm_spec.device_mesh.mesh_shape[0] + total_slices = comm_spec.device_mesh.shape[0] # Get global rank rank = dist.get_rank() @@ -391,7 +388,7 @@ class CommSpec: to determine the buffer shape, and logical_process_axis Argument: - comm_pattern(CollectiveCommPattern): decribe the communication method used in this spec. + comm_pattern(CollectiveCommPattern): describe the communication method used in this spec. sharding_spec(ShardingSpec): This is sharding spec of the tensor which will join the communication action. gather_dim(int, Optional): The gather_dim of the tensor will be gathered. shard_dim(int, Optional): The shard_dim of the tensor will be sharded. @@ -414,7 +411,7 @@ def __init__(self, self.forward_only = forward_only if isinstance(self.logical_process_axis, list): if not mix_gather: - self.device_mesh = self.sharding_spec.device_mesh.flatten_device_mesh + self.device_mesh = self.sharding_spec.device_mesh.flatten() self.logical_process_axis = 0 else: self.device_meshes = self.sharding_spec.device_mesh.flatten_device_meshes diff --git a/colossalai/tensor/compute_spec.py b/colossalai/tensor/compute_spec.py index 73328285ee93..12f8f36bc613 100644 --- a/colossalai/tensor/compute_spec.py +++ b/colossalai/tensor/compute_spec.py @@ -10,7 +10,7 @@ class ComputePattern(Enum): class ComputeSpec(object): """ComputeSpec - The Specification for compuattion pattern + The Specification for computation pattern Args: compute_pattern (ComputePattern): an Enum instance for compute pattern. diff --git a/colossalai/tensor/d_tensor/README.md b/colossalai/tensor/d_tensor/README.md new file mode 100644 index 000000000000..3d862dddbf20 --- /dev/null +++ b/colossalai/tensor/d_tensor/README.md @@ -0,0 +1,103 @@ +# 🔢 Distributed Tensor + +## 📚 Table of Contents + +- [🔢 Distributed Tensor](#-distributed-tensor) + - [📚 Table of Contents](#-table-of-contents) + - [🔗 Introduction](#-introduction) + - [📝 Design](#-design) + - [🔨 Usage](#-usage) + - [🎈 Progress Log](#-progress-log) + +## 🔗 Introduction + +Distributed tensor is a type of tensor that is distributed across multiple devices. It is a wrapper of PyTorch tensor, and it is used to support distributed training. +It can represent the device topology and tensor placement over the devices in the topology. It also provides a set of APIs to manipulate the distributed tensor. + +## 📝 Design + +Our implementation is inspired by the work [Alpa](https://arxiv.org/abs/2201.12023), which unifies data parallelism and tensor parallelism as intra-op parallelism. It uses notations `S` to represent the sharded dimension and `R` to represent the replicated dimension. For example, given a 2D matrix, `[S, R]` represents the tensor is sharded over the first dimension. + +Each sharded dimension will have a subscript to represent its placement over the devices. Assuming we have 4 GPUs and the GPUs are arranged in a 2 x 2 manner. Let's say we have a 2D matrix like below: + + +```text + [1, 2, 3, 4 ] +A = [4, 5, 6, 7 ] + [8, 9, 10, 11] + [12, 13, 14, 15] +``` + +`[S0, R]` would mean that the first dimension is sharded over the rows in the device topology. + +```text +| --------------------—————————————————————-| +| | | +| [1, 2, 3, 4 ] | [1, 2, 3, 4 ] | +| [4, 5, 6, 7 ] | [4, 5, 6, 7 ] | +| | | +| --------------------——————————————————----- +| | | +| [8, 9, 10, 11] | [8, 9, 10, 11] | +| [12, 13, 14, 15] | [12, 13, 14, 15] | +| | | +| --------------------——————————————————----- +``` + +`[S01, R]` would mean that the first dimension is sharded over both the row and column in the device topology. + +```text +| --------------------—————————————————————-| +| | | +| [1, 2, 3, 4 ] | [4, 5, 6, 7 ] | +| | | +| --------------------——————————————————----- +| | | +| [8, 9, 10, 11] | [12, 13, 14, 15] | +| | | +| --------------------——————————————————----- +``` + +## 🔨 Usage + +A sample API usage is given below. + +```python +import torch + +import colossalai +from colossalai.device.device_mesh import DeviceMesh +from colossalai.tensor.d_tensor import DTensor, ShardingSpec + +colossalai.launch_from_torch(config={}) + +# define your device mesh +# assume you have 4 GPUs +physical_mesh_id = torch.arange(0, 4) +mesh_shape = (2, 2) +device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) + +# define a tensor +a = torch.rand(16, 32).cuda() + +# create sharding spec for the tensor +# assume the sharding spec is [S0, R] +dim_partition_dict = {0: [0]} +sharding_spec = ShardingSpec(a.dim(), dim_partition_dict) + +# create a distributed tensor +d_tensor = DTensor(a, device_mesh, sharding_spec) +print(d_tensor) + +global_tensor = d_tensor.to_global() +print(global_tensor) +``` + + +## 🎈 Progress Log + +- [x] Support layout conversion +- [x] Support sharding on 2D device mesh +- [ ] Support sharding on 3D device mesh +- [ ] Support sharding 4D device mesh +- [ ] Support sharding info saving and offline tensor merge (we can save tensor as dtensor and gather the tensors back to the global tensor based on the sharding info in a single process in CPU, useful for distributed training checkpoint load and save.) diff --git a/colossalai/tensor/d_tensor/__init__.py b/colossalai/tensor/d_tensor/__init__.py index e69de29bb2d1..3ae38a12555b 100644 --- a/colossalai/tensor/d_tensor/__init__.py +++ b/colossalai/tensor/d_tensor/__init__.py @@ -0,0 +1,28 @@ +from .api import ( + compute_global_numel, + customized_distributed_tensor_to_param, + distribute_tensor, + distribute_tensor_with_customization, + get_device_mesh, + get_global_shape, + get_layout, + get_sharding_spec, + is_customized_distributed_tensor, + is_distributed_tensor, + is_sharded, + redistribute, + shard_colwise, + shard_rowwise, + sharded_tensor_to_param, + to_global, + to_global_for_customized_distributed_tensor, +) +from .layout import Layout +from .sharding_spec import ShardingSpec + +__all__ = [ + 'is_distributed_tensor', 'distribute_tensor', 'to_global', 'is_sharded', 'shard_rowwise', 'shard_colwise', + 'sharded_tensor_to_param', 'compute_global_numel', 'get_sharding_spec', 'get_global_shape', 'get_device_mesh', + 'redistribute', 'get_layout', 'is_customized_distributed_tensor', 'distribute_tensor_with_customization', + 'to_global_for_customized_distributed_tensor', 'customized_distributed_tensor_to_param', 'Layout', 'ShardingSpec' +] diff --git a/colossalai/tensor/d_tensor/api.py b/colossalai/tensor/d_tensor/api.py new file mode 100644 index 000000000000..95a44e09e16a --- /dev/null +++ b/colossalai/tensor/d_tensor/api.py @@ -0,0 +1,434 @@ +import copy +import operator +from functools import reduce +from typing import Union + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from colossalai.device.device_mesh import DeviceMesh + +from .layout import Layout +from .layout_converter import LayoutConverter +from .sharding_spec import ShardingSpec + +layout_converter = LayoutConverter() + + +def is_distributed_tensor(tensor: torch.Tensor) -> bool: + """ + Check whether the given tensor is a distributed tensor. + + Args: + tensor (torch.Tensor): The tensor to be checked. + + Returns: + bool: Whether the given tensor is a distributed tensor. + """ + return hasattr(tensor, "dist_layout") + + +def is_sharded(dtensor: torch.Tensor) -> bool: + """ + Check if a tensor is sharded. + + Args: + tensor (torch.Tensor): The tensor to be checked. + + Returns: + bool: True if the tensor is sharded, False otherwise. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + return list(dtensor.shape) == list(dtensor.dist_layout.global_shape) + + +def _hijack_detach_and_clone(dtensor: torch.Tensor) -> torch.Tensor: + """ + Hijack the detach and clone methods of the tensor to make sure the dist_layout is copied. + + Args: + tensor (torch.Tensor): The tensor to be hijacked. + + Returns: + torch.Tensor: The hijacked tensor. + """ + dtensor._old_detach = dtensor.detach + dtensor._old_clone = dtensor.clone + + def new_detach(self): + t_ = self._old_detach() + t_.dist_layout = copy.deepcopy(self.dist_layout) + return t_ + + def new_clone(self, *args, **kwargs): + t_ = self._old_clone(*args, **kwargs) + t_.dist_layout = copy.deepcopy(self.dist_layout) + return t_ + + # bind the new methods to the tensor + dtensor.detach = new_detach.__get__(dtensor) + dtensor.clone = new_clone.__get__(dtensor) + return dtensor + + +def _construct_default_sharding_spec(tensor: torch.Tensor,) -> ShardingSpec: + ''' + Construct the default sharding specification for the tensor. + + Args: + tensor (`torch.Tensor`): the tensor to be sharded. + + Returns: + A `ShardingSpec` object without any sharding specified. + ''' + return ShardingSpec(dim_size=tensor.dim(), dim_partition_dict={}) + + +def _apply_layout(tensor, layout): + ''' + Apply the layout to the local tensor during initializing process. + ''' + # layout converter requires a source and target laytout + # we construct the source layer for an unsharded tensor + # and use self.dist_layer as the targer layout for the sharded tensor + source_spec = _construct_default_sharding_spec(tensor) + source_layout = Layout(device_mesh=layout.device_mesh, sharding_spec=source_spec, global_shape=tensor.shape) + sharded_tensor = layout_converter.apply(tensor=tensor, source_layout=source_layout, target_layout=layout) + return sharded_tensor + + +def distribute_tensor(tensor: torch.Tensor, device_mesh: DeviceMesh, sharding_spec: ShardingSpec) -> torch.Tensor: + """ + Convert the given tensor to a distributed tensor. + + Args: + tensor (torch.Tensor): The tensor to be converted. + device_mesh (DeviceMesh): The device mesh for abstraction of the compute devices. + sharding_spec (ShardingSpec): The sharding specification which describes how the tensor will be sharded. + + Returns: + torch.Tensor: The distributed tensor. + """ + assert not is_distributed_tensor(tensor), 'The input tensor is already a distributed tensor.' + dist_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec, global_shape=tensor.shape) + + # shard tensor + sharded_tensor = _apply_layout(tensor, dist_layout) + + # hack some tensor methods + _hijack_detach_and_clone(sharded_tensor) + + return sharded_tensor + + +def redistribute(dtensor: torch.Tensor, device_mesh: DeviceMesh, sharding_spec: ShardingSpec) -> None: + ''' + Convert the layout of the tensor from source_spec to target_spec. + This will update the `local_tensor` and `dist_layout` in place. + + Args: + dtensor (torch.Tensor): the distributed tensor to be converted. + device_mesh (DeviceMesh): the device mesh for abstraction of the compute devices. + target_layout (Layout): the target layout specification. + ''' + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + global_shape = get_global_shape(dtensor) + target_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec, global_shape=global_shape) + resharded_tensor = layout_converter.apply(tensor=dtensor, + source_layout=dtensor.dist_layout, + target_layout=target_layout) + return resharded_tensor + + +def to_global(dtensor: torch.Tensor) -> torch.Tensor: + """ + Convert a distributed tensor to the global tensor with the given layout. + This function returns a native `torch.Tensor` object. + + Args: + dtensor (torch.Tensor): the distributed tensor to be converted. + + Returns: + torch.Tensor: the global tensor. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + layout_converter = LayoutConverter() + + global_sharding_spec = ShardingSpec(dtensor.dim(), {}) + device_mesh = get_device_mesh(dtensor) + global_shape = get_global_shape(dtensor) + global_layout = Layout(device_mesh=device_mesh, sharding_spec=global_sharding_spec, global_shape=global_shape) + + global_tensor = layout_converter.apply(dtensor, dtensor.dist_layout, global_layout) + return global_tensor + + +def shard_rowwise( + tensor: torch.Tensor, + group_or_device_mesh: Union[ProcessGroup, DeviceMesh] = None, +) -> torch.Tensor: + """ + Shard the first dim of the given tensor. + + Args: + tensor (torch.Tensor): The tensor to be sharded. + group_or_device_mesh (Union[ProcessGroup, DeviceMesh], optional): The group or device mesh to shard the tensor. + If None, the tensor will be sharded with respect to the global process group. + Defaults to None. + inplace (bool, optional): Whether to shard the tensor in-place. Defaults to False. + + Returns: + torch.Tensor: The sharded tensor. + """ + # if the group_or_device_mesh is None, we shard the tensor with respect to the global process group + if group_or_device_mesh is None: + group_or_device_mesh = dist.GroupMember.WORLD + + if isinstance(group_or_device_mesh, ProcessGroup): + device_mesh = DeviceMesh.from_process_group(group_or_device_mesh) + else: + assert len(group_or_device_mesh.shape) == 1, 'Only 1D DeviceMesh is accepted for row-wise sharding.' + device_mesh = group_or_device_mesh + + sharding_spec = ShardingSpec(dim_size=tensor.dim(), dim_partition_dict={0: [0]}) + + return distribute_tensor(tensor, device_mesh, sharding_spec) + + +def shard_colwise(tensor: torch.Tensor, group_or_device_mesh: Union[ProcessGroup, DeviceMesh] = None) -> torch.Tensor: + """ + Shard the first dim of the given tensor. + + Args: + tensor (torch.Tensor): The tensor to be sharded. + group_or_device_mesh (Union[ProcessGroup, DeviceMesh], optional): The group or device mesh to shard the tensor. + If None, the tensor will be sharded with respect to the global process group. + Defaults to None. + inplace (bool, optional): Whether to shard the tensor in-place. Defaults to False. + + Returns: + torch.Tensor: The sharded tensor. + """ + # if the group_or_device_mesh is None, we shard the tensor with respect to the global process group + if group_or_device_mesh is None: + group_or_device_mesh = dist.GroupMember.WORLD + + if isinstance(group_or_device_mesh, ProcessGroup): + device_mesh = DeviceMesh.from_process_group(group_or_device_mesh) + else: + assert len(group_or_device_mesh.shape) == 1, 'Only 1D DeviceMesh is accepted for row-wise sharding.' + device_mesh = group_or_device_mesh + sharding_spec = ShardingSpec(dim_size=tensor.dim(), dim_partition_dict={-1: [0]}) + + return distribute_tensor(tensor, device_mesh, sharding_spec) + + +def sharded_tensor_to_param(dtensor: torch.Tensor, requires_grad: bool = True): + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + param = torch.nn.Parameter(dtensor, requires_grad=requires_grad) + + # make it distributed as well + param.dist_layout = dtensor.dist_layout + _hijack_detach_and_clone(param) + + return param + + +def compute_global_numel(dtensor: torch.Tensor) -> int: + """ + Compute the global number of elements in the distributed tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + int: The global number of elements in the distributed tensor. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + numel = reduce(operator.mul, dtensor.dist_layout.global_shape) + return numel + + +def get_layout(dtensor: torch.Tensor) -> Layout: + """ + Get the layout of the distributed tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + Layout: The layout of the distributed tensor. + + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + return dtensor.dist_layout + + +def get_global_shape(dtensor: torch.Tensor) -> torch.Size: + """ + Get the global shape of the distributed tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + torch.Size: The global shape of the distributed tensor. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + return dtensor.dist_layout.global_shape + + +def get_device_mesh(dtensor: torch.Tensor) -> DeviceMesh: + """ + Get the device mesh of the distributed tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + DeviceMesh: The device mesh of the distributed tensor. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + return dtensor.dist_layout.device_mesh + + +def get_sharding_spec(dtensor: torch.Tensor) -> ShardingSpec: + """ + Get the sharding spec of the distributed tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + ShardingSpec: The sharding spec of the distributed tensor. + """ + assert is_distributed_tensor(dtensor), 'The input tensor is not a distributed tensor.' + return dtensor.dist_layout.sharding_spec + + +# ====================================================== +# Some sharding does not obey the SPMD style +# e.g. Fused QKV layer in GPT2 +# we support customize sharding with the following APIs +# ====================================================== +def is_customized_distributed_tensor(tensor: torch.Tensor): + """ + Check whether the given tensor is a customized distributed tensor. + + Args: + tensor (torch.Tensor): The tensor to be checked. + + Returns: + bool: Whether the given tensor is a customized distributed tensor. + """ + return hasattr(tensor, 'shard_fn') and hasattr(tensor, 'gather_fn') + + +def _hijack_detach_and_clone_for_customized_distributed_tensor(dtensor: torch.Tensor) -> torch.Tensor: + """ + Hijack the detach and clone methods of the tensor to make sure the dist_layout is copied. + + Args: + tensor (torch.Tensor): The tensor to be hijacked. + + Returns: + torch.Tensor: The hijacked tensor. + """ + dtensor._old_detach = dtensor.detach + dtensor._old_clone = dtensor.clone + + def new_detach(self): + t_ = self._old_detach() + t_.shard_fn = self.shard_fn + t_.gather_fn = self.gather_fn + return t_ + + def new_clone(self, *args, **kwargs): + t_ = self._old_clone(*args, **kwargs) + t_.shard_fn = self.shard_fn + t_.gather_fn = self.gather_fn + return t_ + + # bind the new methods to the tensor + dtensor.detach = new_detach.__get__(dtensor) + dtensor.clone = new_clone.__get__(dtensor) + return dtensor + + +def distribute_tensor_with_customization(tensor: torch.Tensor, shard_fn, gather_fn: callable): + """ + Distribute the given tensor with the given shard_fn and gather_fn. + + Example: + + ```python + # define shard and gather functions + def shard_fn(tensor): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + return tensor.chunk(world_size, dim=0)[rank] + + def gather_fn(tensor): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + shard_list = [torch.zeros_like(tensor) for _ in range(world_size)] + torch.distributed.all_gather(shard_list, tensor) + return torch.cat(shard_list, dim=0) + + # create a distributed tensor + tensor = torch.rand(4, 4) + dtensor = distribute_tensor_with_customization(tensor, shard_fn, gather_fn) + ``` + + Args: + tensor (torch.Tensor): The tensor to be distributed. + shard_fn (callable): The function to shard the tensor. + gather_fn (callable): The function to gather the tensor. + + Returns: + torch.Tensor: The distributed tensor. + """ + assert callable(shard_fn), 'The shard_fn must be callable.' + assert callable(gather_fn), 'The gather_fn must be callable.' + assert not is_distributed_tensor(tensor), 'The input tensor is already a distributed tensor.' + + sharded_tensor = shard_fn(tensor) + + # set the shard_fn and gather_fn as attributes of the distributed tensor + sharded_tensor.shard_fn = shard_fn + sharded_tensor.gather_fn = gather_fn + + # set the shard_fn and gather_fn as attributes of the distributed tensor + _hijack_detach_and_clone_for_customized_distributed_tensor(sharded_tensor) + + return sharded_tensor + + +def to_global_for_customized_distributed_tensor(dtensor: torch.Tensor) -> torch.Tensor: + """ + Gather the given tensor to the global tensor. + + Args: + dtensor (torch.Tensor): The distributed tensor. + + Returns: + torch.Tensor: The global tensor. + """ + assert is_customized_distributed_tensor(dtensor), 'The input tensor is not a customized distributed tensor.' + return dtensor.gather_fn(dtensor) + + +def customized_distributed_tensor_to_param(dtensor: torch.Tensor, requires_grad: bool = True): + """ + Convert the given customized distributed tensor to a parameter. + """ + assert is_customized_distributed_tensor(dtensor), 'The input tensor is not a customized distributed tensor.' + + param = torch.nn.Parameter(dtensor, requires_grad=requires_grad) + + # make it distributed as well + param.shard_fn = dtensor.shard_fn + param.gather_fn = dtensor.gather_fn + _hijack_detach_and_clone_for_customized_distributed_tensor(param) + return param diff --git a/colossalai/tensor/d_tensor/comm_spec.py b/colossalai/tensor/d_tensor/comm_spec.py index 765d8ec1b01a..79b2e3ef936a 100644 --- a/colossalai/tensor/d_tensor/comm_spec.py +++ b/colossalai/tensor/d_tensor/comm_spec.py @@ -24,12 +24,12 @@ class CommSpec: ''' Communication spec is used to record the communication action. It converts the communication spec to real action which will be used in runtime. It contains comm_pattern to determine the - communication method, process_groups_dict to determine the process groups, gather_dim and shard_dim + communication method, process_group_dict to determine the process groups, gather_dim and shard_dim to determine the buffer shape, and logical_process_axis Argument: comm_pattern(CollectiveCommPattern): decribe the communication method used in this spec. - process_groups_dict(Dict): A dict which contains the process groups used to apply this CommSpec. + process_group_dict(Dict): A dict which contains the process groups used to apply this CommSpec. gather_dim(int, Optional): The gather_dim of the tensor will be gathered. shard_dim(int, Optional): The shard_dim of the tensor will be sharded. logical_process_axis(Union(int, List[int]), Optional): The mesh_dim to implement the communication action. @@ -37,7 +37,7 @@ class CommSpec: def __init__(self, comm_pattern: CollectiveCommPattern, - process_groups_dict: Dict, + process_group_dict: Dict, gather_dim: int = None, shard_dim: int = None, logical_process_axis: int = None): @@ -45,7 +45,7 @@ def __init__(self, self.gather_dim = gather_dim self.shard_dim = shard_dim self.logical_process_axis = logical_process_axis - self.process_groups_dict = process_groups_dict + self.process_group_dict = process_group_dict def __repr__(self): res_list = ["CommSpec:("] @@ -92,68 +92,56 @@ def _all_gather(tensor: torch.Tensor, comm_spec: CommSpec): ''' Implement all gather operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - tensor_list = [ - torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list)) - ] - # without this contiguous operation, the all gather may get some unexpected results. - tensor = tensor.contiguous() - dist.all_gather(tensor_list, tensor, group=process_group) - output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous() - return output + process_group = comm_spec.process_group_dict[comm_spec.logical_process_axis] + world_size = dist.get_world_size(process_group) + tensor_list = [torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) for _ in range(world_size)] + # without this contiguous operation, the all gather may get some unexpected results. + tensor = tensor.contiguous() + dist.all_gather(tensor_list, tensor, group=process_group) + output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous() + return output def _split(tensor: torch.Tensor, comm_spec: CommSpec): ''' Implement shard operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, _ in process_groups_list: - if dist.get_rank() in rank_list: - dim = comm_spec.shard_dim - length = tensor.shape[comm_spec.shard_dim] // len(rank_list) - start = length * rank_list.index(dist.get_rank()) - output = torch.narrow(tensor, dim, start, length).contiguous() - return output + process_group = comm_spec.process_group_dict[comm_spec.logical_process_axis] + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // dist.get_world_size(process_group) + start = length * dist.get_rank(process_group) + output = torch.narrow(tensor, dim, start, length).contiguous() + return output def _all_to_all(tensor: torch.Tensor, comm_spec: CommSpec): ''' Implement all to all operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - new_shape = list(tensor.shape) - new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // len(rank_list) - new_shape = torch.Size(new_shape) - output_tensor_list = [ - torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list)) - ] - dim = comm_spec.shard_dim - length = tensor.shape[comm_spec.shard_dim] // len(rank_list) - input_tensor_list = [ - torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(len(rank_list)) - ] - group = process_group - dist.all_to_all(output_tensor_list, input_tensor_list, group) - output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous() - return output + process_group = comm_spec.process_group_dict[comm_spec.logical_process_axis] + world_size = dist.get_world_size(process_group) + new_shape = list(tensor.shape) + new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // world_size + new_shape = torch.Size(new_shape) + output_tensor_list = [torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(world_size)] + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // world_size + input_tensor_list = [torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(world_size)] + group = process_group + dist.all_to_all(output_tensor_list, input_tensor_list, group) + output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous() + return output def _all_reduce(tensor: torch.Tensor, comm_spec: CommSpec, async_op: bool = False): ''' Implement all reduce operation on device mesh based on information provided by comm_spec. ''' - process_groups_list = comm_spec.process_groups_dict[comm_spec.logical_process_axis] - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - if not tensor.is_contiguous(): - tensor = tensor.contiguous() - dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group, async_op=async_op) - return tensor + process_group = comm_spec.process_group_dict[comm_spec.logical_process_axis] + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group, async_op=async_op) + return tensor class _ReduceGrad(torch.autograd.Function): @@ -269,7 +257,7 @@ def symbolic(graph, input_): def forward(ctx, input_, comm_spec): output = _all_to_all(input_, comm_spec) comm_spec_for_backward = CommSpec(comm_pattern=comm_spec.comm_pattern, - process_groups_dict=comm_spec.process_groups_dict, + process_group_dict=comm_spec.process_group_dict, gather_dim=comm_spec.shard_dim, shard_dim=comm_spec.gather_dim, logical_process_axis=comm_spec.logical_process_axis) diff --git a/colossalai/tensor/d_tensor/d_tensor.py b/colossalai/tensor/d_tensor/d_tensor.py deleted file mode 100644 index c1fe9d50a048..000000000000 --- a/colossalai/tensor/d_tensor/d_tensor.py +++ /dev/null @@ -1,142 +0,0 @@ -from typing import Optional - -import torch -from torch.utils._pytree import tree_map - -from .layout import Layout -from .layout_converter import LayoutConverter, to_global -from .sharding_spec import ShardingSpec - -layout_converter = LayoutConverter() - - -class DTensor(torch.Tensor): - - def __init__(self, local_tensor: torch.Tensor, dist_layout: Layout): - self.local_tensor = local_tensor - self.data_type = local_tensor.dtype - self.entire_shape = local_tensor.shape - self.dist_layout = dist_layout - self._apply_layout() - - @staticmethod - def __new__(cls, local_tensor, layout): - return torch.Tensor._make_subclass(cls, local_tensor, local_tensor.requires_grad) - - def __repr__(self): - return f"DTensor({self.to_global()}, {self.dist_layout})" - - def __str__(self): - return self.__repr__() - - def layout_convert(self, target_layout): - ''' - Convert the layout of the tensor from source_spec to target_spec. - ''' - self.local_tensor = layout_converter.apply(self.local_tensor, self.dist_layout, target_layout) - self.dist_layout = target_layout - - def _apply_layout(self): - ''' - Apply the layout to the local tensor during initializing process. - ''' - source_spec = construct_default_sharding_spec(self.local_tensor) - source_layout = Layout(device_mesh=self.dist_layout.device_mesh, - device_type=self.dist_layout.device_type, - sharding_spec=source_spec, - entire_shape=self.entire_shape) - self.local_tensor = layout_converter.apply(self.local_tensor, source_layout, self.dist_layout) - - @classmethod - def __torch_function__(cls, func, types, args=(), kwargs=None): - if kwargs is None: - kwargs = {} - - def filter_arg(arg): - if isinstance(arg, DTensor): - return arg.local_tensor - else: - return arg - - args = tree_map(filter_arg, args) - kwargs = tree_map(filter_arg, kwargs) - # if we want to convert the result into DTensor, we need to infer the layout of result from the layout of input tensors - # and op type. - - return func(*args, **kwargs) - - @property - def device_mesh(self): - ''' - Return the device mesh of the tensor. - ''' - return self.dist_layout.device_mesh - - @property - def sharding_spec(self): - ''' - Return the sharding specification of the tensor. - ''' - return self.dist_layout.sharding_spec - - def to(self, *args, **kwargs): - ''' - Move the tensor to a new device or convert the tensor to a new dtype. - ''' - self.local_tensor = self.local_tensor.to(*args, **kwargs) - self.data_type = self.local_tensor.dtype - self.dist_layout.device_type = self.local_tensor.device - # TODO: update the device mesh process groups or we should just cache - # both the cpu process groups and the cuda process groups? - return self - - def to_local(self): - ''' - Return the local tensor in this rank. - ''' - return self.local_tensor - - def to_global(self): - ''' - Recover the global tensor from the distributed tensor. - - Note: This function will all_gather the local tensor to the global tensor and it - will not change the layout of the DTensor. This function is mainly used for debugging or - check the correctness of the distributed tensor. - ''' - return to_global(self.local_tensor, self.dist_layout) - - -def distribute_tensor(local_tensor: torch.Tensor, dist_layout: Layout) -> DTensor: - ''' - Distribute the local tensor to the distributed tensor according to the dist_layout specified. - - Args: - local_tensor: tensor to be distributed. - dist_layout: the layout specification of the distributed tensor. - - Returns: - A 'DTensor' object. - ''' - return DTensor(local_tensor, dist_layout) - - -def distribute_module(module: torch.nn.Module, partition_fn: Optional[callable] = None) -> torch.nn.Module: - ''' - This function converts all the parameters in the module to DTensor(DParam). - - Note: This function is subject to future change as the DParam has not been implemented yet. - ''' - for name, param in module.named_parameters(): - if param is not None and not isinstance(param, DTensor): - # TODO: we could convert the parameter to DParam here, - # the type of the parameter could be an optional argument. - setattr(module, name, torch.nn.Parameter(partition_fn(name, param.data))) - return module - - -def construct_default_sharding_spec(tensor: torch.Tensor,) -> ShardingSpec: - ''' - Construct the default sharding specification for the tensor. - ''' - return ShardingSpec(dim_size=tensor.dim(), dim_partition_dict={}) diff --git a/colossalai/tensor/d_tensor/layout.py b/colossalai/tensor/d_tensor/layout.py index 72a2694a1eaf..a35b2f43e44b 100644 --- a/colossalai/tensor/d_tensor/layout.py +++ b/colossalai/tensor/d_tensor/layout.py @@ -1,12 +1,11 @@ import operator -from dataclasses import dataclass from functools import reduce import torch from colossalai.device.device_mesh import DeviceMesh -from .misc import DuplicatedShardingDimensionError, LayoutException, ShardingNotDivisibleError +from .misc import DuplicatedShardingDimensionError, ShardingNotDivisibleError from .sharding_spec import ShardingSpec @@ -14,27 +13,24 @@ class Layout: """Layout of a tensor. Attributes: - device_mesh: the device mesh to store the tensor distributedly. - device_type: the type of the device mesh, e.g. 'cpu' or 'cuda'. + device_mesh: the device mesh to store the tensor distributed. sharding_spec: the sharding specification to describe how the tensor is sharded. - entire_shape: the entire shape of the global tensor. + global_shape: the entire shape of the global tensor. """ - def __init__(self, device_mesh: DeviceMesh, device_type: torch.device, sharding_spec: ShardingSpec, - entire_shape: torch.Size): + def __init__(self, device_mesh: DeviceMesh, sharding_spec: ShardingSpec, global_shape: torch.Size): self.device_mesh = device_mesh - self.device_type = device_type self.sharding_spec = sharding_spec - self.entire_shape = entire_shape + self.global_shape = global_shape self._sanity_check() def __hash__(self) -> int: return hash(f'{self.sharding_spec}') def get_sharded_shape_per_device(self): - sharded_shape = list(self.entire_shape) + sharded_shape = list(self.global_shape) for dim, shard_list in self.sharding_spec.dim_partition_dict.items(): - mesh_list = [self.device_mesh.mesh_shape[mesh_dim] for mesh_dim in shard_list] + mesh_list = [self.device_mesh.shape[mesh_dim] for mesh_dim in shard_list] shard_partitions = reduce(operator.mul, mesh_list, 1) assert sharded_shape[ dim] % shard_partitions == 0, f'Cannot shard dimension {dim} into {shard_partitions} partitions.' @@ -45,22 +41,23 @@ def _sanity_check(self): sharding_spec = self.sharding_spec # make sure all axes in logical device mesh only be used once - dim_check_list = list(range(self.device_mesh.logical_mesh_id.dim())) - for dim, shard_list in sharding_spec.dim_partition_dict.items(): - for element in shard_list: - if element in dim_check_list: - dim_check_list.remove(element) - else: - raise DuplicatedShardingDimensionError( - f"find an invalid sharding axis {element} in dim_partition_dict in tensor dimension {dim}.") + if self.device_mesh.logical_mesh_id is not None: + dim_check_list = list(range(self.device_mesh.logical_mesh_id.dim())) + for dim, shard_list in sharding_spec.dim_partition_dict.items(): + for element in shard_list: + if element in dim_check_list: + dim_check_list.remove(element) + else: + raise DuplicatedShardingDimensionError( + f"find an invalid sharding axis {element} in dim_partition_dict in tensor dimension {dim}.") # make sure that the sharding for a dimension is divisible by the number of devices for dim, shard_list in sharding_spec.dim_partition_dict.items(): - tensor_dim_size = self.entire_shape[dim] + tensor_dim_size = self.global_shape[dim] num_devices = 1 for element in shard_list: - num_devices *= self.device_mesh.mesh_shape[element] + num_devices *= self.device_mesh.shape[element] if tensor_dim_size % num_devices != 0: raise ShardingNotDivisibleError( diff --git a/colossalai/tensor/d_tensor/layout_converter.py b/colossalai/tensor/d_tensor/layout_converter.py index a4f4c9c2dd80..528ed7901c4f 100644 --- a/colossalai/tensor/d_tensor/layout_converter.py +++ b/colossalai/tensor/d_tensor/layout_converter.py @@ -3,14 +3,12 @@ from dataclasses import dataclass from typing import Dict, List, Tuple -import numpy as np import torch -from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, TrainCycleItem from colossalai.context.singleton_meta import SingletonMeta from colossalai.tensor.d_tensor.comm_spec import * from colossalai.tensor.d_tensor.layout import Layout -from colossalai.tensor.sharding_spec import ShardingSpecException +from colossalai.tensor.d_tensor.misc import LayoutException from colossalai.tensor.utils import all_gather_simulator, all_to_all_simulator, shard_simulator from .sharding_spec import ShardingSpec @@ -28,18 +26,6 @@ class LayoutConverterOptions: pass -def to_global(distributed_tensor: torch.Tensor, layout: Layout) -> torch.Tensor: - layout_converter = LayoutConverter() - global_sharding_spec = ShardingSpec(distributed_tensor.dim(), {}) - global_layout = Layout(device_mesh=layout.device_mesh, - device_type=layout.device_type, - sharding_spec=global_sharding_spec, - entire_shape=layout.entire_shape) - with torch.no_grad(): - global_tensor = layout_converter.apply(distributed_tensor, layout, global_layout) - return global_tensor - - def set_layout_converting_options(options: LayoutConverterOptions): """ Configure the shape consistency manager via function call. @@ -49,6 +35,9 @@ def set_layout_converting_options(options: LayoutConverterOptions): class LayoutConverter(metaclass=SingletonMeta): + """ + LayoutConverter is a singleton class which converts the layout of a distributed tensor. + """ def __init__(self): self._options = None @@ -91,15 +80,14 @@ def all_gather_transform_layouts(self, source_layout: Layout) -> Dict[Layout, Co # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - entire_shape = (4, 4, 4) + global_shape = (4, 4, 4) dim_partition_dict = {0: [0], 1: [1]} # [S0,S1,R] sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict) layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec, - entire_shape=entire_shape) + global_shape=global_shape) rst_dict = layout_converter.all_gather_transform_layouts(layout) for layout, comm_spec in rst_dict.items(): @@ -112,7 +100,12 @@ def all_gather_transform_layouts(self, source_layout: Layout) -> Dict[Layout, Co valid_spec_dict = {} comm_pattern = CollectiveCommPattern.GATHER_FWD_SPLIT_BWD source_spec = source_layout.sharding_spec - process_groups_dict = source_layout.device_mesh.process_groups_dict + + # the key of the dict is the axis + # the value is the process group + current_rank = source_layout.device_mesh._global_rank_of_current_process + process_group_dict = source_layout.device_mesh._process_group_dict[current_rank] + for target_pair in source_spec.dim_partition_dict.items(): shard_list = all_gather_simulator(target_pair) index = target_pair[0] @@ -130,7 +123,7 @@ def all_gather_transform_layouts(self, source_layout: Layout) -> Dict[Layout, Co logical_process_axis = target_pair[1][-1] comm_spec = CommSpec( comm_pattern, - process_groups_dict=process_groups_dict, + process_group_dict=process_group_dict, gather_dim=gather_dim, # shard_dim will be used during backward shard_dim=gather_dim, @@ -141,11 +134,10 @@ def all_gather_transform_layouts(self, source_layout: Layout) -> Dict[Layout, Co new_sharding_spec = ShardingSpec(source_spec.dims, dim_partition_dict=new_dim_partition_dict) new_layout = Layout(device_mesh=source_layout.device_mesh, sharding_spec=new_sharding_spec, - device_type=source_layout.device_type, - entire_shape=source_layout.entire_shape) + global_shape=source_layout.global_shape) valid_spec_dict[new_layout] = comm_spec - except ShardingSpecException: + except LayoutException: pass return valid_spec_dict @@ -167,15 +159,14 @@ def all_to_all_transform_layout(self, source_layout: Layout) -> Dict[Layout, Com # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - entire_shape = (4, 4, 4) + global_shape = (4, 4, 4) dim_partition_dict = {0: [0], 1: [1]} # [S0,S1,R] sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict) layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec, - entire_shape=entire_shape) + global_shape=global_shape) rst_dict = layout_converter.all_to_all_transform_layout(layout) for layout, comm_spec in rst_dict.items(): @@ -188,7 +179,12 @@ def all_to_all_transform_layout(self, source_layout: Layout) -> Dict[Layout, Com ''' valid_spec_dict = {} comm_pattern = CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD - process_groups_dict = source_layout.device_mesh.process_groups_dict + + # the key of the dict is the axis + # the value is the process group + current_rank = source_layout.device_mesh._global_rank_of_current_process + process_group_dict = source_layout.device_mesh._process_group_dict[current_rank] + source_spec = source_layout.sharding_spec tensor_dims = source_spec.dims for f_index in range(tensor_dims - 1): @@ -229,7 +225,7 @@ def all_to_all_transform_layout(self, source_layout: Layout) -> Dict[Layout, Com shard_dim = f_index logical_process_axis = b_target_pair[1][-1] comm_spec = CommSpec(comm_pattern, - process_groups_dict, + process_group_dict=process_group_dict, gather_dim=gather_dim, shard_dim=shard_dim, logical_process_axis=logical_process_axis) @@ -252,10 +248,9 @@ def all_to_all_transform_layout(self, source_layout: Layout) -> Dict[Layout, Com new_sharding_spec = ShardingSpec(source_spec.dims, dim_partition_dict=new_dim_partition_dict) new_layout = Layout(device_mesh=source_layout.device_mesh, sharding_spec=new_sharding_spec, - device_type=source_layout.device_type, - entire_shape=source_layout.entire_shape) + global_shape=source_layout.global_shape) valid_spec_dict[new_layout] = comm_spec - except ShardingSpecException: + except LayoutException: pass return valid_spec_dict @@ -278,16 +273,15 @@ def shard_transform_layout(self, source_layout: Layout) -> Dict[Layout, CommSpec # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - entire_shape = (4, 4, 4) + global_shape = (4, 4, 4) dim_partition_dict = {0: [0]} # [S0,R,R] sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict) layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec, - entire_shape=entire_shape) + global_shape=global_shape) rst_dict = layout_converter.shard_transform_layout(layout) for layout, comm_spec in rst_dict.items(): @@ -301,10 +295,14 @@ def shard_transform_layout(self, source_layout: Layout) -> Dict[Layout, CommSpec valid_spec_dict = {} comm_pattern = CollectiveCommPattern.SPLIT_FWD_GATHER_BWD source_spec = source_layout.sharding_spec - process_groups_dict = source_layout.device_mesh.process_groups_dict + + # the key of the dict is the axis + # the value is the process group + current_rank = source_layout.device_mesh._global_rank_of_current_process + process_group_dict = source_layout.device_mesh._process_group_dict[current_rank] # legal sharding dims means the mesh_id is still available to use. - legal_sharding_dims = [i for i in range(len(source_layout.device_mesh.mesh_shape))] + legal_sharding_dims = [i for i in range(len(source_layout.device_mesh.shape))] for dim, shard_list in source_spec.dim_partition_dict.items(): for element in shard_list: legal_sharding_dims.remove(element) @@ -329,7 +327,7 @@ def shard_transform_layout(self, source_layout: Layout) -> Dict[Layout, CommSpec shard_dim = index logical_process_axis = shard_list[-1] comm_spec = CommSpec(comm_pattern, - process_groups_dict, + process_group_dict=process_group_dict, gather_dim=shard_dim, shard_dim=shard_dim, logical_process_axis=logical_process_axis) @@ -340,10 +338,9 @@ def shard_transform_layout(self, source_layout: Layout) -> Dict[Layout, CommSpec dim_partition_dict=new_dim_partition_dict) new_layout = Layout(device_mesh=source_layout.device_mesh, sharding_spec=new_sharding_spec, - device_type=source_layout.device_type, - entire_shape=source_layout.entire_shape) + global_shape=source_layout.global_shape) valid_spec_dict[new_layout] = comm_spec - except ShardingSpecException: + except LayoutException: pass return valid_spec_dict @@ -399,7 +396,7 @@ def layout_converting(self, source_layout: Layout, # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - entire_shape = (4, 4, 4) + global_shape = (4, 4, 4) dim_partition_source = {1: [0, 1]} dim_partition_target = {0: [0, 1]} @@ -407,16 +404,14 @@ def layout_converting(self, source_layout: Layout, # [R,S01,R] sharding_spec_source = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_source) source_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec_source, - entire_shape=entire_shape) + global_shape=global_shape) # [S01,R,R] sharding_spec_target = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_target) target_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec_target, - entire_shape=entire_shape) + global_shape=global_shape) transform_path, comm_action_sequence = layout_converter.layout_converting(source_layout, target_layout) transform_path_str = '->'.join([str(layout.sharding_spec.sharding_sequence) for layout in transform_path]) @@ -505,21 +500,19 @@ def apply(self, tensor: torch.Tensor, source_layout: Layout, target_layout: Layo # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - entire_shape = (4, 4, 4) + global_shape = (4, 4, 4) # [S0,R,R] sharding_spec_source = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_source) source_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec_source, - entire_shape=entire_shape) + global_shape=global_shape) # [R,S0,R] sharding_spec_target = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_target) target_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), sharding_spec=sharding_spec_target, - entire_shape=entire_shape) + global_shape=global_shape) if rank in (0, 1): sharded_tensor_0 = torch.zeros(2, 1) @@ -553,4 +546,5 @@ def apply(self, tensor: torch.Tensor, source_layout: Layout, target_layout: Layo _, comm_action_sequence = self.layout_converting(source_layout, target_layout) for comm_spec in comm_action_sequence: tensor = comm_spec.covert_spec_to_action(tensor) + tensor.dist_layout = target_layout return tensor diff --git a/colossalai/tensor/d_tensor/sharding_spec.py b/colossalai/tensor/d_tensor/sharding_spec.py index 7591f760cb30..565012b58a03 100644 --- a/colossalai/tensor/d_tensor/sharding_spec.py +++ b/colossalai/tensor/d_tensor/sharding_spec.py @@ -14,7 +14,7 @@ class DimSpec: ''' - Sharding spec for single dimension of the sharded tensor decribe the sharding dimension of + Sharding spec for single dimension of the sharded tensor describe the sharding dimension of logical device mesh and give a method to compute the difference between them. This class is used internally in ShardingSpec. @@ -41,7 +41,7 @@ def __repr__(self): def _convert_str_to_shard_list(self, str_spec): ''' - Conver str_spec into shard_list. + Convert str_spec into shard_list. Argument: str_spec(str): dim spec in str type. @@ -58,7 +58,7 @@ def _convert_str_to_shard_list(self, str_spec): def build_difference_2d_dict(self): ''' - Build a difference maping for 2D device mesh case. It will be used to + Build a difference mapping for 2D device mesh case. It will be used to compute the difference between DimSpec pairs. ''' @@ -143,7 +143,7 @@ class ShardingSpec: Argument: dim_partition_dict(Dict[int, List[int]], optional): The key is the dimension of tensor to be sharded, - and the value of the key decribe which logical axis will be sharded in that dimension. + and the value of the key describe which logical axis will be sharded in that dimension. sharding_sequence(List[DimSpec], optional): A straight view of ShardingSpec looks like [R, R, S0, S1]. ''' diff --git a/colossalai/tensor/d_tensor/utils.py b/colossalai/tensor/d_tensor/utils.py index 644bb6306b42..fc22b990d879 100644 --- a/colossalai/tensor/d_tensor/utils.py +++ b/colossalai/tensor/d_tensor/utils.py @@ -29,7 +29,7 @@ def get_comm_cost(layout: Layout, comm_spec: CommSpec, forward_only: bool = Fals # the comm size for all gather is the size of the gathered tensor gather_dim = comm_spec.gather_dim all_gather_axis = layout.sharding_spec.dim_partition_dict[gather_dim][-1] - all_gather_size = device_mesh.mesh_shape[all_gather_axis] + all_gather_size = device_mesh.shape[all_gather_axis] comm_size_for_all_gather = comm_size * all_gather_size forward_communication_cost = device_mesh.all_gather_cost(comm_size_for_all_gather, logical_process_axis) # give a tiny cost to shard diff --git a/colossalai/tensor/dist_spec_mgr.py b/colossalai/tensor/dist_spec_mgr.py index d5c0ce28e9fb..c968050de49d 100644 --- a/colossalai/tensor/dist_spec_mgr.py +++ b/colossalai/tensor/dist_spec_mgr.py @@ -4,10 +4,8 @@ import torch.distributed as dist # from colossalai.nn.layer.utils import divide from numpy import prod -from packaging import version -from colossalai.logging import get_dist_logger -from colossalai.tensor.distspec import _DistSpec +from colossalai.tensor.distspec import DistPlacementPattern, _DistSpec from colossalai.tensor.process_group import ProcessGroup @@ -61,7 +59,7 @@ def _shard_as(tensor: torch.Tensor, old_dist_spec: _DistSpec, dist_spec: _DistSp Args: tensor (torch.Tensor): a global (replicated) tensor before shard dist_spec (_DistSpec): the distributed spec. to be sharded as. - pg (ProcessGrouo): the process group of the corresponding colotensor + pg (ProcessGroup): the process group of the corresponding colotensor Returns: torch.Tensor: a torch tensor after sharded. """ @@ -171,11 +169,21 @@ def handle_trans_spec(tensor: torch.Tensor, old_dist_spec: _DistSpec, dist_spec: pg: ProcessGroup) -> torch.Tensor: assert isinstance(old_dist_spec, _DistSpec), f"{type(old_dist_spec)} should be _DistSpec" assert isinstance(dist_spec, _DistSpec), f"{type(dist_spec)} should be _DistSpec" - forward_trans_handle = getattr(DistSpecManager, f'_{old_dist_spec.placement.value}2{dist_spec.placement.value}') + + trans_func_key = (old_dist_spec.placement, dist_spec.placement) + trans_funcs = { + (DistPlacementPattern.REPLICATE, DistPlacementPattern.REPLICATE): DistSpecManager._r2r, + (DistPlacementPattern.REPLICATE, DistPlacementPattern.SHARD): DistSpecManager._r2s, + (DistPlacementPattern.SHARD, DistPlacementPattern.REPLICATE): DistSpecManager._s2r, + (DistPlacementPattern.SHARD, DistPlacementPattern.SHARD): DistSpecManager._s2s + } + + forward_trans_handle = trans_funcs[trans_func_key] if not DistSpecManager._use_autograd_function: return forward_trans_handle(tensor, old_dist_spec, dist_spec, pg) - backward_trans_handle = getattr(DistSpecManager, - f'_{dist_spec.placement.value}2{old_dist_spec.placement.value}') + + backward_trans_handle = trans_funcs[(dist_spec.placement, old_dist_spec.placement)] + return TransformDistSpec.apply(tensor, old_dist_spec, dist_spec, pg, forward_trans_handle, backward_trans_handle) diff --git a/colossalai/tensor/distspec.py b/colossalai/tensor/distspec.py index 8dd0d8791537..3a09f1426e31 100644 --- a/colossalai/tensor/distspec.py +++ b/colossalai/tensor/distspec.py @@ -15,7 +15,7 @@ class _DistSpec: A class indicates Distributed Specification. The DistSpec is only works for the tensor parallel process groups. Because the dist spec of data parallel process group can be automatically deduced. - This is an internal data structrue. + This is an internal data structure. The API for users should be `ShardSpec` and `ReplicaSpec`. Args: diff --git a/colossalai/tensor/param_op_hook.py b/colossalai/tensor/param_op_hook.py index ed705da0eb0d..8ed8176d996a 100644 --- a/colossalai/tensor/param_op_hook.py +++ b/colossalai/tensor/param_op_hook.py @@ -164,16 +164,16 @@ def _get_grad_args(*args): for obj in args: if _is_grad_tensor(obj): return args, None - # otherwise, the first arguement should be a tuple of grad tensors + # otherwise, the first argument should be a tuple of grad tensors # if there is no grad tensor, the backward of PreFwdPostBwd can't be triggered arg_zero = args[0] if not isinstance(arg_zero, tuple): - raise NotImplementedError("Some torch function is incompatible because of its complcated inputs.") + raise NotImplementedError("Some torch function is incompatible because of its complicated inputs.") check_grad_flag = False for obj in arg_zero: check_grad_flag |= _is_grad_tensor(obj) if not check_grad_flag: - raise NotImplementedError("Some torch function is incompatible because of its complcated inputs.") + raise NotImplementedError("Some torch function is incompatible because of its complicated inputs.") return arg_zero, args[1:] diff --git a/colossalai/tensor/process_group.py b/colossalai/tensor/process_group.py index f108bdc247f5..8d2e9a616d76 100644 --- a/colossalai/tensor/process_group.py +++ b/colossalai/tensor/process_group.py @@ -130,7 +130,7 @@ def set_cpu_groups(self): @property def has_cpu_groups(self) -> bool: """has_cpu_groups - If cpu groups have been initailized. + If cpu groups have been initialized. Returns: bool: cpu process groups have been initialized or not. diff --git a/colossalai/tensor/shape_consistency.py b/colossalai/tensor/shape_consistency.py index 2831b10a3c57..99d782c3f6e8 100644 --- a/colossalai/tensor/shape_consistency.py +++ b/colossalai/tensor/shape_consistency.py @@ -73,7 +73,7 @@ def get_all_all_gather_spec(self, source_spec: ShardingSpec, orig_cost_dict: Dict[str, float]) -> Dict[ShardingSpec, float]: ''' Get all valid sharding specs from source_spec with single all-gather operation, and - accumulate commucation cost on origin cost which will finally be used in auto sharding solver. + accumulate communication cost on origin cost which will finally be used in auto sharding solver. For the all-gather operation, we just care about the S dimension. Argument: @@ -145,7 +145,7 @@ def get_all_all_to_all_spec(self, source_spec: ShardingSpec, orig_cost_dict: Dict[str, float]) -> Dict[ShardingSpec, float]: ''' Get all valid sharding specs from source_spec with single all-to-all operation, and - accumulate commucation cost on origin cost which will finally be used in auto sharding solver. + accumulate communication cost on origin cost which will finally be used in auto sharding solver. For the all-to-all operation, we just care about the pairs containing S dimension. Argument: @@ -252,7 +252,7 @@ def get_all_all_to_all_spec(self, source_spec: ShardingSpec, def get_all_shard_spec(self, source_spec: ShardingSpec, orig_cost_dict): ''' Get all valid sharding specs from source_spec with single shard operation, and - accumulate commucation cost on origin cost which will finally be used in auto sharding solver. + accumulate communication cost on origin cost which will finally be used in auto sharding solver. For the sharding operation, we just care about legal sharding dimensions. Argument: @@ -285,7 +285,7 @@ def get_all_shard_spec(self, source_spec: ShardingSpec, orig_cost_dict): comm_pattern = CollectiveCommPattern.SPLIT_FWD_GATHER_BWD # legal sharding dims means the mesh_id is still available to use. - legal_sharding_dims = [i for i in range(len(source_spec.device_mesh.mesh_shape))] + legal_sharding_dims = [i for i in range(len(source_spec.device_mesh.shape))] for dim, shard_list in source_spec.dim_partition_dict.items(): for element in shard_list: legal_sharding_dims.remove(element) @@ -386,7 +386,7 @@ def get_all_mix_gather_spec(self, source_spec: ShardingSpec, def get_all_one_step_transform_spec(self, source_spec: ShardingSpec, orig_cost_dict) -> Dict[ShardingSpec, float]: ''' Get all valid sharding specs from source_spec with one step transform, and - accumulate commucation cost on origin cost which will finally be used in auto sharding solver. + accumulate communication cost on origin cost which will finally be used in auto sharding solver. Note: all-gather will eliminate a sharding dimension, all-to-all will keep sharding dimension same as before, and shard will add a sharding dimension. Therefore, the result of above operations are mutual exclusive, @@ -435,7 +435,7 @@ def gather_analysis(comm_spec: CommSpec, discard_input: bool, alloc_numel: int, """ input_shape = compute_shape(comm_spec.sharding_spec) input_numel = np.prod(input_shape) - output_numel = input_numel * comm_spec.device_mesh.mesh_shape[comm_spec.logical_process_axis] + output_numel = input_numel * comm_spec.device_mesh.shape[comm_spec.logical_process_axis] peak_numel = max(peak_numel, alloc_numel + output_numel * 2) alloc_numel += output_numel if discard_input: @@ -461,7 +461,7 @@ def split_analysis(comm_spec: CommSpec, discard_input: bool, alloc_numel: int, p # generate a new tensor input_shape = compute_shape(comm_spec.sharding_spec) input_numel = np.prod(input_shape) - output_numel = input_numel // comm_spec.device_mesh.mesh_shape[comm_spec.logical_process_axis] + output_numel = input_numel // comm_spec.device_mesh.shape[comm_spec.logical_process_axis] alloc_numel += output_numel peak_numel = max(peak_numel, alloc_numel) if discard_input: @@ -577,7 +577,7 @@ def shape_consistency(self, source_spec: ShardingSpec, Step3: Repeat above steps until the source spec transform to target spec. - During finding the transform path, commucation cost will be accumulated, and it + During finding the transform path, communication cost will be accumulated, and it will be finally used in auto parallel solver. Additionally, to avoid repeating the path search in runtime, we cached all solved path diff --git a/colossalai/tensor/sharding_spec.py b/colossalai/tensor/sharding_spec.py index cdd0338850cf..e594fd297dc4 100644 --- a/colossalai/tensor/sharding_spec.py +++ b/colossalai/tensor/sharding_spec.py @@ -18,7 +18,7 @@ class _DimSpec: ''' - Sharding spec for single dimension of the sharded tensor decribe the sharding dimension of + Sharding spec for single dimension of the sharded tensor describe the sharding dimension of logical device mesh and give a method to compute the difference between them. This class is used internally in ShardingSpec. @@ -45,7 +45,7 @@ def __repr__(self): def _convert_str_to_shard_list(self, str_spec): ''' - Conver str_spec into shard_list. + Convert str_spec into shard_list. Argument: str_spec(str): dim spec in str type. @@ -62,7 +62,7 @@ def _convert_str_to_shard_list(self, str_spec): def build_difference_2d_dict(self): ''' - Build a difference maping for 2D device mesh case. It will be used to + Build a difference mapping for 2D device mesh case. It will be used to compute the difference between DimSpec pairs. ''' @@ -166,7 +166,7 @@ class ShardingSpec: device_mesh(DeviceMesh): A logical view of a physical mesh. entire_shape(torch.Size): The entire shape of tensor before sharded. dim_partition_dict(Dict[int, List[int]]， optional): The key is the dimension of tensor to be sharded, - and the value of the key decribe which logical axis will be sharded in that dimension. + and the value of the key describe which logical axis will be sharded in that dimension. sharding_sequence(List[_DimSpec], optional): A straight view of ShardingSpec looks like [R, R, S0, S1]. ''' @@ -195,7 +195,7 @@ def __init__(self, def __repr__(self): res_list = ["DistSpec:"] res_list.append(f"\n\tshard_sequence: " + ",".join(str(dimspec) for dimspec in self.sharding_sequence)) - res_list.append(f"\n\tdevice_mesh_shape: {self.device_mesh.mesh_shape}") + res_list.append(f"\n\tdevice_mesh_shape: {self.device_mesh.shape}") return ' '.join(res_list) def _sanity_check(self): @@ -222,7 +222,7 @@ def _sanity_check(self): num_devices = 1 for element in shard_list: - num_devices *= self.device_mesh.mesh_shape[element] + num_devices *= self.device_mesh.shape[element] if tensor_dim_size % num_devices != 0: raise ShardingNotDivisibleError( @@ -288,7 +288,7 @@ def get_sharded_shape_per_device(self): sharded_shape = list(self.entire_shape) for dim, shard_list in self.dim_partition_dict.items(): - mesh_list = [self.device_mesh.mesh_shape[mesh_dim] for mesh_dim in shard_list] + mesh_list = [self.device_mesh.shape[mesh_dim] for mesh_dim in shard_list] shard_partitions = reduce(operator.mul, mesh_list, 1) assert sharded_shape[ dim] % shard_partitions == 0, f'Cannot shard dimension {dim} into {shard_partitions} partitions.' diff --git a/colossalai/tensor/utils.py b/colossalai/tensor/utils.py index 0c2ead630d59..e7d51d099e02 100644 --- a/colossalai/tensor/utils.py +++ b/colossalai/tensor/utils.py @@ -18,7 +18,7 @@ def all_gather_simulator(target_pair): Argument: target_pair(Tuple[int, List[int]]): The first element is the dimension of tensor to be sharded, - and the second element decribes which logical axis will be sharded in that dimension. + and the second element describes which logical axis will be sharded in that dimension. ''' _, shard_list = target_pair new_shard_list = shard_list[:-1] @@ -36,7 +36,7 @@ def all_to_all_simulator(f_target_pair, b_target_pair): Therefore, if the behind shard_list is not None, we just extend it to the front shard_list. Argument: target_pair(Tuple[int, List[int]]): The first element is the dimension of tensor to be sharded, - and the second element decribes which logical axis will be sharded in that dimension. + and the second element describes which logical axis will be sharded in that dimension. e.g.: all-to-all(S0, S1) -> [S01, R] all-to-all(S0, R) -> [R, S0] @@ -46,7 +46,7 @@ def all_to_all_simulator(f_target_pair, b_target_pair): Argument: target_pair(Tuple[int, List[int]]): The first element is the dimension of tensor to be sharded, - and the second element decribes which logical axis will be sharded in that dimension. + and the second element describes which logical axis will be sharded in that dimension. ''' _, f_shard_list = f_target_pair _, b_shard_list = b_target_pair @@ -77,7 +77,7 @@ def shard_simulator(target_pair, legal_sharding_dims): Argument: target_pair(Tuple[int, List[int]]): The first element is the dimension of tensor to be sharded, - and the second element decribes which logical axis will be sharded in that dimension. + and the second element describes which logical axis will be sharded in that dimension. ''' _, shard_list = target_pair shard_list_list = [] diff --git a/colossalai/testing/__init__.py b/colossalai/testing/__init__.py index e3dd500dea8e..0db33361c6a0 100644 --- a/colossalai/testing/__init__.py +++ b/colossalai/testing/__init__.py @@ -1,7 +1,25 @@ -from .comparison import assert_equal, assert_not_equal, assert_close, assert_close_loose, assert_equal_in_group -from .utils import parameterize, rerun_on_exception, rerun_if_address_is_in_use, skip_if_not_enough_gpus +from .comparison import ( + assert_close, + assert_close_loose, + assert_equal, + assert_equal_in_group, + assert_hf_output_close, + assert_not_equal, + check_state_dict_equal, +) +from .pytest_wrapper import run_on_environment_flag +from .utils import ( + clear_cache_before_run, + free_port, + parameterize, + rerun_if_address_is_in_use, + rerun_on_exception, + skip_if_not_enough_gpus, + spawn, +) __all__ = [ 'assert_equal', 'assert_not_equal', 'assert_close', 'assert_close_loose', 'assert_equal_in_group', 'parameterize', - 'rerun_on_exception', 'rerun_if_address_is_in_use', 'skip_if_not_enough_gpus' + 'rerun_on_exception', 'rerun_if_address_is_in_use', 'skip_if_not_enough_gpus', 'free_port', 'spawn', + 'clear_cache_before_run', 'run_on_environment_flag', 'check_state_dict_equal', 'assert_hf_output_close' ] diff --git a/colossalai/testing/comparison.py b/colossalai/testing/comparison.py index e00d0da168c7..8d9ec8ab5f35 100644 --- a/colossalai/testing/comparison.py +++ b/colossalai/testing/comparison.py @@ -1,8 +1,11 @@ +from typing import Any, List, OrderedDict + import torch import torch.distributed as dist from torch import Tensor from torch.distributed import ProcessGroup from torch.testing import assert_close +from torch.utils._pytree import tree_flatten def assert_equal(a: Tensor, b: Tensor): @@ -14,7 +17,12 @@ def assert_not_equal(a: Tensor, b: Tensor): def assert_close_loose(a: Tensor, b: Tensor, rtol: float = 1e-3, atol: float = 1e-3): - assert_close(a, b, rtol=rtol, atol=atol) + assert_close(a, + b, + rtol=rtol, + atol=atol, + msg=f"Tensor not close, shape: {a.shape} vs {b.shape}, \ + dtype: {a.dtype} vs {b.dtype}") def assert_equal_in_group(tensor: Tensor, process_group: ProcessGroup = None): @@ -28,3 +36,100 @@ def assert_equal_in_group(tensor: Tensor, process_group: ProcessGroup = None): a = tensor_list[i] b = tensor_list[i + 1] assert torch.all(a == b), f'expected tensors on rank {i} and {i + 1} to be equal but they are not, {a} vs {b}' + + +def check_state_dict_equal(d1: OrderedDict, d2: OrderedDict, ignore_device: bool = True): + assert len(list(d1.keys())) == len(list(d2.keys())), \ + f"Number of keys unequal: {len(list(d1.keys()))} vs {len(list(d2.keys()))}" + for k, v1 in d1.items(): + assert k in d2 + v2 = d2[k] + if isinstance(v1, dict): + assert isinstance(v2, dict) + check_state_dict_equal(v1, v2, ignore_device) + elif isinstance(v1, list): + assert isinstance(v2, list) + for v1_i, v2_i in zip(v1, v2): + if isinstance(v1_i, torch.Tensor): + assert isinstance(v2_i, torch.Tensor) + if not ignore_device: + v1_i = v1_i.to("cpu") + v2_i = v2_i.to("cpu") + assert_close_loose(v1_i, v2_i) + elif isinstance(v1_i, dict): + assert isinstance(v2_i, dict) + check_state_dict_equal(v1_i, v2_i, ignore_device) + else: + assert v1_i == v2_i, f"{v1_i} not equals to {v2_i}" + elif isinstance(v1, torch.Tensor): + assert isinstance(v2, torch.Tensor) + if not ignore_device: + v1 = v1.to("cpu") + v2 = v2.to("cpu") + assert_close_loose(v1, v2) + else: + assert v1 == v2, f"{v1} not equals to {v2}" + + +def check_state_dict_equal_pytree(d1: OrderedDict, d2: OrderedDict, ignore_device: bool = True): + flat_d1, _ = tree_flatten(d1) + flat_d2, _ = tree_flatten(d2) + assert len(flat_d1) == len(flat_d2) + for v1, v2 in zip(flat_d1, flat_d2): + if isinstance(v1, torch.Tensor): + assert isinstance(v2, torch.Tensor) + if not ignore_device: + v1 = v1.to("cpu") + v2 = v2.to("cpu") + assert_close_loose(v1, v2) + else: + assert v1 == v2, f"{v1} not equals to {v2}" + + +def assert_hf_output_close(out1: Any, + out2: Any, + ignore_keys: List[str] = None, + track_name: str = "", + atol=1e-5, + rtol=1e-5): + """ + Check if two outputs from huggingface are equal. + + Args: + out1 (Any): the first output + out2 (Any): the second output + ignore_keys (List[str]): the keys to ignore when comparing two dicts + track_name (str): the name of the value compared, used to track the path + """ + if isinstance(out1, dict) and isinstance(out2, dict): + # if two values are dict + # we recursively check the keys + assert set(out1.keys()) == set(out2.keys()) + for k in out1.keys(): + if ignore_keys is not None and k in ignore_keys: + continue + assert_hf_output_close(out1[k], + out2[k], + track_name=f"{track_name}.{k}", + ignore_keys=ignore_keys, + atol=atol, + rtol=rtol) + elif isinstance(out1, (list, tuple)) and isinstance(out2, (list, tuple)): + # if two values are list + # we recursively check the elements + assert len(out1) == len(out2) + for i in range(len(out1)): + assert_hf_output_close(out1[i], + out2[i], + track_name=f"{track_name}.{i}", + ignore_keys=ignore_keys, + atol=atol, + rtol=rtol) + elif isinstance(out1, Tensor) and isinstance(out2, Tensor): + if out1.shape != out2.shape: + raise AssertionError(f"{track_name}: shape mismatch: {out1.shape} vs {out2.shape}") + assert torch.allclose( + out1, out2, atol=atol, rtol=rtol + ), f"{track_name}: tensor value mismatch\nvalue 1: {out1}\nvalue 2: {out2}, \nmean error: {torch.abs(out1 - out2).mean()}" + else: + assert out1 == out2, f"{track_name}: value mismatch.\nout1: {out1}\nout2: {out2}" diff --git a/colossalai/testing/pytest_wrapper.py b/colossalai/testing/pytest_wrapper.py index a472eb3723ec..6a80e1dcc548 100644 --- a/colossalai/testing/pytest_wrapper.py +++ b/colossalai/testing/pytest_wrapper.py @@ -1,10 +1,9 @@ """ This file will not be automatically imported by `colossalai.testing` -as this file has a dependency on `pytest`. Therefore, you need to +as this file has a dependency on `pytest`. Therefore, you need to explicitly import this file `from colossalai.testing.pytest_wrapper import `.from """ -import pytest import os @@ -30,10 +29,16 @@ def test_for_something(): pytest test_for_something.py """ + try: + import pytest + except ImportError: + raise ImportError( + 'This function requires `pytest` to be installed, please do `pip install pytest` and try again.') + assert isinstance(name, str) flag = os.environ.get(name.upper(), '0') - reason = f'Environment varialbe {name} is {flag}' + reason = f'Environment variable {name} is {flag}' if flag == '1': return pytest.mark.skipif(False, reason=reason) else: diff --git a/colossalai/testing/utils.py b/colossalai/testing/utils.py index 64c1d6e7bcd0..a4370a8d4933 100644 --- a/colossalai/testing/utils.py +++ b/colossalai/testing/utils.py @@ -1,8 +1,13 @@ +import gc +import random import re -import torch -from typing import Callable, List, Any +import socket from functools import partial from inspect import signature +from typing import Any, Callable, List + +import torch +import torch.multiprocessing as mp from packaging import version @@ -12,10 +17,10 @@ def parameterize(argument: str, values: List[Any]) -> Callable: we want to avoid the number of distributed network initialization, we need to have this extra decorator on the function launched by torch.multiprocessing. - If a function is wrapped with this wrapper, non-paramterized arguments must be keyword arguments, - positioanl arguments are not allowed. + If a function is wrapped with this wrapper, non-parametrized arguments must be keyword arguments, + positional arguments are not allowed. - Usgae:: + Usage:: # Example 1: @parameterize('person', ['xavier', 'davis']) @@ -28,7 +33,7 @@ def say_something(person, msg): # > xavier: hello # > davis: hello - # Exampel 2: + # Example 2: @parameterize('person', ['xavier', 'davis']) @parameterize('msg', ['hello', 'bye', 'stop']) def say_something(person, msg): @@ -43,7 +48,7 @@ def say_something(person, msg): # > davis: hello # > davis: bye # > davis: stop - + Args: argument (str): the name of the argument to parameterize values (List[Any]): a list of values to iterate for this argument @@ -85,13 +90,13 @@ def test_method(): def test_method(): print('hey') raise RuntimeError('Address already in use') - + # rerun for infinite times if Runtime error occurs @rerun_on_exception(exception_type=RuntimeError, max_try=None) def test_method(): print('hey') raise RuntimeError('Address already in use') - + # rerun only the exception message is matched with pattern # for infinite times if Runtime error occurs @rerun_on_exception(exception_type=RuntimeError, pattern="^Address.*$") @@ -101,11 +106,11 @@ def test_method(): Args: exception_type (Exception, Optional): The type of exception to detect for rerun - pattern (str, Optional): The pattern to match the exception message. + pattern (str, Optional): The pattern to match the exception message. If the pattern is not None and matches the exception message, the exception will be detected for rerun - max_try (int, Optional): Maximum reruns for this function. The default value is 5. - If max_try is None, it will rerun foreven if exception keeps occurings + max_try (int, Optional): Maximum reruns for this function. The default value is 5. + If max_try is None, it will rerun forever if exception keeps occurring """ def _match_lines(lines, pattern): @@ -139,7 +144,7 @@ def _run_until_success(*args, **kwargs): # Override signature # otherwise pytest.mark.parameterize will raise the following error: - # function does not use argumetn xxx + # function does not use argument xxx sig = signature(func) _run_until_success.__signature__ = sig @@ -162,10 +167,10 @@ def test_something(): """ # check version torch_version = version.parse(torch.__version__) - assert torch_version.major == 1 + assert torch_version.major >= 1 # only torch >= 1.8 has ProcessRaisedException - if torch_version.minor >= 8: + if torch_version >= version.parse("1.8.0"): exception = torch.multiprocessing.ProcessRaisedException else: exception = Exception @@ -202,3 +207,72 @@ def _execute_by_gpu_num(*args, **kwargs): return _execute_by_gpu_num return _wrap_func + + +def free_port() -> int: + """Get a free port on localhost. + + Returns: + int: A free port on localhost. + """ + while True: + port = random.randint(20000, 65000) + try: + with socket.socket() as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("localhost", port)) + return port + except OSError: + continue + + +def spawn(func, nprocs=1, **kwargs): + """ + This function is used to spawn processes for testing. + + Usage: + # must contains arguments rank, world_size, port + def do_something(rank, world_size, port): + ... + + spawn(do_something, nprocs=8) + + # can also pass other arguments + def do_something(rank, world_size, port, arg1, arg2): + ... + + spawn(do_something, nprocs=8, arg1=1, arg2=2) + + Args: + func (Callable): The function to be spawned. + nprocs (int, optional): The number of processes to spawn. Defaults to 1. + """ + port = free_port() + wrapped_func = partial(func, world_size=nprocs, port=port, **kwargs) + mp.spawn(wrapped_func, nprocs=nprocs) + + +def clear_cache_before_run(): + """ + This function is a wrapper to clear CUDA and python cache before executing the function. + + Usage: + @clear_cache_before_run() + def test_something(): + ... + """ + + def _wrap_func(f): + + def _clear_cache(*args, **kwargs): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.reset_max_memory_allocated() + torch.cuda.reset_max_memory_cached() + torch.cuda.synchronize() + gc.collect() + f(*args, **kwargs) + + return _clear_cache + + return _wrap_func diff --git a/colossalai/trainer/_trainer.py b/colossalai/trainer/_trainer.py index 60bbc4eeee32..bfe1c403fd48 100644 --- a/colossalai/trainer/_trainer.py +++ b/colossalai/trainer/_trainer.py @@ -31,9 +31,9 @@ class Trainer: >>> # Initialize your engine, train_dataloader, test_dataloader, lr_scheduler >>> engine, train_dataloader, _, _ = colossalai.initialize(model, optimizer, criterion) >>> # Beginning training progress - >>> timier = ... + >>> timer = ... >>> logger = ... - >>> trainer = Trainer(engine=engine, logger=logger, timer=timier) + >>> trainer = Trainer(engine=engine, logger=logger, timer=timer) >>> # add hooks you would like to use here. >>> hook_list = [] >>> trainer.fit( @@ -56,7 +56,7 @@ def __init__( timer: MultiTimer = None, logger: DistributedLogger = None, ): - # training-ralated params + # training-related params self._engine = engine self._max_epochs = 0 self._cur_epoch = 0 @@ -118,7 +118,7 @@ def _set_current_step(self, epoch: int): self._cur_step = epoch * self._steps_per_epoch def _call_timer(self, action: str, item: str, *args, **kwargs) -> None: - """Call timer funciton with a given timer name. + """Call timer function with a given timer name. Args: action (str): Function to be called on timer. diff --git a/colossalai/utils/__init__.py b/colossalai/utils/__init__.py index 3f16bd91e5fe..7b2e8480c66c 100644 --- a/colossalai/utils/__init__.py +++ b/colossalai/utils/__init__.py @@ -7,7 +7,6 @@ count_zeros_fp32, disposable, ensure_path_exists, - free_port, is_ddp_ignored, is_dp_rank_0, is_model_parallel_parameter, @@ -37,7 +36,6 @@ __all__ = [ 'checkpoint', - 'free_port', 'print_rank_0', 'sync_model_param', 'is_ddp_ignored', diff --git a/colossalai/utils/checkpoint/module_checkpoint.py b/colossalai/utils/checkpoint/module_checkpoint.py index a109b3702577..d390da864cd3 100644 --- a/colossalai/utils/checkpoint/module_checkpoint.py +++ b/colossalai/utils/checkpoint/module_checkpoint.py @@ -89,7 +89,7 @@ def load_checkpoint(path: str, torch_load_kwargs: (dict, optional): The kwargs of torch.load inside the function load_state_dict_kwargs (dict, optional): The kwargs of load_state_dict inside the function """ - # initialize the default paramters + # initialize the default parameters if not torch_load_kwargs: torch_load_kwargs = dict() if not load_state_dict_kwargs: diff --git a/colossalai/utils/checkpoint/utils.py b/colossalai/utils/checkpoint/utils.py index 5652600ffd9b..682cd0903d5b 100644 --- a/colossalai/utils/checkpoint/utils.py +++ b/colossalai/utils/checkpoint/utils.py @@ -34,7 +34,7 @@ def gather_tensor(colo_tensor: ColoTensor) -> None: dist.barrier() if dist.get_rank() == 0: - setattr(colo_tensor, 'save_ready', True) # set saving signitrue + setattr(colo_tensor, 'save_ready', True) # set saving signature def scatter_tensor(colo_tensor: ColoTensor, dist_spec: _DistSpec) -> None: diff --git a/colossalai/utils/common.py b/colossalai/utils/common.py index e15981140be1..8022e84dc24b 100644 --- a/colossalai/utils/common.py +++ b/colossalai/utils/common.py @@ -50,23 +50,6 @@ def ensure_path_exists(filename: str): Path(dirpath).mkdir(parents=True, exist_ok=True) -def free_port() -> int: - """Get a free port on localhost. - - Returns: - int: A free port on localhost. - """ - while True: - port = random.randint(20000, 65000) - try: - with socket.socket() as sock: - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind(("localhost", port)) - return port - except OSError: - continue - - def sync_model_param(model, parallel_mode): r"""Make sure data parameters are consistent during Data Parallel Mode. @@ -341,7 +324,7 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2): norm_type = float(norm_type) # Parameters can be on CPU or CUDA - # If parameters are on CPU, disable CUDA kernerls + # If parameters are on CPU, disable CUDA kernels # Calculate norm. if norm_type == inf: diff --git a/colossalai/utils/data_sampler/data_parallel_sampler.py b/colossalai/utils/data_sampler/data_parallel_sampler.py index 945dc54b397a..2318e07a7f8d 100644 --- a/colossalai/utils/data_sampler/data_parallel_sampler.py +++ b/colossalai/utils/data_sampler/data_parallel_sampler.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -# adpated from torch.utils.data.DistributedSampler +# adapted from torch.utils.data.DistributedSampler import math import random diff --git a/colossalai/utils/model/lazy_init_context.py b/colossalai/utils/model/lazy_init_context.py deleted file mode 100644 index cf05f966089d..000000000000 --- a/colossalai/utils/model/lazy_init_context.py +++ /dev/null @@ -1,242 +0,0 @@ -#!/usr/bin/env python -# coding: utf-8 - -import inspect -import types -from typing import Callable, List - -import torch -import torch.nn as nn - -from colossalai.tensor import ColoParameter, ColoTensor -from colossalai.utils.model.utils import substitute_init_recursively - - -class LazyInitContext(): - """ - A context to allow for lazy weight initialization of PyTorch modules. It intercepts the tensor - initialization functions for lazy initialization - - Note: - This API is only experimental and subject to future changes. - - Usage: - with LazyInitContext() as ctx: - model = nn.Linear(10, 10) - model.weight.zero_() - - # make sure the weight is a meta tensor - assert model.weight.is_meta - - # initialize weights - ctx.lazy_init_parameters(model) - - # make sure the weight is not a meta tensor - # and initialized correctly - assert not model.weight.is_meta and torch.all(model.weight == 0) - - Args: - to_meta (bool): optional, whether to initialize the model with meta tensors, default is True. This - argument exists for now because some corner cases such as self.weight = torch.zeros(...) cannot be captured yet. - extra_torch_tensor_func (List[str]): extra torch tensor functions related - to value setting, such as `zero_` and `triu_`. `zero_` is pre-added by default. - """ - - tensor_set_value_func = ['zero_', 'fill_'] - - def __init__(self, to_meta: bool = True, extra_torch_tensor_func: List[str] = None): - # TODO: hijack the torch constructor functions as well - self._to_meta = to_meta - self._intercepted_nn_init_func_cache = {} - self._nn_init_methods = self._get_nn_init_methods() - self._torch_mod_cls = torch.nn.modules.module.Module - - if extra_torch_tensor_func: - # use tuple to remove duplicates - self._torch_tensor_funcs = tuple(self.tensor_set_value_func + extra_torch_tensor_func) - else: - self._torch_tensor_funcs = self.tensor_set_value_func - - @property - def to_meta(self): - return self._to_meta - - def _cache_init_func(self, func): - """ - This method wraps the ``torch.nn.init`` method and torch tensor value-setting functions - so that the function call is cached instead of being executed. - """ - - def wrapped_init_func(tensor, *args, **kwargs): - if tensor not in self._intercepted_nn_init_func_cache: - self._intercepted_nn_init_func_cache[tensor] = [] - self._intercepted_nn_init_func_cache[tensor].append((func, args, kwargs)) - - return wrapped_init_func - - def _get_nn_init_methods(self): - """ - This method looks for all available functions in the ``torch.nn.init`` - module. - """ - nn_init_method_names = dir(torch.nn.init) - nn_init_methods = [] - - # look for all methods in ``torch.nn.init`` module - for name in nn_init_method_names: - nn_init_methods.append((name, getattr(torch.nn.init, name))) - - def _is_init_method(item): - name, func = item - - if (not isinstance(func, types.FunctionType) or name.startswith('_') or not name.endswith('_')): - return False - else: - return True - - # remove methods which are not init functions - nn_init_methods = list(filter(_is_init_method, nn_init_methods)) - return nn_init_methods - - def _wrap_module_init(self, func): - """ - This method wraps the calls to the `__init__` of ``torch.nn.Module`` and replaces - the argument device with value 'meta' so that all modules are created as meta tensors. - """ - has_device = 'device' in inspect.signature(func).parameters - - def layer_lazy_init(module, *args, **kwargs): - # if this module contains device argument - # we set it to meta to initialize as meta backend - if has_device: - kwargs['device'] = 'meta' - func(module, *args, **kwargs) - - # if device is not found, we intialize it and convert to meta - if not has_device: - module.to('meta') - - return layer_lazy_init - - def _get_tmp_origin_func_ref(self, name): - """ - Generate a function name for consistency during caching and retrieving. - """ - return f'_orig_{name}' - - def _patch_nn_init_funcs(self): - # patch nn.init functions - for name, func in self._nn_init_methods: - setattr(torch.nn.init, name, self._cache_init_func(func)) - - def _unpatch_nn_init_funcs(self): - # unpatch nn.init functions - for name, func in self._nn_init_methods: - setattr(torch.nn.init, name, func) - - def _patch_submodule_init(self): - # patch classes __init__ methods - def _activate_wrap_init(cls): - cls.__orig_init__ = cls.__init__ - cls.__init__ = self._wrap_module_init(cls.__init__) - - substitute_init_recursively(self._torch_mod_cls, _activate_wrap_init, set()) - - def _unpatch_submodule_init(self): - - def _recover_orig_init(cls): - cls.__init__ = cls.__orig_init__ - - substitute_init_recursively(self._torch_mod_cls, _recover_orig_init, set()) - - def _patch_torch_tensor_funcs(self): - # patch tensor value-setting functions - for func_name in self._torch_tensor_funcs: - origin_func_name = self._get_tmp_origin_func_ref(func_name) - origin_func = getattr(torch.Tensor, func_name) - setattr(torch.Tensor, origin_func_name, origin_func) - setattr(torch.Tensor, func_name, self._cache_init_func(origin_func)) - - def _unpatch_torch_tensor_funcs(self): - for func_name in self._torch_tensor_funcs: - origin_func_name = self._get_tmp_origin_func_ref(func_name) - origin_func = getattr(torch.Tensor, origin_func_name) - setattr(torch.Tensor, func_name, origin_func) - - def __enter__(self): - self._patch_torch_tensor_funcs() - self._patch_nn_init_funcs() - - if self._to_meta: - self._patch_submodule_init() - return self - - def __exit__(self, *args, **kwargs): - if self._to_meta: - self._unpatch_submodule_init() - self._unpatch_nn_init_funcs() - self._unpatch_torch_tensor_funcs() - - def lazy_init_parameters(self, model: torch.nn.Module, device='cpu'): - """ - Initialize the weights of the meta-tensor model. - - Args: - model (`torch.nn.Module`): the model instantiated under the context. - device (str): the device on which weights are initialized - - """ - - def _init_recursively(module: nn.Module): - # recursively initialize the module - for mod in module.children(): - _init_recursively(mod) - - # initialize and shard tensors directly attached to the current module - for name, param in module.named_parameters(recurse=False): - _init_and_shard(module, name, param) - - for name, buf in module.named_buffers(recurse=False): - _init_and_shard(module, name, buf) - - @torch.no_grad() - def _init_and_shard(module, name, tensor): - # check whether the tensor is a buffer or parameter - is_param = isinstance(tensor, nn.parameter.Parameter) - - # get sharding spec - dist_spec = getattr(tensor, 'dist_spec', None) - pg = getattr(tensor, 'pg', None) - comp_spec = getattr(tensor, 'comp_spec', None) - - # convert the tensor from meta to materialized one - if tensor.is_meta: - materialized_tensor = torch.empty_like(tensor, device=device) - # if this tensor is a meta tensor, it must have an init function - assert tensor in self._intercepted_nn_init_func_cache - else: - materialized_tensor = tensor - - # apply init function - if tensor in self._intercepted_nn_init_func_cache: - init_func, args, kwargs = self._intercepted_nn_init_func_cache[tensor][-1] - init_func(materialized_tensor, *args, **kwargs) - - # convert it to ColoTensor or ColoParameter - if is_param: - tensor = ColoParameter.from_torch_tensor(materialized_tensor, requires_grad=tensor.requires_grad) - else: - tensor = ColoTensor.from_torch_tensor(materialized_tensor) - - # override the original tensor - with torch.no_grad(): - setattr(module, name, tensor) - - # apply sharding - if dist_spec: - tensor.process_group = pg - tensor.set_tensor_spec(dist_spec, comp_spec) - - _init_recursively(model) - - return model diff --git a/colossalai/utils/model/utils.py b/colossalai/utils/model/utils.py index f49607376439..21bc530934d3 100644 --- a/colossalai/utils/model/utils.py +++ b/colossalai/utils/model/utils.py @@ -70,7 +70,7 @@ def _init_subclass(cls, **kwargs): cls.__init__ = preprocess_after(cls.__init__) # Replace .__init__() for all existing subclasses of torch.nn.Module - # Excution self._post_init_method after the default init function. + # Execution self._post_init_method after the default init function. substitute_init_recursively(torch.nn.modules.module.Module, _enable_class, set()) # holding on to the current __init__subclass__ for exit diff --git a/colossalai/utils/moe.py b/colossalai/utils/moe.py index 90783e5d9b8e..86d04c11958b 100644 --- a/colossalai/utils/moe.py +++ b/colossalai/utils/moe.py @@ -38,7 +38,7 @@ def sync_moe_model_param(model: nn.Module): param_dict = get_moe_epsize_param_dict(model) - # synchrosize the parameters whose dp_group is the whole world + # synchronize the parameters whose dp_group is the whole world if 1 in param_dict: src_rank = gpc.get_ranks_in_group(ParallelMode.DATA)[0] for param in param_dict[1]: diff --git a/colossalai/utils/profiler/legacy/comm_profiler.py b/colossalai/utils/profiler/legacy/comm_profiler.py index a4f5729c97ec..334f0113ee90 100644 --- a/colossalai/utils/profiler/legacy/comm_profiler.py +++ b/colossalai/utils/profiler/legacy/comm_profiler.py @@ -111,7 +111,7 @@ def append(s: str = None): res.append(sep) if self.warn_flag: - append("Warnning: there exists multiple communication operations in the same time. As a result, " + append("Warning: there exists multiple communication operations in the same time. As a result, " "the profiling result is not accurate.") if self.total_cuda_time == 0: @@ -123,12 +123,12 @@ def append(s: str = None): append("total number of calls: {}".format(self.total_count)) append("All events:") - seperation = '-' * 74 + separation = '-' * 74 row_format = '{:^10}' + '{:^12}' * 2 + '{:^16}' + '{:^12}' * 2 - append(seperation) + append(separation) append(row_format.format('Location', 'GPU time', 'Percentage', 'Comm volume', 'Bandwidth', 'Num of calls')) - append(seperation) + append(separation) show_list = sorted(self.ops_record.items(), key=lambda kv: -kv[1].self_cuda_time) for location, event in show_list: diff --git a/colossalai/utils/profiler/legacy/pcie_profiler.py b/colossalai/utils/profiler/legacy/pcie_profiler.py index 526222941ef9..8f812f5cfc7b 100644 --- a/colossalai/utils/profiler/legacy/pcie_profiler.py +++ b/colossalai/utils/profiler/legacy/pcie_profiler.py @@ -130,12 +130,12 @@ def append(s: str = None): append("Possible data transmission events in PCIE:") - seperation = '-' * 62 + separation = '-' * 62 row_format = '{:^10}' + '{:^12}' + '{:^16}' + '{:^12}' * 2 - append(seperation) + append(separation) append(row_format.format('Location', 'GPU time', 'Trans volume', 'Bandwidth', 'Num of calls')) - append(seperation) + append(separation) show_list = sorted(self.ops_record.items(), key=lambda kv: -kv[1].cuda_time) for location, event in show_list: diff --git a/colossalai/utils/profiler/legacy/prof_utils.py b/colossalai/utils/profiler/legacy/prof_utils.py index 87ad644a7ecc..2f7eee827651 100644 --- a/colossalai/utils/profiler/legacy/prof_utils.py +++ b/colossalai/utils/profiler/legacy/prof_utils.py @@ -32,9 +32,9 @@ def _format_memory(nbytes): return str(nbytes) + ' B' -def _format_bandwidth(volme: float or int, time_us: int): +def _format_bandwidth(volume: float or int, time_us: int): sec_div_mb = (1000.0 / 1024.0)**2 - mb_per_sec = volme / time_us * sec_div_mb + mb_per_sec = volume / time_us * sec_div_mb if mb_per_sec >= 1024.0: return '{:.3f} GB/s'.format(mb_per_sec / 1024.0) diff --git a/colossalai/utils/rank_recorder/README.md b/colossalai/utils/rank_recorder/README.md index e30a925d2a92..da8a6039d543 100644 --- a/colossalai/utils/rank_recorder/README.md +++ b/colossalai/utils/rank_recorder/README.md @@ -1,5 +1,5 @@ # Rank Recorder -This is a useful tool to get the records of certain functions in each rank. The records of each rank will dump into a json file after the end of multiple process program. You can parse and visualise the json file easily. +This is a useful tool to get the records of certain functions in each rank. The records of each rank will dump into a json file after the end of multiple process program. You can parse and visualize the json file easily. Before using the tool, you should ensure dist.is_initialized() return true before exit of program. @@ -20,7 +20,7 @@ with recorder(record_name, current_rank) as r: ``` ## Example -This is a demo to display kernel select in cuda and visualise the cost of several procedures in each rank. +This is a demo to display kernel select in cuda and visualize the cost of several procedures in each rank. ```python import time diff --git a/colossalai/utils/rank_recorder/rank_recorder.py b/colossalai/utils/rank_recorder/rank_recorder.py index c088ceeb2e87..40bb7e184a12 100644 --- a/colossalai/utils/rank_recorder/rank_recorder.py +++ b/colossalai/utils/rank_recorder/rank_recorder.py @@ -133,7 +133,7 @@ def merge_recode(self): with open(self.export_name + '.json', 'w', encoding='utf-8') as f: json.dump(recoders, f, ensure_ascii=False) - def visualise_record(self): + def visualize_record(self): with open(self.export_name + '.json', 'r', encoding='utf-8') as f: records = json.load(f) records = dict(records) @@ -171,7 +171,7 @@ def exit_worker(self): if rank == 1: # take the base time of rank 0 as standard self.merge_recode() - self.visualise_record() + self.visualize_record() recorder = Recorder() diff --git a/colossalai/utils/tensor_detector/readme.md b/colossalai/utils/tensor_detector/readme.md index 840dc8f4eca6..d6852ea55b54 100644 --- a/colossalai/utils/tensor_detector/readme.md +++ b/colossalai/utils/tensor_detector/readme.md @@ -46,7 +46,7 @@ detector.detect() I have made some comments on the right of the output for your understanding. -Note that the total `Mem` of all the tensors and parameters is not equal to `Total GPU Memery Allocated`. PyTorch's memory management is really complicated, and for models of a large scale, it's impossible to figure out clearly. +Note that the total `Mem` of all the tensors and parameters is not equal to `Total GPU Memory Allocated`. PyTorch's memory management is really complicated, and for models of a large scale, it's impossible to figure out clearly. **The order of print is not equal to the order the tensor creates, but they are really close.** @@ -61,7 +61,7 @@ Note that the total `Mem` of all the tensors and parameters is not equal to `Tot + mlp.2.bias cuda:0 (32,) True torch.float32 128 B ------------------------------------------------------------------------------------------------------------ Detect Location: "test_tensor_detector.py" line 27 -Totle GPU Memery Allocated on cuda:0 is 4.5 KB +Total GPU Memory Allocated on cuda:0 is 4.5 KB ------------------------------------------------------------------------------------------------------------ @@ -72,7 +72,7 @@ Totle GPU Memery Allocated on cuda:0 is 4.5 KB + Tensor cuda:0 (32,) True torch.float32 128 B # output ------------------------------------------------------------------------------------------------------------ Detect Location: "test_tensor_detector.py" line 30 -Totle GPU Memery Allocated on cuda:0 is 5.5 KB +Total GPU Memory Allocated on cuda:0 is 5.5 KB ------------------------------------------------------------------------------------------------------------ @@ -82,7 +82,7 @@ Totle GPU Memery Allocated on cuda:0 is 5.5 KB + Tensor cuda:0 () True torch.float32 4 B # loss ------------------------------------------------------------------------------------------------------------ Detect Location: "test_tensor_detector.py" line 32 -Totle GPU Memery Allocated on cuda:0 is 6.0 KB +Total GPU Memory Allocated on cuda:0 is 6.0 KB ------------------------------------------------------------------------------------------------------------ @@ -103,7 +103,7 @@ Totle GPU Memery Allocated on cuda:0 is 6.0 KB - Tensor cuda:0 (8,) True torch.float32 32 B # deleted activation ------------------------------------------------------------------------------------------------------------ Detect Location: "test_tensor_detector.py" line 34 -Totle GPU Memery Allocated on cuda:0 is 10.0 KB +Total GPU Memory Allocated on cuda:0 is 10.0 KB ------------------------------------------------------------------------------------------------------------ @@ -117,7 +117,7 @@ Totle GPU Memery Allocated on cuda:0 is 10.0 KB + Tensor cuda:0 (32,) False torch.float32 128 B ------------------------------------------------------------------------------------------------------------ Detect Location: "test_tensor_detector.py" line 36 -Totle GPU Memery Allocated on cuda:0 is 14.0 KB +Total GPU Memory Allocated on cuda:0 is 14.0 KB ------------------------------------------------------------------------------------------------------------ ``` diff --git a/colossalai/utils/tensor_detector/tensor_detector.py b/colossalai/utils/tensor_detector/tensor_detector.py index a8186f76834c..cfcd4e47b4cb 100644 --- a/colossalai/utils/tensor_detector/tensor_detector.py +++ b/colossalai/utils/tensor_detector/tensor_detector.py @@ -55,7 +55,7 @@ def get_tensor_mem(self, tensor): return self.mem_format(memory_size) def mem_format(self, real_memory_size): - # format the tensor memory into a reasonal magnitude + # format the tensor memory into a reasonable magnitude if real_memory_size >= 2**30: return str(real_memory_size / (2**30)) + ' GB' if real_memory_size >= 2**20: @@ -71,7 +71,7 @@ def collect_tensors_state(self): if (not self.include_cpu) and obj.device == torch.device('cpu'): continue self.detected.append(id(obj)) - # skip paramters we had added in __init__ when module is an instance of nn.Module for the first epoch + # skip parameters we had added in __init__ when module is an instance of nn.Module for the first epoch if id(obj) not in self.tensor_info: name = type(obj).__name__ @@ -84,7 +84,7 @@ def collect_tensors_state(self): name = par_name + ' (with grad)' else: # with no grad attached - # there will be no new paramters created during running + # there will be no new parameters created during running # so it must be in saved_tensor_info continue # we can also marked common tensors as tensor(with grad) @@ -155,7 +155,7 @@ def print_tensors_state(self): if device == torch.device('cpu'): continue gpu_mem_alloc = self.mem_format(torch.cuda.memory_allocated(device)) - self.info += f"Totle GPU Memery Allocated on {device} is {gpu_mem_alloc}\n" + self.info += f"Total GPU Memory Allocated on {device} is {gpu_mem_alloc}\n" self.info += LINE self.info += '\n\n' if self.show_info: diff --git a/colossalai/zero/__init__.py b/colossalai/zero/__init__.py index 098ccbb45c5a..3465079e4fbb 100644 --- a/colossalai/zero/__init__.py +++ b/colossalai/zero/__init__.py @@ -1,41 +1,16 @@ -from typing import Tuple - -import torch -import torch.nn as nn - -from colossalai.logging import get_dist_logger -from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 -from colossalai.zero.sharded_optim import LowLevelZeroOptimizer, ShardedOptimizerV2 - -from ..nn.optimizer.zero_optimizer import ZeroOptimizer - - -def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model_config, - optimizer_config) -> Tuple[ShardedModelV2, ShardedOptimizerV2]: - """ - A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading - - :param model: Your model object - :type model: :class:`torch.nn.Module` - :param optimizer_config: Your optimizer object - :type optimizer_config: :class:`dict` - - :return: (model, optimizer) - :rtype: Tuple - """ - - logger = get_dist_logger('convert_to_zero_v2') - - logger.info(f'optimizer_config is {optimizer_config}', ranks=[0]) - if optimizer_config is None: - optimizer_config = dict() - logger.info(f'model_config is {model_config}', ranks=[0]) - if model_config is None: - model_config = dict() - - zero_model = ShardedModelV2(model, **model_config) - zero_optimizer = ShardedOptimizerV2(zero_model, optimizer, **optimizer_config) - return zero_model, zero_optimizer - - -__all__ = ['convert_to_zero_v2', 'LowLevelZeroOptimizer', 'ShardedModelV2', 'ShardedOptimizerV2', 'ZeroOptimizer'] +from .gemini import ( + ColoInitContext, + GeminiAdamOptimizer, + GeminiDDP, + ZeroDDP, + ZeroOptimizer, + get_static_torch_model, + post_process_colo_init_ctx, +) +from .low_level import LowLevelZeroOptimizer +from .wrapper import zero_model_wrapper, zero_optim_wrapper + +__all__ = [ + 'ZeroDDP', 'GeminiDDP', 'ZeroOptimizer', 'GeminiAdamOptimizer', 'zero_model_wrapper', 'zero_optim_wrapper', + 'LowLevelZeroOptimizer', 'ColoInitContext', 'post_process_colo_init_ctx', 'get_static_torch_model' +] diff --git a/colossalai/zero/gemini/__init__.py b/colossalai/zero/gemini/__init__.py new file mode 100644 index 000000000000..60f85ca2f540 --- /dev/null +++ b/colossalai/zero/gemini/__init__.py @@ -0,0 +1,11 @@ +from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration +from .colo_init_context import ColoInitContext, post_process_colo_init_ctx +from .gemini_ddp import GeminiDDP, ZeroDDP +from .gemini_mgr import GeminiManager +from .gemini_optimizer import GeminiAdamOptimizer, ZeroOptimizer +from .utils import get_static_torch_model + +__all__ = [ + 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager', 'search_chunk_configuration', 'ZeroDDP', 'GeminiDDP', + 'get_static_torch_model', 'GeminiAdamOptimizer', 'ZeroOptimizer', 'ColoInitContext', 'post_process_colo_init_ctx' +] diff --git a/colossalai/gemini/chunk/__init__.py b/colossalai/zero/gemini/chunk/__init__.py similarity index 100% rename from colossalai/gemini/chunk/__init__.py rename to colossalai/zero/gemini/chunk/__init__.py diff --git a/colossalai/gemini/chunk/chunk.py b/colossalai/zero/gemini/chunk/chunk.py similarity index 99% rename from colossalai/gemini/chunk/chunk.py rename to colossalai/zero/gemini/chunk/chunk.py index a7682eaf62e9..51da9be2b1f8 100644 --- a/colossalai/gemini/chunk/chunk.py +++ b/colossalai/zero/gemini/chunk/chunk.py @@ -416,7 +416,7 @@ def copy_tensor_to_chunk_slice(self, tensor: torch.Tensor, data_slice: torch.Ten Copy data slice to the memory space indexed by the input tensor in the chunk. Args: - tensor (torch.Tensor): the tensor used to retrive meta information + tensor (torch.Tensor): the tensor used to retrieve meta information data_slice (torch.Tensor): the tensor to be copied to the chunk """ # sanity check diff --git a/colossalai/gemini/chunk/manager.py b/colossalai/zero/gemini/chunk/manager.py similarity index 94% rename from colossalai/gemini/chunk/manager.py rename to colossalai/zero/gemini/chunk/manager.py index 30ac4d354647..38d34f14863e 100644 --- a/colossalai/gemini/chunk/manager.py +++ b/colossalai/zero/gemini/chunk/manager.py @@ -3,10 +3,11 @@ import torch -from colossalai.gemini.chunk import Chunk, ChunkFullError, TensorState from colossalai.tensor import ColoTensor from colossalai.utils import get_current_device +from .chunk import Chunk, ChunkFullError, TensorState + class ChunkManager: """ @@ -72,7 +73,7 @@ def register_tensor(self, if tensor.numel() > chunk_size: chunk_size = tensor.numel() - dp_size = tensor.process_group.dp_world_size() + dp_size = tensor.get_dp_world_size() chunk_size = chunk_size + (-chunk_size % dp_size) chunk = Chunk( @@ -101,7 +102,7 @@ def access_chunk(self, chunk: Chunk) -> None: """ if chunk in self.accessed_chunks: return - self.__sub_memroy_usage(chunk.memory_usage) + self.__sub_memory_usage(chunk.memory_usage) if chunk.device_type == 'cpu': chunk.shard_move(get_current_device()) self.__add_accessed_chunk(chunk) @@ -113,7 +114,7 @@ def release_chunk(self, chunk: Chunk) -> None: if chunk not in self.accessed_chunks: return if chunk.can_release: - self.__sub_memroy_usage(chunk.memory_usage) + self.__sub_memory_usage(chunk.memory_usage) self.__sub_accessed_chunk(chunk) self.__add_memory_usage(chunk.memory_usage) @@ -122,7 +123,7 @@ def move_chunk(self, chunk: Chunk, device: torch.device, force_copy: bool = Fals """ if not chunk.can_move or chunk.device_type == device.type: return - self.__sub_memroy_usage(chunk.memory_usage) + self.__sub_memory_usage(chunk.memory_usage) chunk.shard_move(device, force_copy) self.__add_memory_usage(chunk.memory_usage) @@ -137,7 +138,7 @@ def reduce_chunk(self, chunk: Chunk) -> bool: """ if not chunk.can_reduce: return False - self.__sub_memroy_usage(chunk.memory_usage) + self.__sub_memory_usage(chunk.memory_usage) chunk.reduce() self.__sub_accessed_chunk(chunk) self.__add_memory_usage(chunk.memory_usage) @@ -156,7 +157,7 @@ def copy_tensor_to_chunk_slice(self, tensor: torch.Tensor, data: torch.Tensor) - Copy data to the chunk. Args: - tensor (torch.Tensor): the tensor used to retrive meta information + tensor (torch.Tensor): the tensor used to retrieve meta information data (torch.Tensor): the tensor to be copied to the chunk """ chunk = self.tensor_chunk_map[tensor] @@ -227,11 +228,11 @@ def __get_chunk_group(self, group_name: str) -> Deque: return self.chunk_groups[group_name] def __close_one_chunk(self, chunk: Chunk): - self.__sub_memroy_usage(chunk.memory_usage) + self.__sub_memory_usage(chunk.memory_usage) chunk.close_chunk() self.__add_memory_usage(chunk.memory_usage) - def __sub_memroy_usage(self, usage: Dict[str, int]): + def __sub_memory_usage(self, usage: Dict[str, int]): for k, v in usage.items(): self.total_mem[k] -= v diff --git a/colossalai/gemini/chunk/search_utils.py b/colossalai/zero/gemini/chunk/search_utils.py similarity index 70% rename from colossalai/gemini/chunk/search_utils.py rename to colossalai/zero/gemini/chunk/search_utils.py index fe9650721d74..6c3d4f9a1b41 100644 --- a/colossalai/gemini/chunk/search_utils.py +++ b/colossalai/zero/gemini/chunk/search_utils.py @@ -5,14 +5,19 @@ import torch.distributed as dist import torch.nn as nn -from colossalai.gemini.memory_tracer import MemStats, OrderedParamGenerator from colossalai.tensor import ColoParameter from colossalai.utils import is_ddp_ignored +from colossalai.zero.gemini.memory_tracer import MemStats, OrderedParamGenerator def _filter_exlarge_params(model: nn.Module, size_dict: Dict[int, List[int]]) -> None: - """ + """_filter_exlarge_params + Filter those parameters whose size is too large (more than 3x standard deviations) from others. + + Args: + model (nn.Module): the model. + size_dict (Dict[int, List[int]]): the size dict of parameters. """ agg_size_list = [] for key in size_dict: @@ -33,7 +38,16 @@ def _filter_exlarge_params(model: nn.Module, size_dict: Dict[int, List[int]]) -> def _get_unused_byte(size_list: List[int], chunk_size: int) -> int: - """Get unused byte for a certain chunk size. + """_get_unused_byte + + Get unused byte for a certain chunk size. + + Args: + size_list (List[int]): the size list of parameters. + chunk_size (int): the chunk size. + + Returns: + int: the unused byte. """ acc = 0 left = 0 @@ -45,10 +59,22 @@ def _get_unused_byte(size_list: List[int], chunk_size: int) -> int: return left + acc -def _tensor_numel(local_param: ColoParameter, strict_ddp_flag: bool): - if strict_ddp_flag: +def _tensor_numel(local_param: ColoParameter, strict_ddp_flag: bool) -> int: + """_tensor_numel + + Get the number of elements of a tensor. + + Args: + local_param (ColoParameter): The local parameter. + strict_ddp_flag (bool): whether to enable the strict ddp mode. + + Returns: + int: the number of elements. + """ + if strict_ddp_flag and type(local_param) is ColoParameter: return local_param.numel_global() else: + # if local_param is not ColoParameter, we assume it's replicated return local_param.numel() @@ -59,7 +85,8 @@ def classify_params_by_dp_degree(param_order: OrderedParamGenerator, Classify the parameters by their dp degree Args: - param_order (OrderedParamGenerator): the order of param be visied + param_order (OrderedParamGenerator): the order of param be vised + strict_ddp_flag (bool, optional): whether to enable the strict ddp mode. Defaults to False. Returns: Dict[int, List[ColoParameter]]: a dict contains the classification results. @@ -67,11 +94,13 @@ def classify_params_by_dp_degree(param_order: OrderedParamGenerator, """ params_dict: Dict[int, List[ColoParameter]] = dict() for param in param_order.generate(): - assert isinstance(param, ColoParameter), "please init model in the ColoInitContext" + # assert isinstance(param, ColoParameter), "please init model in the ColoInitContext" if is_ddp_ignored(param): continue - if strict_ddp_flag: + if strict_ddp_flag or type(param) is not ColoParameter: + # if model is not initialized with ColoInitContext, we assume it's replicated + # TODO(ver217): integrate DTensor param_key = dist.get_world_size() else: param_key = param.process_group.dp_world_size() @@ -85,19 +114,21 @@ def classify_params_by_dp_degree(param_order: OrderedParamGenerator, def search_chunk_configuration( model: nn.Module, - search_range_mb: float, - search_interval_byte: int, # hidden size is the best value for the interval - min_chunk_size_mb: float = 32, + search_range_m: float, + search_interval: int, # hidden size is the best value for the interval + min_chunk_size_m: float = 32, filter_exlarge_params: bool = True, strict_ddp_flag: bool = False, memstas: Optional[MemStats] = None) -> Tuple[Dict, int, int]: """search_chunk_configuration + Search the chunk configuration for a model. + Args: model (nn.Module): torch module - search_range_mb (float): searching range in mega byte. - search_interval_byte (int): searching interval in byte. - min_chunk_size_mb (float, optional): the minimum size of a distributed chunk. + search_range_m (float): searching range divided by 2^20. + search_interval (int): searching interval. + min_chunk_size_m (float, optional): the minimum size of a distributed chunk, divided by 2^20.. filter_exlarge_params (bool, optional): filter extreme large parameters. Defaults to True. strict_ddp_flag (bool, optional): whether to enable the strict ddp mode. all parameters keep replicated in this mode. @@ -114,9 +145,9 @@ def search_chunk_configuration( for p in model.parameters(): param_order.append(p) - search_range_byte = round(search_range_mb * 1024**2) - min_chunk_size_byte = round(min_chunk_size_mb * 1024**2) - assert search_range_byte >= 0 + search_range = round(search_range_m * 1024**2) + min_chunk_size = round(min_chunk_size_m * 1024**2) + assert search_range >= 0 params_dict = classify_params_by_dp_degree(param_order, strict_ddp_flag) size_lcm = np.lcm.reduce(list(params_dict.keys())) @@ -131,7 +162,7 @@ def search_chunk_configuration( total_param_size += group_acc_size # let small parameters keep gathered in CUDA all the time - if group_acc_size < min_chunk_size_byte: + if group_acc_size < min_chunk_size: config_dict[dp_degree] = dict(chunk_size=group_acc_size, keep_gathered=True) else: size_dict[dp_degree] = size_list @@ -139,15 +170,15 @@ def search_chunk_configuration( if filter_exlarge_params: _filter_exlarge_params(model, size_dict) - max_size = min_chunk_size_byte + max_size = min_chunk_size for key in size_dict: max_size = max(max_size, max(size_dict[key])) - start_size = int(math.ceil(max_size / search_interval_byte) * search_interval_byte) + start_size = int(math.ceil(max_size / search_interval) * search_interval) min_chunk_waste = float('+inf') best_chunk_size = start_size - for chunk_size in range(start_size, start_size + search_range_byte + 1, search_interval_byte): + for chunk_size in range(start_size, start_size + search_range + 1, search_interval): temp_waste = 0 for key in size_dict: temp_waste += _get_unused_byte(size_dict[key], chunk_size) diff --git a/colossalai/gemini/chunk/utils.py b/colossalai/zero/gemini/chunk/utils.py similarity index 67% rename from colossalai/gemini/chunk/utils.py rename to colossalai/zero/gemini/chunk/utils.py index 83512b8e0ee5..e98e9cf9c314 100644 --- a/colossalai/gemini/chunk/utils.py +++ b/colossalai/zero/gemini/chunk/utils.py @@ -5,10 +5,11 @@ import torch.distributed as dist import torch.nn as nn -from colossalai.gemini.chunk import ChunkManager -from colossalai.gemini.chunk.search_utils import search_chunk_configuration from colossalai.utils import is_ddp_ignored +from .manager import ChunkManager +from .search_utils import search_chunk_configuration + def safe_div(a, b): if a == 0: @@ -19,12 +20,13 @@ def safe_div(a, b): def init_chunk_manager(model: nn.Module, init_device: Optional[torch.device] = None, hidden_dim: Optional[int] = None, + verbose: bool = False, **kwargs) -> ChunkManager: if hidden_dim: - search_interval_byte = hidden_dim + search_interval = hidden_dim else: - search_interval_byte = 1024 # defaults to 1kb - kwargs["search_interval_byte"] = search_interval_byte + search_interval = 1024 # defaults to 1024 + kwargs["search_interval"] = search_interval dist.barrier() begin = time() @@ -34,13 +36,13 @@ def init_chunk_manager(model: nn.Module, dist.barrier() end = time() span_s = end - begin - mb_size = 1024**2 - total_size /= mb_size - wasted_size /= mb_size + mega_unit = 1024**2 + total_size /= mega_unit + wasted_size /= mega_unit - if dist.get_rank() == 0: + if verbose and dist.get_rank() == 0: print("searching chunk configuration is completed in {:.2f} s.\n".format(span_s), - "used number: {:.2f} MB, wasted number: {:.2f} MB\n".format(total_size, wasted_size), + "used number: {:.2f} * 2^20, wasted number: {:.2f} * 2^20\n".format(total_size, wasted_size), "total wasted percentage is {:.2f}%".format(100 * safe_div(wasted_size, total_size + wasted_size)), sep='', flush=True) diff --git a/colossalai/utils/model/colo_init_context.py b/colossalai/zero/gemini/colo_init_context.py similarity index 96% rename from colossalai/utils/model/colo_init_context.py rename to colossalai/zero/gemini/colo_init_context.py index 87ae413a2a8a..75f8576ca477 100644 --- a/colossalai/utils/model/colo_init_context.py +++ b/colossalai/zero/gemini/colo_init_context.py @@ -3,10 +3,8 @@ import torch from torch import nn -from colossalai.nn.parallel.layers import ColoEmbedding, ColoLinear, register_colo_module from colossalai.tensor import ColoParameter, ColoTensor, ProcessGroup - -from .utils import InsertPostInitMethodToModuleSubClasses +from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses # find named_params includes replica @@ -76,7 +74,7 @@ def __init__(self, """ Args: device (torch.device): the device where parameters initialized are resident. Defaults to torch.device('cpu'). - dtype (torch.dtype): the dtype of parameters initialized. Defults to torch.float. + dtype (torch.dtype): the dtype of parameters initialized. Defaults to torch.float. default_pg (ProcessGroup): the default process group for all initialized parameters. default_dist_spec: the default distributed specifications. """ @@ -89,6 +87,7 @@ def __init__(self, self._default_dist_spec = default_dist_spec def _register_colo_modules(self): + from colossalai.nn.parallel.layers import ColoEmbedding, ColoLinear, register_colo_module register_colo_module(torch.nn.Linear, ColoLinear()) register_colo_module(torch.nn.Embedding, ColoEmbedding()) @@ -165,7 +164,7 @@ def post_process_colo_init_ctx(model: torch.nn.Module, model (torch.nn.module): the model device (torch.device, optional): device type of the model params. Defaults to torch.device('cpu'). dtype (torch.dtype, optional): dtype of the model params. Defaults to torch.float. - default_pg (Optional[ProcessGroup], optional): default process group. Defaults to None. Inidicates a DP-only process group. + default_pg (Optional[ProcessGroup], optional): default process group. Defaults to None. Indicates a DP-only process group. default_dist_spec (Any, optional): default dist spec of params. Defaults to None. Raises: diff --git a/colossalai/zero/gemini/gemini_ddp.py b/colossalai/zero/gemini/gemini_ddp.py new file mode 100644 index 000000000000..08384ee82d0b --- /dev/null +++ b/colossalai/zero/gemini/gemini_ddp.py @@ -0,0 +1,793 @@ +import itertools +from collections import OrderedDict +from contextlib import nullcontext +from functools import partial +from typing import Dict, Iterator, List, Optional, Set, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn as nn + +from colossalai.checkpoint_io.utils import calculate_tensor_size +from colossalai.lazy import LazyTensor +from colossalai.logging import get_dist_logger +from colossalai.nn.parallel.data_parallel import ColoDDP, _cast_float, free_storage +from colossalai.tensor import ProcessGroup as ColoProcessGroup +from colossalai.tensor import ReplicaSpec +from colossalai.tensor.colo_parameter import ColoParameter, ColoTensor, ColoTensorSpec +from colossalai.tensor.param_op_hook import ColoParamOpHookManager +from colossalai.utils import get_current_device, is_ddp_ignored + +from .chunk import Chunk, ChunkManager, TensorState, init_chunk_manager +from .gemini_hook import GeminiZeROHook +from .gemini_mgr import GeminiManager +from .memory_tracer import MemStats, OrderedParamGenerator +from .utils import get_temp_total_chunk_on_cuda + +try: + from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, _IncompatibleKeys +except ImportError: + _EXTRA_STATE_KEY_SUFFIX = '_extra_state' + +__all__ = [ + 'ZeroDDP', + 'GeminiDDP', +] + + +class ZeroDDP(ColoDDP): + """ZeRO DDP for ColoTensor. + Warning: Nested ZeroDDP is not supported now. + It is designed to be used with ChunkManager and GeminiManager. + For more details, see the API reference of ``ChunkManager`` and ``GeminiManager``. + + Args: + module (torch.nn.Module): Module to apply ZeRO-DP. + gemini_manager (GeminiManager): Manages the chunk manager and heterogeneous memory space. + For more details, see the API reference of ``GeminiManager``. + pin_memory (bool): Chunks on CPU Memory use pin-memory. + force_outputs_fp32 (bool): If set to True, outputs will be fp32. Otherwise, outputs will be fp16. + Defaults to False. + strict_ddp_mode (bool): If set to True, there is no tensor sharding, each tensor is replicated. + Defaults to False. Users can set it to True, when they clearly know that they only need DDP. + scatter_after_inference (bool): If set to True, the model will be scattered after inference. This will save memory but slow down the consecutive inference. + mixed_precision (torch.dtype): If set to torch.float16, the model will be trained in fp16. Otherwise, the model will be trained in bf16. Defaults to torch.float16. + """ + + def __init__(self, + module: torch.nn.Module, + gemini_manager: GeminiManager, + pin_memory: bool = False, + force_outputs_fp32: bool = False, + strict_ddp_mode: bool = False, + scatter_after_inference: bool = True, + mixed_precision: torch.dtype = torch.float16) -> None: + assert mixed_precision in (torch.float16, torch.bfloat16) + self.gemini_manager = gemini_manager + self.chunk_manager: ChunkManager = gemini_manager.chunk_manager + self.force_outputs_fp32 = force_outputs_fp32 + self.param_op_hook = GeminiZeROHook(gemini_manager) + self.fp32_params: List[ColoTensor] = list() + self.fp16_params: List[ColoParameter] = list() + self.overflow_counter = 0 + self.grads_device: Dict[torch.Tensor, torch.device] = dict() + self.param2name: Dict[nn.Parameter, str] = dict() + self.name2param: Dict[str, nn.Parameter] = dict() + self.scatter_after_inference = scatter_after_inference + self.mixed_precision = mixed_precision + + self._logger = get_dist_logger() + + if self.gemini_manager._premade_memstats_: + # build chunk in param runtime visited order. + param_order = self.gemini_manager.memstats()._param_runtime_order + else: + # build chunk in param initialized order. + # Note: in this way, it can not get filter unused params during runtime. + param_order = OrderedParamGenerator() + for p in module.parameters(): + param_order.append(p) + + self._init_chunks(param_order=param_order, + strict_ddp_mode=strict_ddp_mode, + cpu_offload=self.gemini_manager.policy_name != 'cuda', + pin_memory=pin_memory) + + for name, param in module.named_parameters(): + self.param2name[param] = name + for m_name, m_var in module.named_modules(): + for p_name, p_var in m_var.named_parameters(recurse=False): + param_name = m_name + '.' + p_name if m_name else p_name + self.name2param[param_name] = p_var + super().__init__(module, process_group=ColoProcessGroup()) + self._non_persistent_buffers_set = self._get_non_persistent_buffers_set(module) + self._cast_buffers() + + def _get_non_persistent_buffers_set(self, + module, + memo: Optional[Set[nn.Module]] = None, + prefix: str = '', + remove_duplicate: bool = True): + r""" + Args: + memo: a memo to store the set of modules already added to the result + prefix: a prefix that will be added to the name of the module + remove_duplicate: whether to remove the duplicated module instances in the result + or not + """ + + if memo is None: + memo = set() + self_non_persistent_set = set() + if module not in memo: + if remove_duplicate: + memo.add(module) + self_non_persistent_set = set( + map(lambda key: prefix + ('.' if prefix else '') + key, module._non_persistent_buffers_set)) + for name, sub_module in module._modules.items(): + if sub_module is None: + continue + submodule_prefix = prefix + ('.' if prefix else '') + name + child_non_persistent_set = self._get_non_persistent_buffers_set(sub_module, memo, submodule_prefix, + remove_duplicate) + self_non_persistent_set = set.union(self_non_persistent_set, child_non_persistent_set) + return self_non_persistent_set + + def _post_forward(self): + """This function is only triggered for inference. + """ + access_list = list(self.chunk_manager.accessed_chunks) + # we need to scatter all accessed chunks and move them to their original places + for chunk in access_list: + if chunk.keep_gathered: + self.chunk_manager.fake_release_chunk(chunk) + else: + assert chunk.can_release + self.chunk_manager.release_chunk(chunk) + first_param = next(iter(chunk.tensors_info)) + self.chunk_manager.move_chunk(chunk, self.grads_device[first_param]) + assert self.chunk_manager.accessed_mem == 0 + + def forward(self, *args, **kwargs): + # check whether we are in a inference mode + grad_flag = torch.is_grad_enabled() + if not grad_flag: + assert not self.gemini_manager.need_warmup or not self.gemini_manager.is_warmup( + ), "You should run a completed iteration as your warmup iter" + + args, kwargs = _cast_float(args, self.mixed_precision), _cast_float(kwargs, self.mixed_precision) + self.module.zero_grad(set_to_none=True) + if not grad_flag: + outputs = self._inference_forward(*args, **kwargs) + else: + self.gemini_manager.pre_iter(*args) + with ColoParamOpHookManager.use_hooks(self.param_op_hook): + outputs = self.module(*args, **kwargs) + + if self.force_outputs_fp32: + return _cast_float(outputs, torch.float) + return outputs + + def _inference_forward(self, *args, **kwargs): + """This function is only triggered for inference. + """ + fwd_ctx = ColoParamOpHookManager.use_hooks(self.param_op_hook) + if not self.scatter_after_inference: + # gather all chunks + for chunk in self.chunk_manager.get_chunks(self.fp16_params): + self.chunk_manager.access_chunk(chunk) + fwd_ctx = nullcontext() + with fwd_ctx: + outputs = self.module(*args, **kwargs) + if self.scatter_after_inference: + # scatter chunks + self._post_forward() + # reset all recorded attributes + self.gemini_manager.reset_attributes() + return outputs + + def _setup_grads_ptr(self): + for p in self.module.parameters(): + if is_ddp_ignored(p): + continue + p.grad = None + + def _pre_backward(self): + # set a visit label for all parameters + # the label is used to check whether the parameter is correctly reduced + for param in self.param2name: + if not is_ddp_ignored(param): + setattr(param, "_gemini_reduced", False) + + def _post_backward(self): + if self.chunk_manager.accessed_mem != 0: + error_params = ["Reduction failed at followed parameters:"] + for param in self.param2name: + if not is_ddp_ignored(param) and not getattr(param, "_gemini_reduced"): + error_params.append(self.param2name[param]) + error_str = "\n\t".join(error_params) + raise RuntimeError("ZERO DDP error: the synchronization of gradients doesn't exit properly.", + "The most possible reason is that the model is not compatible with ZeroDDP.\n", + f"{error_str}") + self._setup_grads_ptr() + self._logger.debug( + f'comp cuda demand time: {self.gemini_manager._comp_cuda_demand_time}, layout time: {self.gemini_manager._layout_time}, evict time: {self.gemini_manager._evict_time}, CPU->CUDA vol: {self.gemini_manager._h2d_volume}B, CUDA->CPU vol: {self.gemini_manager._d2h_volume}' + ) + self.gemini_manager.post_iter() + + def backward(self, loss: torch.Tensor): + self._pre_backward() + with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook): + loss.backward() + self._post_backward() + + def backward_by_grad(self, tensor, grad): + with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook): + torch.autograd.backward(tensor, grad) + self._post_backward() + + def grad_handle(self, p, grad): + empty_grad = torch.empty_like(grad) + free_storage(empty_grad) + with torch._C.DisableTorchFunction(): + chunk = self.chunk_manager.get_chunk(p) + if chunk.tensors_info[p].state != TensorState.HOLD_AFTER_BWD: + raise RuntimeError(f"Parameter `{self.param2name[p]}` failed at the gradient reduction. " + "Some unsupported torch function is operated upon this parameter.") + self.chunk_manager.trans_tensor_state(p, TensorState.READY_FOR_REDUCE) + chunk.copy_tensor_to_chunk_slice(p, grad) + reduced = self.chunk_manager.reduce_chunk(chunk) + if reduced: + if chunk.is_gathered: + chunk.cuda_global_chunk.div_(chunk.pg_size) + else: + chunk.cuda_shard.div_(chunk.pg_size) + # check overflow elements + self.overflow_counter += chunk.has_inf_or_nan + # record l2 norm for gradient clipping + if chunk.l2_norm_flag: + chunk.set_l2_norm() + self.chunk_manager.move_chunk(chunk, self.grads_device[p], force_copy=True) + return empty_grad + + def zero_grad(self, set_to_none: bool = False) -> None: + self.module.zero_grad(set_to_none=True) + + def set_chunk_grad_device(self, chunk: Chunk, device: torch.device) -> None: + for tensor in chunk.get_tensors(): + self.grads_device[tensor] = device + + def state_dict(self, + destination=None, + prefix='', + keep_vars=False, + only_rank_0: bool = True, + dtype: torch.dtype = torch.float16): + """Returns a dictionary containing a whole state of the module. + + Both parameters and persistent buffers (e.g. running averages) are included. + Keys are corresponding parameter and buffer names. + Parameters and buffers set to ``None`` are not included. + + Warning: The non strict state dict would ignore the parameters if the tensors of the parameters + are shared with other parameters which have been included in the dictionary. + When you need to load the state dict, you should set the argument `strict` to False. + + Returns: + dict: + a dictionary containing a whole state of the module + """ + if destination is None: + destination = OrderedDict() + destination._metadata = OrderedDict() + destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version) + self._save_to_state_dict(destination, prefix, keep_vars, only_rank_0, dtype) + + for hook in self._state_dict_hooks.values(): + hook_result = hook(self, destination, prefix, local_metadata) + if hook_result is not None: + destination = hook_result + return destination + + def _get_chunk_to_save_data(self, chunk: Chunk, only_rank_0: bool, dtype: torch.dtype = torch.float16) -> Dict: + """ + get gathered chunk content. + + Args: + chunk (Chunk): a chunk + only_rank_0 (bool): whether to only save data on rank 0 + + Returns: + Dict: a dict whose key is param name and value is param with correct payload + """ + # save parameters + chunk_to_save_data = dict() + temp_chunk = get_temp_total_chunk_on_cuda(chunk) + if torch.is_floating_point(temp_chunk): + temp_chunk = temp_chunk.to(dtype) + for tensor, tensor_info in chunk.tensors_info.items(): + record_tensor = torch.empty([0]) + record_flag = (not only_rank_0) | (dist.get_rank(chunk.torch_pg) == 0) + if record_flag: + record_tensor = temp_chunk[tensor_info.offset:tensor_info.end].view(tensor.shape).cpu() + + assert tensor not in chunk_to_save_data + chunk_to_save_data[tensor] = record_tensor + + del temp_chunk + return chunk_to_save_data + + def _get_param_to_save_data(self, param_list: List[torch.nn.Parameter], only_rank_0: bool, + dtype: torch.dtype) -> Dict: + """ + get param content from chunks. + + Args: + param_list (_type_): a list of torch.nn.Parameters + only_rank_0 (_type_): _description_ + + Returns: + Dict: a dict whose key is param name and value is param with correct payload + """ + # save parameters + param_to_save_data = dict() + chunk_list = self.chunk_manager.get_chunks(param_list) + for chunk in chunk_list: + param_to_save_data.update(self._get_chunk_to_save_data(chunk, only_rank_0, dtype)) + return param_to_save_data + + def _save_to_state_dict(self, destination, prefix, keep_vars, only_rank_0=True, dtype=torch.float16): + r"""Saves module state to `destination` dictionary, containing a state + of the module, but not its descendants. This is called on every + submodule in :meth:`~torch.nn.Module.state_dict`. + + In rare cases, subclasses can achieve class-specific behavior by + overriding this method with custom logic. + + Args: + destination (dict): a dict where state will be stored + prefix (str): the prefix for parameters and buffers used in this + module + """ + assert keep_vars is False, "`state_dict` with parameter, `keep_vars=True`, is not supported now." + + # get copies of fp32 parameters in CPU + # as memory of fp16_params may be reused by grad, it's not reliable, we should use fp32_params and convert to fp16 + param_to_save_data = self._get_param_to_save_data(self.fp32_params, only_rank_0, dtype) + # get the mapping between copies and fp16 parameters + p_mapping = dict() + for p, fp32_p in zip(self.fp16_params, self.fp32_params): + name = self.param2name[p] + assert fp32_p in param_to_save_data, "Parameter '{}' is neglected in the chunk list".format(name) + record_parameter = param_to_save_data[fp32_p] + p_mapping[p] = record_parameter + for name, param in self.name2param.items(): + if param is not None: + if is_ddp_ignored(param): + # deal with ddp ignored parameters + destination[prefix + name] = param if keep_vars else param.detach() + else: + destination[prefix + name] = p_mapping[param] + del p_mapping + del param_to_save_data + + # save all buffers + for name, buf in self.named_buffers(): + if buf is not None and name not in self._non_persistent_buffers_set: + destination[prefix + name] = buf if keep_vars else buf.detach() + # save extra states + extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX + if getattr(self.__class__, "get_extra_state", + torch.nn.Module.get_extra_state) is not torch.nn.Module.get_extra_state: + destination[extra_state_key] = self.get_extra_state() + + def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict: bool = True): + r"""Copies parameters and buffers from :attr:`state_dict` into + this module and its descendants. If :attr:`strict` is ``True``, then + the keys of :attr:`state_dict` must exactly match the keys returned + by this module's :meth:`~torch.nn.Module.state_dict` function. + + Args: + state_dict (dict): a dict containing parameters and + persistent buffers. + strict (bool, optional): whether to strictly enforce that the keys + in :attr:`state_dict` match the keys returned by this module's + :meth:`~torch.nn.Module.state_dict` function. Default: ``True`` + + Returns: + ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: + * **missing_keys** is a list of str containing the missing keys + * **unexpected_keys** is a list of str containing the unexpected keys + + Note: + If a parameter or buffer is registered as ``None`` and its corresponding key + exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a + ``RuntimeError``. + """ + missing_keys: List[str] = [] + unexpected_keys: List[str] = [] + error_msgs: List[str] = [] + + # copy state_dict so _load_from_state_dict can modify it + metadata = getattr(state_dict, '_metadata', None) + state_dict = state_dict.copy() + if metadata is not None: + # mypy isn't aware that "_metadata" exists in state_dict + state_dict._metadata = metadata # type: ignore[attr-defined] + + prefix = '' + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + self._load_from_state_dict(state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs) + + if strict: + if len(unexpected_keys) > 0: + error_msgs.insert( + 0, 'Unexpected key(s) in state_dict: {}. '.format(', '.join( + '"{}"'.format(k) for k in unexpected_keys))) + if len(missing_keys) > 0: + error_msgs.insert( + 0, 'Missing key(s) in state_dict: {}. '.format(', '.join('"{}"'.format(k) for k in missing_keys))) + + if len(error_msgs) > 0: + raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( + self.__class__.__name__, "\n\t".join(error_msgs))) + return _IncompatibleKeys(missing_keys, unexpected_keys) + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, + error_msgs): + r"""Copies parameters and buffers from :attr:`state_dict` into only + this module, but not its descendants. This is called on every submodule + in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this + module in input :attr:`state_dict` is provided as :attr:`local_metadata`. + For state dicts without metadata, :attr:`local_metadata` is empty. + Subclasses can achieve class-specific backward compatible loading using + the version number at `local_metadata.get("version", None)`. + + .. note:: + :attr:`state_dict` is not the same object as the input + :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So + it can be modified. + + Args: + state_dict (dict): a dict containing parameters and + persistent buffers. + prefix (str): the prefix for parameters and buffers used in this + module + local_metadata (dict): a dict containing the metadata for this module. + See + strict (bool): whether to strictly enforce that the keys in + :attr:`state_dict` with :attr:`prefix` match the names of + parameters and buffers in this module + missing_keys (list of str): if ``strict=True``, add missing keys to + this list + unexpected_keys (list of str): if ``strict=True``, add unexpected + keys to this list + error_msgs (list of str): error messages should be added to this + list, and will be reported together in + :meth:`~torch.nn.Module.load_state_dict` + """ + for hook in self._load_state_dict_pre_hooks.values(): + hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) + + persistent_buffers = {k: v for k, v in self.named_buffers() if k not in self._non_persistent_buffers_set} + local_name_params = itertools.chain(self.named_parameters(), persistent_buffers.items()) + local_state = {k: v for k, v in local_name_params if v is not None} + + def load(param_name, dest_tensor, copy_func): + state_key = prefix + param_name + if state_key in state_dict: + input_param = state_dict[state_key] + # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+ + if len(dest_tensor.shape) == 0 and len(input_param.shape) == 1: + input_param = input_param[0] + if input_param.shape != dest_tensor.shape: + # local shape should match the one in checkpoint + error_msgs.append('size mismatch for {}: copying a param with shape {} from checkpoint, ' + 'the shape in current model is {}.'.format(state_key, input_param.shape, + dest_tensor.shape)) + return + try: + with torch.no_grad(): + copy_func(input_param) + except Exception as ex: + error_msgs.append('While copying the parameter named "{}", ' + 'whose dimensions in the model are {} and ' + 'whose dimensions in the checkpoint are {}, ' + 'an exception occurred : {}.'.format(state_key, dest_tensor.size(), + input_param.size(), ex.args)) + elif strict: + missing_keys.append(state_key) + + def load_fp32_parameter(chunk_slice, data): + chunk_slice.copy_(data.flatten()) + + for name, param in self.named_parameters(): + if is_ddp_ignored(param): + # deal with ddp ignored parameters + load(name, param, param.copy_) + + fp32_to_name = dict() + for p, fp32_p in zip(self.fp16_params, self.fp32_params): + if p is not None: + name = self.param2name[p] + fp32_to_name[fp32_p] = name + + chunk_list = self.chunk_manager.get_chunks(self.fp32_params) + for chunk in chunk_list: + temp_chunk = get_temp_total_chunk_on_cuda(chunk) + + for tensor, tensor_info in chunk.tensors_info.items(): + parameter_name = fp32_to_name[tensor] + parameter_slice = temp_chunk[tensor_info.offset:tensor_info.end] + load(parameter_name, tensor, partial(load_fp32_parameter, parameter_slice)) + + if chunk.is_gathered: + chunk.cuda_global_chunk.copy_(temp_chunk) + elif chunk.cuda_shard is not None: + chunk.cuda_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end]) + else: + chunk.cpu_shard.copy_(temp_chunk[chunk.shard_begin:chunk.shard_end]) + + del temp_chunk + + for chunk_32 in chunk_list: + chunk_16 = chunk_32.paired_chunk + assert chunk_16 is not None + chunk_16.optim_update() + + for name, buf in persistent_buffers.items(): + if buf is not None: + load(name, buf, buf.copy_) + + extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX + if getattr(self.__class__, "set_extra_state", + torch.nn.Module.set_extra_state) is not torch.nn.Module.set_extra_state: + if extra_state_key in state_dict: + self.set_extra_state(state_dict[extra_state_key]) + elif strict: + missing_keys.append(extra_state_key) + elif strict and (extra_state_key in state_dict): + unexpected_keys.append(extra_state_key) + + if strict: + for key in state_dict.keys(): + if key.startswith(prefix) and key != extra_state_key: + input_name = key[len(prefix):] + if input_name not in local_state: + unexpected_keys.append(key) + + def _init_chunks(self, param_order, strict_ddp_mode: bool, cpu_offload: bool, pin_memory: bool): + ddp_pg = ColoProcessGroup() + for p in param_order.generate(): + self._preprocess_param(p) + assert type(p) is ColoParameter + + # gather sharded parameters in the strict ddp mode + if strict_ddp_mode: + if not p.is_replicate(): + p.set_dist_spec(ReplicaSpec()) + p.set_process_group(pg=ddp_pg) + + # ignore the parameters with no gradient + if not p.requires_grad: + self.set_params_to_ignore([p]) + + # move ignored parameters to CUDA + if is_ddp_ignored(p): + p.data = p.data.to(device=get_current_device(), dtype=self.mixed_precision) + continue + + # create a fp32 parameter + fp32_data = p.data.float() + fp32_p = ColoTensor(fp32_data, spec=ColoTensorSpec(p.process_group)) + # create a fp16 parameter + p.data = p.data.to(self.mixed_precision) + + # register the fp16 parameter and fp32 parameter in the chunk manager + dp_world_size = p.process_group.dp_world_size() + self.chunk_manager.register_tensor(tensor=p, + group_type='fp16_param', + config_key=dp_world_size, + cpu_offload=cpu_offload, + pin_memory=pin_memory) + self.chunk_manager.register_tensor(tensor=fp32_p, + group_type='fp32_param', + config_key=dp_world_size, + cpu_offload=cpu_offload, + pin_memory=pin_memory) + + self.fp16_params.append(p) + self.fp32_params.append(fp32_p) + self.grads_device[p] = self.gemini_manager.default_device + + self.chunk_manager.close_all_groups() + + for p, fp32_p in zip(self.fp16_params, self.fp32_params): + chunk_16 = self.chunk_manager.get_chunk(p) + chunk_32 = self.chunk_manager.get_chunk(fp32_p) + chunk_32.init_pair(chunk_16) + + # keep gathered chunks are in CUDA + if chunk_16.keep_gathered: + self.grads_device[p] = get_current_device() + + def _cast_buffers(self): + for buffer in self.module.buffers(): + if isinstance(buffer, LazyTensor): + buffer.materialize() + buffer.data = buffer.cuda() + if torch.is_floating_point(buffer): + buffer.data = buffer.to(self.mixed_precision) + + def _preprocess_param(self, p: Union[nn.Parameter, ColoParameter, 'LazyTensor']) -> None: + """Convert parameter to ColoParameter in-place. + Args: + p (Union[nn.Parameter, ColoParameter, LazyTensor]): parameter to be converted + """ + if type(p) is ColoParameter: + # model is initialized with ColoInitContext + return + requires_grad = p.requires_grad + if isinstance(p, LazyTensor): + # model is initialized with LazyInitContext + p.materialize() + p.__class__ = ColoParameter + p.__init__(p, requires_grad=requires_grad) + + def state_dict_shard(self, + prefix: str = '', + keep_vars: bool = False, + max_shard_size: int = 1024, + only_rank_0: bool = True, + dtype: torch.dtype = torch.float16) -> Iterator[Tuple[OrderedDict, int]]: + """Returns dictionaries containing a whole state of the module one by one. The max size of dictionary shard is specified by ``max_shard_size``. + + Both parameters and persistent buffers (e.g. running averages) are included. + Keys are corresponding parameter and buffer names. + Parameters and buffers set to ``None`` are not included. + + Args: + prefix (str, optional): the prefix for parameters and buffers used in this + module. Defaults to ''. + keep_vars (bool, optional): whether to keep variables. Defaults to False. + max_shard_size (int, optional): max size of state dict shard (in MB). Defaults to 1024. + only_rank_0 (bool, optional): only get data on rank0. Defaults to True. + + + Yields: + Iterator[OrderedDict]: A generator of state dict shard + """ + sharder = _StateDictSharder(max_shard_size) + + # get the mapping between copies and fp16 parameters + fp16_to_fp32 = dict() + for p, fp32_p in zip(self.fp16_params, self.fp32_params): + fp16_to_fp32[p] = fp32_p + + # key is fp32 param, and value is gathered param on CPU + gathered_param_buffer = dict() + for name, param in self.name2param.items(): + if param is not None: + if is_ddp_ignored(param): + # deal with ddp ignored parameters + gathered_param = param if keep_vars else param.detach() + else: + # as memory of fp16 param may be reused, we should use fp32 param and then convert to fp16 + fp32_param = fp16_to_fp32[param] + if fp32_param not in gathered_param_buffer: + chunk = self.chunk_manager.get_chunk(fp32_param) + gathered_param_buffer.update(self._get_chunk_to_save_data(chunk, only_rank_0, dtype)) + gathered_param = gathered_param_buffer.pop(fp32_param) + + block, block_size = sharder.append(prefix + name, gathered_param) + if block is not None: + yield block, block_size + + del fp16_to_fp32 + del gathered_param_buffer + + # save all buffers + for name, buf in self.named_buffers(): + if buf is not None and name not in self._non_persistent_buffers_set: + buffer = buf if keep_vars else buf.detach() + block, block_size = sharder.append(prefix + name, buffer) + if block is not None: + yield block, block_size + # save extra states + extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX + if getattr(self.__class__, "get_extra_state", + torch.nn.Module.get_extra_state) is not torch.nn.Module.get_extra_state: + extra_state = self.get_extra_state() + block, block_size = sharder.append(extra_state_key, extra_state) + if block is not None: + yield block, block_size + + yield sharder.current_block, sharder.current_block_size + + +class _StateDictSharder: + + def __init__(self, max_shard_size: int) -> None: + self.max_shard_size = max_shard_size + self.current_block = OrderedDict() + self.current_block_size = 0 + + def append(self, name: str, tensor: torch.Tensor) -> Tuple[Optional[OrderedDict], int]: + tensor_size = calculate_tensor_size(tensor) + ret_block = None + ret_block_size = 0 + + # before we return the current block and create a new block, + # we need to ensure that the current block is not empty + if self.current_block_size + tensor_size > self.max_shard_size and self.current_block_size > 0: + ret_block = self.current_block + ret_block_size = self.current_block_size + self.current_block = OrderedDict() + self.current_block_size = 0 + self.current_block[name] = tensor + self.current_block_size += tensor_size + return ret_block, ret_block_size + + +class GeminiDDP(ZeroDDP): + + def __init__(self, + module: torch.nn.Module, + device: torch.device, + placement_policy: str = "cpu", + pin_memory: bool = False, + force_outputs_fp32: bool = False, + strict_ddp_mode: bool = False, + scatter_after_inference: bool = True, + search_range_m: int = 32, + hidden_dim: Optional[int] = None, + min_chunk_size_m: float = 32, + memstats: Optional[MemStats] = None, + mixed_precision: torch.dtype = torch.float16, + verbose: bool = False) -> None: + """ + A torch.Module wrapper using ZeRO-DP and Gemini. + ZeRO is for parallel. Gemini is for memory management. + WARNING: The class will modify the module inline! + + Example: + model is initialized under the context of ColoInitContext + >>> model = GeminiDDP(model, torch.cuda.current_device(), "cuda") + >>> logits = model(x) + >>> loss = criterion(logits, labels) + >>> model.backward(loss) + + Args: + module (torch.nn.Module): the model to be wrapped. + device (torch.device): device to place the model. + placement_policy (str, optional): "cpu", "cuda", "auto". Defaults to "cpu". + pin_memory (bool, optional): use pin memory on CPU. Defaults to False. + force_outputs_fp32 (bool, optional): force outputs are fp32. Defaults to False. + search_range_m (int, optional): chunk size searching range divided by 2^20. Defaults to 32. + hidden_dim (int, optional): the hidden dimension of DNN. + Users can provide this argument to speed up searching. + If users do not know this argument before training, it is ok. We will use a default value 1024. + min_chunk_size_m (float, optional): the minimum chunk size divided by 2^20. + If the aggregate size of parameters is still smaller than the minimum chunk size, + all parameters will be compacted into one small chunk. + memstats (MemStats, optional) the memory statistics collector by a runtime memory tracer. + """ + # some ugly hotfix for the compatibility with Lightning + if search_range_m is None: + search_range_m = 32 + + chunk_manager = init_chunk_manager(model=module, + init_device=device, + hidden_dim=hidden_dim, + search_range_m=search_range_m, + min_chunk_size_m=min_chunk_size_m, + strict_ddp_flag=strict_ddp_mode, + verbose=verbose) + gemini_manager = GeminiManager(placement_policy, chunk_manager, memstats) + super().__init__(module, + gemini_manager, + pin_memory, + force_outputs_fp32, + strict_ddp_mode, + scatter_after_inference, + mixed_precision=mixed_precision) diff --git a/colossalai/zero/utils/gemini_hook.py b/colossalai/zero/gemini/gemini_hook.py similarity index 95% rename from colossalai/zero/utils/gemini_hook.py rename to colossalai/zero/gemini/gemini_hook.py index bddc307a0504..dbc2924858e6 100644 --- a/colossalai/zero/utils/gemini_hook.py +++ b/colossalai/zero/gemini/gemini_hook.py @@ -5,10 +5,10 @@ import torch -from colossalai.gemini import TensorState -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.tensor.param_op_hook import ColoParamOpHook from colossalai.utils import is_ddp_ignored +from colossalai.zero.gemini import TensorState +from colossalai.zero.gemini.gemini_mgr import GeminiManager class TrainingPhase(Enum): diff --git a/colossalai/gemini/gemini_mgr.py b/colossalai/zero/gemini/gemini_mgr.py similarity index 97% rename from colossalai/gemini/gemini_mgr.py rename to colossalai/zero/gemini/gemini_mgr.py index 72a5e4a7f19b..c38e6eff840d 100644 --- a/colossalai/gemini/gemini_mgr.py +++ b/colossalai/zero/gemini/gemini_mgr.py @@ -4,10 +4,8 @@ import torch -from colossalai.gemini.chunk import Chunk, ChunkManager -from colossalai.gemini.memory_tracer import MemStats - -from .memory_tracer import ChunkMemStatsCollector +from .chunk import Chunk, ChunkManager +from .memory_tracer import ChunkMemStatsCollector, MemStats from .placement_policy import PlacementPolicyFactory diff --git a/colossalai/zero/gemini/gemini_optimizer.py b/colossalai/zero/gemini/gemini_optimizer.py new file mode 100644 index 000000000000..7d0db6b1fa23 --- /dev/null +++ b/colossalai/zero/gemini/gemini_optimizer.py @@ -0,0 +1,749 @@ +# this code is inspired by the DeepSpeed library and implemented with our own design from scratch +import copy +import gc +import math +import warnings +from typing import Any, Dict, Iterator, OrderedDict, Set, Tuple + +import torch +import torch.distributed as dist +from torch.nn import Parameter +from torch.optim import Optimizer + +from colossalai.amp.naive_amp.mixed_precision_mixin import BF16MixedPrecisionMixin, FP16MixedPrecisionMixin +from colossalai.checkpoint_io.utils import calculate_tensor_size +from colossalai.logging import get_dist_logger +from colossalai.nn.optimizer import ColossalaiOptimizer, CPUAdam, FusedAdam, HybridAdam +from colossalai.tensor.d_tensor import is_distributed_tensor +from colossalai.utils import disposable, get_current_device, is_ddp_ignored + +from .chunk import Chunk, ChunkManager +from .gemini_ddp import ZeroDDP + +__all__ = ['ZeroOptimizer', 'GeminiAdamOptimizer'] + +_AVAIL_OPTIM_LIST = {FusedAdam, CPUAdam, HybridAdam} + + +class GeminiFP16MixedPrecisionMixin(FP16MixedPrecisionMixin): + + def __init__(self, + module: ZeroDDP, + initial_scale: float = 2**16, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32) -> None: + super().__init__(initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, + max_scale) + self.module = module + + def check_local_overflow(self) -> bool: + return self.module.overflow_counter > 0 + + def pre_zero_grad(self) -> None: + self.module.overflow_counter = 0 + + +class ZeroOptimizer(ColossalaiOptimizer): + """A wrapper for optimizer. ``ZeroDDP`` and ``ZeroOptimizer`` implement Zero Redundancy Optimizer (ZeRO state-3). + + Note: + You must use ``ZeroDDP`` with ``ZeroOptimizer``. + + Note: + Make sure you set ``placement_policy`` of ``GeminiManager`` to `"auto"`, + if you set ``gpu_margin_mem_ratio > 0``. + + Args: + optim (Optimizer): An Optimizer instance. + module (ZeroDDP): A ``ZeroDDP`` instance. + gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward) + which will be used when using hybrid CPU optimizer. + This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto". + Defaults to 0.0. + initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32. + min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. + growth_factor (float, optional): Growth_factor used by DynamicGradScaler. Defaults to 2. + backoff_factor (float, optional): Backoff_factor used by DynamicGradScaler. Defaults to 0.5. + growth_interval (float, optional): Growth_interval used by DynamicGradScaler. Defaults to 1000. + hysteresis (float, optional): Hysteresis used by DynamicGradScaler. Defaults to 2. + max_scale (int, optional): Max_scale used by DynamicGradScaler. Defaults to 2**32. + clipping_norm (float, optional): The norm value used to clip gradient. Defaults to 0.0. + norm_type (float, optional): The type of norm used for gradient clipping. Currently, only L2-norm (norm_type=2.0) + is supported in ZeroOptimizer. Defaults to 2.0. + verbose (bool, optional): Whether to print verbose information, including grad overflow info. Defaults to False. + """ + + def __init__(self, + optim: Optimizer, + module: ZeroDDP, + gpu_margin_mem_ratio: float = 0.0, + initial_scale: float = 2**32, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32, + clipping_norm: float = 0.0, + norm_type: float = 2.0, + verbose: bool = False, + **defaults: Any): + super().__init__(optim) + assert isinstance(module, ZeroDDP) + assert type(optim) in _AVAIL_OPTIM_LIST, "You should use an optimizer in the available list:\n" \ + f"{_AVAIL_OPTIM_LIST}" + self.module = module + self.gemini_manager = module.gemini_manager + self.chunk_manager: ChunkManager = self.gemini_manager.chunk_manager + self.param_to_range: Dict[Parameter, Tuple[int, int]] = dict() + self.param_to_chunk32: Dict[Parameter, Chunk] = dict() + self.chunk16_set: Set[Chunk] = set() + self.clipping_flag = clipping_norm > 0.0 + self.max_norm = clipping_norm + self.verbose = verbose + self.param_groups_backup = list() + + # Mapping from integer id to real/fake param tensor, used for checkpointing. + self.id_to_real_params: Dict[int, Parameter] = dict() + self.id_to_fake_params: Dict[int, Parameter] = dict() + + if self.clipping_flag: + assert norm_type == 2.0, "ZeroOptimizer only supports L2 norm now" + + ddp_param_list = [] + for name, param in module.named_parameters(): + if is_ddp_ignored(param): + if param.requires_grad: + warnings.warn(f"Parameter `{name}` is ignored by DDP but requires gradient! " + "You should handle its optimizer update by yourself!") + else: + ddp_param_list.append(param) + + for p, fp32_p in zip(ddp_param_list, module.fp32_params): + chunk_16 = self.chunk_manager.get_chunk(p) + if chunk_16 not in self.chunk16_set: + chunk_16.l2_norm_flag = self.clipping_flag + self.chunk16_set.add(chunk_16) + + self.__init__optimizer() + + if module.mixed_precision is torch.float16: + self.mix_precision_mixin = GeminiFP16MixedPrecisionMixin(module, + initial_scale=initial_scale, + min_scale=min_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + max_scale=max_scale) + elif module.mixed_precision is torch.bfloat16: + self.mix_precision_mixin = BF16MixedPrecisionMixin() + else: + raise RuntimeError(f"Unsupported mixed precision type: {module.mixed_precision}") + + self._logger = get_dist_logger() + + self.gpu_margin_mem_ratio: float = float(gpu_margin_mem_ratio) + assert 0.0 <= self.gpu_margin_mem_ratio <= 1.0, f'gpu_margin_mem_ratio must >=0.0 and <=1.0' + # Only move fp32 shards from CPU to GPU when user allows and inner optimizer is valid + # Inner optimizer must support optimizing hybrid (CPU and CUDA) tensors, + # and it must set `num_fp32_shards_per_param` correctly + self._should_move_fp32_params_h2d: bool = self.gemini_manager.is_cuda_margin_mem_avail and self.gpu_margin_mem_ratio > 0.0 and getattr( + optim, 'num_fp32_shards_per_param', 0) >= 2 + if self.gpu_margin_mem_ratio > 0.0 and not self.gemini_manager.is_cuda_margin_mem_avail: + self._logger.warning(f'gpu_margin_mem_ratio is meaningless when placement_policy is not "auto"', ranks=[0]) + + self._register_states = disposable(self._register_states_) + + def _set_grad_ptr(self): + for group in self.param_groups: + for fake_param in group['params']: + chunk32 = self.param_to_chunk32[fake_param] + begin, end = self.param_to_range[fake_param] + chunk16 = chunk32.paired_chunk + + fake_param.data = chunk16.payload[begin:end] + fake_param.grad = fake_param.data + fake_param.data = chunk32.payload[begin:end] + + def _update_fp16_params(self): + none_tensor = torch.empty([0]) + for group in self.param_groups: + for fake_param in group['params']: + assert fake_param.grad is None + fake_param.data = none_tensor.to(fake_param.device) + + for chunk16 in self.chunk16_set: + chunk16.optim_update() + + def _clear_global_norm(self) -> None: + for c16 in self.chunk16_set: + c16.l2_norm = None + + def _calc_global_norm(self) -> float: + norm_sqr: float = 0.0 + group_to_norm = dict() + for c16 in self.chunk16_set: + assert c16.l2_norm is not None + + if c16.is_gathered: + norm_sqr += c16.l2_norm + else: + # this chunk is sharded, use communication to collect total norm + if c16.torch_pg not in group_to_norm: + group_to_norm[c16.torch_pg] = 0.0 + group_to_norm[c16.torch_pg] += c16.l2_norm + + c16.l2_norm = None # clear l2 norm + + comm_buffer = torch.zeros(1, dtype=torch.float, device=get_current_device()) + for group, part_norm in group_to_norm.items(): + comm_buffer.fill_(part_norm) + dist.all_reduce(comm_buffer, group=group) + norm_sqr += comm_buffer.item() + + global_norm = math.sqrt(norm_sqr) + return global_norm + + def _get_combined_scale(self): + div_scale = self.mix_precision_mixin.get_grad_div_scale() + + if self.clipping_flag: + total_norm = self._calc_global_norm() + clip = ((total_norm / div_scale) + 1e-6) / self.max_norm + if clip > 1: + div_scale = clip * div_scale + + return -1 if div_scale == 1.0 else div_scale + + def zero_grad(self, *args, **kwargs): + self.mix_precision_mixin.pre_zero_grad() + return self.optim.zero_grad(set_to_none=True) + + def step(self, *args, **kwargs): + self._maybe_move_fp32_params() + self._set_grad_ptr() + + if self.mix_precision_mixin.should_skip_step(): + if self.verbose: + self._logger.info(f'Found overflow. Skip step') + self._clear_global_norm() # clear recorded norm + self.zero_grad() # reset all gradients + self._update_fp16_params() + return + + # get combined scale. combined scale = loss scale * clipping norm + # so that gradient = gradient / combined scale + combined_scale = self._get_combined_scale() + + ret = self.optim.step(div_scale=combined_scale, *args, **kwargs) + self._register_states() + self.zero_grad() + self._update_fp16_params() + return ret + + def clip_grad_norm(self, model: torch.nn.Module, max_norm: float, norm_type: float = 2.0): + raise NotImplementedError + + def backward(self, loss: torch.Tensor): + loss = self.mix_precision_mixin.pre_backward(loss) + self.module.backward(loss) + + def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tensor): + # This function is called except the last stage of pipeline parallel + # It receives the scaled grad from the previous rank + # No need to scale the grad again + # Need to unscale when optimizing + grad = self.mix_precision_mixin.pre_backward_by_grad(grad) + self.module.backward_by_grad(tensor, grad) + + def _maybe_move_fp32_params(self): + if self._should_move_fp32_params_h2d: + self._should_move_fp32_params_h2d = False + available_cuda_margin_mem = self.gemini_manager.cuda_margin_mem * self.gpu_margin_mem_ratio + fp32_params_available_cuda_margin_mem = available_cuda_margin_mem / self.optim.num_fp32_shards_per_param + fp32_params_used_cuda_margin_mem = 0 + + for group in self.param_groups: + for fake_param in group['params']: + chunk32 = self.param_to_chunk32[fake_param] + chunk16 = chunk32.paired_chunk + + if chunk32.device_type == 'cuda': + continue + + if fp32_params_used_cuda_margin_mem + chunk32.payload_mem < fp32_params_available_cuda_margin_mem: + self.chunk_manager.move_chunk(chunk32, get_current_device()) + # stores grad now + self.chunk_manager.move_chunk(chunk16, get_current_device()) + self.module.set_chunk_grad_device(chunk16, get_current_device()) + fp32_params_used_cuda_margin_mem += chunk32.payload_mem + + for group in self.param_groups: + for fake_param in group['params']: + chunk32 = self.param_to_chunk32[fake_param] + if chunk32.device_type == 'cuda': + state = self.optim.state[fake_param] + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(get_current_device()) + + def _register_states_(self): + for group in self.optim.param_groups: + for p in group['params']: + state = self.optim.state[p] + for val in state.values(): + if isinstance(val, torch.Tensor): + self.chunk_manager.add_extern_static_tensor(val) + + def __init__optimizer(self): + + def get_range_pair(local_chunk: Chunk, local_param: Parameter): + param_info = local_chunk.tensors_info[local_param] + if local_chunk.keep_gathered: + return param_info.offset, param_info.end + begin = max(0, param_info.offset - local_chunk.shard_begin) + end = min(local_chunk.shard_size, param_info.end - local_chunk.shard_begin) + return begin, end + + param_id = -1 + for group in self.optim.param_groups: + fake_params_list = list() + group_backup = {k: v for k, v in group.items() if k != 'params'} + group_ids = [] + for param in group['params']: + + # Record the mapping of id to current param. + param_id += 1 + self.id_to_real_params[param_id] = param + group_ids.append(param_id) + + # If current param is controlled by current process, add it to fake_param. + if is_ddp_ignored(param): + continue + chunk16 = self.chunk_manager.get_chunk(param) + range_pair = get_range_pair(chunk16, param) + if range_pair[0] >= range_pair[1]: + continue + grad_device = self.module.grads_device[param] + fake_param = torch.nn.Parameter(torch.empty([0], device=grad_device)) + self.param_to_chunk32[fake_param] = chunk16.paired_chunk + self.param_to_range[fake_param] = range_pair + self.id_to_fake_params[param_id] = fake_param + fake_params_list.append(fake_param) + + # Update self.optim.param_groups as well as backup group. + group['params'] = fake_params_list + group_backup['params'] = group_ids + self.param_groups_backup.append(group_backup) + + def get_offsets(self, param_id: int) -> tuple: + ''' + Args: + param_id(int): The id of parameter. + + Returns: + chunk_offset(int): Offset of parameter inside the chunk. + shard_offset(int): Offset of its optimizer state shard + relative to the whole optimizer state. + shard_size(int): Length of parameter shard owned by current process. + ''' + + if param_id not in self.id_to_fake_params: + return -1, -1, -1 + fake_param = self.id_to_fake_params[param_id] + chunk = self.param_to_chunk32[fake_param].paired_chunk + param = self.id_to_real_params[param_id] + param_info = chunk.tensors_info[param] + + begin_in_chunk, end_in_chunk = self.param_to_range[fake_param] + chunk_offset = begin_in_chunk + if chunk.keep_gathered: + shard_offset = 0 + else: + shard_offset = begin_in_chunk + chunk.shard_begin - param_info.offset + shard_size = end_in_chunk - begin_in_chunk + assert chunk_offset >= 0 and shard_offset >= 0 + return chunk_offset, shard_offset, shard_size + + def collect_states(self, param_id: int, only_rank_0: bool = True) -> dict: + """ + Args: + param_id (int): id of the parameter whose state is to be gathered at master rank. + only_rank_0(bool): if True, states will be collected only on master rank, otherwise collected on every rank. + + Returns: + collected_states(dict): the gathered optimzier state of parameter with given id + if this method is called by master rank, otherwise an empty dict. + + This method can work only when called by all processes simultaneously. + """ + + # Get param & chunk & process group. + param = self.id_to_real_params[param_id] + fake_param = self.id_to_fake_params.get(param_id, None) + chunk = self.chunk_manager.get_chunk(param) + process_group = chunk.torch_pg + rank = dist.get_rank(process_group) + master_rank = 0 + collected_states = {} + + # Fetch names of states through all_gather. + local_state_names = None + if fake_param is not None: + local_state_names = list(self.optim.state[fake_param].keys()) + gathered_state_names = [None for _ in range(dist.get_world_size(process_group))] + dist.barrier() + dist.all_gather_object(gathered_state_names, local_state_names) + state_names = None + for names in gathered_state_names: + if names is not None: + # Assume different devices share the same set of state names if they have. + state_names = copy.deepcopy(names) + break + + # Directly return if this parameter doesn't have optimizer states. + # e.g. parameter freezed/layer dropped + if state_names is None: + return collected_states + + # Boolean variable is_collector indicates that whether the current rank + # needs to gather the whole optimizer states. + # Only master rank is collector when only_rank_0 is True. + # Every rank is collector when only_rank_0 is False. + is_collector = (rank == master_rank) or (not only_rank_0) + + # If the chunk is kept gathered, + # the parameteres are treated the same as that of those in strict DDP during training. + # So states can be directly fetched from current device. + if chunk.keep_gathered: + assert param_id in self.id_to_fake_params + if is_collector: + states = self.optim.state[fake_param] + for state_name in state_names: + if state_name == 'step': + # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32. + collected_states[state_name] = torch.tensor(states['step'], + dtype=torch.float32, + requires_grad=False).cpu() + else: + state_tensor = states[state_name].detach().clone().to(torch.float32).cpu() + collected_states[state_name] = torch.reshape(state_tensor, param.shape) + return collected_states + + # Check whether the param with given id is managed by current process. + own_param = param_id in self.id_to_fake_params + + # Collector gets prepared for state collecting. + if is_collector: + for state_name in state_names: + if state_name == 'step': + # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32. + collected_states[state_name] = torch.tensor(0.0, dtype=torch.float32, requires_grad=False).cpu() + else: + collected_states[state_name] = torch.zeros(param.numel(), dtype=torch.float32, + requires_grad=False).cpu() + + # Materials for gathering, including compacted state tensors, and the offset of shard inside each state. + compacted_states = self.pack_optimizer_states_to_tensor(param_id, state_names) if own_param else None + _, shard_offset, shard_size = self.get_offsets(param_id) + + # Collectors gather state shards through all_gathering. + gathered_state_shards = [None for _ in range(dist.get_world_size(process_group))] + + dist.barrier() + dist.all_gather_object(gathered_state_shards, [compacted_states, shard_offset, shard_size]) + + if is_collector: + for state_shard in gathered_state_shards: + compacted_states = state_shard[0] + shard_offset = state_shard[1] + shard_size = state_shard[2] + if compacted_states is None: + continue + self.load_from_compacted_states(compacted_states, collected_states, state_names, shard_offset, + shard_size) + + # Clean gathered states + for state_shard in gathered_state_shards: + del state_shard[0] + gc.collect() + + # Reshape tensors + if is_collector: + for state_name, state_tensor in collected_states.items(): + if state_tensor.numel() == param.numel(): + collected_states[state_name] = torch.reshape(state_tensor, param.shape) + + return collected_states + + def pack_optimizer_states_to_tensor(self, + param_id: int, + state_names: list, + device: torch.device = torch.device('cuda'), + dtype: torch.dtype = torch.float32) -> torch.Tensor: + ''' + With param id given, pack its optimizer states into a compact tensor and return. + ''' + if param_id not in self.id_to_fake_params: + return None + + fake_param = self.id_to_fake_params[param_id] + param_range = self.param_to_range[fake_param] + states = self.optim.state[fake_param] + shard_size = param_range[1] - param_range[0] + compacted_size = 0 + for name in state_names: + if name == 'step': + compacted_size += 1 + else: + compacted_size += shard_size + compacted_states = torch.zeros(compacted_size, dtype=dtype, device=device, requires_grad=False) + + next_state_offset = 0 + for state_name, state_tensor in states.items(): + # State 'step' needs special operation. + if state_name == 'step': + if isinstance(state_tensor, torch.Tensor): + compacted_states[next_state_offset] = state_tensor[0].item() + else: + assert isinstance(state_tensor, int) + compacted_states[next_state_offset] = state_tensor + next_state_offset += 1 + else: + assert state_tensor.numel() == shard_size + compacted_states[next_state_offset:next_state_offset + shard_size].copy_(state_tensor) + next_state_offset += shard_size + + return compacted_states + + def load_from_compacted_states(self, compacted_states: torch.Tensor, collected_states: dict, state_names: list, + shard_start: int, shard_size: int): + ''' + Given a tensor carrying compacted optimizer states, + update these states to collected_states. + ''' + shard_end = shard_start + shard_size + next_state_offset = 0 + + for state_name in state_names: + if state_name == 'step': + collected_states['step'].data = torch.tensor(compacted_states[next_state_offset].item(), + dtype=torch.float32, + requires_grad=False).cpu() + next_state_offset += 1 + else: + target_segment = collected_states[state_name][shard_start:shard_end] + target_segment.copy_(compacted_states[next_state_offset:next_state_offset + shard_size]) + next_state_offset += shard_size + + def get_param_groups_for_saving(self) -> list: + ''' + Return the param_groups in Pytorch format when saving to checkpoint. + ''' + + param_groups = copy.deepcopy(self.param_groups_backup) + + # To be compatible with pytorch checkpointing, + # store extra hyperparameters used by pytorch Adam optimizer. + torch_special_hyperparameters = { + 'amsgrad': False, + 'maximize': False, + 'foreach': None, + 'capturable': False, + 'differentiable': False, + 'fused': False + } + + for group in param_groups: + for k, v in torch_special_hyperparameters.items(): + if k not in group: + group[k] = v + + return param_groups + + def state_dict(self, only_rank_0: bool = True) -> dict: + """ + Args: + only_rank_0 (bool): a boolean value indicating whether the state_dict is collected + only on rank 0, dafault to True. + + Returns: + The complete state of the optimizer as a :class:`dict`. + It contains two entries: + + * state - a dict holding current optimization state. Its content + differs between optimizer classes. + * param_groups - a list containing all parameter groups where each + parameter group is a dict. + + Warning: This method will gather and return the whole optimizer state_dict, + so it should be called only when memory resources are abundant. + """ + state_dict = {} + state_dict['param_groups'] = self.get_param_groups_for_saving() + + # Collect optimizer states. + state_dict['state'] = dict() + for param_id in self.id_to_real_params.keys(): + dist.barrier() + state_dict['state'][param_id] = self.collect_states(param_id=param_id, only_rank_0=only_rank_0) + return state_dict + + def load_param_groups(self, saved_param_groups: list): + """ + Load saved_param_groups into + self.param_groups and self.param_groups_backup + """ + self.param_groups_backup = copy.deepcopy(saved_param_groups) + + # discard the older param_groups + self.optim.param_groups = [] + + for group in saved_param_groups: + fake_params_list = list() + updated_group = {k: v for k, v in group.items() if k != 'params'} + for param_id in group['params']: + if param_id not in self.id_to_fake_params: + continue + fake_param = self.id_to_fake_params[param_id] + fake_params_list.append(fake_param) + updated_group['params'] = fake_params_list + self.optim.param_groups.append(updated_group) + + def load_single_param_states(self, param_id: int, saved_states: dict): + """ + Load saved optimizer states into parameter with given id. + """ + + def cast(param, state_range, value, key=None): + """ + Make a copy of the needed segment of value and cast it to device of param. + """ + assert isinstance(value, torch.Tensor) + ret_val = value + if (key == "step"): + assert value.numel() == 1 + ret_val = int(value.item()) + else: + state_start, state_end = state_range + ret_val = torch.zeros(state_end - state_start, + dtype=torch.float32, + device=param.device, + requires_grad=False) + ret_val.copy_(value.flatten()[state_start:state_end]) + return ret_val + + assert param_id in self.id_to_fake_params + fake_param = self.id_to_fake_params[param_id] + _, state_offset, param_size = self.get_offsets(param_id) + state_range = (state_offset, state_offset + param_size) + + # Copy states assigned to param (and cast tensors to appropriate types). + updated_states = dict() + for k, v in saved_states.items(): + updated_states[k] = cast(fake_param, state_range, v, k) + del v # clean loaded states + self.optim.state[fake_param].update(updated_states) + + def load_param_states(self, param_states: dict): + """Loads param states from a state_dict. The param_states can be complete or sharded. + During loading, filter out the part of states not considered by current process. + + Args: + param_states (dict): A mapping from param_id to its states. + """ + for param_id, states in param_states.items(): + if param_id in self.id_to_fake_params: + self.load_single_param_states(param_id, states) + + def optimizer_loading_epilogue(self): + # Epilogue when loading state_dict to pytorch optimizer. + self.optim._hook_for_profile() # To support multiprocessing pickle/unpickle. + self.optim.defaults.setdefault('differentiable', False) + + def load_state_dict(self, state_dict: dict): + """Loads optimizer state from complete optimizer state_dict. + During loading, filter out the part of states not considered by current process. + + Args: + state_dict (dict): optimizer state. Should be an object returned + from a call to :meth:`state_dict`. + """ + assert 'param_groups' in state_dict + assert 'state' in state_dict + self.load_param_groups(state_dict['param_groups']) + self.load_param_states(state_dict['state']) + self.optimizer_loading_epilogue() + + def state_shard(self, + prefix: str = '', + max_shard_size: int = 1024, + only_rank_0: bool = True) -> Iterator[Tuple[OrderedDict, int]]: + """Returns dictionaries containing shards of optimizer states one by one. + The max size of each dictionary shard is specified by ``max_shard_size``. + + Args: + prefix (str, optional): the prefix for states. Default to ''. + max_shard_size (int, optional): max size of state dict shard (in MB). Defaults to 1024. + only_rank_0 (bool, optional): a boolean value indicating whether the state_dict is collected + only on rank 0, dafault to True. + + Yields: + Iterator[OrderedDict]: A generator of state dict shard of optimizer states. + """ + + current_block = {} + current_block_size = 0 + + for param_id in self.id_to_real_params.keys(): + + dist.barrier() + state = self.collect_states(param_id=param_id, only_rank_0=only_rank_0) + + ret_block = None + ret_block_size = 0 + + # A state might contain more than one tensors. + # e.g. each Adam state includes: 'step', 'exp_avg', 'exp_avg_sq' + state_size = 0 + isDTensor = False + for state_tensor in state.values(): + + # When state_tensor is not of Tensor class, + # e.g., a SGD optimizer with momentum set to 0 can have None as state + # The calculation of tensor size should be skipped to avoid error. + if not isinstance(state_tensor, torch.Tensor): + continue + + # If the states are stored as DTensors, mark isDTensor as true. + if is_distributed_tensor(state_tensor): + isDTensor = True + state_size += calculate_tensor_size(state_tensor) + + if not isDTensor: + + if current_block_size + state_size > max_shard_size and current_block_size > 0: + ret_block = current_block + ret_block_size = current_block_size + current_block = {} + current_block_size = 0 + + current_block[param_id] = state + current_block_size += state_size + + if ret_block != None: + yield ret_block, ret_block_size + + yield current_block, current_block_size + + +class GeminiAdamOptimizer(ZeroOptimizer): + + def __init__(self, model: torch.nn.Module, **defaults: Any) -> None: + optimizer = HybridAdam(model.parameters(), **defaults) + super().__init__(optimizer, model, **defaults) diff --git a/colossalai/gemini/memory_tracer/__init__.py b/colossalai/zero/gemini/memory_tracer/__init__.py similarity index 100% rename from colossalai/gemini/memory_tracer/__init__.py rename to colossalai/zero/gemini/memory_tracer/__init__.py diff --git a/colossalai/gemini/memory_tracer/chunk_memstats_collector.py b/colossalai/zero/gemini/memory_tracer/chunk_memstats_collector.py similarity index 87% rename from colossalai/gemini/memory_tracer/chunk_memstats_collector.py rename to colossalai/zero/gemini/memory_tracer/chunk_memstats_collector.py index 1a5b6bf525be..83903bbf4023 100644 --- a/colossalai/gemini/memory_tracer/chunk_memstats_collector.py +++ b/colossalai/zero/gemini/memory_tracer/chunk_memstats_collector.py @@ -1,10 +1,10 @@ from typing import Optional -from colossalai.gemini.chunk import ChunkManager -from colossalai.gemini.memory_tracer import MemStats from colossalai.utils import get_current_device from colossalai.utils.memory import colo_device_memory_capacity +from colossalai.zero.gemini.chunk import ChunkManager +from .memory_stats import MemStats from .memstats_collector import MemStatsCollector @@ -25,7 +25,7 @@ def __init__(self, chunk_manager: ChunkManager, memstats: Optional[MemStats] = N # override def record_model_data_volume(self) -> None: """ - record model data volumn on cuda and cpu. + record model data volume on cuda and cpu. """ if self._start_flag and not self.use_outside_memstats: cuda_mem = self._chunk_manager.total_mem['cuda'] diff --git a/colossalai/gemini/memory_tracer/memory_monitor.py b/colossalai/zero/gemini/memory_tracer/memory_monitor.py similarity index 95% rename from colossalai/gemini/memory_tracer/memory_monitor.py rename to colossalai/zero/gemini/memory_tracer/memory_monitor.py index f8d99dbce7a4..4bb585677d5b 100644 --- a/colossalai/gemini/memory_tracer/memory_monitor.py +++ b/colossalai/zero/gemini/memory_tracer/memory_monitor.py @@ -45,7 +45,7 @@ def clear(self): class AsyncMemoryMonitor(MemoryMonitor): """ - An Async Memory Monitor runing during computing. Sampling memory usage of the current GPU + An Async Memory Monitor running during computing. Sampling memory usage of the current GPU at interval of `1/(10**power)` sec. The idea comes from Runtime Memory Tracer of PatrickStar @@ -67,7 +67,7 @@ class AsyncMemoryMonitor(MemoryMonitor): async_mem_monitor.save('log.pkl') Args: - power (int, optional): the power of time interva. Defaults to 10. + power (int, optional): the power of time interval. Defaults to 10. .. _PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management: https://arxiv.org/abs/2108.05818 diff --git a/colossalai/gemini/memory_tracer/memory_stats.py b/colossalai/zero/gemini/memory_tracer/memory_stats.py similarity index 98% rename from colossalai/gemini/memory_tracer/memory_stats.py rename to colossalai/zero/gemini/memory_tracer/memory_stats.py index 84fa00fb9361..41d7e5754e96 100644 --- a/colossalai/gemini/memory_tracer/memory_stats.py +++ b/colossalai/zero/gemini/memory_tracer/memory_stats.py @@ -2,7 +2,7 @@ import torch -from colossalai.gemini.memory_tracer import OrderedParamGenerator +from .param_runtime_order import OrderedParamGenerator class MemStats(object): @@ -59,7 +59,7 @@ def increase_preop_step(self, param_list: List[torch.nn.Parameter]): time step. Args: - param_list (List[torch.nn.Parameter]): a list of torch paramters. + param_list (List[torch.nn.Parameter]): a list of torch parameters. """ for p in param_list: if p not in self._param_step_dict: diff --git a/colossalai/gemini/memory_tracer/memstats_collector.py b/colossalai/zero/gemini/memory_tracer/memstats_collector.py similarity index 92% rename from colossalai/gemini/memory_tracer/memstats_collector.py rename to colossalai/zero/gemini/memory_tracer/memstats_collector.py index d939da6eb4cf..0694be48550a 100644 --- a/colossalai/gemini/memory_tracer/memstats_collector.py +++ b/colossalai/zero/gemini/memory_tracer/memstats_collector.py @@ -1,12 +1,7 @@ import time -from typing import List, Optional - -import torch - -from colossalai.gemini.memory_tracer import SyncCudaMemoryMonitor -from colossalai.gemini.stateful_tensor import StatefulTensor -from colossalai.utils.memory import colo_device_memory_used +from typing import Optional +from .memory_monitor import SyncCudaMemoryMonitor from .memory_stats import MemStats @@ -49,7 +44,7 @@ def next_period_non_model_data_usage(self, device_type: str) -> int: assert self._step_total > 0, 'Cannot get mem stats info before collection phase.' assert len(self._memstats.non_model_data_list(device_type)) > self._step_idx, \ f"{len(self._memstats.non_model_data_list(device_type))} should be > than step idx {self._step_idx}, "\ - f"step total {self._step_total}" + f"step total {self._step_total}" next_non_model_data = self._memstats.non_model_data_list(device_type)[self._step_idx] self._step_idx = (self._step_idx + 1) % self._step_total return next_non_model_data @@ -75,6 +70,8 @@ def record_model_data_volume(self) -> None: Sampling model data statistics. """ if self._start_flag and not self.use_outside_memstats: + from colossalai.zero.legacy.gemini import StatefulTensor + # The following code work for ZeroInitContext, which is deprecated in v0.1.12 cuda_mem = StatefulTensor.GST_MGR.total_mem['cuda'] self._memstats.record_max_cuda_model_data(cuda_mem) diff --git a/colossalai/gemini/memory_tracer/param_runtime_order.py b/colossalai/zero/gemini/memory_tracer/param_runtime_order.py similarity index 100% rename from colossalai/gemini/memory_tracer/param_runtime_order.py rename to colossalai/zero/gemini/memory_tracer/param_runtime_order.py diff --git a/colossalai/gemini/memory_tracer/runtime_mem_tracer.py b/colossalai/zero/gemini/memory_tracer/runtime_mem_tracer.py similarity index 95% rename from colossalai/gemini/memory_tracer/runtime_mem_tracer.py rename to colossalai/zero/gemini/memory_tracer/runtime_mem_tracer.py index a643751da7e2..0c9eac8b63e3 100644 --- a/colossalai/gemini/memory_tracer/runtime_mem_tracer.py +++ b/colossalai/zero/gemini/memory_tracer/runtime_mem_tracer.py @@ -1,9 +1,14 @@ import torch.nn -from colossalai.gemini.memory_tracer import MemStats -from colossalai.gemini.ophooks.runtime_mem_tracer_hook import GradMemStats, GradMemTracerHook, ParamMemTracerHook from colossalai.nn.parallel.data_parallel import _cast_float from colossalai.tensor.param_op_hook import ColoParamOpHookManager +from colossalai.zero.legacy.gemini.ophooks.runtime_mem_tracer_hook import ( + GradMemStats, + GradMemTracerHook, + ParamMemTracerHook, +) + +from .memory_stats import MemStats __all__ = ['RuntimeMemTracer'] diff --git a/colossalai/gemini/memory_tracer/static_memstats_collector.py b/colossalai/zero/gemini/memory_tracer/static_memstats_collector.py similarity index 98% rename from colossalai/gemini/memory_tracer/static_memstats_collector.py rename to colossalai/zero/gemini/memory_tracer/static_memstats_collector.py index 3209881e100c..b8f9a095f422 100644 --- a/colossalai/gemini/memory_tracer/static_memstats_collector.py +++ b/colossalai/zero/gemini/memory_tracer/static_memstats_collector.py @@ -6,7 +6,7 @@ from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.fx.profiler import calculate_fwd_out, calculate_fwd_tmp, is_compatible_with_meta -from colossalai.gemini.chunk import ChunkManager +from colossalai.zero.gemini.chunk import ChunkManager if is_compatible_with_meta(): from colossalai.fx.profiler import MetaTensor diff --git a/colossalai/gemini/memory_tracer/utils.py b/colossalai/zero/gemini/memory_tracer/utils.py similarity index 96% rename from colossalai/gemini/memory_tracer/utils.py rename to colossalai/zero/gemini/memory_tracer/utils.py index 6962c058110e..65f6ba775139 100644 --- a/colossalai/gemini/memory_tracer/utils.py +++ b/colossalai/zero/gemini/memory_tracer/utils.py @@ -7,7 +7,7 @@ def colo_model_optimizer_usage(optim) -> Tuple[int, int]: """Trace the optimizer memory usage Args: - optim (ShardedOptimV2): an instance of ShardedOptimver + optim (ShardedOptimV2): an instance of ShardedOptimizer Returns: Tuple[int, int]: cuda/cpu memory usage in Byte diff --git a/colossalai/gemini/placement_policy.py b/colossalai/zero/gemini/placement_policy.py similarity index 98% rename from colossalai/gemini/placement_policy.py rename to colossalai/zero/gemini/placement_policy.py index fed1cc2985ff..84a868872f88 100644 --- a/colossalai/gemini/placement_policy.py +++ b/colossalai/zero/gemini/placement_policy.py @@ -5,11 +5,12 @@ import torch -from colossalai.gemini.chunk import Chunk, ChunkManager -from colossalai.gemini.memory_tracer import ChunkMemStatsCollector from colossalai.utils import get_current_device from colossalai.utils.memory import colo_device_memory_capacity +from .chunk import Chunk, ChunkManager +from .memory_tracer import ChunkMemStatsCollector + class PlacementPolicy(ABC): need_mem_stats: bool = False diff --git a/colossalai/nn/parallel/utils.py b/colossalai/zero/gemini/utils.py similarity index 96% rename from colossalai/nn/parallel/utils.py rename to colossalai/zero/gemini/utils.py index 08fdb6026e38..6f4a253b504b 100644 --- a/colossalai/nn/parallel/utils.py +++ b/colossalai/zero/gemini/utils.py @@ -6,9 +6,10 @@ import torch.distributed as dist import torch.nn as nn -from colossalai.gemini.chunk import Chunk from colossalai.utils import get_current_device +from .chunk import Chunk + def get_temp_total_chunk_on_cuda(chunk: Chunk): if chunk.is_gathered: @@ -72,12 +73,12 @@ def get_static_torch_model(zero_ddp_model, zero_ddp_model (ZeroDDP): a zero ddp model device (torch.device): the device of the final torch model dtype (torch.dtype): the dtype of the final torch model - only_rank_0 (bool): if True, only rank0 has the coverted torch model + only_rank_0 (bool): if True, only rank0 has the converted torch model Returns: torch.nn.Module: a static torch model used for saving checkpoints or numeric checks """ - from colossalai.nn.parallel import ZeroDDP + from colossalai.zero.gemini.gemini_ddp import ZeroDDP assert isinstance(zero_ddp_model, ZeroDDP) state_dict = zero_ddp_model.state_dict(only_rank_0=only_rank_0) diff --git a/colossalai/zero/legacy/__init__.py b/colossalai/zero/legacy/__init__.py new file mode 100644 index 000000000000..3783d38e61b2 --- /dev/null +++ b/colossalai/zero/legacy/__init__.py @@ -0,0 +1,45 @@ +from typing import Tuple + +import torch +import torch.nn as nn + +from colossalai.logging import get_dist_logger + +from .init_ctx import ZeroInitContext, no_shard_zero_context, no_shard_zero_decrator +from .shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from .sharded_model import ShardedModelV2 +from .sharded_optim import ShardedOptimizerV2 + + +def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model_config, + optimizer_config) -> Tuple[ShardedModelV2, ShardedOptimizerV2]: + """ + A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading + + :param model: Your model object + :type model: :class:`torch.nn.Module` + :param optimizer_config: Your optimizer object + :type optimizer_config: :class:`dict` + + :return: (model, optimizer) + :rtype: Tuple + """ + + logger = get_dist_logger('convert_to_zero_v2') + + logger.info(f'optimizer_config is {optimizer_config}', ranks=[0]) + if optimizer_config is None: + optimizer_config = dict() + logger.info(f'model_config is {model_config}', ranks=[0]) + if model_config is None: + model_config = dict() + + zero_model = ShardedModelV2(model, **model_config) + zero_optimizer = ShardedOptimizerV2(zero_model, optimizer, **optimizer_config) + return zero_model, zero_optimizer + + +__all__ = [ + 'convert_to_zero_v2', 'ShardedModelV2', 'ShardedOptimizerV2', 'ZeroInitContext', 'no_shard_zero_context', + 'no_shard_zero_decrator', 'TensorShardStrategy', 'BucketTensorShardStrategy' +] diff --git a/colossalai/zero/legacy/gemini/__init__.py b/colossalai/zero/legacy/gemini/__init__.py new file mode 100644 index 000000000000..754ae9bc0044 --- /dev/null +++ b/colossalai/zero/legacy/gemini/__init__.py @@ -0,0 +1,9 @@ +from .ophooks import BaseOpHook, register_ophooks_recursively +from .stateful_tensor import StatefulTensor +from .stateful_tensor_mgr import StatefulTensorMgr +from .tensor_placement_policy import AutoTensorPlacementPolicy, CPUTensorPlacementPolicy, CUDATensorPlacementPolicy + +__all__ = [ + 'StatefulTensorMgr', 'StatefulTensor', 'CPUTensorPlacementPolicy', 'CUDATensorPlacementPolicy', + 'AutoTensorPlacementPolicy', 'register_ophooks_recursively', 'BaseOpHook' +] diff --git a/colossalai/gemini/gemini_context.py b/colossalai/zero/legacy/gemini/gemini_context.py similarity index 100% rename from colossalai/gemini/gemini_context.py rename to colossalai/zero/legacy/gemini/gemini_context.py diff --git a/colossalai/gemini/ophooks/__init__.py b/colossalai/zero/legacy/gemini/ophooks/__init__.py similarity index 100% rename from colossalai/gemini/ophooks/__init__.py rename to colossalai/zero/legacy/gemini/ophooks/__init__.py diff --git a/colossalai/gemini/ophooks/_shard_grad_ophook.py b/colossalai/zero/legacy/gemini/ophooks/_shard_grad_ophook.py similarity index 88% rename from colossalai/gemini/ophooks/_shard_grad_ophook.py rename to colossalai/zero/legacy/gemini/ophooks/_shard_grad_ophook.py index 5115ff74da16..8f8fec64924e 100644 --- a/colossalai/gemini/ophooks/_shard_grad_ophook.py +++ b/colossalai/zero/legacy/gemini/ophooks/_shard_grad_ophook.py @@ -8,7 +8,7 @@ @OPHOOKS.register_module class ShardGradMemTracerHook(BaseOpHook): """ - A hook to process sharded param before and afther FWD and BWD operator executing. + A hook to process sharded param before and after FWD and BWD operator executing. """ def __init__(self): diff --git a/colossalai/gemini/ophooks/_shard_param_ophook.py b/colossalai/zero/legacy/gemini/ophooks/_shard_param_ophook.py similarity index 93% rename from colossalai/gemini/ophooks/_shard_param_ophook.py rename to colossalai/zero/legacy/gemini/ophooks/_shard_param_ophook.py index 57f76970cc86..a2a62fb9788a 100644 --- a/colossalai/gemini/ophooks/_shard_param_ophook.py +++ b/colossalai/zero/legacy/gemini/ophooks/_shard_param_ophook.py @@ -1,4 +1,5 @@ import torch + from colossalai.registry import OPHOOKS from . import BaseOpHook @@ -7,7 +8,7 @@ @OPHOOKS.register_module class ShardParamHook(BaseOpHook): """ - A hook to process sharded param before and afther FWD and BWD operator executing. + A hook to process sharded param before and after FWD and BWD operator executing. """ def __init__(self): diff --git a/colossalai/gemini/ophooks/runtime_mem_tracer_hook.py b/colossalai/zero/legacy/gemini/ophooks/runtime_mem_tracer_hook.py similarity index 96% rename from colossalai/gemini/ophooks/runtime_mem_tracer_hook.py rename to colossalai/zero/legacy/gemini/ophooks/runtime_mem_tracer_hook.py index 6d0df4e615ca..f40d6ced1ee0 100644 --- a/colossalai/gemini/ophooks/runtime_mem_tracer_hook.py +++ b/colossalai/zero/legacy/gemini/ophooks/runtime_mem_tracer_hook.py @@ -5,9 +5,9 @@ import torch -from colossalai.gemini.memory_tracer import MemStats, SyncCudaMemoryMonitor -from colossalai.gemini.tensor_utils import alloc_storage, free_storage from colossalai.tensor.param_op_hook import ColoParamOpHook +from colossalai.zero.gemini.memory_tracer import MemStats, SyncCudaMemoryMonitor +from colossalai.zero.legacy.gemini.tensor_utils import alloc_storage, free_storage class TrainingPhase(Enum): diff --git a/colossalai/gemini/ophooks/utils.py b/colossalai/zero/legacy/gemini/ophooks/utils.py similarity index 97% rename from colossalai/gemini/ophooks/utils.py rename to colossalai/zero/legacy/gemini/ophooks/utils.py index 84e8298c1d51..f88ad2b00e9e 100644 --- a/colossalai/gemini/ophooks/utils.py +++ b/colossalai/zero/legacy/gemini/ophooks/utils.py @@ -88,7 +88,7 @@ def register_ophooks_recursively(module: torch.nn.Module, ophook_list: List[BaseOpHook], name: str = "", filter_fn: Optional[Callable] = None): - r"""Recursilvely register pre/post hooks for all submodules in the module in FWD and BWD.""" + r"""Recursively register pre/post hooks for all submodules in the module in FWD and BWD.""" assert isinstance(module, torch.nn.Module) assert isinstance(ophook_list, (list, tuple)) assert len(ophook_list) > 0, 'expected at least 1 hook in the argument ophook_list but found 0' @@ -103,7 +103,7 @@ def register_ophooks_recursively(module: torch.nn.Module, if len(list(module.parameters(recurse=False))) == 0: return - # return from flitered module + # return from filtered module if filter_fn is not None and filter_fn(module): return diff --git a/colossalai/gemini/paramhooks/__init__.py b/colossalai/zero/legacy/gemini/paramhooks/__init__.py similarity index 100% rename from colossalai/gemini/paramhooks/__init__.py rename to colossalai/zero/legacy/gemini/paramhooks/__init__.py diff --git a/colossalai/gemini/paramhooks/_param_hookmgr.py b/colossalai/zero/legacy/gemini/paramhooks/_param_hookmgr.py similarity index 99% rename from colossalai/gemini/paramhooks/_param_hookmgr.py rename to colossalai/zero/legacy/gemini/paramhooks/_param_hookmgr.py index ee57cb46a90d..84f32be358e3 100644 --- a/colossalai/gemini/paramhooks/_param_hookmgr.py +++ b/colossalai/zero/legacy/gemini/paramhooks/_param_hookmgr.py @@ -1,6 +1,7 @@ +import functools from typing import Callable, List + import torch -import functools class BaseParamHookMgr(object): diff --git a/colossalai/gemini/stateful_tensor.py b/colossalai/zero/legacy/gemini/stateful_tensor.py similarity index 97% rename from colossalai/gemini/stateful_tensor.py rename to colossalai/zero/legacy/gemini/stateful_tensor.py index 18fc8fd14d3c..1619ae40798d 100644 --- a/colossalai/gemini/stateful_tensor.py +++ b/colossalai/zero/legacy/gemini/stateful_tensor.py @@ -1,9 +1,9 @@ from enum import Enum -from typing import Optional +from typing import Optional, Union + import torch -from typing import Union -from colossalai.gemini.gemini_context import GeminiMemoryManager +from .gemini_context import GeminiMemoryManager def sizeof_tensor(tensor: torch.Tensor): @@ -19,7 +19,7 @@ class TensorState(Enum): class StatefulTensor(object): - """A Structure stores a Torch Tensor and labeled states. + """A Structure stores a Torch Tensor and labeled states. Inspired from the paper: PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management diff --git a/colossalai/gemini/stateful_tensor_mgr.py b/colossalai/zero/legacy/gemini/stateful_tensor_mgr.py similarity index 92% rename from colossalai/gemini/stateful_tensor_mgr.py rename to colossalai/zero/legacy/gemini/stateful_tensor_mgr.py index c300f9bffc89..4f9ea7c6d520 100644 --- a/colossalai/gemini/stateful_tensor_mgr.py +++ b/colossalai/zero/legacy/gemini/stateful_tensor_mgr.py @@ -1,13 +1,16 @@ import functools -import torch import types -from colossalai.utils.cuda import get_current_device -from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage -from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState -from colossalai.gemini.tensor_placement_policy import TensorPlacementPolicy +from time import time from typing import List + +import torch + from colossalai.logging import get_dist_logger -from time import time +from colossalai.utils.cuda import get_current_device + +from .stateful_tensor import StatefulTensor, TensorState +from .tensor_placement_policy import TensorPlacementPolicy +from .tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage class StatefulTensorMgr(object): @@ -50,7 +53,7 @@ def finish_iter(self): self._evict_time = 0 def adjust_layout(self) -> None: - """ Adjust the layout of statefuil tensor according to the information provided + """ Adjust the layout of stateful tensor according to the information provided by mem_stats_collector, which should belongs to a Sharded Model. """ # find stateful tensor in state COMPUTE diff --git a/colossalai/gemini/tensor_placement_policy.py b/colossalai/zero/legacy/gemini/tensor_placement_policy.py similarity index 95% rename from colossalai/gemini/tensor_placement_policy.py rename to colossalai/zero/legacy/gemini/tensor_placement_policy.py index cfcfb385667c..165ae51fee60 100644 --- a/colossalai/gemini/tensor_placement_policy.py +++ b/colossalai/zero/legacy/gemini/tensor_placement_policy.py @@ -1,15 +1,16 @@ +import functools from abc import ABC, abstractmethod from time import time -from typing import List, Optional +from typing import List, Optional, Type + import torch + from colossalai.utils import get_current_device from colossalai.utils.memory import colo_device_memory_capacity +from colossalai.zero.gemini.memory_tracer import MemStatsCollector -from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage -from colossalai.gemini.stateful_tensor import StatefulTensor -from colossalai.gemini.memory_tracer import MemStatsCollector -from typing import Type -import functools +from .stateful_tensor import StatefulTensor +from .tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage class TensorPlacementPolicy(ABC): diff --git a/colossalai/gemini/tensor_utils.py b/colossalai/zero/legacy/gemini/tensor_utils.py similarity index 96% rename from colossalai/gemini/tensor_utils.py rename to colossalai/zero/legacy/gemini/tensor_utils.py index bcc159f9954a..843e330ee2c6 100644 --- a/colossalai/gemini/tensor_utils.py +++ b/colossalai/zero/legacy/gemini/tensor_utils.py @@ -1,6 +1,8 @@ +from typing import Tuple, Union + import torch -from colossalai.gemini.stateful_tensor import StatefulTensor -from typing import Union, Tuple + +from .stateful_tensor import StatefulTensor def is_storage_empty(tensor: torch.Tensor) -> bool: @@ -75,7 +77,7 @@ def colo_model_data_tensor_move_inline(t: Union[StatefulTensor, torch.Tensor], t move a tensor to the target_device Args: t (Union[StatefulTensor, torch.Tensor]): the tensor be moved - target_device: a traget device, if type is int, it the index of cuda card. + target_device: a target device, if type is int, it the index of cuda card. """ if not isinstance(target_device, torch.device): target_device = torch.device(f'cuda:{target_device}') diff --git a/colossalai/zero/init_ctx/__init__.py b/colossalai/zero/legacy/init_ctx/__init__.py similarity index 100% rename from colossalai/zero/init_ctx/__init__.py rename to colossalai/zero/legacy/init_ctx/__init__.py diff --git a/colossalai/zero/init_ctx/init_context.py b/colossalai/zero/legacy/init_ctx/init_context.py similarity index 85% rename from colossalai/zero/init_ctx/init_context.py rename to colossalai/zero/legacy/init_ctx/init_context.py index 572ddd9e4e3f..84e2d2f4f8e1 100644 --- a/colossalai/zero/init_ctx/init_context.py +++ b/colossalai/zero/legacy/init_ctx/init_context.py @@ -1,53 +1,52 @@ import contextlib import functools -from typing import Optional from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Optional import torch -import torch.nn as nn import torch.distributed as dist +import torch.nn as nn from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc from colossalai.context.singleton_meta import SingletonMeta +from colossalai.core import global_context as gpc from colossalai.logging import get_dist_logger -from colossalai.zero.shard_utils import BaseShardStrategy -from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16 -from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 -from colossalai.zero.sharded_param import ShardedParamV2 from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses +from colossalai.zero.legacy.shard_utils import BaseShardStrategy +from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_bf16, cast_tensor_to_fp16 +from colossalai.zero.legacy.sharded_model.sharded_model_v2 import ShardedModelV2 +from colossalai.zero.legacy.sharded_param import ShardedParamV2 -class ZeroContextConfig(object): +@dataclass +class ZeroContextConfig: """The configuration used to control zero context initialization. Args: target_device (torch.device): The device where param data are after exiting the context. - replicated (bool, optional): Whether the param is replicated across data parallel group. + is_replicated (bool, optional): Whether the param is replicated across data parallel group. Some parameters are not replicated, e.g. parameters in MOE experts. shard_param (bool, optional): Is param sharded after exiting the context. Defaults to False. """ - def __init__(self, target_device: torch.device, replicated: bool = True, shard_param: bool = False): - super().__init__() + target_device: torch.device + is_replicated: bool = True + shard_param: bool = False - if shard_param: - assert replicated, "Non-replicated parameters can't be sharded." + def __post_init__(self): + if self.shard_param: + assert self.is_replicated, "Non-replicated parameters can't be sharded." - # replicated no-shard parameters should locate in cuda, since we will broadcast them soon - if replicated and not shard_param: - assert target_device.type == 'cuda', "Replicated no-shard paramters should locate in cuda." - - self.target_device = target_device - self.is_replicated: bool = replicated - self.shard_param: bool = shard_param + if self.is_replicated and not self.shard_param: + assert self.target_device.type == 'cuda', "Replicated no-shard parameters should be located in cuda." class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): """A context to initialize model. 1. Convert the model to fp16. - 2. The paramaters of the module are adapted to type ShardedParameter. + 2. The parameters of the module are adapted to type ShardedParameter. 3. Shard the param and grad according to flags. Args: @@ -56,6 +55,7 @@ class ZeroInitContext(InsertPostInitMethodToModuleSubClasses): seed (int, optional): Random seed for weight initialization shard_param (bool, optional): Is param sharded after exiting the context. Defaults to False. default_dtype (torch.dtype, optional): If it's not None, parameters will be initialized as ``default_dtype`` then converted to fp16. + bf16 (bool, optional): If it's True, parameters will be initialized as ``torch.bfloat16``. Otherwise, parameters will be initialized as ``torch.float16``. Defaults to False. model_numel_tensor (torch.Tensor, optional): A tensor which will store the number of elements of model. Defaults to torch.zeros(1, dtype=torch.int). """ @@ -65,6 +65,7 @@ def __init__(self, seed: int = 2**10 - 1, shard_param: bool = False, default_dtype: Optional[torch.dtype] = None, + bf16: bool = False, model_numel_tensor: torch.Tensor = torch.zeros(1, dtype=torch.long)): super().__init__(default_dtype=default_dtype) @@ -72,9 +73,10 @@ def __init__(self, self.param_list = [] self.model_numel_tensor = model_numel_tensor self.seed = seed + self.bf16 = bf16 self.dp_process_group = gpc.get_group(ParallelMode.DATA) - self.config = ZeroContextConfig(target_device=target_device, replicated=True, shard_param=shard_param) + self.config = ZeroContextConfig(target_device=target_device, is_replicated=True, shard_param=shard_param) ZeroContextMgr().current_context = self @@ -98,7 +100,7 @@ def calc_fanin_fanout(tensor: torch.Tensor): """We use this function to substitute fan-in and fan-out calculation in torch.nn.init. This can help us get correct fan-in and fan-out for sharded tensor. """ - assert isinstance(tensor, nn.Parameter), "Sharded tensor initilization is only allowed for paramters" + assert isinstance(tensor, nn.Parameter), "Sharded tensor initialization is only allowed for parameters" # get correct shape of input tensor if not hasattr(tensor, 'colo_attr') or not tensor.colo_attr.param_is_sharded: @@ -124,7 +126,7 @@ def calc_fanin_fanout(tensor: torch.Tensor): return fan_in, fan_out def _pre_context_exec(self): - """ + """ The Callback function when entering the context """ self.logger = get_dist_logger("ZeroInitContext") @@ -184,9 +186,10 @@ def _post_init_method(self, module: torch.nn.Module, *args, **kwargs): NOTE() The module may be passed to this function multiple times. """ self.top_module = module + half_dtype = torch.float16 if not self.bf16 else torch.bfloat16 def half_fn(t: torch.Tensor): - return t.half() if t.is_floating_point() else t + return t.to(half_dtype) if t.is_floating_point() else t for param in module.parameters(recurse=False): # avoid adapting a param to ShardedParam twice @@ -227,9 +230,10 @@ def half_fn(t: torch.Tensor): # We must cast buffers # If we use BN, buffers may be on CPU and Float # We must cast them + cast_fn = cast_tensor_to_fp16 if not self.bf16 else cast_tensor_to_bf16 for buffer in module.buffers(recurse=False): buffer.data = buffer.data.to(device=torch.cuda.current_device()) - buffer.data = cast_tensor_to_fp16(buffer.data) + buffer.data = cast_fn(buffer.data) class ZeroContextMgr(metaclass=SingletonMeta): @@ -248,7 +252,7 @@ def hijack_context_config(self, **kwargs): def no_shard_zero_context(is_replicated: bool = True) -> AbstractContextManager: return ZeroContextMgr().hijack_context_config(target_device=torch.device('cuda', torch.cuda.current_device()), - replicated=is_replicated, + is_replicated=is_replicated, shard_param=False) diff --git a/colossalai/zero/shard_utils/__init__.py b/colossalai/zero/legacy/shard_utils/__init__.py similarity index 100% rename from colossalai/zero/shard_utils/__init__.py rename to colossalai/zero/legacy/shard_utils/__init__.py diff --git a/colossalai/zero/shard_utils/base_shard_strategy.py b/colossalai/zero/legacy/shard_utils/base_shard_strategy.py similarity index 87% rename from colossalai/zero/shard_utils/base_shard_strategy.py rename to colossalai/zero/legacy/shard_utils/base_shard_strategy.py index 7c2f4c9f6659..7ca951091640 100644 --- a/colossalai/zero/shard_utils/base_shard_strategy.py +++ b/colossalai/zero/legacy/shard_utils/base_shard_strategy.py @@ -2,7 +2,8 @@ from typing import List, Optional import torch.distributed as dist -from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor + +from colossalai.zero.legacy.sharded_param.sharded_tensor import ShardedTensor class BaseShardStrategy(ABC): diff --git a/colossalai/zero/shard_utils/bucket_tensor_shard_strategy.py b/colossalai/zero/legacy/shard_utils/bucket_tensor_shard_strategy.py similarity index 87% rename from colossalai/zero/shard_utils/bucket_tensor_shard_strategy.py rename to colossalai/zero/legacy/shard_utils/bucket_tensor_shard_strategy.py index a7bd7cf538e7..d663104831ce 100644 --- a/colossalai/zero/shard_utils/bucket_tensor_shard_strategy.py +++ b/colossalai/zero/legacy/shard_utils/bucket_tensor_shard_strategy.py @@ -2,18 +2,19 @@ import torch import torch.distributed as dist -from colossalai.utils import get_current_device -from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor from torch._utils import _flatten_dense_tensors as flatten +from colossalai.utils import get_current_device +from colossalai.zero.legacy.sharded_param.sharded_tensor import ShardedTensor + from .tensor_shard_strategy import TensorShardStrategy class BucketTensorShardStrategy(TensorShardStrategy): - """Use the same shard scheme as `TensorShardStrategy`'s, but it gathers tensors of a sub-module together, - which will fully utilize network bandwidth. - It is especially useful when sub-module contains bias, - since we cannot utilize network bandwidth well if we only gather a bias tensor (bias is usaully small). + """Use the same shard scheme as `TensorShardStrategy`'s, but it gathers tensors of a sub-module together, + which will fully utilize network bandwidth. + It is especially useful when sub-module contains bias, + since we cannot utilize network bandwidth well if we only gather a bias tensor (bias is usually small). """ def gather(self, tensor_list: List[ShardedTensor], process_group: Optional[dist.ProcessGroup] = None): diff --git a/colossalai/zero/shard_utils/commons.py b/colossalai/zero/legacy/shard_utils/commons.py similarity index 95% rename from colossalai/zero/shard_utils/commons.py rename to colossalai/zero/legacy/shard_utils/commons.py index 71cef44c177f..bf5ae325caf4 100644 --- a/colossalai/zero/shard_utils/commons.py +++ b/colossalai/zero/legacy/shard_utils/commons.py @@ -1,7 +1,7 @@ -import torch -import torch.nn.functional as F from typing import Tuple +import torch + def get_shard(tensor: torch.Tensor, rank: int, world_size: int) -> Tuple[torch.Tensor, int]: """Return the local shard of a full tensor.""" diff --git a/colossalai/zero/shard_utils/tensor_shard_strategy.py b/colossalai/zero/legacy/shard_utils/tensor_shard_strategy.py similarity index 86% rename from colossalai/zero/shard_utils/tensor_shard_strategy.py rename to colossalai/zero/legacy/shard_utils/tensor_shard_strategy.py index 5bdd95400d82..d1df4803b820 100644 --- a/colossalai/zero/shard_utils/tensor_shard_strategy.py +++ b/colossalai/zero/legacy/shard_utils/tensor_shard_strategy.py @@ -2,11 +2,12 @@ import torch import torch.distributed as dist + from colossalai.utils import get_current_device -from colossalai.zero.shard_utils import BaseShardStrategy -from colossalai.zero.shard_utils.commons import get_shard -from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor -from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline +from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move_inline +from colossalai.zero.legacy.shard_utils import BaseShardStrategy +from colossalai.zero.legacy.shard_utils.commons import get_shard +from colossalai.zero.legacy.sharded_param.sharded_tensor import ShardedTensor class TensorShardStrategy(BaseShardStrategy): @@ -27,7 +28,7 @@ def _shard_tensor(self, t: ShardedTensor, process_group: Optional[dist.ProcessGr Args: t (ShardedTensor): a tensor to be sharded. - process_group (Optional[dist.ProcessGroup], optional): the process group among which tensor shards. + process_group (Optional[dist.ProcessGroup], optional): the process group among which tensor shards. Defaults to None. """ if t.is_sharded: diff --git a/colossalai/zero/sharded_model/__init__.py b/colossalai/zero/legacy/sharded_model/__init__.py similarity index 61% rename from colossalai/zero/sharded_model/__init__.py rename to colossalai/zero/legacy/sharded_model/__init__.py index 725179295c60..93120bdc34b4 100644 --- a/colossalai/zero/sharded_model/__init__.py +++ b/colossalai/zero/legacy/sharded_model/__init__.py @@ -1,3 +1,3 @@ from .sharded_model_v2 import ShardedModelV2 -__all__ = ['ShardedModelV2'] \ No newline at end of file +__all__ = ['ShardedModelV2'] diff --git a/colossalai/zero/sharded_model/_utils.py b/colossalai/zero/legacy/sharded_model/_utils.py similarity index 86% rename from colossalai/zero/sharded_model/_utils.py rename to colossalai/zero/legacy/sharded_model/_utils.py index 85a3ab73dd1b..f1d642cf3f13 100644 --- a/colossalai/zero/sharded_model/_utils.py +++ b/colossalai/zero/legacy/sharded_model/_utils.py @@ -1,9 +1,9 @@ -from typing import Any, Callable, List, Tuple +from typing import Any, Callable, List, Tuple, Union import torch import torch.nn.functional as F -from typing import Union -from colossalai.gemini.stateful_tensor import StatefulTensor + +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor def get_gradient_predivide_factor(world_size: int) -> float: @@ -43,11 +43,19 @@ def cast_tensor_to_fp32(tensor: Union[torch.Tensor, StatefulTensor]) -> torch.Te if isinstance(tensor, StatefulTensor): tensor = tensor.payload - if torch.is_floating_point(tensor) and tensor.dtype is torch.float16: + if torch.is_floating_point(tensor) and tensor.dtype in (torch.float16, torch.bfloat16): return tensor.float() return tensor +def cast_tensor_to_bf16(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, StatefulTensor): + tensor = tensor.payload + if torch.is_floating_point(tensor) and tensor.dtype is torch.float32: + return tensor.bfloat16() + return tensor + + def apply_to_tensors(x: Any, fn: Callable): if torch.is_tensor(x): return fn(x) diff --git a/colossalai/zero/sharded_model/reduce_scatter.py b/colossalai/zero/legacy/sharded_model/reduce_scatter.py similarity index 100% rename from colossalai/zero/sharded_model/reduce_scatter.py rename to colossalai/zero/legacy/sharded_model/reduce_scatter.py diff --git a/colossalai/zero/sharded_model/sharded_model_v2.py b/colossalai/zero/legacy/sharded_model/sharded_model_v2.py similarity index 94% rename from colossalai/zero/sharded_model/sharded_model_v2.py rename to colossalai/zero/legacy/sharded_model/sharded_model_v2.py index 094f7d76a86d..e7064277fb3c 100644 --- a/colossalai/zero/sharded_model/sharded_model_v2.py +++ b/colossalai/zero/legacy/sharded_model/sharded_model_v2.py @@ -13,28 +13,29 @@ from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc -from colossalai.gemini.memory_tracer import MemStatsCollector, StaticMemStatsCollector -from colossalai.gemini.ophooks import register_ophooks_recursively -from colossalai.gemini.paramhooks import BaseParamHookMgr -from colossalai.gemini.stateful_tensor import TensorState -from colossalai.gemini.stateful_tensor_mgr import StatefulTensorMgr -from colossalai.gemini.tensor_placement_policy import TensorPlacementPolicy, TensorPlacementPolicyFactory -from colossalai.gemini.tensor_utils import colo_model_data_move_to_cpu from colossalai.logging import get_dist_logger from colossalai.utils import disposable, get_current_device from colossalai.utils.memory import colo_device_memory_capacity -from colossalai.zero.shard_utils import BaseShardStrategy -from colossalai.zero.sharded_model.reduce_scatter import ReduceScatterBucketer -from colossalai.zero.utils import ZeroHook +from colossalai.zero.gemini.memory_tracer import MemStatsCollector, StaticMemStatsCollector +from colossalai.zero.legacy.gemini.ophooks import register_ophooks_recursively +from colossalai.zero.legacy.gemini.paramhooks import BaseParamHookMgr +from colossalai.zero.legacy.gemini.stateful_tensor import TensorState +from colossalai.zero.legacy.gemini.stateful_tensor_mgr import StatefulTensorMgr +from colossalai.zero.legacy.gemini.tensor_placement_policy import TensorPlacementPolicy, TensorPlacementPolicyFactory +from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_move_to_cpu +from colossalai.zero.legacy.shard_utils import BaseShardStrategy +from colossalai.zero.legacy.sharded_model.reduce_scatter import ReduceScatterBucketer from ._utils import ( cast_float_arguments, + cast_tensor_to_bf16, cast_tensor_to_fp16, cast_tensor_to_fp32, chunk_and_pad, free_storage, get_gradient_predivide_factor, ) +from .zero_hook import ZeroHook try: from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX @@ -68,12 +69,13 @@ class ShardedModelV2(nn.Module): If it's 'auto', they are moving dynamically based on CPU and CUDA memory usage. It will utilize heterogeneous memory space evenly and well. Note that 'auto' policy can only work well when no other processes use CUDA during your training. Defaults to 'cuda'. - gradient_predivide_factor (Optional[float], optional): Gradient is divived by this value before reduce-scatter. Defaults to 1.0. + gradient_predivide_factor (Optional[float], optional): Gradient is divided by this value before reduce-scatter. Defaults to 1.0. reuse_fp16_shard (bool, optional): Whether to reuse fp16 shard for param and grad. Enabling this can reduce GPU memory usage, but you have to make sure you disable it when using gradient accumulation. In this mode, grad will be fp16. Make sure your optimizer supports mixed precision (fp32 param and fp16 grad). We find that PyTorch's optimizers don't support mixed precision, so we recommend you enable this only when using our CPUAdam with CPU offload. Defaults to False. + bf16 (bool, optional): Whether to use bfloat16 for param and grad. Defaults to False. """ def __init__(self, @@ -86,11 +88,13 @@ def __init__(self, tensor_placement_policy: str = 'cuda', gradient_predivide_factor: Optional[float] = 1.0, reuse_fp16_shard: bool = False, + bf16: bool = False, *args, **kwargs): assert not isinstance(module, ShardedModelV2), 'Nested ShardedModelV2 is not supported.' super().__init__() self.logger = get_dist_logger() + self.bf16 = bf16 # We force users to use ZeroInitContext for submodule in module.modules(): @@ -192,7 +196,7 @@ def cpu_offload(self): def dump_memory_stats(self, filename: Optional[str] = 'dump_mem_stats.log') -> None: """ - dummy memory tracer collected infomation to a file. + dummy memory tracer collected information to a file. try: # forward: model(inputs) # backward: optimizer.backward() @@ -201,7 +205,7 @@ def dump_memory_stats(self, filename: Optional[str] = 'dump_mem_stats.log') -> N exit(0) """ if self._use_memory_tracer: - self.logger.error(f'dump memort tracer collected infomation to a {filename}', ranks=[0]) + self.logger.error(f'dump memory tracer collected information to a {filename}', ranks=[0]) if gpc.get_global_rank() == 0: with open(filename, 'w+') as f: f.write(f'cuda reserved {torch.cuda.memory_reserved(get_current_device()) / 1e9} GB\n') @@ -232,7 +236,8 @@ def _post_forward_operations(self): def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: self._pre_forward_operations(*args) - args, kwargs = cast_float_arguments(cast_tensor_to_fp16, *args, **kwargs) + cast_fn = cast_tensor_to_bf16 if self.bf16 else cast_tensor_to_fp16 + args, kwargs = cast_float_arguments(cast_fn, *args, **kwargs) outputs = self.module(*args, **kwargs) self._post_forward_operations() return outputs @@ -293,7 +298,7 @@ def _post_backward_operations(self) -> None: if not p.requires_grad: continue # Leave the gradient accumulation state (_require_backward_grad_sync) as-is if not synchronizing this pass. - # NOTE() (no-sync)/sync pass: (not conduct)/conduct gradient allreducing between process group. + # NOTE() (no-sync)/sync pass: (not conduct)/conduct gradient all reducing between process group. # If _require_backward_grad_sync is True, # p.grad remains the accumulated unsharded gradient from prior no-sync passes. # We also allows to interleave no-sync pass with sync passes, if desired. @@ -380,12 +385,12 @@ def _save_grad(self, param: Parameter, grad: torch.Tensor): # make parameters point to gradient assert param.colo_attr.saved_grad.is_null( - ), 'Gradien accumulation is not supported when reuse_fp16_shard=True' + ), 'Gradient accumulation is not supported when reuse_fp16_shard=True' param.colo_attr.grad_payload_reset(grad.data) # release the memory of param # we set a false None for parameter's payload - # so we can get paramter's device and dtype later in optimizer + # so we can get parameter's device and dtype later in optimizer param.colo_attr.data_payload_reset(torch.empty(0, device=grad.device, dtype=grad.dtype)) if param.colo_attr.is_replicated: @@ -494,6 +499,7 @@ def _colo_load_from_state_dict(self, error_msgs (list of str): error messages should be added to this list, and will be reported together in :meth:`~torch.nn.Module.load_state_dict` + shard_strategy (Optional[BaseShardStrategy], optional): A shard strategy to manage shard behavior. Defaults to None. """ for hook in self._load_state_dict_pre_hooks.values(): hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) diff --git a/colossalai/zero/sharded_model/utils.py b/colossalai/zero/legacy/sharded_model/utils.py similarity index 91% rename from colossalai/zero/sharded_model/utils.py rename to colossalai/zero/legacy/sharded_model/utils.py index 69f5a23ac920..08806e78ea3b 100644 --- a/colossalai/zero/sharded_model/utils.py +++ b/colossalai/zero/legacy/sharded_model/utils.py @@ -1,7 +1,8 @@ +import copy + import torch -from colossalai.zero.sharded_model import ShardedModelV2 -import copy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 def col_model_deepcopy(sharded_model: ShardedModelV2, other_model: torch.nn.Module): diff --git a/colossalai/zero/utils/zero_hook.py b/colossalai/zero/legacy/sharded_model/zero_hook.py similarity index 92% rename from colossalai/zero/utils/zero_hook.py rename to colossalai/zero/legacy/sharded_model/zero_hook.py index 87bf2c0f5086..50f4bdfc775d 100644 --- a/colossalai/zero/utils/zero_hook.py +++ b/colossalai/zero/legacy/sharded_model/zero_hook.py @@ -3,14 +3,14 @@ import torch import torch.distributed as dist -from colossalai.gemini.memory_tracer import MemStatsCollector -from colossalai.gemini.ophooks import BaseOpHook -from colossalai.gemini.stateful_tensor import TensorState -from colossalai.gemini.stateful_tensor_mgr import StatefulTensorMgr from colossalai.logging import get_dist_logger from colossalai.registry import OPHOOKS from colossalai.utils import get_current_device -from colossalai.zero.shard_utils import BaseShardStrategy +from colossalai.zero.gemini.memory_tracer import MemStatsCollector +from colossalai.zero.legacy.gemini.ophooks import BaseOpHook +from colossalai.zero.legacy.gemini.stateful_tensor import TensorState +from colossalai.zero.legacy.gemini.stateful_tensor_mgr import StatefulTensorMgr +from colossalai.zero.legacy.shard_utils import BaseShardStrategy @OPHOOKS.register_module diff --git a/colossalai/zero/legacy/sharded_optim/__init__.py b/colossalai/zero/legacy/sharded_optim/__init__.py new file mode 100644 index 000000000000..b71a70aeffa4 --- /dev/null +++ b/colossalai/zero/legacy/sharded_optim/__init__.py @@ -0,0 +1,3 @@ +from .sharded_optim_v2 import ShardedOptimizerV2 + +__all__ = ['ShardedOptimizerV2'] diff --git a/colossalai/zero/sharded_optim/sharded_optim_v2.py b/colossalai/zero/legacy/sharded_optim/sharded_optim_v2.py similarity index 90% rename from colossalai/zero/sharded_optim/sharded_optim_v2.py rename to colossalai/zero/legacy/sharded_optim/sharded_optim_v2.py index 43a0b7d76107..41dd174cb65a 100644 --- a/colossalai/zero/sharded_optim/sharded_optim_v2.py +++ b/colossalai/zero/legacy/sharded_optim/sharded_optim_v2.py @@ -14,13 +14,13 @@ from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc -from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState -from colossalai.gemini.tensor_placement_policy import AutoTensorPlacementPolicy -from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage from colossalai.logging import get_dist_logger from colossalai.nn.optimizer import ColossalaiOptimizer -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model._utils import cast_tensor_to_fp32 +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState +from colossalai.zero.legacy.gemini.tensor_placement_policy import AutoTensorPlacementPolicy +from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp32 class OptimState(Enum): @@ -67,8 +67,8 @@ class ShardedOptimizerV2(ColossalaiOptimizer): growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000. hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2. max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32. - dp_process_group (Optional[ProcessGroup], optional): data paralle process group. Defaults to None. - mp_process_group (Optional[ProcessGroup], optional): model paralle process group. Defaults to None. + dp_process_group (Optional[ProcessGroup], optional): data parallel process group. Defaults to None. + mp_process_group (Optional[ProcessGroup], optional): model parallel process group. Defaults to None. .. _PatrickStar\: Parallel Training of Pre-trained Models via Chunk-based Memory Management: https://arxiv.org/abs/2108.05818 @@ -94,6 +94,7 @@ def __init__(self, super().__init__(optimizer) self.shard_strategy = sharded_model.shard_strategy self.model: ShardedModelV2 = sharded_model + self.bf16 = sharded_model.bf16 self.gpu_margin_mem_ratio: float = float(gpu_margin_mem_ratio) assert 0.0 <= self.gpu_margin_mem_ratio <= 1.0, f'gpu_margin_mem_ratio must >=0.0 and <=1.0' @@ -117,6 +118,7 @@ def __init__(self, self._found_overflow: Tensor = torch.IntTensor([0]).to(torch.cuda.current_device()) self._logger = get_dist_logger("ShardedOptimizerV2") self._verbose = verbose + self._grad_prepared: bool = False # this should be set to true when _prepare_grads() and reset to false when backward # Store fp32 param shards self._register_master_weight() @@ -166,8 +168,10 @@ def zero_grad(self, *args, **kwargs): self._zero_grad() def backward(self, loss: Tensor) -> None: - loss = self.loss_scale * loss - self.optim_state = OptimState.SCALED + if not self.bf16: + loss = self.loss_scale * loss + self.optim_state = OptimState.SCALED + self._grad_prepared = False self.model.backward(loss) def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None: @@ -175,30 +179,33 @@ def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None: # It receives the scaled grad from the previous rank # No need to scale the grad again # Need to unscale when optimizing - self.optim_state = OptimState.SCALED + if not self.bf16: + self.optim_state = OptimState.SCALED + self._grad_prepared = False self.model.backward_by_grad(tensor, grad) def clip_grad_norm(self, model: nn.Module, max_norm: float): - if self.optim_state == OptimState.SCALED: - self._prepare_grads() + self._prepare_grads() + if not self.bf16 and self.optim_state == OptimState.SCALED: self._unscale_grads() return super().clip_grad_norm(model, max_norm) def step(self, *args, **kwargs): + self._prepare_grads() # unscale grads if scaled - if self.optim_state == OptimState.SCALED: - self._prepare_grads() + if not self.bf16 and self.optim_state == OptimState.SCALED: self._unscale_grads() self._maybe_move_fp32_shards() - found_inf = self._check_overflow() - self.grad_scaler.update(found_inf) + if not self.bf16: + found_inf = self._check_overflow() + self.grad_scaler.update(found_inf) - if found_inf: - self._logger.warning('found inf during ShardedOptimV2 step') - self._zero_grad(recover_data=True) - return + if found_inf: + self._logger.warning('found inf during ShardedOptimV2 step') + self._zero_grad(recover_data=True) + return self._point_param_fp16_to_master_param() @@ -274,7 +281,7 @@ def _register_master_weight(self): assert hasattr(p, 'colo_attr'), 'The parameter must be wrapped with ShardedParam' shard_flag = not p.colo_attr.sharded_data_tensor.is_sharded and p.colo_attr.is_replicated if shard_flag: - # we always shard replicated paramters + # we always shard replicated parameters self.shard_strategy.shard([p.colo_attr.sharded_data_tensor], self.dp_process_group) self.master_params[p] = StatefulTensor(cast_tensor_to_fp32(p.colo_attr.data_payload.to(self.device))) if shard_flag: @@ -304,6 +311,8 @@ def _maybe_move_fp32_shards(self): state[k] = v.cuda() def _prepare_grads(self): + if self._grad_prepared: + return for group in self.optim.param_groups: for p in group['params']: if p.colo_attr.saved_grad.is_null(): @@ -312,7 +321,7 @@ def _prepare_grads(self): # If reuse_fp16_shard, grad fp16 which wasn't be offloaded may be evicted to CPU if not p.colo_attr.offload_grad: colo_model_data_tensor_move_inline(p.colo_attr.saved_grad, torch.cuda.current_device()) - # FIXME(ver217): p.data here is an empty tensor on CUDA and has no useful infomation + # FIXME(ver217): p.data here is an empty tensor on CUDA and has no useful information # If we change p.grad directly # it may raise error because of different shape/dtype/device of p.data and p.grad # We just set p.data = p.colo_attr.saved_grad.payload here @@ -320,6 +329,7 @@ def _prepare_grads(self): p.grad = p.colo_attr.grad_payload # Set p.data to empty tensor, in case of memory leaking p.colo_attr.set_data_none() + self._grad_prepared = True def _point_param_fp16_to_master_param(self): # assign master param pointers to p.data. @@ -333,7 +343,7 @@ def _point_param_fp16_to_master_param(self): def _copy_master_model_to_model_fp16(self): # Copy master param data (fp32) to payload of colo_attr (fp16) - # TODO() improve efficiency by gathering tensors into a chunk and transfering + # TODO() improve efficiency by gathering tensors into a chunk and transferring # a chunk. for group in self.optim.param_groups: for p in group['params']: @@ -350,14 +360,15 @@ def _copy_master_param_to_param_fp16(self, p): p.data = self.master_params[p].payload - # we need to allocate new memory for keep_not_shard paramters + # we need to allocate new memory for keep_not_shard parameters # in order to use copy, otherwise, the sizes of tensor is not compatible if p.colo_attr.data_payload.numel() != p.data.numel(): p.colo_attr.data_payload_reset( torch.empty(p.data.shape, dtype=p.colo_attr.data_payload.dtype, device=p.colo_attr.data_payload.device)) # TODO() optimize this line CPU (fp32) -> GPU (fp16) - p.colo_attr.sharded_data_tensor.payload_copy(p.half().detach()) + half_dtype = torch.bfloat16 if self.bf16 else torch.float16 + p.colo_attr.sharded_data_tensor.payload_copy(p.to(half_dtype).detach()) p.colo_attr.set_data_none() if p.colo_attr.keep_not_shard and p.colo_attr.is_replicated: diff --git a/colossalai/zero/legacy/sharded_param/__init__.py b/colossalai/zero/legacy/sharded_param/__init__.py new file mode 100644 index 000000000000..47e2ce2fa0e0 --- /dev/null +++ b/colossalai/zero/legacy/sharded_param/__init__.py @@ -0,0 +1,4 @@ +from .sharded_param import ShardedParamV2 +from .sharded_tensor import ShardedTensor + +__all__ = ['ShardedTensor', 'ShardedParamV2'] diff --git a/colossalai/zero/sharded_param/sharded_param.py b/colossalai/zero/legacy/sharded_param/sharded_param.py similarity index 93% rename from colossalai/zero/sharded_param/sharded_param.py rename to colossalai/zero/legacy/sharded_param/sharded_param.py index db0f2d149431..4bcc4b62104a 100644 --- a/colossalai/zero/sharded_param/sharded_param.py +++ b/colossalai/zero/legacy/sharded_param/sharded_param.py @@ -1,9 +1,11 @@ +from typing import List, Optional, Tuple + import torch -from typing import Optional, Tuple -from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor -from colossalai.gemini.tensor_utils import colo_tensor_mem_usage -from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState -from typing import List + +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState +from colossalai.zero.legacy.gemini.tensor_utils import colo_tensor_mem_usage + +from .sharded_tensor import ShardedTensor EMPTY_TENSOR_DICT = {} diff --git a/colossalai/zero/sharded_param/sharded_tensor.py b/colossalai/zero/legacy/sharded_param/sharded_tensor.py similarity index 92% rename from colossalai/zero/sharded_param/sharded_tensor.py rename to colossalai/zero/legacy/sharded_param/sharded_tensor.py index 77f4aec30f32..af60312600f2 100644 --- a/colossalai/zero/sharded_param/sharded_tensor.py +++ b/colossalai/zero/legacy/sharded_param/sharded_tensor.py @@ -1,5 +1,6 @@ import torch -from colossalai.gemini.stateful_tensor import StatefulTensor, TensorState + +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState class ShardedTensor(StatefulTensor): diff --git a/colossalai/zero/low_level/__init__.py b/colossalai/zero/low_level/__init__.py new file mode 100644 index 000000000000..ae3c1de3a5bc --- /dev/null +++ b/colossalai/zero/low_level/__init__.py @@ -0,0 +1,3 @@ +from .low_level_optim import LowLevelZeroOptimizer + +__all__ = ['LowLevelZeroOptimizer'] diff --git a/colossalai/zero/sharded_optim/_utils.py b/colossalai/zero/low_level/_utils.py similarity index 95% rename from colossalai/zero/sharded_optim/_utils.py rename to colossalai/zero/low_level/_utils.py index 9ca2fdf5aa06..218f7603bc54 100644 --- a/colossalai/zero/sharded_optim/_utils.py +++ b/colossalai/zero/low_level/_utils.py @@ -91,10 +91,18 @@ def get_grad_accumulate_object(tensor): return grad_acc_obj -def split_half_float_double(tensor_list): +def split_by_dtype(tensor_list): + """ + Splits a list of PyTorch tensors into sublists based on their data type. + + :param tensor_list: A list of PyTorch tensors. + :type tensor_list: list[torch.Tensor] + :return: A list of sublists, where each sublist contains tensors of a specific data type. + :rtype: list[list[torch.Tensor]] + """ dtypes = ["torch.cuda.HalfTensor", "torch.cuda.FloatTensor", "torch.cuda.DoubleTensor", "torch.cuda.BFloat16Tensor"] buckets = [] - for i, dtype in enumerate(dtypes): + for _, dtype in enumerate(dtypes): bucket = [t for t in tensor_list if t.type() == dtype] if bucket: buckets.append(bucket) @@ -253,7 +261,7 @@ def sync_param(flat_tensor, tensor_list): share the same memory space. This function will update the tensor list so that they point to the same value. - :param flat_tensor: A flat tensor obtained by calling `torch._utils._unflatten_dense_tensors` on a tensor lsit + :param flat_tensor: A flat tensor obtained by calling `torch._utils._unflatten_dense_tensors` on a tensor list :param tensor_list: A list of tensors corresponding to the flattened tensor :type flat_tensor: torch.Tensor :type tensor_list: List[torch.Tensor] diff --git a/colossalai/zero/sharded_optim/bookkeeping/__init__.py b/colossalai/zero/low_level/bookkeeping/__init__.py similarity index 100% rename from colossalai/zero/sharded_optim/bookkeeping/__init__.py rename to colossalai/zero/low_level/bookkeeping/__init__.py diff --git a/colossalai/zero/sharded_optim/bookkeeping/base_store.py b/colossalai/zero/low_level/bookkeeping/base_store.py similarity index 100% rename from colossalai/zero/sharded_optim/bookkeeping/base_store.py rename to colossalai/zero/low_level/bookkeeping/base_store.py diff --git a/colossalai/zero/sharded_optim/bookkeeping/bucket_store.py b/colossalai/zero/low_level/bookkeeping/bucket_store.py similarity index 100% rename from colossalai/zero/sharded_optim/bookkeeping/bucket_store.py rename to colossalai/zero/low_level/bookkeeping/bucket_store.py diff --git a/colossalai/zero/sharded_optim/bookkeeping/gradient_store.py b/colossalai/zero/low_level/bookkeeping/gradient_store.py similarity index 100% rename from colossalai/zero/sharded_optim/bookkeeping/gradient_store.py rename to colossalai/zero/low_level/bookkeeping/gradient_store.py diff --git a/colossalai/zero/sharded_optim/bookkeeping/parameter_store.py b/colossalai/zero/low_level/bookkeeping/parameter_store.py similarity index 61% rename from colossalai/zero/sharded_optim/bookkeeping/parameter_store.py rename to colossalai/zero/low_level/bookkeeping/parameter_store.py index cbf708b3471f..1f3ba7cbc3bc 100644 --- a/colossalai/zero/sharded_optim/bookkeeping/parameter_store.py +++ b/colossalai/zero/low_level/bookkeeping/parameter_store.py @@ -11,9 +11,9 @@ class ParameterStore(BaseStore): def __init__(self, torch_pg: ProcessGroup): super().__init__(torch_pg) # param partitioning data structures - self._fp16_param_to_rank = dict() - self._rank_groupid_to_fp16_param_list = dict() - self._rank_group_id_to_flat_fp16_param = dict() + self._param_to_rank = dict() + self._rank_group_id_to_param_list = dict() + self._rank_group_id_to_flat_param = dict() # param reduction data structures self._is_param_reduced = dict() @@ -29,7 +29,7 @@ def set_param_to_rank(self, tensor: Tensor, rank: int) -> None: :type rank: int """ - self._fp16_param_to_rank[tensor] = rank + self._param_to_rank[tensor] = rank def get_param_rank(self, tensor: Tensor) -> int: """ @@ -38,7 +38,7 @@ def get_param_rank(self, tensor: Tensor) -> int: :param tensor: A :class:`torch.Tensor` object :type tensor: torch.Tensor """ - return self._fp16_param_to_rank[tensor] + return self._param_to_rank[tensor] def belongs_to_current_rank(self, tensor) -> bool: """ @@ -51,29 +51,29 @@ def belongs_to_current_rank(self, tensor) -> bool: :rtype: bool """ - tensor_rank = self._fp16_param_to_rank[tensor] + tensor_rank = self._param_to_rank[tensor] return tensor_rank == self._local_rank - def add_fp16_param_list_by_rank_group(self, rank, group_id, tensor_list) -> None: - if rank not in self._rank_groupid_to_fp16_param_list: - self._rank_groupid_to_fp16_param_list[rank] = dict() + def add_param_list_by_rank_group(self, rank, group_id, tensor_list) -> None: + if rank not in self._rank_group_id_to_param_list: + self._rank_group_id_to_param_list[rank] = dict() - if group_id not in self._rank_groupid_to_fp16_param_list[rank]: - self._rank_groupid_to_fp16_param_list[rank][group_id] = [] + if group_id not in self._rank_group_id_to_param_list[rank]: + self._rank_group_id_to_param_list[rank][group_id] = [] - self._rank_groupid_to_fp16_param_list[rank][group_id].extend(tensor_list) + self._rank_group_id_to_param_list[rank][group_id].extend(tensor_list) - def get_fp16_params_by_rank_group(self, rank, group_id) -> List[Tensor]: - return self._rank_groupid_to_fp16_param_list[rank][group_id] + def get_params_by_rank_group(self, rank, group_id) -> List[Tensor]: + return self._rank_group_id_to_param_list[rank][group_id] - def add_flat_fp16_param_by_rank_group(self, rank, group_id, tensor) -> None: - if rank not in self._rank_group_id_to_flat_fp16_param: - self._rank_group_id_to_flat_fp16_param[rank] = dict() + def add_flat_param_by_rank_group(self, rank, group_id, tensor) -> None: + if rank not in self._rank_group_id_to_flat_param: + self._rank_group_id_to_flat_param[rank] = dict() - self._rank_group_id_to_flat_fp16_param[rank][group_id] = tensor + self._rank_group_id_to_flat_param[rank][group_id] = tensor - def get_flat_fp16_param_by_rank_group(self, rank, group_id) -> Tensor: - return self._rank_group_id_to_flat_fp16_param[rank][group_id] + def get_flat_param_by_rank_group(self, rank, group_id) -> Tensor: + return self._rank_group_id_to_flat_param[rank][group_id] def is_param_reduced(self, tensor): return self._is_param_reduced[tensor] diff --git a/colossalai/zero/sharded_optim/bookkeeping/tensor_bucket.py b/colossalai/zero/low_level/bookkeeping/tensor_bucket.py similarity index 100% rename from colossalai/zero/sharded_optim/bookkeeping/tensor_bucket.py rename to colossalai/zero/low_level/bookkeeping/tensor_bucket.py diff --git a/colossalai/zero/sharded_optim/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py similarity index 73% rename from colossalai/zero/sharded_optim/low_level_optim.py rename to colossalai/zero/low_level/low_level_optim.py index 49fb8b54b7d2..ee03c0f0ae15 100644 --- a/colossalai/zero/sharded_optim/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -6,7 +6,11 @@ import torch.distributed as dist from torch.optim import Optimizer -from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler +from colossalai.amp.naive_amp.mixed_precision_mixin import ( + BF16MixedPrecisionMixin, + FP16MixedPrecisionMixin, + MixedPrecisionMixin, +) from colossalai.context import ParallelMode from colossalai.core import global_context as gpc from colossalai.logging import get_dist_logger @@ -21,12 +25,37 @@ has_inf_or_nan, reduce_tensor_dp_group, release_param_grad, - split_half_float_double, + split_by_dtype, sync_param, ) from .bookkeeping import BucketStore, GradientStore, ParameterStore, TensorBucket +class LowLevelZeroFP16MixedPrecisionMixin(FP16MixedPrecisionMixin): + + def __init__(self, + num_working_param_groups: int, + grad_store: GradientStore, + initial_scale: float = 2**16, + min_scale: float = 1, + growth_factor: float = 2, + backoff_factor: float = 0.5, + growth_interval: int = 1000, + hysteresis: int = 2, + max_scale: float = 2**32) -> None: + super().__init__(initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, + max_scale) + self.num_working_param_groups = num_working_param_groups + self.grad_store = grad_store + + def check_local_overflow(self) -> bool: + for group_id in range(self.num_working_param_groups): + for avg_grad in self.grad_store.get_averaged_gradients_by_group(group_id): + if avg_grad is not None and has_inf_or_nan(avg_grad): + return True + return False + + class LowLevelZeroOptimizer(ColossalaiOptimizer): """Optimizer used for ZeRO-1 and ZeRO-2. """ @@ -55,6 +84,7 @@ def __init__( # 2. contiguous gradients # 3. cpu offload # 4. support when some parameters requires_grad = False + # 5. support layer drop super(LowLevelZeroOptimizer, self).__init__(optim=optimizer) self._dtype = self.optim.param_groups[0]['params'][0].dtype self._logger = get_dist_logger() @@ -89,26 +119,16 @@ def __init__( self._mp_torch_group = gpc.get_group(mp_parallel_mode) else: raise NotImplementedError - # fp16 and fp32 params for mixed precision training - self._fp16_param_groups = dict() - self._fp32_flat_param_groups_of_current_rank = dict() + + # working and master params for mixed precision training + self._working_param_groups = dict() + self._master_flat_param_groups_of_current_rank = dict() # communication params self._overlap_communication = overlap_communication self._reduce_bucket_size = reduce_bucket_size self._communication_dtype = communication_dtype - # gradient scaler - self.grad_scaler = DynamicGradScaler(initial_scale=initial_scale, - min_scale=min_scale, - growth_factor=growth_factor, - backoff_factor=backoff_factor, - growth_interval=growth_interval, - hysteresis=hysteresis, - max_scale=max_scale, - verbose=verbose) - self._found_overflow = torch.FloatTensor([0]).to(get_current_device()) - # gradient clipping self._clip_grad_norm = clip_grad_norm @@ -137,8 +157,8 @@ def __init__( if param.requires_grad: group_params.append(param) - # add the fp16 params to fp16_param_groups for bookkeeping - self._fp16_param_groups[group_id] = group_params + # add the working params to working_param_groups for bookkeeping + self._working_param_groups[group_id] = group_params # assign parameters to ranks # the params in the list are sorted @@ -147,7 +167,7 @@ def __init__( # store the mapping between param to rank # each param should belong to only one rank for rank, params in enumerate(params_per_rank): - self._param_store.add_fp16_param_list_by_rank_group(rank, group_id, params) + self._param_store.add_param_list_by_rank_group(rank, group_id, params) for param in params: self._param_store.set_param_to_rank(param, rank) @@ -158,37 +178,37 @@ def __init__( # flatten the reordered tensors for rank in range(self._world_size): - tensor_list = self._param_store.get_fp16_params_by_rank_group(rank, group_id) + tensor_list = self._param_store.get_params_by_rank_group(rank, group_id) with torch.no_grad(): flat_tensor = flatten(tensor_list) flat_tensor = flat_tensor.data.cuda() - self._param_store.add_flat_fp16_param_by_rank_group(rank, group_id, flat_tensor) + self._param_store.add_flat_param_by_rank_group(rank, group_id, flat_tensor) # sync parameters for rank in range(self._world_size): - flat_tensor = self._param_store.get_flat_fp16_param_by_rank_group(rank, group_id) - tensor_list = self._param_store.get_fp16_params_by_rank_group(rank, group_id) + flat_tensor = self._param_store.get_flat_param_by_rank_group(rank, group_id) + tensor_list = self._param_store.get_params_by_rank_group(rank, group_id) sync_param(flat_tensor=flat_tensor, tensor_list=tensor_list) - # create a copy of fp32 weights of the parameters for which this rank is responsible - fp16_flat_current_rank = self._param_store.get_flat_fp16_param_by_rank_group(self._local_rank, group_id) - fp32_flat_current_rank = fp16_flat_current_rank.float() + # create a copy of fp32 master weights of the parameters for which this rank is responsible + working_flat_current_rank = self._param_store.get_flat_param_by_rank_group(self._local_rank, group_id) + master_flat_current_rank = working_flat_current_rank.float() device = 'cpu' if self._cpu_offload else get_current_device() - fp32_flat_current_rank = fp32_flat_current_rank.to(device) - fp32_flat_current_rank.requires_grad = True - self._fp32_flat_param_groups_of_current_rank[group_id] = fp32_flat_current_rank + master_flat_current_rank = master_flat_current_rank.to(device) + master_flat_current_rank.requires_grad = True + self._master_flat_param_groups_of_current_rank[group_id] = master_flat_current_rank # need to replace the params in the `params` field in the optimizer # so that when the optimizer calls step(), it only updates the tensors # managed by this data parallel rank - param_group['params'] = [fp32_flat_current_rank] + param_group['params'] = [master_flat_current_rank] # set reduction state - for param in self._fp16_param_groups[group_id]: + for param in self._working_param_groups[group_id]: self._param_store.set_param_reduction_state(param, False) - # intialize communication stream for - # communication-compuation overlapping + # initialize communication stream for + # communication-computation overlapping if self._overlap_communication: self._comm_stream = torch.cuda.Stream() @@ -198,17 +218,28 @@ def __init__( if self._overlap_communication or self._partition_grads: self._attach_reduction_hook() + # initialize mixed precision mixin + self.mixed_precision_mixin: Optional[MixedPrecisionMixin] = None + if self._dtype is torch.float16: + self.mixed_precision_mixin = LowLevelZeroFP16MixedPrecisionMixin(self.num_param_groups, + self._grad_store, + initial_scale=initial_scale, + min_scale=min_scale, + growth_factor=growth_factor, + backoff_factor=backoff_factor, + growth_interval=growth_interval, + hysteresis=hysteresis, + max_scale=max_scale) + elif self._dtype is torch.bfloat16: + self.mixed_precision_mixin = BF16MixedPrecisionMixin() + @property def dtype(self): return self._dtype - @property - def loss_scale(self): - return self.grad_scaler.scale - @property def num_param_groups(self): - return len(self._fp16_param_groups) + return len(self._working_param_groups) def _sanity_checks(self): assert torch.cuda.is_available(), 'CUDA is required' @@ -238,7 +269,7 @@ def _partition_param_list(self, param_list): params_per_rank = [[] for _ in range(self._world_size)] numel_per_rank = [0 for _ in range(self._world_size)] - # partititon the parameters in a greedy fashion + # partition the parameters in a greedy fashion sorted_params = sorted(param_list, key=lambda x: x.numel(), reverse=True) for param in sorted_params: # allocate this parameter to the rank with @@ -260,13 +291,13 @@ def _grad_handler(self, param, grad, reduce_rank): return grad def _attach_reduction_hook(self): - # we iterate over the fp16 params + # we iterate over the working params # on each param, we register a hook to its AccumulateGrad object for group_id in range(self.num_param_groups): - param_group = self._fp16_param_groups[group_id] + param_group = self._working_param_groups[group_id] for param in param_group: if param.requires_grad: - # determines the reduction destionation rank + # determines the reduction destination rank # this is only valid for stage 2 # dst_rank = None means using all-reduce # else using reduce @@ -314,7 +345,7 @@ def _reduce_tensor_list_with_one_dtype(self, tensor_list, bucket_size, reduce_ra self._reduce_tensor_bucket(bucket=param_bucket, reduce_rank=reduce_rank) def _reduce_grads(self, reduce_rank, grads, bucket_size): - grad_buckets_by_dtype = split_half_float_double(grads) + grad_buckets_by_dtype = split_by_dtype(grads) for tensor_list in grad_buckets_by_dtype: self._reduce_tensor_list_with_one_dtype(tensor_list=tensor_list, @@ -390,7 +421,8 @@ def _add_to_reduction_bucket(self, param, reduce_rank=None): ################################ def backward(self, loss, retain_graph=False, sync_grad=True): - loss = self.loss_scale * loss + if self.mixed_precision_mixin is not None: + loss = self.mixed_precision_mixin.pre_backward(loss) loss.backward(retain_graph=retain_graph) # finish gradient reduction @@ -417,7 +449,9 @@ def zero_grad(self, set_to_none=True): :param set_to_none: Whether set the gradient to None. Default value is True. :type set_to_none: bool """ - for _, param_group in self._fp16_param_groups.items(): + if self.mixed_precision_mixin is not None: + self.mixed_precision_mixin.pre_zero_grad() + for _, param_group in self._working_param_groups.items(): for param in param_group: if set_to_none: param.grad = None @@ -433,43 +467,40 @@ def zero_grad(self, set_to_none=True): def step(self, closure=None): assert closure is None, 'closure is not supported by step()' - # check for overflow - found_inf = self._check_overflow() - self.grad_scaler.update(found_inf) - - # update loss scale if overflow occurs - if found_inf: + if self.mixed_precision_mixin is not None and self.mixed_precision_mixin.should_skip_step(): self._grad_store.reset_all_average_gradients() + if self._verbose: + self._logger.info(f'Found overflow. Skip step') self.zero_grad() return - # copy the grad of fp16 param to fp32 param + # copy the grad of working param to master param single_grad_partition_groups = [] norm_groups = [] for group_id in range(self.num_param_groups): # compute norm norm_group = compute_norm(gradients=self._grad_store.get_averaged_gradients_by_group(group_id), - params=self._param_store.get_fp16_params_by_rank_group(group_id=group_id, - rank=self._local_rank), + params=self._param_store.get_params_by_rank_group(group_id=group_id, + rank=self._local_rank), dp_group=self._dp_torch_group, mp_group=self._mp_torch_group) norm_groups.append(norm_group) - # create flat gradient for the flat fp32 params - fp16_avg_grads = self._grad_store.get_averaged_gradients_by_group(group_id) - flat_fp16_avg_grads = flatten(fp16_avg_grads) + # create flat gradient for the flat fp32 master params + working_avg_grads = self._grad_store.get_averaged_gradients_by_group(group_id) + flat_working_avg_grads = flatten(working_avg_grads) - dtype = self._fp32_flat_param_groups_of_current_rank[group_id].dtype - flat_fp32_avg_grads = flat_fp16_avg_grads.to(dtype) + dtype = self._master_flat_param_groups_of_current_rank[group_id].dtype + flat_master_avg_grads = flat_working_avg_grads.to(dtype) - param_shape = self._fp32_flat_param_groups_of_current_rank[group_id].shape - assert param_shape == flat_fp32_avg_grads.shape, \ - f'fp32 param and grad have different shape {param_shape} vs {flat_fp32_avg_grads.shape}' + param_shape = self._master_flat_param_groups_of_current_rank[group_id].shape + assert param_shape == flat_master_avg_grads.shape, \ + f'fp32 param and grad have different shape {param_shape} vs {flat_master_avg_grads.shape}' - single_grad_partition_groups.append(flat_fp32_avg_grads) - device = self._fp32_flat_param_groups_of_current_rank[group_id].device - self._fp32_flat_param_groups_of_current_rank[group_id].grad = flat_fp32_avg_grads.to(device) + single_grad_partition_groups.append(flat_master_avg_grads) + device = self._master_flat_param_groups_of_current_rank[group_id].device + self._master_flat_param_groups_of_current_rank[group_id].grad = flat_master_avg_grads.to(device) self._grad_store.reset_average_gradients_by_group(group_id) # unscale and clip grads @@ -478,66 +509,45 @@ def step(self, closure=None): # update the parameters self.optim.step() - # release the fp32 grad - release_param_grad(self._fp32_flat_param_groups_of_current_rank.values()) + # release the master grad + release_param_grad(self._master_flat_param_groups_of_current_rank.values()) - # update fp16 partition updated by the current rank - for group_id in range(len(self._fp16_param_groups)): - fp16_param = self._param_store.get_flat_fp16_param_by_rank_group(rank=self._local_rank, group_id=group_id) - fp32_param = self._fp32_flat_param_groups_of_current_rank[group_id] - fp16_param.data.copy_(fp32_param) + # update working partition updated by the current rank + for group_id in range(len(self._working_param_groups)): + working_param = self._param_store.get_flat_param_by_rank_group(rank=self._local_rank, group_id=group_id) + master_param = self._master_flat_param_groups_of_current_rank[group_id] + working_param.data.copy_(master_param) # broadcast the updated model weights handles = [] for group_id in range(self.num_param_groups): for index in range(self._world_size): rank = self._dp_global_ranks[index] - fp16_param = self._param_store.get_flat_fp16_param_by_rank_group(rank=index, group_id=group_id) - handle = dist.broadcast(fp16_param, src=rank, group=self._dp_torch_group, async_op=True) + working_param = self._param_store.get_flat_param_by_rank_group(rank=index, group_id=group_id) + handle = dist.broadcast(working_param, src=rank, group=self._dp_torch_group, async_op=True) handles.append(handle) for handle in handles: handle.wait() - ################## - # FP16 Utilities # - ################## - - def _check_overflow(self): - # clear previous overflow record - self._found_overflow.fill_(0.0) - - # check for overflow - for group_id in range(len(self._fp16_param_groups)): - for avg_grad in self._grad_store.get_averaged_gradients_by_group(group_id): - if avg_grad is not None and has_inf_or_nan(avg_grad): - self._found_overflow.fill_(1.0) - break - - # all-reduce across dp group - dist.all_reduce(self._found_overflow, op=dist.ReduceOp.MAX, group=self._dp_torch_group) - - # all-reduce over model parallel group - if self._mp_torch_group: - dist.all_reduce(self._found_overflow, op=dist.ReduceOp.MAX, group=self._mp_torch_group) - - if self._found_overflow.item() > 0: - return True - else: - return False + ############################# + # Mixed Precision Utilities # + ############################# def _unscale_and_clip_grads(self, grad_groups_flat, total_norm): # compute combined scale factor for this group - combined_scale = self.loss_scale + div_scale = 1.0 + if self.mixed_precision_mixin is not None: + div_scale = self.mixed_precision_mixin.get_grad_div_scale() if self._clip_grad_norm > 0.: # norm is in fact norm*scale - clip = ((total_norm / self.loss_scale) + 1e-6) / self._clip_grad_norm + clip = ((total_norm / div_scale) + 1e-6) / self._clip_grad_norm if clip > 1: - combined_scale = clip * self.loss_scale + div_scale = clip * div_scale for grad in grad_groups_flat: - grad.data.mul_(1. / combined_scale) + grad.data.mul_(1. / div_scale) ############################ # Gradient Synchronization # @@ -551,7 +561,7 @@ def _sync_grad(self): # accumulate gradient for group_id in range(self.num_param_groups): - param_group = self._param_store.get_fp16_params_by_rank_group(self._local_rank, group_id) + param_group = self._param_store.get_params_by_rank_group(self._local_rank, group_id) avg_gradients_group = self._grad_store.get_averaged_gradients_by_group(group_id) @@ -572,8 +582,8 @@ def _reduce_grad_stage1(self): # if not overlapping communication (no reduction hook is attached) # we need to manually reduce these gradients if not self._overlap_communication: - for group_id in range(len(self._fp16_param_groups)): - param_group = self._fp16_param_groups[group_id] + for group_id in range(len(self._working_param_groups)): + param_group = self._working_param_groups[group_id] for param in param_group: if param.grad is not None: self._add_to_reduction_bucket(param) diff --git a/colossalai/zero/sharded_optim/__init__.py b/colossalai/zero/sharded_optim/__init__.py deleted file mode 100644 index 30c26fb75f30..000000000000 --- a/colossalai/zero/sharded_optim/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .low_level_optim import LowLevelZeroOptimizer -from .sharded_optim_v2 import ShardedOptimizerV2 - -__all__ = ['ShardedOptimizerV2', 'LowLevelZeroOptimizer'] diff --git a/colossalai/zero/sharded_param/__init__.py b/colossalai/zero/sharded_param/__init__.py deleted file mode 100644 index 5642a504acf7..000000000000 --- a/colossalai/zero/sharded_param/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from colossalai.zero.sharded_param.sharded_tensor import ShardedTensor -from colossalai.zero.sharded_param.sharded_param import ShardedParamV2 - -__all__ = ['ShardedTensor', 'ShardedParamV2'] diff --git a/colossalai/zero/utils/__init__.py b/colossalai/zero/utils/__init__.py deleted file mode 100644 index c4e687228957..000000000000 --- a/colossalai/zero/utils/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .zero_hook import ZeroHook - -__all__ = ['ZeroHook'] \ No newline at end of file diff --git a/colossalai/nn/parallel/zero_wrapper.py b/colossalai/zero/wrapper.py similarity index 82% rename from colossalai/nn/parallel/zero_wrapper.py rename to colossalai/zero/wrapper.py index be8d1da7c24e..3e48f49fa305 100644 --- a/colossalai/nn/parallel/zero_wrapper.py +++ b/colossalai/zero/wrapper.py @@ -4,10 +4,13 @@ import torch import torch.nn as nn -from .gemini_parallel import GeminiDDP +from .gemini import GeminiDDP -def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None): +def zero_model_wrapper(model: nn.Module, + zero_stage: int = 1, + gemini_config: Optional[Dict] = None, + verbose: bool = False): """This wrapper function is used to wrap your training model for ZeRO DDP. Example: @@ -23,7 +26,7 @@ def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Opt zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper. https://arxiv.org/abs/1910.02054 gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled - when the stage is set to 3. You can set the arguemnts of `GeminiDDP` in the gemini_config. + when the stage is set to 3. You can set the arguments of `GeminiDDP` in the gemini_config. Here is an example where we set the device of the model, the placement policy of Gemini, and the size of hidden dimension to help Gemini find out a unified chunk size. @@ -40,7 +43,7 @@ def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Opt if zero_stage in [1, 2]: wrapped_model = model else: - wrapped_model = GeminiDDP(model, **gemini_config) + wrapped_model = GeminiDDP(model, **gemini_config, verbose=verbose) setattr(wrapped_model, "_colo_zero_stage", zero_stage) @@ -58,7 +61,8 @@ def zero_optim_wrapper(model: nn.Module, max_scale: float = 2**32, max_norm: float = 0.0, norm_type: float = 2.0, - optim_config: Optional[Dict] = None): + optim_config: Optional[Dict] = None, + verbose: bool = False): """This wrapper function is used to wrap your training optimizer for ZeRO DDP. Args: @@ -74,11 +78,12 @@ def zero_optim_wrapper(model: nn.Module, max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm. norm_type (float, optional): norm_type used for `clip_grad_norm`. - optim_config (dict, optinoal): The configuration used for the ZeRO optimizer. + optim_config (dict, optional): The configuration used for the ZeRO optimizer. Example: >>> zero2_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True) >>> optim = zero_optim_wrapper(model, optim, optim_config=zero2_config) + verbose (bool, optional): Whether to print the verbose info. """ assert hasattr(model, "_colo_zero_stage"), "You should use `zero_ddp_wrapper` first" zero_stage = getattr(model, "_colo_zero_stage") @@ -99,11 +104,11 @@ def zero_optim_wrapper(model: nn.Module, config_dict['max_scale'] = max_scale if zero_stage in [1, 2]: - from colossalai.zero.sharded_optim.low_level_optim import LowLevelZeroOptimizer + from colossalai.zero.low_level import LowLevelZeroOptimizer config_dict['partition_grad'] = zero_stage == 2 config_dict['clip_grad_norm'] = max_norm - return LowLevelZeroOptimizer(optimizer, **config_dict) + return LowLevelZeroOptimizer(optimizer, **config_dict, verbose=verbose) else: - from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer + from colossalai.zero.gemini.gemini_optimizer import ZeroOptimizer config_dict['clipping_norm'] = max_norm - return ZeroOptimizer(optimizer, model, **config_dict) + return ZeroOptimizer(optimizer, model, **config_dict, verbose=verbose) diff --git a/docker/Dockerfile b/docker/Dockerfile index 49ff9b344268..a1e136ee58a5 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -5,17 +5,37 @@ LABEL org.opencontainers.image.source = "https://github.com/hpcaitech/ColossalAI LABEL org.opencontainers.image.licenses = "Apache License 2.0" LABEL org.opencontainers.image.base.name = "docker.io/library/hpcaitech/cuda-conda:11.3" +# enable passwordless ssh +RUN mkdir ~/.ssh && \ + printf "Host * \n ForwardAgent yes\nHost *\n StrictHostKeyChecking no" > ~/.ssh/config && \ + ssh-keygen -t rsa -N "" -f ~/.ssh/id_rsa && \ + cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys + +# enable RDMA support +RUN apt-get update && \ + apt-get install -y infiniband-diags perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + # install torch RUN conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch +# install ninja +RUN apt-get update && \ + apt-get install -y --no-install-recommends ninja-build && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + # install apex RUN git clone https://github.com/NVIDIA/apex && \ cd apex && \ + git checkout 91fcaa && \ pip install packaging && \ pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" ./ # install colossalai -RUN git clone https://github.com/hpcaitech/ColossalAI.git \ +ARG VERSION=main +RUN git clone -b ${VERSION} https://github.com/hpcaitech/ColossalAI.git \ && cd ./ColossalAI \ && CUDA_EXT=1 pip install -v --no-cache-dir . diff --git a/README-zh-Hans.md b/docs/README-zh-Hans.md similarity index 86% rename from README-zh-Hans.md rename to docs/README-zh-Hans.md index 81c45abfd833..e229c65d890c 100644 --- a/README-zh-Hans.md +++ b/docs/README-zh-Hans.md @@ -24,8 +24,11 @@
## 新闻 +* [2023/07] [65B Model Pretraining Accelerated by 38%, Best Practices for Building LLaMA-Like Base Models Open-Source](https://www.hpc-ai.tech/blog/large-model-pretraining) +* [2023/03] [ColossalChat: An Open-Source Solution for Cloning ChatGPT With a Complete RLHF Pipeline](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +* [2023/03] [Intel and Colossal-AI Partner to Deliver Cost-Efficient Open-Source Solution for Protein Folding Structure Prediction](https://www.hpc-ai.tech/blog/intel-habana) * [2023/03] [AWS and Google Fund Colossal-AI with Startup Cloud Programs](https://www.hpc-ai.tech/blog/aws-and-google-fund-colossal-ai-with-startup-cloud-programs) -* [2023/02] [Open source solution replicates ChatGPT training process! Ready to go with only 1.6GB GPU memory](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) +* [2023/02] [Open Source Solution Replicates ChatGPT Training Process! Ready to go with only 1.6GB GPU Memory](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) * [2023/01] [Hardware Savings Up to 46 Times for AIGC and Automatic Parallelism](https://medium.com/pytorch/latest-colossal-ai-boasts-novel-automatic-parallelism-and-offers-savings-up-to-46x-for-stable-1453b48f3f02) * [2022/11] [Diffusion Pretraining and Hardware Fine-Tuning Can Be Almost 7X Cheaper](https://www.hpc-ai.tech/blog/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper) * [2022/10] [Use a Laptop to Analyze 90% of Proteins, With a Single-GPU Inference Sequence Exceeding 10,000](https://www.hpc-ai.tech/blog/use-a-laptop-to-analyze-90-of-proteins-with-a-single-gpu-inference-sequence-exceeding) @@ -36,9 +39,18 @@
  • 为何选择 Colossal-AI
  • 特点
    • +
  • + Colossal-AI 成功案例 + +
  • 并行训练样例展示
    • -
  • - Colossal-AI 成功案例 - -
  • 安装
      @@ -115,8 +119,106 @@ Colossal-AI 为您提供了一系列并行组件。我们的目标是让您的

      (返回顶端)

      +## Colossal-AI 成功案例 +### ColossalChat + + + +[ColossalChat](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat): 完整RLHF流程0门槛克隆 [ChatGPT](https://openai.com/blog/chatgpt/) +[[代码]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat) +[[博客]](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +[[在线样例]](https://www.youtube.com/watch?v=HcTiHzApHm0) +[[教程]](https://www.youtube.com/watch?v=-qFBZFmOJfg) + +

      + +

      + +- 最高可提升RLHF PPO阶段3训练速度10倍 + +

      + +

      + +- 最高可提升单机训练速度7.73倍，单卡推理速度1.42倍 + +

      + +

      + +- 单卡模型容量最多提升10.3倍 +- 最小demo训练流程最低仅需1.62GB显存 (任意消费级GPU) + +

      + +

      + +- 提升单卡的微调模型容量3.7倍 +- 同时保持高速运行 + +

      (back to top)

      + +### AIGC +加速AIGC(AI内容生成)模型，如[Stable Diffusion v1](https://github.com/CompVis/stable-diffusion) 和 [Stable Diffusion v2](https://github.com/Stability-AI/stablediffusion) + +

      + +

      + +- [训练](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): 减少5.6倍显存消耗，硬件成本最高降低46倍(从A100到RTX3060) + +

      + +

      + +- [DreamBooth微调](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/dreambooth): 仅需3-5张目标主题图像个性化微调 + +

      + +

      + +- [推理](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): GPU推理显存消耗降低2.5倍 + + +

      (返回顶端)

      + +### 生物医药 + +加速 [AlphaFold](https://alphafold.ebi.ac.uk/) 蛋白质结构预测 + +

      + +

      + +- [FastFold](https://github.com/hpcaitech/FastFold): 加速AlphaFold训练与推理、数据前处理、推理序列长度超过10000残基 + +

      + +

      + +- [FastFold with Intel](https://github.com/hpcaitech/FastFold): 3倍推理加速和39%成本节省 + +

      + +

      + +- [xTrimoMultimer](https://github.com/biomap-research/xTrimoMultimer): 11倍加速蛋白质单体与复合物结构预测 + +

      (返回顶端)

      + ## 并行训练样例展示 +### LLaMA +

      + +

      +- 650亿参数大模型预训练加速38% +[[代码]](https://github.com/hpcaitech/ColossalAI/tree/example/llama/examples/language/llama) +[[博客]](https://www.hpc-ai.tech/blog/large-model-pretraining) ### GPT-3

      @@ -211,79 +313,6 @@ Colossal-AI 为您提供了一系列并行组件。我们的目标是让您的

      (返回顶端)

      -## Colossal-AI 成功案例 -### ChatGPT -低成本复现[ChatGPT](https://openai.com/blog/chatgpt/)完整流程 [[代码]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/ChatGPT) [[博客]](https://www.hpc-ai.tech/blog/colossal-ai-chatgpt) -

      - -

      - -- 最高可提升单机训练速度7.73倍，单卡推理速度1.42倍 - -

      - -

      - -- 单卡模型容量最多提升10.3倍 -- 最小demo训练流程最低仅需1.62GB显存 (任意消费级GPU) - -

      - -

      - -- 提升单卡的微调模型容量3.7倍 -- 同时保持高速运行 - -

      (back to top)

      - -### AIGC -加速AIGC(AI内容生成)模型，如[Stable Diffusion v1](https://github.com/CompVis/stable-diffusion) 和 [Stable Diffusion v2](https://github.com/Stability-AI/stablediffusion) - -

      - -

      - -- [训练](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): 减少5.6倍显存消耗，硬件成本最高降低46倍(从A100到RTX3060) - -

      - -

      - -- [DreamBooth微调](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/dreambooth): 仅需3-5张目标主题图像个性化微调 - -

      - -

      - -- [推理](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion): GPU推理显存消耗降低2.5倍 - - -

      (返回顶端)

      - -### 生物医药 - -加速 [AlphaFold](https://alphafold.ebi.ac.uk/) 蛋白质结构预测 - -

      - -

      - -- [FastFold](https://github.com/hpcaitech/FastFold): 加速AlphaFold训练与推理、数据前处理、推理序列长度超过10000残基 - -

      - -

      - -- [FastFold with Intel](https://github.com/hpcaitech/FastFold): 3倍推理加速和39%成本节省 - -

      - -

      - -- [xTrimoMultimer](https://github.com/biomap-research/xTrimoMultimer): 11倍加速蛋白质单体与复合物结构预测 - -

      (返回顶端)

      - ## 安装 环境要求: @@ -291,6 +320,8 @@ Colossal-AI 为您提供了一系列并行组件。我们的目标是让您的 - PyTorch >= 1.11 (PyTorch 2.x 正在适配中) - Python >= 3.7 - CUDA >= 11.0 +- [NVIDIA GPU Compute Capability](https://developer.nvidia.com/cuda-gpus) >= 7.0 (V100/RTX20 and higher) +- Linux OS 如果你遇到安装问题，可以向本项目 [反馈](https://github.com/hpcaitech/ColossalAI/issues/new/choose)。 @@ -386,16 +417,16 @@ docker run -ti --gpus all --rm --ipc=host colossalai bash 真诚感谢所有贡献者！ - - -*贡献者头像的展示顺序是随机的。* + + +

      (返回顶端)

      ## CI/CD -我们使用[GitHub Actions](https://github.com/features/actions)来自动化大部分开发以及部署流程。如果想了解这些工作流是如何运行的，请查看这个[文档](.github/workflows/README.md). +我们使用[GitHub Actions](https://github.com/features/actions)来自动化大部分开发以及部署流程。如果想了解这些工作流是如何运行的，请查看这个[文档](https://github.com/hpcaitech/ColossalAI/blob/main/.github/workflows/README.md). ## 引用我们 diff --git a/docs/README.md b/docs/README.md index f520608d552c..f0cb50ffe217 100644 --- a/docs/README.md +++ b/docs/README.md @@ -98,7 +98,7 @@ Lastly, if you want to skip some code, you just need to add the following annota ``` -If you have any dependency required, please add it to `requriements-doc-test.txt` for pip and `conda-doc-test-deps.yml` for Conda. +If you have any dependency required, please add it to `requirements-doc-test.txt` for pip and `conda-doc-test-deps.yml` for Conda. ### 💉 Auto Documentation diff --git a/REFERENCE.md b/docs/REFERENCE.md similarity index 98% rename from REFERENCE.md rename to docs/REFERENCE.md index 2681198191cb..0984b2dc3f28 100644 --- a/REFERENCE.md +++ b/docs/REFERENCE.md @@ -1,6 +1,6 @@ # References -The Colossal-AI project aims to provide a wide array of parallelism techniques for the machine learning community in the big-model era. This project is inspired by quite a few reserach works, some are conducted by some of our developers and the others are research projects open-sourced by other organizations. We would like to credit these amazing projects below in the IEEE citation format. +The Colossal-AI project aims to provide a wide array of parallelism techniques for the machine learning community in the big-model era. This project is inspired by quite a few research works, some are conducted by some of our developers and the others are research projects open-sourced by other organizations. We would like to credit these amazing projects below in the IEEE citation format. ## By Our Team diff --git a/docs/requirements-doc-test.txt b/docs/requirements-doc-test.txt index 6a6bb3bee9b0..79e04bd5615d 100644 --- a/docs/requirements-doc-test.txt +++ b/docs/requirements-doc-test.txt @@ -4,3 +4,4 @@ packaging tensornvme psutil transformers +pytest diff --git a/docs/sidebars.json b/docs/sidebars.json index 44287c17eadf..8be40e4512f9 100644 --- a/docs/sidebars.json +++ b/docs/sidebars.json @@ -26,8 +26,11 @@ "collapsed": true, "items": [ "basics/command_line_tool", - "basics/define_your_config", "basics/launch_colossalai", + "basics/booster_api", + "basics/booster_plugins", + "basics/booster_checkpoint", + "basics/define_your_config", "basics/initialize_features", "basics/engine_trainer", "basics/configure_parallelization", @@ -40,8 +43,11 @@ "label": "Features", "collapsed": true, "items": [ + "features/mixed_precision_training_with_booster", "features/mixed_precision_training", + "features/gradient_accumulation_with_booster", "features/gradient_accumulation", + "features/gradient_clipping_with_booster", "features/gradient_clipping", "features/gradient_handler", "features/zero_with_chunk", @@ -57,7 +63,8 @@ ] }, "features/pipeline_parallel", - "features/nvme_offload" + "features/nvme_offload", + "features/cluster_utils" ] }, { diff --git a/docs/source/en/Colossal-Auto/get_started/run_demo.md b/docs/source/en/Colossal-Auto/get_started/run_demo.md index 6f7a82966f20..34872e399c81 100644 --- a/docs/source/en/Colossal-Auto/get_started/run_demo.md +++ b/docs/source/en/Colossal-Auto/get_started/run_demo.md @@ -4,7 +4,7 @@ Colossal-Auto simplifies the process of deploying large-scale machine learning m ### 1. Basic usage -Colossal-Auto can be used to find a hybrid SPMD parallel strategy includes data, tensor(i.e., 1D, 2D, sequencial) for each operation. You can follow the [GPT example](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt/experiments/auto_parallel). +Colossal-Auto can be used to find a hybrid SPMD parallel strategy includes data, tensor(i.e., 1D, 2D, sequential) for each operation. You can follow the [GPT example](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt/experiments/auto_parallel). Detailed instructions can be found in its `README.md`. ### 2. Integration with activation checkpoint diff --git a/docs/source/en/advanced_tutorials/add_your_parallel.md b/docs/source/en/advanced_tutorials/add_your_parallel.md index be7284a7ab64..1caf58c8734e 100644 --- a/docs/source/en/advanced_tutorials/add_your_parallel.md +++ b/docs/source/en/advanced_tutorials/add_your_parallel.md @@ -56,7 +56,7 @@ follow the steps below to create a new distributed initialization. world_size: int, config: Config, data_parallel_size: int, - pipeline_parlalel_size: int, + pipeline_parallel_size: int, tensor_parallel_size: int, arg1, arg2): diff --git a/docs/source/en/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md b/docs/source/en/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md index e01caf76d2b3..bfa5539fe3a6 100644 --- a/docs/source/en/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md +++ b/docs/source/en/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md @@ -121,7 +121,7 @@ Inside the initialization of Experts, the local expert number of each GPU will b ## Train Your Model -Do not to forget to use `colossalai.initialize` function in `colosalai` to add gradient handler for the engine. +Do not to forget to use `colossalai.initialize` function in `colossalai` to add gradient handler for the engine. We handle the back-propagation of MoE models for you. In `colossalai.initialize`, we will automatically create a `MoeGradientHandler` object to process gradients. You can find more information about the handler `MoeGradientHandler` in colossal directory. @@ -137,3 +137,4 @@ criterion = MoeLoss( Finally, just use trainer or engine in `colossalai` to do your training. Otherwise, you should take care of gradient by yourself. + diff --git a/docs/source/en/advanced_tutorials/meet_gemini.md b/docs/source/en/advanced_tutorials/meet_gemini.md index 4889b30a6cf8..e94e3fea3710 100644 --- a/docs/source/en/advanced_tutorials/meet_gemini.md +++ b/docs/source/en/advanced_tutorials/meet_gemini.md @@ -9,16 +9,21 @@ When you only have a few GPUs for large model training tasks, **heterogeneous tr ## Usage -At present, Gemini supports compatibility with ZeRO parallel mode, and it is really simple to use Gemini. Set attribute of zero model_config, i.e., tensor_placement_policy='auto'. - -``` -zero = dict( - model_config=dict( - tensor_placement_policy='auto', - shard_strategy=BucketTensorShardStrategy() - ), - optimizer_config=dict( - ...) +At present, Gemini supports compatibility with ZeRO parallel mode, and it is really simple to use Gemini: Inject the features of `GeminiPlugin` into training components with `booster`. More instructions of `booster` please refer to [**usage of booster**](../basics/booster_api.md). + +```python +from torchvision.models import resnet18 +from colossalai.booster import Booster +from colossalai.zero import ColoInitContext +from colossalai.booster.plugin import GeminiPlugin +plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, max_norm=1.0, initial_scale=2**5) +booster = Booster(plugin=plugin) +ctx = ColoInitContext() +with ctx: + model = resnet18() +optimizer = HybridAdam(model.parameters(), lr=1e-3) +criterion = lambda x: x.mean() +model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) ) ``` @@ -44,7 +49,7 @@ In some solutions, the [Zero-offload](https://arxiv.org/abs/2101.06840) adopted -Colossal-AI designed Gemini, just like two-stars, which manages the memory space of CPU and GPU efficiently. It can make the tensor dynamically distributed in the storage space of CPU-GPU during training, so that the model training can break through the memory wall of GPU. The memory manager consists of two parts: **MemStatsCollector (MSC)** and **StatefuleTensorMgr (STM)**. +Colossal-AI designed Gemini, just like two-stars, which manages the memory space of CPU and GPU efficiently. It can make the tensor dynamically distributed in the storage space of CPU-GPU during training, so that the model training can break through the memory wall of GPU. The memory manager consists of two parts: **MemStatsCollector (MSC)** and **StatefulTensorMgr (STM)**. We take advantage of the iterative characteristics of the deep learning network training process. We divide iterations into two stages: warmup and non-warmup. One or several iterative steps at the beginning belong to the warmup stage, and the other iterative steps belong to the non-warmup stage. In the warmup stage, we collect information for the MSC, while in the non-warmup stage, STM gets the information collected by the MSC to move the tensor, so as to minimize the CPU-GPU data movement volume. @@ -86,3 +91,5 @@ The important duty of MSC is to adjust the tensor layout position. For example, In the warmup stage, since we haven't finished a complete iteration yet, we don't know actual memory occupation. At this time, we limit the upper bound of memory usage of the model data. For example, only 30% of the GPU memory can be used. This ensures that we can successfully complete the warmup state. In the non-warmup stage, we need to use the memory information of non-model data collected in the warm-up stage to reserve the peak memory required by the computing device for the next Period, which requires us to move some model tensors. In order to avoid frequent replacement of the same tensor in and out of the CPU-GPU, causing a phenomenon similar to [cache thrashing](https://en.wikipedia.org/wiki/Thrashing_(computer_science)). Using the iterative characteristics of DNN training, we design the OPT cache swap out strategy. Specifically, in the warmup stage, we record the sampling time required by each tensor computing device. If we need to expel some HOLD tensors, we will choose the latest tensor needed on this device as the victim. + + diff --git a/docs/source/en/advanced_tutorials/opt_service.md b/docs/source/en/advanced_tutorials/opt_service.md index b317de91bbdd..eccfa12f9389 100644 --- a/docs/source/en/advanced_tutorials/opt_service.md +++ b/docs/source/en/advanced_tutorials/opt_service.md @@ -20,7 +20,7 @@ To launch the distributed inference service quickly, you can download the OPT-12 2. Prepare a prebuilt service image -Pull a docker image from dockerhub installed with Colossal-AI inference. +Pull a docker image from docker hub installed with Colossal-AI inference. ```bash docker pull hpcaitech/energon-ai:latest @@ -53,7 +53,7 @@ export CHECKPOINT_DIR="your_opt_checkpoint_path" # the ${CONFIG_DIR} must contain a server.sh file as the entry of service export CONFIG_DIR="config_file_path" -docker run --gpus all --rm -it -p 8020:8020 -v ${CHECKPOINT_DIR}:/model_checkpoint -v ${CONFIG_DIR}:/config --ipc=host energonai:lastest +docker run --gpus all --rm -it -p 8020:8020 -v ${CHECKPOINT_DIR}:/model_checkpoint -v ${CONFIG_DIR}:/config --ipc=host energonai:latest ``` Then open `https://[IP-ADDRESS]:8020/docs#` in your browser to try out! diff --git a/docs/source/en/advanced_tutorials/parallelize_your_training_like_Megatron.md b/docs/source/en/advanced_tutorials/parallelize_your_training_like_Megatron.md index e7698e5e9d1b..281fd47554ca 100644 --- a/docs/source/en/advanced_tutorials/parallelize_your_training_like_Megatron.md +++ b/docs/source/en/advanced_tutorials/parallelize_your_training_like_Megatron.md @@ -69,7 +69,7 @@ After the forward operation of the embedding module, each word in all sequences
      The embedding module
      -Each transformer layer contains two blocks. The self-attention operation is called in the first block and a two-layer percepton is located in the second block. +Each transformer layer contains two blocks. The self-attention operation is called in the first block and a two-layer perception is located in the second block.
      @@ -141,16 +141,16 @@ for mn, module in model.named_modules(): if 'mlp.c_fc' in mn: if 'weight' in pn or 'bias' in pn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice # keep the shape of the output from c_fc param.compute_spec.set_output_replicate(False) elif 'mlp.c_proj' in mn: if 'weight' in pn: split_param_row_tp1d(param, pg) # row slice elif 'wte' in mn or 'wpe' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif 'c_attn' in mn or 'c_proj' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice ``` The modified model is illustrated below. @@ -175,13 +175,13 @@ In this way, users can train their models as usual. In our latest example, a Gemini + ZeRO DDP model is also defined to reduce overhead and improve efficiency.For the details of this part, please refer to [ZeRO](../features/zero_with_chunk.md). You can combine these two parts to understand our entire training process: ```python -def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: str = "auto"): +def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placement_policy: str = "auto"): from colossalai.nn.parallel import GeminiDDP model = GeminiDDP(model, device=get_current_device(), - placement_policy=placememt_policy, + placement_policy=placement_policy, pin_memory=True, - search_range_mb=32) + search_range_m=32) return model ``` @@ -190,3 +190,5 @@ def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: The above optimization we made allows us to pretrain the GPT-2 model on a single GPU. We only need to set the parameter `GPUNUM`=1 in `run.sh`, and then we can complete the model training on a single GPU when running the file. The GPT-2 example is accessible at [Train GPT with Colossal-AI](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). + + diff --git a/docs/source/en/advanced_tutorials/train_vit_using_pipeline_parallelism.md b/docs/source/en/advanced_tutorials/train_vit_using_pipeline_parallelism.md index b26599740c5f..6adfe4f113da 100644 --- a/docs/source/en/advanced_tutorials/train_vit_using_pipeline_parallelism.md +++ b/docs/source/en/advanced_tutorials/train_vit_using_pipeline_parallelism.md @@ -195,7 +195,7 @@ def build_cifar(batch_size): ## Training ViT using pipeline -You can set the size of pipeline parallel and number of microbatches in config. `NUM_CHUNKS` is useful when using interleved-pipeline (for more details see [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://arxiv.org/abs/2104.04473) ). The original batch will be split into `num_microbatches`, and each stage will load a micro batch each time. Then we will generate an approriate schedule for you to execute the pipeline training. If you don't need the output and label of model, you can set `return_output_label` to `False` when calling `trainer.fit()` which can further reduce GPU memory usage. +You can set the size of pipeline parallel and number of microbatches in config. `NUM_CHUNKS` is useful when using interleaved-pipeline (for more details see [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://arxiv.org/abs/2104.04473) ). The original batch will be split into `num_microbatches`, and each stage will load a micro batch each time. Then we will generate an appropriate schedule for you to execute the pipeline training. If you don't need the output and label of model, you can set `return_output_label` to `False` when calling `trainer.fit()` which can further reduce GPU memory usage. You should `export DATA=/path/to/cifar`. diff --git a/docs/source/en/advanced_tutorials/train_vit_with_hybrid_parallelism.md b/docs/source/en/advanced_tutorials/train_vit_with_hybrid_parallelism.md index 1f3086559939..a2deaeb88893 100644 --- a/docs/source/en/advanced_tutorials/train_vit_with_hybrid_parallelism.md +++ b/docs/source/en/advanced_tutorials/train_vit_with_hybrid_parallelism.md @@ -12,18 +12,18 @@ Author: Yuxuan Lou ## Introduction -In this example for ViT model, Colossal-AI provides three different parallelism techniques which acclerate model training: data parallelism, pipeline parallelism and tensor parallelism. +In this example for ViT model, Colossal-AI provides three different parallelism techniques which accelerate model training: data parallelism, pipeline parallelism and tensor parallelism. We will show you how to train ViT on CIFAR-10 dataset with these parallelism techniques. To run this example, you will need 2-4 GPUs. -## Tabel of Contents +## Table of Contents 1. Colossal-AI installation 2. Steps to train ViT with data parallelism 3. Steps to train ViT with pipeline parallelism 4. Steps to train ViT with tensor parallelism or hybrid parallelism ## Colossal-AI Installation -You can install Colossal-AI pacakage and its dependencies with PyPI. +You can install Colossal-AI package and its dependencies with PyPI. ```bash pip install colossalai ``` @@ -31,7 +31,7 @@ pip install colossalai ## Data Parallelism -Data parallism is one basic way to accelerate model training process. You can apply data parallism to training by only two steps: +Data parallelism is one basic way to accelerate model training process. You can apply data parallelism to training by only two steps: 1. Define a configuration file 2. Change a few lines of code in train script @@ -94,7 +94,7 @@ from torchvision import transforms from torchvision.datasets import CIFAR10 ``` -#### Lauch Colossal-AI +#### Launch Colossal-AI In train script, you need to initialize the distributed environment for Colossal-AI after your config file is prepared. We call this process `launch`. In Colossal-AI, we provided several launch methods to initialize the distributed backend. In most cases, you can use `colossalai.launch` and `colossalai.get_default_parser` to pass the parameters via command line. Besides, Colossal-AI can utilize the existing launch tool provided by PyTorch as many users are familiar with by using `colossalai.launch_from_torch`. For more details, you can view the related [documents](https://www.colossalai.org/docs/basics/launch_colossalai). @@ -108,7 +108,7 @@ disable_existing_loggers() logger = get_dist_logger() ``` -After initialization, you can acess the variables in the config file by using `colossalai.core.global_context`. +After initialization, you can access the variables in the config file by using `colossalai.core.global_context`. ```python #access parameters @@ -162,7 +162,7 @@ optimizer = colossalai.nn.Lamb(model.parameters(), lr=1.8e-2, weight_decay=0.1) # build loss criterion = torch.nn.CrossEntropyLoss() -# lr_scheduelr +# lr_scheduler lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=gpc.config.NUM_EPOCHS) ``` @@ -230,10 +230,10 @@ torchrun --standalone --nproc_per_node train_dp.py --config ./config ## Pipeline Parallelism -Aside from data parallelism, Colossal-AI also support pipleline parallelism. In specific, Colossal-AI uses 1F1B pipeline introduced by NVIDIA. For more details, you can view the related [documents](https://www.colossalai.org/tutorials/features/pipeline_parallel). +Aside from data parallelism, Colossal-AI also support pipeline parallelism. In specific, Colossal-AI uses 1F1B pipeline introduced by NVIDIA. For more details, you can view the related [documents](https://www.colossalai.org/tutorials/features/pipeline_parallel). ### Define your configuration file(`hybrid_parallel/configs/vit_pipeline.py`) -To apply pipleline parallel on the data parallel basis, you only need to add a **parallel dict** +To apply pipeline parallel on the data parallel basis, you only need to add a **parallel dict** ```python from colossalai.amp import AMP_TYPE @@ -250,7 +250,7 @@ clip_grad_norm = 1.0 Other configs： ```python -# hyperparameters +# hyper parameters # BATCH_SIZE is as per GPU # global batch size = BATCH_SIZE x data parallel size BATCH_SIZE = 256 @@ -276,7 +276,7 @@ Colossal-AI provides two methods to build a pipeline model from the existing mod - `colossalai.builder.build_pipeline_model_from_cfg` - `colossalai.builder.build_pipeline_model` -Besides, you can also build a pipeline model from scrath with Colossal-AI. +Besides, you can also build a pipeline model from scratch with Colossal-AI. ```python import math from typing import Callable @@ -521,7 +521,7 @@ def build_cifar(batch_size): return train_dataloader, test_dataloader -# craete dataloaders +# create dataloaders train_dataloader , test_dataloader = build_cifar() # create loss function @@ -539,7 +539,7 @@ lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, #### Start Colossal-AI engine ```python -# intiailize +# initialize engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model=model, optimizer=optimizer, criterion=criterion, @@ -613,9 +613,9 @@ NUM_MICRO_BATCHES = parallel['pipeline'] TENSOR_SHAPE = (BATCH_SIZE // NUM_MICRO_BATCHES, SEQ_LENGTH, HIDDEN_SIZE) ``` -Ohter configs: +Other configs: ```python -# hyperparameters +# hyper parameters # BATCH_SIZE is as per GPU # global batch size = BATCH_SIZE x data parallel size BATCH_SIZE = 256 diff --git a/docs/source/en/basics/booster_api.md b/docs/source/en/basics/booster_api.md new file mode 100644 index 000000000000..1e75c343c14f --- /dev/null +++ b/docs/source/en/basics/booster_api.md @@ -0,0 +1,78 @@ +# Booster API + +Author: [Mingyan Jiang](https://github.com/jiangmingyan) [Jianghai Chen](https://github.com/CjhHa1) + +**Prerequisite:** + +- [Distributed Training](../concepts/distributed_training.md) +- [Colossal-AI Overview](../concepts/colossalai_overview.md) + +**Example Code** + +- [Train with Booster](https://github.com/hpcaitech/ColossalAI/blob/main/examples/tutorial/new_api/cifar_resnet/README.md) + +## Introduction + +In our new design, `colossalai.booster` replaces the role of `colossalai.initialize` to inject features into your training components (e.g. model, optimizer, dataloader) seamlessly. With these new APIs, you can integrate your model with our parallelism features more friendly. Also calling `colossalai.booster` is the standard procedure before you run into your training loops. In the sections below, I will cover how `colossalai.booster` works and what we should take note of. + +### Plugin + +Plugin is an important component that manages parallel configuration (eg: The gemini plugin encapsulates the gemini acceleration solution). Currently supported plugins are as follows: + +**_GeminiPlugin:_** This plugin wraps the Gemini acceleration solution, that ZeRO with chunk-based memory management. + +**_TorchDDPPlugin:_** This plugin wraps the DDP acceleration solution of Pytorch. It implements data parallelism at the module level which can run across multiple machines. + +**_LowLevelZeroPlugin:_** This plugin wraps the 1/2 stage of Zero Redundancy Optimizer. Stage 1 : Shards optimizer states across data parallel workers/GPUs. Stage 2 : Shards optimizer states + gradients across data parallel workers/GPUs. + + +**_TorchFSDPPlugin:_** This plugin wraps the FSDP acceleration solution of Pytorch and can be used to train models with zero-dp. + +### API of booster + +{{ autodoc:colossalai.booster.Booster }} + +## Usage + +In a typical workflow, you should launch distributed environment at the beginning of training script and create objects needed (such as models, optimizers, loss function, data loaders etc.) firstly, then call `colossalai.booster` to inject features into these objects, After that, you can use our booster APIs and these returned objects to continue the rest of your training processes. + +A pseudo-code example is like below: + +```python +import torch +from torch.optim import SGD +from torchvision.models import resnet18 + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin + +def train(): + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + model = resnet18() + criterion = lambda x: x.mean() + optimizer = SGD((model.parameters()), lr=0.001) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1) + model, optimizer, criterion, _, scheduler = booster.boost(model, optimizer, criterion, lr_scheduler=scheduler) + + x = torch.randn(4, 3, 224, 224) + x = x.to('cuda') + output = model(x) + loss = criterion(output) + booster.backward(loss, optimizer) + optimizer.clip_grad_by_norm(1.0) + optimizer.step() + scheduler.step() + + save_path = "./model" + booster.save_model(model, save_path, True, True, "", 10, use_safetensors=use_safetensors) + + new_model = resnet18() + booster.load_model(new_model, save_path) +``` + +[more design details](https://github.com/hpcaitech/ColossalAI/discussions/3046) + + diff --git a/docs/source/en/basics/booster_checkpoint.md b/docs/source/en/basics/booster_checkpoint.md new file mode 100644 index 000000000000..b2840fe87441 --- /dev/null +++ b/docs/source/en/basics/booster_checkpoint.md @@ -0,0 +1,46 @@ +# Booster Checkpoint + +Author: [Hongxin Liu](https://github.com/ver217) + +**Prerequisite:** +- [Booster API](./booster_api.md) + +## Introduction + +We've introduced the [Booster API](./booster_api.md) in the previous tutorial. In this tutorial, we will introduce how to save and load checkpoints using booster. + +## Model Checkpoint + +{{ autodoc:colossalai.booster.Booster.save_model }} + +Model must be boosted by `colossalai.booster.Booster` before saving. `checkpoint` is the path to saved checkpoint. It can be a file, if `shard=False`. Otherwise, it should be a directory. If `shard=True`, the checkpoint will be saved in a sharded way. This is useful when the checkpoint is too large to be saved in a single file. Our sharded checkpoint format is compatible with [huggingface/transformers](https://github.com/huggingface/transformers). + +{{ autodoc:colossalai.booster.Booster.load_model }} + +Model must be boosted by `colossalai.booster.Booster` before loading. It will detect the checkpoint format automatically, and load in corresponding way. + +## Optimizer Checkpoint + +{{ autodoc:colossalai.booster.Booster.save_optimizer }} + +Optimizer must be boosted by `colossalai.booster.Booster` before saving. + +{{ autodoc:colossalai.booster.Booster.load_optimizer }} + +Optimizer must be boosted by `colossalai.booster.Booster` before loading. + +## LR Scheduler Checkpoint + +{{ autodoc:colossalai.booster.Booster.save_lr_scheduler }} + +LR scheduler must be boosted by `colossalai.booster.Booster` before saving. `checkpoint` is the local path to checkpoint file. + +{{ autodoc:colossalai.booster.Booster.load_lr_scheduler }} + +LR scheduler must be boosted by `colossalai.booster.Booster` before loading. `checkpoint` is the local path to checkpoint file. + +## Checkpoint design + +More details about checkpoint design can be found in our discussion [A Unified Checkpoint System Design](https://github.com/hpcaitech/ColossalAI/discussions/3339). + + diff --git a/docs/source/en/basics/booster_plugins.md b/docs/source/en/basics/booster_plugins.md new file mode 100644 index 000000000000..c5c45abce8f7 --- /dev/null +++ b/docs/source/en/basics/booster_plugins.md @@ -0,0 +1,72 @@ +# Booster Plugins + +Author: [Hongxin Liu](https://github.com/ver217) + +**Prerequisite:** +- [Booster API](./booster_api.md) + +## Introduction + +As mentioned in [Booster API](./booster_api.md), we can use booster plugins to customize the parallel training. In this tutorial, we will introduce how to use booster plugins. + +We currently provide the following plugins: + +- [Low Level Zero Plugin](#low-level-zero-plugin): It wraps the `colossalai.zero.low_level.LowLevelZeroOptimizer` and can be used to train models with zero-dp. It only supports zero stage-1 and stage-2. +- [Gemini Plugin](#gemini-plugin): It wraps the [Gemini](../features/zero_with_chunk.md) which implements Zero-3 with chunk-based and heterogeneous memory management. +- [Torch DDP Plugin](#torch-ddp-plugin): It is a wrapper of `torch.nn.parallel.DistributedDataParallel` and can be used to train models with data parallelism. +- [Torch FSDP Plugin](#torch-fsdp-plugin): It is a wrapper of `torch.distributed.fsdp.FullyShardedDataParallel` and can be used to train models with zero-dp. + +More plugins are coming soon. + +## Plugins + +### Low Level Zero Plugin + +This plugin implements Zero-1 and Zero-2 (w/wo CPU offload), using `reduce` and `gather` to synchronize gradients and weights. + +Zero-1 can be regarded as a better substitute of Torch DDP, which is more memory efficient and faster. It can be easily used in hybrid parallelism. + +Zero-2 does not support local gradient accumulation. Though you can accumulate gradient if you insist, it cannot reduce communication cost. That is to say, it's not a good idea to use Zero-2 with pipeline parallelism. + +{{ autodoc:colossalai.booster.plugin.LowLevelZeroPlugin }} + +We've tested compatibility on some famous models, following models may not be supported: + +- `timm.models.convit_base` +- dlrm and deepfm models in `torchrec` +- `diffusers.VQModel` +- `transformers.AlbertModel` +- `transformers.AlbertForPreTraining` +- `transformers.BertModel` +- `transformers.BertForPreTraining` +- `transformers.GPT2DoubleHeadsModel` + +Compatibility problems will be fixed in the future. + +> ⚠ This plugin can only load optimizer checkpoint saved by itself with the same number of processes now. This will be fixed in the future. + +### Gemini Plugin + +This plugin implements Zero-3 with chunk-based and heterogeneous memory management. It can train large models without much loss in speed. It also does not support local gradient accumulation. More details can be found in [Gemini Doc](../features/zero_with_chunk.md). + +{{ autodoc:colossalai.booster.plugin.GeminiPlugin }} + +### Torch DDP Plugin + +More details can be found in [Pytorch Docs](https://pytorch.org/docs/main/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel). + +{{ autodoc:colossalai.booster.plugin.TorchDDPPlugin }} + +### Torch FSDP Plugin + +> ⚠ This plugin is not available when torch version is lower than 1.12.0. + +> ⚠ This plugin does not support save/load sharded model checkpoint now. + +> ⚠ This plugin does not support optimizer that use multi params group. + +More details can be found in [Pytorch Docs](https://pytorch.org/docs/main/fsdp.html). + +{{ autodoc:colossalai.booster.plugin.TorchFSDPPlugin }} + + diff --git a/docs/source/en/basics/colotensor_concept.md b/docs/source/en/basics/colotensor_concept.md index 2d8acd88dfd4..abe470fe0794 100644 --- a/docs/source/en/basics/colotensor_concept.md +++ b/docs/source/en/basics/colotensor_concept.md @@ -2,6 +2,8 @@ Author: [Jiarui Fang](https://github.com/feifeibear), [Hongxin Liu](https://github.com/ver217) and [Haichen Huang](https://github.com/1SAA) +> ⚠️ The information on this page is outdated and will be deprecated. + **Prerequisite:** - [Colossal-AI Overview](../concepts/colossalai_overview.md) - [Distributed Training](../concepts/distributed_training.md) @@ -42,7 +44,7 @@ Therefore, when using Distributed Spec, we only need to describe the way that th ## Compute Spec -An instance of class [ComputeSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.compute_spec.html#colossalai.tensor.compute_spec.ComputeSpec) describes how a Coloensor be used in DNN training. Currently, we will set the correct Compute Pattern for the ColoTensor as the parameters of the module. The specific application scenarios will be shown in the next document. +An instance of class [ComputeSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.compute_spec.html#colossalai.tensor.compute_spec.ComputeSpec) describes how a Colotensor be used in DNN training. Currently, we will set the correct Compute Pattern for the ColoTensor as the parameters of the module. The specific application scenarios will be shown in the next document. ## ColoParameter @@ -50,18 +52,18 @@ An instance of class [ComputeSpec](https://colossalai.readthedocs.io/en/latest/c ## Example -Let's see an example. A ColoTensor is initialized and sharded on 8 GPUs using tp_degree=4, dp_dgree=2. And then the tensor is sharded along the last dim among the TP process groups. Finally, we reshard it along the first dim (0 dim) among the TP process groups. We encourage users to run the code and observe the shape of each tensor. +Let's see an example. A ColoTensor is initialized and sharded on 8 GPUs using tp_degree=4, dp_degree=2. And then the tensor is sharded along the last dim among the TP process groups. Finally, we reshard it along the first dim (0 dim) among the TP process groups. We encourage users to run the code and observe the shape of each tensor. ```python import torch import torch.multiprocessing as mp -from colossalai.utils import free_port, print_rank_0 +from colossalai.utils import print_rank_0 from functools import partial import colossalai from colossalai.tensor import ProcessGroup, ColoTensor, ColoTensorSpec, ShardSpec, ComputeSpec, ComputePattern -from colossalai.utils import free_port +from colossalai.testing import spawn import torch @@ -83,8 +85,7 @@ def run_dist_tests(rank, world_size, port): print_rank_0(f"shape {t1.shape}, {t1.data}") def test_dist_cases(world_size): - run_func = partial(run_dist_tests, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist_tests, world_size) if __name__ == '__main__': test_dist_cases(4) diff --git a/docs/source/en/basics/configure_parallelization.md b/docs/source/en/basics/configure_parallelization.md index 4ac0299eac14..fd1e72ccd45a 100644 --- a/docs/source/en/basics/configure_parallelization.md +++ b/docs/source/en/basics/configure_parallelization.md @@ -2,6 +2,8 @@ Author: Shenggui Li, Siqi Mai +> ⚠️ The information on this page is outdated and will be deprecated. Please check [Booster Plugins](../basics/booster_plugins.md) for more information. + **Prerequisite:** - [Distributed Training](../concepts/distributed_training.md) - [Paradigms of Parallelism](../concepts/paradigms_of_parallelism.md) diff --git a/docs/source/en/basics/define_your_config.md b/docs/source/en/basics/define_your_config.md index d2569691b7dc..048ffcacbb8f 100644 --- a/docs/source/en/basics/define_your_config.md +++ b/docs/source/en/basics/define_your_config.md @@ -2,6 +2,9 @@ Author: Guangyang Lu, Shenggui Li, Siqi Mai +> ⚠️ The information on this page is outdated and will be deprecated. Please check [Booster API](../basics/booster_api.md) for more information. + + **Prerequisite:** - [Distributed Training](../concepts/distributed_training.md) - [Colossal-AI Overview](../concepts/colossalai_overview.md) diff --git a/docs/source/en/basics/engine_trainer.md b/docs/source/en/basics/engine_trainer.md index 39792f622aa9..d2f99563f042 100644 --- a/docs/source/en/basics/engine_trainer.md +++ b/docs/source/en/basics/engine_trainer.md @@ -2,6 +2,8 @@ Author: Shenggui Li, Siqi Mai +> ⚠️ The information on this page is outdated and will be deprecated. Please check [Booster API](../basics/booster_api.md) for more information. + **Prerequisite:** - [Initialize Features](./initialize_features.md) @@ -172,7 +174,7 @@ In this config file, we specify that we want to use batch size 128 per GPU and r #### Step 2. Initialize Distributed Environment We need to initialize the distributed training environment. This has been introduced in the tutorial on how to -[launch Colossal-AI](./launch_colossalai.md). For this demostration, we use `launch_from_torch` and PyTorch launch utility. +[launch Colossal-AI](./launch_colossalai.md). For this demonstration, we use `launch_from_torch` and PyTorch launch utility. ```python import colossalai diff --git a/docs/source/en/basics/initialize_features.md b/docs/source/en/basics/initialize_features.md index e768d2022ad8..b89017427476 100644 --- a/docs/source/en/basics/initialize_features.md +++ b/docs/source/en/basics/initialize_features.md @@ -2,6 +2,8 @@ Author: Shenggui Li, Siqi Mai +> ⚠️ The information on this page is outdated and will be deprecated. Please check [Booster API](../basics/booster_api.md) for more information. + **Prerequisite:** - [Distributed Training](../concepts/distributed_training.md) - [Colossal-AI Overview](../concepts/colossalai_overview.md) diff --git a/docs/source/en/basics/launch_colossalai.md b/docs/source/en/basics/launch_colossalai.md index be487f8539a5..334757ea75af 100644 --- a/docs/source/en/basics/launch_colossalai.md +++ b/docs/source/en/basics/launch_colossalai.md @@ -87,14 +87,13 @@ import colossalai args = colossalai.get_default_parser().parse_args() # launch distributed environment -colossalai.launch(config=, +colossalai.launch(config=args.config, rank=args.rank, world_size=args.world_size, host=args.host, port=args.port, backend=args.backend ) - ``` @@ -107,12 +106,21 @@ First, we need to set the launch method in our code. As this is a wrapper of the use `colossalai.launch_from_torch`. The arguments required for distributed environment such as rank, world size, host and port are all set by the PyTorch launcher and can be read from the environment variable directly. +config.py +```python +BATCH_SIZE = 512 +LEARNING_RATE = 3e-3 +WEIGHT_DECAY = 0.3 +NUM_EPOCHS = 2 +``` +train.py ```python import colossalai colossalai.launch_from_torch( - config=, + config="./config.py", ) +... ``` Next, we can easily start multiple processes with `colossalai run` in your terminal. Below is an example to run the code diff --git a/docs/source/en/basics/model_checkpoint.md b/docs/source/en/basics/model_checkpoint.md index 09d44e7c2709..70334f1c41e7 100644 --- a/docs/source/en/basics/model_checkpoint.md +++ b/docs/source/en/basics/model_checkpoint.md @@ -2,6 +2,8 @@ Author : Guangyang Lu +> ⚠️ The information on this page is outdated and will be deprecated. Please check [Booster Checkpoint](../basics/booster_checkpoint.md) for more information. + **Prerequisite:** - [Launch Colossal-AI](./launch_colossalai.md) - [Initialize Colossal-AI](./initialize_features.md) diff --git a/docs/source/en/concepts/colossalai_overview.md b/docs/source/en/concepts/colossalai_overview.md index d75d20196b08..7617c62a4e00 100644 --- a/docs/source/en/concepts/colossalai_overview.md +++ b/docs/source/en/concepts/colossalai_overview.md @@ -6,20 +6,20 @@ Author: Shenggui Li, Siqi Mai With the development of deep learning model size, it is important to shift to a new training paradigm. The traditional training method with no parallelism and optimization became a thing of the past and new training methods are the key to make training large-scale models efficient and cost-effective. -Colossal-AI is designed to be a unfied system to provide an integrated set of training skills and utilities to the user. You can find the common training utilities such as mixed precision training and gradient accumulation. Besides, we provide an array of parallelism including data, tensor and pipeline parallelism. We optimize tensor parallelism with different multi-dimensional distributed matrix-matrix multiplication algorithm. We also provided different pipeline parallelism methods to allow the user to scale their model across nodes efficiently. More advanced features such as offloading can be found in this tutorial documentation in detail as well. +Colossal-AI is designed to be a unified system to provide an integrated set of training skills and utilities to the user. You can find the common training utilities such as mixed precision training and gradient accumulation. Besides, we provide an array of parallelism including data, tensor and pipeline parallelism. We optimize tensor parallelism with different multi-dimensional distributed matrix-matrix multiplication algorithm. We also provided different pipeline parallelism methods to allow the user to scale their model across nodes efficiently. More advanced features such as offloading can be found in this tutorial documentation in detail as well. ## General Usage -We aim to make Colossal-AI easy to use and non-instrusive to user code. There is a simple general workflow if you want to use Colossal-AI. +We aim to make Colossal-AI easy to use and non-intrusive to user code. There is a simple general workflow if you want to use Colossal-AI.
      Workflow
      -1. Prepare a configiguration file where specifies the features you want to use and your parameters. +1. Prepare a configuration file where specifies the features you want to use and your parameters. 2. Initialize distributed backend with `colossalai.launch` -3. Inject the training features into your training components (e.g. model, optimizer) with `colossalai.initialize`. +3. Inject the training features into your training components (e.g. model, optimizer) with `colossalai.booster`. 4. Run training and testing We will cover the whole workflow in the `basic tutorials` section. @@ -34,3 +34,5 @@ The Colossal-AI system will be expanded to include more training skills, these n 4. expansion of existing parallelism methods We welcome ideas and contribution from the community and you can post your idea for future development in our forum. + + diff --git a/docs/source/en/features/1D_tensor_parallel.md b/docs/source/en/features/1D_tensor_parallel.md index 530c2e7b64bc..7157af210bc5 100644 --- a/docs/source/en/features/1D_tensor_parallel.md +++ b/docs/source/en/features/1D_tensor_parallel.md @@ -7,7 +7,7 @@ Author: Zhengda Bian, Yongbin Li - [Configure Parallelization](../basics/configure_parallelization.md) **Example Code** -- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_1d.py) +- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **Related Paper** - [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://deepakn94.github.io/assets/papers/megatron-sc21.pdf) @@ -19,9 +19,16 @@ An efficient 1D tensor parallelism implementation was introduced by [Megatron-LM Let's take a linear layer as an example, which consists of a GEMM $Y = XA$. Given 2 processors, we split the columns of $A$ into $[A_1 ~ A_2]$, and calculate $Y_i = XA_i$ on each processor, which then forms $[Y_1 ~ Y_2] = [XA_1 ~ XA_2]$. This is called a column-parallel fashion. -When a second linear layer $Z=YB$ follows the column-parallel one, we split $B$ into $\left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right]$, +When a second linear layer $Z=YB$ follows the column-parallel one, we split $B$ into +$$ +\left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right] +$$ which is called a row-parallel fashion. -To calculate $Z = [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right]$, we first calculate $Y_iB_i$ on each processor, then use an all-reduce to aggregate the results as $Z=Y_1B_1+Y_2B_2$. +To calculate +$$ +Z = [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right] +$$ +we first calculate $Y_iB_i$ on each processor, then use an all-reduce to aggregate the results as $Z=Y_1B_1+Y_2B_2$. We also need to note that in the backward pass, the column-parallel linear layer needs to aggregate the gradients of the input tensor $X$, because on each processor $i$ we only have $\dot{X_i}=\dot{Y_i}A_i^T$. Thus, we apply an all-reduce across the processors to get $\dot{X}=\dot{Y}A^T=\dot{Y_1}A_1^T+\dot{Y_2}A_2^T$. @@ -35,7 +42,7 @@ Given $P$ processors, we present the theoretical computation and memory cost, as ## Usage -To enable 1D tensor parallelism for our model, e.g. on 2 GPUs, we need to configure the parallism setting as below. +To enable 1D tensor parallelism for our model, e.g. on 2 GPUs, we need to configure the parallelism setting as below. ```python CONFIG = dict(parallel=dict( data=1, diff --git a/docs/source/en/features/2D_tensor_parallel.md b/docs/source/en/features/2D_tensor_parallel.md index 582614c2f2f4..aae8cc9eef97 100644 --- a/docs/source/en/features/2D_tensor_parallel.md +++ b/docs/source/en/features/2D_tensor_parallel.md @@ -8,7 +8,7 @@ Author: Zhengda Bian, Yongbin Li - [1D Tensor Parallelism](./1D_tensor_parallel.md) **Example Code** -- [ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_2d.py) +- [ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **Related Paper** - [An Efficient 2D Method for Training Super-Large Deep Learning Models](https://arxiv.org/pdf/2104.05343.pdf) @@ -22,33 +22,33 @@ Let's still take a linear layer $Y = XA$ as an example. Given $P=q\times q$ processors (necessary condition), e.g. $q=2$, we split both the input $X$ and weight $A$ into $$ -\left[\begin{matrix} X_{10} & X_{11} \\ X_{00} & X_{01} \end{matrix} \right] +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \end{matrix} \right] \text{~and~} -\left[\begin{matrix} A_{10} & A_{11} \\ A_{00} & A_{01} \end{matrix} \right]. +\left[\begin{matrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{matrix} \right]. $$ The calculation includes $q$ steps. When $t=1$, $X_{i0}$ is broadcasted in its row, and $A_{0j}$ is broadcasted in its column. So, we have $$ -\left[\begin{matrix} X_{10},A_{00} & X_{10},A_{01} \\ X_{00},A_{00} & X_{00},A_{01} \end{matrix} \right]. +\left[\begin{matrix} X_{00},A_{00} & X_{00},A_{01} \\ X_{10},A_{00} & X_{10},A_{01} \end{matrix} \right]. $$ Then we multiply $X_{i0}$ and $A_{0j}$ on each processor $(i, j)$ as $$ -\left[\begin{matrix} X_{10}A_{00} & X_{10}A_{01} \\ X_{00}A_{00} & X_{00}A_{01} \end{matrix} \right] (1). +\left[\begin{matrix} X_{00}A_{00} & X_{00}A_{01} \\ X_{10}A_{00} & X_{10}A_{01} \end{matrix} \right] (1). $$ Similarly, when $t=2$, $X_{i1}$ is broadcasted in its row, $A_{1j}$ is broadcasted in its column, and we multiply them as $$ -\left[\begin{matrix} X_{11}A_{10} & X_{11}A_{11} \\ X_{01}A_{10} & X_{01}A_{11} \end{matrix} \right] (2). +\left[\begin{matrix} X_{01}A_{10} & X_{01}A_{11} \\ X_{11}A_{10} & X_{11}A_{11} \end{matrix} \right] (2). $$ By adding $(1)$ and $(2)$ up, we have $$ -Y = XA = \left[\begin{matrix} X_{10}A_{00}+X_{11}A_{10} & X_{10}A_{01}+X_{11}A_{11} \\ X_{00}A_{00}+X_{01}A_{10} & X_{00}A_{01}+X_{01}A_{11} \end{matrix} \right]. +Y = XA = \left[\begin{matrix} X_{00}A_{00}+X_{01}A_{10} & X_{00}A_{01}+X_{01}A_{11} \\ X_{10}A_{00}+X_{11}A_{10} & X_{10}A_{01}+X_{11}A_{11} \end{matrix} \right]. $$ ## Efficiency @@ -60,7 +60,7 @@ Given $P=q\times q$ processors, we present the theoretical computation and memor ## Usage -To enable 2D tensor parallelism for our model, e.g. on 4 GPUs, we need to configure the parallism setting as below. +To enable 2D tensor parallelism for our model, e.g. on 4 GPUs, we need to configure the parallelism setting as below. ```python CONFIG = dict(parallel=dict( data=1, diff --git a/docs/source/en/features/2p5D_tensor_parallel.md b/docs/source/en/features/2p5D_tensor_parallel.md index 34a261ea0aa0..a81d14f10627 100644 --- a/docs/source/en/features/2p5D_tensor_parallel.md +++ b/docs/source/en/features/2p5D_tensor_parallel.md @@ -9,7 +9,7 @@ Author: Zhengda Bian, Yongbin Li - [2D Tensor Parallelism](./2D_tensor_parallel.md) **Example Code** -- [ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_2p5d.py) +- [ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **Related Paper** - [2.5-dimensional distributed model training](https://arxiv.org/pdf/2105.14500.pdf) @@ -23,29 +23,30 @@ Let's still take a linear layer $Y = XA$ as an example. Given $P=q \times q \times d$ processors (necessary condition), e.g. $q=d=2$, we split the input $X$ into $d\times q$ rows and $q$ columns as $$ -\left[\begin{matrix} X_{30} & X_{31} \\ X_{20} & X_{21} \\ X_{10} & X_{11} \\ X_{00} & X_{01}\end{matrix} \right], +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \\ X_{20} & X_{21} \\ X_{30} & X_{31}\end{matrix} \right], $$ + which can be reshaped into $d$ layers as $$ -\left[\begin{matrix} X_{10} & X_{11} \\ X_{00} & X_{01} \end{matrix} \right] \text{~and~}\left[\begin{matrix} X_{30} & X_{31} \\ X_{20} & X_{21} \end{matrix} \right]. +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \end{matrix} \right] \text{~and~}\left[\begin{matrix} X_{20} & X_{21} \\ X_{30} & X_{31} \end{matrix} \right]. $$ Also, the weight $A$ is split into $$ -\left[\begin{matrix} A_{10} & A_{11} \\ A_{00} & A_{01} \end{matrix} \right]. +\left[\begin{matrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{matrix} \right]. $$ For each layer of $X$, we use the SUMMA algorithm to multiply $X$ and $A$. Then, we have the output $$ -\left[\begin{matrix} Y_{10}=X_{10}A_{00}+X_{11}A_{10} & Y_{11}=X_{10}A_{01}+X_{11}A_{11} \\ Y_{00}=X_{00}A_{00}+X_{01}A_{10} & Y_{01}=X_{00}A_{01}+X_{01}A_{11} \end{matrix} \right] +\left[\begin{matrix} Y_{00}=X_{00}A_{00}+X_{01}A_{10} & Y_{01}=X_{00}A_{01}+X_{01}A_{11} \\ Y_{10}=X_{10}A_{00}+X_{11}A_{10} & Y_{11}=X_{10}A_{01}+X_{11}A_{11} \end{matrix} \right] \text{~and~} $$ $$ -\left[\begin{matrix} Y_{30}=X_{30}A_{00}+X_{31}A_{10} & Y_{31}=X_{30}A_{01}+X_{31}A_{11} \\ Y_{20}=X_{20}A_{00}+X_{21}A_{10} & Y_{21}=X_{20}A_{01}+X_{21}A_{11} \end{matrix} \right]. +\left[\begin{matrix} Y_{20}=X_{20}A_{00}+X_{21}A_{10} & Y_{21}=X_{20}A_{01}+X_{21}A_{11} \\ Y_{30}=X_{30}A_{00}+X_{31}A_{10} & Y_{31}=X_{30}A_{01}+X_{31}A_{11} \end{matrix} \right]. $$ ## Efficiency @@ -57,7 +58,7 @@ Given $P=q \times q \times d$ processors, we present the theoretical computation ## Usage -To enable 2.5D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallism setting as below. +To enable 2.5D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallelism setting as below. ```python CONFIG = dict(parallel=dict( data=1, diff --git a/docs/source/en/features/3D_tensor_parallel.md b/docs/source/en/features/3D_tensor_parallel.md index 1207376335ce..0e28f08b23c9 100644 --- a/docs/source/en/features/3D_tensor_parallel.md +++ b/docs/source/en/features/3D_tensor_parallel.md @@ -9,7 +9,7 @@ Author: Zhengda Bian, Yongbin Li - [2D Tensor Parallelism](./2D_tensor_parallel.md) **Example Code** -- [ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_3d.py) +- [ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **Related Paper** - [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://arxiv.org/pdf/2105.14450.pdf) @@ -67,7 +67,7 @@ Given $P=q \times q \times q$ processors, we present the theoretical computation ## Usage -To enable 3D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallism setting as below. +To enable 3D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallelism setting as below. ```python CONFIG = dict(parallel=dict( data=1, diff --git a/docs/source/en/features/cluster_utils.md b/docs/source/en/features/cluster_utils.md new file mode 100644 index 000000000000..7331d5e73ae0 --- /dev/null +++ b/docs/source/en/features/cluster_utils.md @@ -0,0 +1,16 @@ +# Cluster Utilities + +Author: [Hongxin Liu](https://github.com/ver217) + +**Prerequisite:** +- [Distributed Training](../concepts/distributed_training.md) + +## Introduction + +We provide a utility class `colossalai.cluster.DistCoordinator` to coordinate distributed training. It's useful to get various information about the cluster, such as the number of nodes, the number of processes per node, etc. + +## API Reference + +{{ autodoc:colossalai.cluster.DistCoordinator }} + + diff --git a/docs/source/en/features/gradient_accumulation.md b/docs/source/en/features/gradient_accumulation.md index d8781ee691bc..91d89b815bf7 100644 --- a/docs/source/en/features/gradient_accumulation.md +++ b/docs/source/en/features/gradient_accumulation.md @@ -1,4 +1,4 @@ -# Gradient Accumulation +# Gradient Accumulation (Outdated) Author: Shenggui Li, Yongbin Li @@ -28,7 +28,7 @@ gradient_accumulation = ## Hands-on Practice We provide a [runnable example](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/gradient_accumulation) -to demonstrate gradient accumulation. In this example, we set the gradinet accumulation size to be 4. You can run the script using this command: +to demonstrate gradient accumulation. In this example, we set the gradient accumulation size to be 4. You can run the script using this command: ```shell python -m torch.distributed.launch --nproc_per_node 1 --master_addr localhost --master_port 29500 run_resnet_cifar10_with_engine.py @@ -43,3 +43,5 @@ iteration 1, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0 iteration 2, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) iteration 3, first 10 elements of param: tensor([-0.0141, 0.0464, 0.0507, 0.0321, 0.0356, -0.0150, 0.0172, -0.0118, 0.0222, 0.0473], device='cuda:0', grad_fn=) ``` + + diff --git a/docs/source/en/features/gradient_accumulation_with_booster.md b/docs/source/en/features/gradient_accumulation_with_booster.md new file mode 100644 index 000000000000..201e3bc2b643 --- /dev/null +++ b/docs/source/en/features/gradient_accumulation_with_booster.md @@ -0,0 +1,144 @@ +# Gradient Accumulation (Latest) + +Author: [Mingyan Jiang](https://github.com/jiangmingyan) + +**Prerequisite** +- [Define Your Configuration](../basics/define_your_config.md) +- [Training Booster](../basics/booster_api.md) + +## Introduction + +Gradient accumulation is a common way to enlarge your batch size for training. When training large-scale models, memory can easily become the bottleneck and the batch size can be very small, (e.g. 2), leading to unsatisfactory convergence. Gradient accumulation works by adding up the gradients calculated in multiple iterations, and only update the parameters in the preset iteration. + +## Usage + +It is simple to use gradient accumulation in Colossal-AI. Just call `booster.no_sync()` which returns a context manager. It accumulate gradients without synchronization, meanwhile you should not update the weights. + +## Hands-on Practice + +We now demonstrate gradient accumulation. In this example, we let the gradient accumulation size to be 4. + +### Step 1. Import libraries in train.py +Create a `train.py` and import the necessary dependencies. The version of `torch` should not be lower than 1.8.1. + +```python +import os +from pathlib import Path + +import torch +from torchvision import transforms +from torchvision.datasets import CIFAR10 +from torchvision.models import resnet18 +from torch.utils.data import DataLoader + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.cluster.dist_coordinator import priority_execution +``` + +### Step 2. Initialize Distributed Environment +We then need to initialize distributed environment. For demo purpose, we uses `launch_from_torch`. You can refer to [Launch Colossal-AI](../basics/launch_colossalai.md) for other initialization methods. + +```python +# initialize distributed setting +parser = colossalai.get_default_parser() +args = parser.parse_args() +# launch from torch +colossalai.launch_from_torch(config=dict()) +``` + +### Step 3. Create training components +Build your model, optimizer, loss function, lr scheduler and dataloaders. Note that the root path of the dataset is obtained from the environment variable `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` to a path on your machine. Data will be automatically downloaded to the root path. + +```python +# define the training hyperparameters +BATCH_SIZE = 128 +GRADIENT_ACCUMULATION = 4 + +# build resnet +model = resnet18(num_classes=10) + +# build dataloaders +with priority_execution(): + train_dataset = CIFAR10(root=Path(os.environ.get('DATA', './data')), + download=True, + transform=transforms.Compose([ + transforms.RandomCrop(size=32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ])) + +# build criterion +criterion = torch.nn.CrossEntropyLoss() + +# optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) +``` + +### Step 4. Inject Feature +Create a `TorchDDPPlugin` object to instantiate a `Booster`, and boost these training components. + +```python +plugin = TorchDDPPlugin() +booster = Booster(plugin=plugin) +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +model, optimizer, criterion, train_dataloader, _ = booster.boost(model=model, + optimizer=optimizer, + criterion=criterion, + dataloader=train_dataloader) +``` + +### Step 5. Train with Booster +Use booster in a normal training loops, and verify gradient accumulation. `param_by_iter` is to record the distributed training information. +```python +optimizer.zero_grad() +for idx, (img, label) in enumerate(train_dataloader): + sync_context = booster.no_sync(model) + img = img.cuda() + label = label.cuda() + if idx % (GRADIENT_ACCUMULATION - 1) != 0: + with sync_context: + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + else: + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + ele_1st = next(model.parameters()).flatten()[0] + param_by_iter.append(str(ele_1st.item())) + + if idx != 0 and idx % (GRADIENT_ACCUMULATION - 1) == 0: + break + + for iteration, val in enumerate(param_by_iter): + print(f'iteration {iteration} - value: {val}') + + if param_by_iter[-1] != param_by_iter[0]: + print('The parameter is only updated in the last iteration') + +``` + +### Step 6. Invoke Training Scripts +To verify gradient accumulation, we can just check the change of parameter values. When gradient accumulation is set, parameters are only updated in the last step. You can run the script using this command: +```shell +colossalai run --nproc_per_node 1 train.py +``` + +You will see output similar to the text below. This shows gradient is indeed accumulated as the parameter is not updated +in the first 3 steps, but only updated in the last step. + +```text +iteration 0, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 1, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 2, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 3, first 10 elements of param: tensor([-0.0141, 0.0464, 0.0507, 0.0321, 0.0356, -0.0150, 0.0172, -0.0118, 0.0222, 0.0473], device='cuda:0', grad_fn=) +``` + + diff --git a/docs/source/en/features/gradient_clipping.md b/docs/source/en/features/gradient_clipping.md index f606dde6c393..5a23c68e3e27 100644 --- a/docs/source/en/features/gradient_clipping.md +++ b/docs/source/en/features/gradient_clipping.md @@ -1,4 +1,4 @@ -# Gradient Clipping +# Gradient Clipping (Outdated) Author: Boxiang Wang, Haichen Huang, Yongbin Li @@ -60,3 +60,5 @@ to demonstrate gradient clipping. In this example, we set the gradient clipping ```shell python -m torch.distributed.launch --nproc_per_node 1 --master_addr localhost --master_port 29500 train_with_engine.py ``` + + diff --git a/docs/source/en/features/gradient_clipping_with_booster.md b/docs/source/en/features/gradient_clipping_with_booster.md new file mode 100644 index 000000000000..341a608a5c7b --- /dev/null +++ b/docs/source/en/features/gradient_clipping_with_booster.md @@ -0,0 +1,142 @@ +# Gradient Clipping (Latest) + +Author: [Mingyan Jiang](https://github.com/jiangmingyan) + +**Prerequisite** +- [Define Your Configuration](../basics/define_your_config.md) +- [Training Booster](../basics/booster_api.md) + +**Related Paper** +- [On the difficulty of training Recurrent Neural Networks](https://arxiv.org/abs/1211.5063) + +## Introduction + +In order to speed up training process and seek global optimum for better performance, more and more learning rate schedulers have been proposed. People turn to control learning rate to adjust descent pace during training, which makes gradient vector better to be uniformed in every step. In that case, the descent pace can be controlled as expected. As a result, gradient clipping, a technique which can normalize the gradient vector to circumscribe it in a uniformed length, becomes indispensable for those who desire their better performance of their models. + +You do not have to worry about implementing gradient clipping when using Colossal-AI, we support gradient clipping in a powerful and convenient way. All you need is just an additional command in your configuration file. + +## Why you should use gradient clipping provided by Colossal-AI + +The reason of why we do not recommend users to write gradient clipping by themselves is that naive gradient clipping may fail when applying tensor parallelism, pipeline parallelism or MoE. + +According to the illustration below, each GPU only owns a portion of parameters of the weight in a linear layer. To get correct norm of gradient vector of the weight of the linear layer, the norm of every gradient vector in each GPU should be summed together. More complicated thing is that the distribution of bias is different from the distribution of the weight. The communication group is different in the sum operation. + +(PS: This situation is an old version of 2D parallelism, the implementation in the code is not the same. But it is a good example about the difficulty to unify all communication in gradient clipping.) + +
      + +
      Layout of parameters
      +
      + +Do not worry about it, since Colossal-AI have handled it for you. + +## Usage +To use gradient clipping, you can just add the following code to your configuration file, and after boosted, you can call `clip_grad_by_norm` or `clip_grad_by_value` method of optimizer, if it support clip gradients. + +## Hands-On Practice + +We now demonstrate how to use gradient clipping. In this example, we set the gradient clipping vector norm to be 1.0. + +### step 1. Import libraries in train.py +Create a `train.py` and import the necessary dependencies. + +```python +import os +from pathlib import Path + +import torch +from torchvision import transforms +from torchvision.datasets import CIFAR10 +from torchvision.models import resnet34 +from tqdm import tqdm + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.nn.lr_scheduler import CosineAnnealingLR +``` + +### Step 2. Initialize Distributed Environment +We then need to initialize distributed environment. For demo purpose, we uses `launch_from_torch`. You can refer to [Launch Colossal-AI](../basics/launch_colossalai.md) +for other initialization methods. + +```python +colossalai.launch_from_torch(config=dict()) +logger = get_dist_logger() +``` + + +### Step 3. Create training components + +Build your model, optimizer, loss function, lr scheduler and dataloaders. Note that the root path of the dataset is obtained from the environment variable `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` to a path on your machine. Data will be automatically downloaded to the root path. +```python +# define training hyperparameters +NUM_EPOCHS = 200 +BATCH_SIZE = 128 +GRADIENT_CLIPPING = 0.1 +# build resnet +model = resnet34(num_classes=10) +# build dataloaders +train_dataset = CIFAR10(root=Path(os.environ.get('DATA', './data')), + download=True, + transform=transforms.Compose([ + transforms.RandomCrop(size=32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ])) +# build criterion +criterion = torch.nn.CrossEntropyLoss() + +# optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) + +# lr_scheduler +lr_scheduler = CosineAnnealingLR(optimizer, total_steps=NUM_EPOCHS) + +``` +### Step 4. Inject Gradient Clipping Feature + +Create a `TorchDDPPlugin` object and `Booster` object, get a data loader from plugin, then boost all training components. +```python +plugin = TorchDDPPlugin() +booster = Booster(mixed_precision='fp16', plugin=plugin) +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +model, optimizer, criterion, train_dataloader, lr_scheduler = booster.boost(model,optimizer, criterion,train_dataloader, lr_scheduler) + +``` + +### Step 5. Train with Booster +Use booster in a normal training loops. +```python +# verify gradient clipping +model.train() +for idx, (img, label) in enumerate(train_dataloader): + img = img.cuda() + label = label.cuda() + + model.zero_grad() + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + optimizer.clip_grad_by_norm(max_norm=GRADIENT_CLIPPING) + optimizer.step() + lr_scheduler.step() + + ele_1st = next(model.parameters()).flatten()[0] + logger.info(f'iteration {idx}, loss: {train_loss}, 1st element of parameters: {ele_1st.item()}') + + # only run for 4 iterations + if idx == 3: + break +``` + +### Step 6. Invoke Training Scripts +You can run the script using this command: + +```shell +colossalai run --nproc_per_node 1 train.py +``` + + diff --git a/docs/source/en/features/mixed_precision_training.md b/docs/source/en/features/mixed_precision_training.md index 71cb6971d346..8579d586ed5f 100644 --- a/docs/source/en/features/mixed_precision_training.md +++ b/docs/source/en/features/mixed_precision_training.md @@ -1,4 +1,4 @@ -# Auto Mixed Precision Training +# Auto Mixed Precision Training (Outdated) Author: Chuanrui Wang, Shenggui Li, Yongbin Li @@ -101,7 +101,7 @@ you can use `colossalai.amp.convert_to_amp`. ```python from colossalai.amp import AMP_TYPE -# exmaple of using torch amp +# example of using torch amp model, optimizer, criterion = colossalai.amp.convert_to_amp(model, optimizer, criterion, @@ -220,7 +220,7 @@ The default parameters of Naive AMP: - initial_scale(int): initial scale of gradient scaler - growth_factor(int): the growth rate of loss scale - backoff_factor(float): the decrease rate of loss scale -- hysterisis(int): delay shift in dynamic loss scaling +- hysteresis(int): delay shift in dynamic loss scaling - max_scale(int): maximum loss scale allowed - verbose(bool): if set to `True`, will print debug info @@ -292,7 +292,7 @@ colossalai.launch_from_torch(config=args.config) ### Step 4. Create training components Build your model, optimizer, loss function, lr scheduler and dataloaders. Note that the root path of the dataset is -obtained from the environment varialbe `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` +obtained from the environment variable `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` to a path on your machine. Data will be automatically downloaded to the root path. ```python @@ -326,7 +326,7 @@ to a path on your machine. Data will be automatically downloaded to the root pat # build loss criterion = torch.nn.CrossEntropyLoss() - # lr_scheduelr + # lr_scheduler lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=gpc.config.NUM_EPOCHS) ``` @@ -362,6 +362,7 @@ for epoch in range(gpc.config.NUM_EPOCHS): Use the following command to start the training scripts. You can change `--nproc_per_node` to use a different number of GPUs. -```python +```shell python -m torch.distributed.launch --nproc_per_node 4 --master_addr localhost --master_port 29500 train_with_engine.py --config config/config_AMP_torch.py ``` + diff --git a/docs/source/en/features/mixed_precision_training_with_booster.md b/docs/source/en/features/mixed_precision_training_with_booster.md new file mode 100644 index 000000000000..1240b47d5d2e --- /dev/null +++ b/docs/source/en/features/mixed_precision_training_with_booster.md @@ -0,0 +1,259 @@ +# Auto Mixed Precision Training (Latest) + +Author: [Mingyan Jiang](https://github.com/jiangmingyan) + +**Prerequisite** + +- [Define Your Configuration](../basics/define_your_config.md) +- [Training Booster](../basics/booster_api.md) + +**Related Paper** + +- [Accelerating Scientific Computations with Mixed Precision Algorithms](https://arxiv.org/abs/0808.2794) + +## Introduction + +AMP stands for automatic mixed precision training. +In Colossal-AI, we have incorporated different implementations of mixed precision training: + +1. torch.cuda.amp +2. apex.amp +3. naive amp + +| Colossal-AI | support tensor parallel | support pipeline parallel | fp16 extent | +| -------------- | ----------------------- | ------------------------- | ---------------------------------------------------------------------------------------------------- | +| AMP_TYPE.TORCH | ✅ | ❌ | Model parameters, activation, gradients are downcast to fp16 during forward and backward propagation | +| AMP_TYPE.APEX | ❌ | ❌ | More fine-grained, we can choose opt_level O0, O1, O2, O3 | +| AMP_TYPE.NAIVE | ✅ | ✅ | Model parameters, forward and backward operations are all downcast to fp16 | + +The first two rely on the original implementation of PyTorch (version 1.6 and above) and NVIDIA Apex. +The last method is similar to Apex O2 level. +Among these methods, apex AMP is not compatible with tensor parallelism. +This is because that tensors are split across devices in tensor parallelism, thus, it is required to communicate among different processes to check if inf or nan occurs in the whole model weights. +We modified the torch amp implementation so that it is compatible with tensor parallelism now. + +> ❌️ fp16 and zero are not compatible +> +> ⚠️ Pipeline only support naive AMP currently + +We recommend you to use torch AMP as it generally gives better accuracy than naive AMP if no pipeline is used. + +## Table of Contents + +In this tutorial we will cover: + +1. [AMP introduction](#amp-introduction) +2. [AMP in Colossal-AI](#amp-in-colossal-ai) +3. [Hands-on Practice](#hands-on-practice) + +## AMP Introduction + +Automatic Mixed Precision training is a mixture of FP16 and FP32 training. + +Half-precision float point format (FP16) has lower arithmetic complexity and higher compute efficiency. Besides, fp16 requires half of the storage needed by fp32 and saves memory & network bandwidth, which makes more memory available for large batch size and model size. + +However, there are other operations, like reductions, which require the dynamic range of fp32 to avoid numeric overflow/underflow. That's the reason why we introduce automatic mixed precision, attempting to match each operation to its appropriate data type, which can reduce the memory footprint and augment training efficiency. + +
      + +
      Illustration of an ordinary AMP (figure from PatrickStar paper)
      +
      + +## AMP in Colossal-AI + +We supported three AMP training methods and allowed the user to train with AMP with no code. If you want to train with amp, just assign `mixed_precision` with `fp16` when you instantiate the `Booster`. Now booster support torch amp, the other two(apex amp, naive amp) are still started by `colossalai.initialize`, if needed, please refer to [this](./mixed_precision_training.md). Next we will support `bf16`, `fp8`. + +### Start with Booster + +instantiate `Booster` with `mixed_precision="fp16"`, then you can train with torch amp. + + + +```python +""" + Mapping: + 'fp16': torch amp + 'fp16_apex': apex amp, + 'bf16': bf16, + 'fp8': fp8, + 'fp16_naive': naive amp +""" +from colossalai import Booster +booster = Booster(mixed_precision='fp16',...) +``` + + + +or you can create a `FP16TorchMixedPrecision` object, such as: + + + +```python +from colossalai.mixed_precision import FP16TorchMixedPrecision +mixed_precision = FP16TorchMixedPrecision( + init_scale=2.**16, + growth_factor=2.0, + backoff_factor=0.5, + growth_interval=2000) +booster = Booster(mixed_precision=mixed_precision,...) +``` + + + +The same goes for other types of amps. + +### Torch AMP Configuration + +{{ autodoc:colossalai.booster.mixed_precision.FP16TorchMixedPrecision }} + +### Apex AMP Configuration + +For this mode, we rely on the Apex implementation for mixed precision training. +We support this plugin because it allows for finer control on the granularity of mixed precision. +For example, O2 level (optimization level 2) will keep batch normalization in fp32. + +If you look for more details, please refer to [Apex Documentation](https://nvidia.github.io/apex/). + +{{ autodoc:colossalai.booster.mixed_precision.FP16ApexMixedPrecision }} + +### Naive AMP Configuration + +In Naive AMP mode, we achieved mixed precision training while maintaining compatibility with complex tensor and pipeline parallelism. +This AMP mode will cast all operations into fp16. +The following code block shows the mixed precision api for this mode. + +{{ autodoc:colossalai.booster.mixed_precision.FP16NaiveMixedPrecision }} + +When using `colossalai.booster`, you are required to first instantiate a model, an optimizer and a criterion. +The output model is converted to AMP model of smaller memory consumption. +If your input model is already too large to fit in a GPU, please instantiate your model weights in `dtype=torch.float16`. +Otherwise, try smaller models or checkout more parallelization training techniques! + +## Hands-on Practice + +Now we will introduce the use of AMP with Colossal-AI. In this practice, we will use Torch AMP as an example. + +### Step 1. Import libraries in train.py + +Create a `train.py` and import the necessary dependencies. Remember to install `scipy` and `timm` by running +`pip install timm scipy`. + +```python +import os +from pathlib import Path + +import torch +from timm.models import vit_base_patch16_224 +from titans.utils import barrier_context +from torchvision import datasets, transforms + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.nn.lr_scheduler import LinearWarmupLR +``` + +### Step 2. Initialize Distributed Environment + +We then need to initialize distributed environment. For demo purpose, we uses `launch_from_torch`. You can refer to [Launch Colossal-AI](../basics/launch_colossalai.md) +for other initialization methods. + +```python +# initialize distributed setting +parser = colossalai.get_default_parser() +args = parser.parse_args() + +# launch from torch +colossalai.launch_from_torch(config=dict()) + +``` + +### Step 3. Create training components + +Build your model, optimizer, loss function, lr scheduler and dataloaders. Note that the root path of the dataset is +obtained from the environment variable `DATA`. You may `export DATA=/path/to/data` or change `Path(os.environ['DATA'])` +to a path on your machine. Data will be automatically downloaded to the root path. + +```python +# define the constants +NUM_EPOCHS = 2 +BATCH_SIZE = 128 + +# build model +model = vit_base_patch16_224(drop_rate=0.1) + +# build dataloader +train_dataset = datasets.Caltech101( + root=Path(os.environ['DATA']), + download=True, + transform=transforms.Compose([ + transforms.Resize(256), + transforms.RandomResizedCrop(224), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + Gray2RGB(), + transforms.Normalize([0.5, 0.5, 0.5], + [0.5, 0.5, 0.5]) + ])) + +# build optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, weight_decay=0.1) + +# build loss +criterion = torch.nn.CrossEntropyLoss() + +# lr_scheduler +lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=NUM_EPOCHS) +``` + +### Step 4. Inject AMP Feature + +Create a `MixedPrecision`(if needed) and `TorchDDPPlugin` object, call `colossalai.boost` convert the training components to be running with FP16. + +```python +plugin = TorchDDPPlugin() +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +booster = Booster(mixed_precision='fp16', plugin=plugin) + +# if you need to customize the config, do like this +# >>> from colossalai.mixed_precision import FP16TorchMixedPrecision +# >>> mixed_precision = FP16TorchMixedPrecision( +# >>> init_scale=2.**16, +# >>> growth_factor=2.0, +# >>> backoff_factor=0.5, +# >>> growth_interval=2000) +# >>> plugin = TorchDDPPlugin() +# >>> booster = Booster(mixed_precision=mixed_precision, plugin=plugin) + +# boost model, optimizer, criterion, dataloader, lr_scheduler +model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model, optimizer, criterion, dataloader, lr_scheduler) +``` + +### Step 5. Train with Booster + +Use booster in a normal training loops. + +```python +model.train() +for epoch in range(NUM_EPOCHS): + for img, label in enumerate(train_dataloader): + img = img.cuda() + label = label.cuda() + optimizer.zero_grad() + output = model(img) + loss = criterion(output, label) + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() +``` + +### Step 6. Invoke Training Scripts + +Use the following command to start the training scripts. You can change `--nproc_per_node` to use a different number of GPUs. + +```shell +colossalai run --nproc_per_node 1 train.py +``` + + diff --git a/docs/source/en/features/nvme_offload.md b/docs/source/en/features/nvme_offload.md index 2933c3db6c58..6ed6f2dee5d6 100644 --- a/docs/source/en/features/nvme_offload.md +++ b/docs/source/en/features/nvme_offload.md @@ -53,11 +53,11 @@ It's compatible with all parallel methods in ColossalAI. > ⚠ It only offloads optimizer states on CPU. This means it only affects CPU training or Zero/Gemini with offloading. -## Exampls +## Examples Let's start from two simple examples -- training GPT with different methods. These examples relies on `transformers`. -We should install denpendencies first: +We should install dependencies first: ```shell pip install psutil transformers @@ -78,8 +78,9 @@ from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel import colossalai from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin ``` Then we define a loss function: @@ -99,7 +100,7 @@ class GPTLMLoss(nn.Module): shift_labels.view(-1)) ``` -And we define some utility functions, which generates random data, computes the number of paramters of a model and get memory usage of current process: +And we define some utility functions, which generates random data, computes the number of parameters of a model and get memory usage of current process: ```python def get_data(batch_size: int, seq_len: int, @@ -192,17 +193,23 @@ def train_gemini_cpu(nvme_offload_fraction: float = 0.0): optimizer = HybridAdam(model.parameters(), nvme_offload_fraction=nvme_offload_fraction) print(f'Model numel: {get_model_numel(model) / 1024**3:.3f} B') - gemini_config = dict(strict_ddp_mode=True, device=torch.cuda.current_device(), - placement_policy='cpu', pin_memory=True, hidden_dim=config.n_embd) - model = zero_model_wrapper(model, zero_stage=3, gemini_config=gemini_config) - optimizer = zero_optim_wrapper(model, optimizer, initial_scale=2**5) + plugin = GeminiPlugin( + strict_ddp_mode=True, + device=torch.cuda.current_device(), + placement_policy='cpu', + pin_memory=True, + hidden_dim=config.n_embd, + initial_scale=2**5 + ) + booster = Booster(plugin) + model, optimizer, criterion, _* = booster.boost(model, optimizer, criterion) start = time.time() for step in range(3): data = get_data(4, 128, config.vocab_size) outputs = model(**data) loss = criterion(outputs.logits, data['input_ids']) - optimizer.backward(loss) + booster.backward(loss, optimizer) optimizer.step() optimizer.zero_grad() print(f'[{step}] loss: {loss.item():.3f}') @@ -251,7 +258,7 @@ Time: 3.691 s Mem usage: 5298.344 MB ``` -NVME offload saves about 294 MB memory. Note that enabling `pin_memory` of Gemini can accelerate training but increase memory usage. So this result also meets our expectation. If we disable `pin_memory`, we can aslo observe a memory usage drop about 900 MB. +NVME offload saves about 294 MB memory. Note that enabling `pin_memory` of Gemini can accelerate training but increase memory usage. So this result also meets our expectation. If we disable `pin_memory`, we can also observe a memory usage drop about 900 MB. ## API Reference diff --git a/docs/source/en/features/pipeline_parallel.md b/docs/source/en/features/pipeline_parallel.md index ac49863b3c71..30654b0b0195 100644 --- a/docs/source/en/features/pipeline_parallel.md +++ b/docs/source/en/features/pipeline_parallel.md @@ -156,4 +156,4 @@ trainer.fit(train_dataloader=train_dataloader, display_progress=True) ``` -We use `2` pipeline stages and the batch will be splitted into `4` micro batches. +We use `2` pipeline stages and the batch will be split into `4` micro batches. diff --git a/docs/source/en/features/zero_with_chunk.md b/docs/source/en/features/zero_with_chunk.md index 6b0a9585af85..b50d2d02217b 100644 --- a/docs/source/en/features/zero_with_chunk.md +++ b/docs/source/en/features/zero_with_chunk.md @@ -3,7 +3,7 @@ Author: [Hongxiu Liu](https://github.com/ver217), [Jiarui Fang](https://github.com/feifeibear), [Zijian Ye](https://github.com/ZijianYY) **Prerequisite:** -- [Define Your Configuration](../basics/define_your_config.md) +- [Train with booster](../basics/booster_api.md) **Example Code** @@ -32,11 +32,11 @@ and the first and second momentum estimates) are partitioned across the processe 3. **Shard Parameter**: The 16-bit model parameters are partitioned across the processes of a data parallel group. -4. **[Gemini](../advanced_tutorials/meet_gemini.md)**: Dynamic heterogeneous memory space manager for paramters, gradients and optimizer states. +4. **[Gemini](../advanced_tutorials/meet_gemini.md)**: Dynamic heterogeneous memory space manager for parameters, gradients and optimizer states. Besides, this article will introduce the Zero Redundancy Optimizer with chunk-based memory management. -When using ZeRO, we distributed the model by sharding the parameters. The advantage of this method is that the memory of each node is load balanced. But this approach has two significiant disadvantages. First, during communication, a temporary memory buffer needs to be allocated and released afterwards, leading to the memory fragmentation problem. Secondly, using tensor as the granularity for communication will cause the network bandwidth underutilized. Generally, the longer the transmitted message length, the higher the bandwidth utilization. +When using ZeRO, we distributed the model by sharding the parameters. The advantage of this method is that the memory of each node is load balanced. But this approach has two significant disadvantages. First, during communication, a temporary memory buffer needs to be allocated and released afterwards, leading to the memory fragmentation problem. Secondly, using tensor as the granularity for communication will cause the network bandwidth underutilized. Generally, the longer the transmitted message length, the higher the bandwidth utilization. Using the Chunk mechanism introduced in ColossalAI v0.1.8, we can improve the efficiency of ZeRO. We store a continuous set of parameters in initialization order into a Chunk (a chunk is a continuous memory space), and each Chunk has the same size. Organizing memory in chunks can lead to efficient use of network bandwidth between PCI-e and GPU-GPU, reduce the number of communications, and avoid potential memory fragmentation. @@ -67,12 +67,12 @@ Define the model parameters as follows: chunk_manager = init_chunk_manager(model=module, init_device=device, hidden_dim=hidden_dim, - search_range_mb=search_range_mb, - min_chunk_size_mb=min_chunk_size_mb) + search_range_m=search_range_m, + min_chunk_size_m=min_chunk_size_m) gemini_manager = GeminiManager(placement_policy, chunk_manager) ``` -`hidden_dim` is the hidden dimension of DNN. Users can provide this argument to speed up searching. If users do not know this argument before training, it is ok. We will use a default value 1024. `min_chunk_size_mb` is the the minimum chunk size in MegaByte. If the aggregate size of parameters is still samller than the minimum chunk size, all parameters will be compacted into one small chunk. +`hidden_dim` is the hidden dimension of DNN. Users can provide this argument to speed up searching. If users do not know this argument before training, it is ok. We will use a default value 1024. `min_chunk_size_m` is a floating point, being the minimum chunk size divided by 2^20 (e.g., if min_chunk_size_m=2.5, then the minimum chunk size should be 2.5*(2^20)).If the aggregate size of parameters is still smaller than the minimum chunk size, all parameters will be compacted into one small chunk. Initialization of the optimizer. ```python @@ -97,6 +97,7 @@ For simplicity, we just use randomly generated data here. First we only need to import `GPT2LMHeadModel` from `Huggingface transformers` to define our model, which does not require users to define or modify the model, so that users can use it more conveniently. +Define a GPT model: ```python class GPTLMModel(nn.Module): @@ -182,34 +183,6 @@ def split_param_col_tp1d(param: ColoParameter, pg: ProcessGroup): split_param_single_dim_tp1d(-1, param, pg) ``` -Define a model which uses Gemini + ZeRO DDP: - -```python -def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: str = "auto"): - cai_version = colossalai.__version__ - if version.parse(cai_version) > version.parse("0.1.10"): - from colossalai.nn.parallel import GeminiDDP - model = GeminiDDP(model, - device=get_current_device(), - placement_policy=placememt_policy, - pin_memory=True, - search_range_mb=32) - elif version.parse(cai_version) <= version.parse("0.1.10") and version.parse(cai_version) >= version.parse("0.1.9"): - from colossalai.gemini import ChunkManager, GeminiManager - chunk_size = ChunkManager.search_chunk_size(model, 64 * 1024**2, 32) - gemini_manager = GeminiManager(placememt_policy, chunk_manager) - chunk_manager = ChunkManager(chunk_size, - pg, - enable_distributed_storage=True, - init_device=GeminiManager.get_default_device(placememt_policy)) - model = ZeroDDP(model, gemini_manager) - else: - raise NotImplemented(f"CAI version {cai_version} is not supported") - return model -``` - -As we pre-train GPT in this example, we just use a simple language model loss. - Write a function to get random inputs: ```python @@ -219,9 +192,15 @@ def get_data(batch_size, seq_len, vocab_size): return input_ids, attention_mask ``` -Finally, we can define our training loop: +Finally, we define a model which uses Gemini + ZeRO DDP and define our training loop, As we pre-train GPT in this example, we just use a simple language model loss: ```python +from colossalai.nn.optimizer import HybridAdam + +from colossalai.booster import Booster +from colossalai.zero import ColoInitContext +from colossalai.booster.plugin import GeminiPlugin + def main(): args = parse_args() BATCH_SIZE = 8 @@ -232,22 +211,23 @@ def main(): # build criterion criterion = GPTLMLoss() + optimizer = HybridAdam(model.parameters(), lr=0.001) torch.manual_seed(123) default_pg = ProcessGroup(tp_degree=args.tp_degree) - default_dist_spec = ShardSpec([-1], [args.tp_degree]) if args.shardinit else None + default_dist_spec = ShardSpec([-1], [args.tp_degree]) # build GPT model with ColoInitContext(device='cpu', default_dist_spec=default_dist_spec, default_pg=default_pg): model = gpt2_medium(checkpoint=True) pg = default_pg # Tensor Parallelism (TP) tensor_parallelize(model, pg) + # Gemini + ZeRO DP, Note it must be used after TP - model = gemini_zero_dpp(model, pg, args.placement) - # build optimizer - optimizer = GeminiAdamOptimizer(model, lr=1e-3, initial_scale=2**5) - numel = sum([p.numel() for p in model.parameters()]) - get_tflops_func = partial(get_tflops, numel, BATCH_SIZE, SEQ_LEN) + plugin = GeminiPlugin(placement_policy='cuda', max_norm=1.0, initial_scale=2**5) + booster = Booster(plugin=plugin) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + torch.cuda.synchronize() model.train() for n in range(NUM_STEPS): @@ -256,10 +236,12 @@ def main(): optimizer.zero_grad() outputs = model(input_ids, attn_mask) loss = criterion(outputs, input_ids) - optimizer.backward(loss) + booster.backward(loss, optimizer) optimizer.step() torch.cuda.synchronize() ``` > ⚠️ Note: If you want to use the Gemini module, please do not use the [Gradient Accumulation](../features/gradient_accumulation.md) we mentioned before。 The complete example can be found on [Train GPT with Colossal-AI](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). + + diff --git a/docs/source/en/get_started/installation.md b/docs/source/en/get_started/installation.md index 672fd8ae03a4..6fc4ce2c922a 100644 --- a/docs/source/en/get_started/installation.md +++ b/docs/source/en/get_started/installation.md @@ -4,6 +4,8 @@ Requirements: - PyTorch >= 1.11 (PyTorch 2.x in progress) - Python >= 3.7 - CUDA >= 11.0 +- [NVIDIA GPU Compute Capability](https://developer.nvidia.com/cuda-gpus) >= 7.0 (V100/RTX20 and higher) +- Linux OS If you encounter any problem about installation, you may want to raise an [issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose) in this repository. @@ -27,7 +29,7 @@ CUDA_EXT=1 pip install colossalai ## Download From Source -> The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. :) +> The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. ```shell git clone https://github.com/hpcaitech/ColossalAI.git @@ -37,14 +39,29 @@ cd ColossalAI pip install -r requirements/requirements.txt # install colossalai -pip install . +CUDA_EXT=1 pip install . ``` -If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer): +If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer), just don't specify the `CUDA_EXT`: ```shell -CUDA_EXT=1 pip install . +pip install . ``` +For Users with CUDA 10.2, you can still build ColossalAI from source. However, you need to manually download the cub library and copy it to the corresponding directory. + +```bash +# clone the repository +git clone https://github.com/hpcaitech/ColossalAI.git +cd ColossalAI + +# download the cub library +wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip +unzip 1.8.0.zip +cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ + +# install +CUDA_EXT=1 pip install . +``` diff --git a/docs/source/en/get_started/run_demo.md b/docs/source/en/get_started/run_demo.md index f47bdbbd62fc..1ce185e26db0 100644 --- a/docs/source/en/get_started/run_demo.md +++ b/docs/source/en/get_started/run_demo.md @@ -7,19 +7,18 @@ can also run on systems with only one GPU. Quick demos showing how to use Coloss ## Single GPU Colossal-AI can be used to train deep learning models on systems with only one GPU and achieve baseline -performances. We provided an example to [train ResNet on CIFAR10 dataset](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/resnet) -with only one GPU. You can find the example in [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples). +performances. We provided an example to [train ResNet on CIFAR10 dataset](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/resnet) +with only one GPU. You can find the example in [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples). Detailed instructions can be found in its `README.md`. ## Multiple GPUs Colossal-AI can be used to train deep learning models on distributed systems with multiple GPUs and accelerate the -training process drastically by applying efficient parallelization techniques. When we have several parallelism for you -to try out. +training process drastically by applying efficient parallelization techniques. When we have several parallelism for you to try out. #### 1. data parallel -You can use the same [ResNet example](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/resnet) as the +You can use the same [ResNet example](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/resnet) as the single-GPU demo above. By setting `--nproc_per_node` to be the number of GPUs you have on your machine, the example is turned into a data parallel example. @@ -27,17 +26,19 @@ is turned into a data parallel example. Hybrid parallel includes data, tensor, and pipeline parallelism. In Colossal-AI, we support different types of tensor parallelism (i.e. 1D, 2D, 2.5D and 3D). You can switch between different tensor parallelism by simply changing the configuration -in the `config.py`. You can follow the [GPT example](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/language/gpt). +in the `config.py`. You can follow the [GPT example](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). Detailed instructions can be found in its `README.md`. #### 3. MoE parallel -We provided [an example of WideNet](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/widenet) to demonstrate +We provided [an example of ViT-MoE](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/moe) to demonstrate MoE parallelism. WideNet uses mixture of experts (MoE) to achieve better performance. More details can be found in [Tutorial: Integrate Mixture-of-Experts Into Your Model](../advanced_tutorials/integrate_mixture_of_experts_into_your_model.md) #### 4. sequence parallel Sequence parallel is designed to tackle memory efficiency and sequence length limit problems in NLP tasks. We provided -[an example of BERT](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/language/bert/sequene_parallel) in -[ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples). You can follow the `README.md` to execute the code. +[an example of BERT](https://github.com/hpcaitech/ColossalAI/tree/main/examples/tutorial/sequence_parallel) in +[ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples). You can follow the `README.md` to execute the code. + + diff --git a/docs/source/zh-Hans/advanced_tutorials/add_your_parallel.md b/docs/source/zh-Hans/advanced_tutorials/add_your_parallel.md index 4825a6fa1d6c..059eb014affd 100644 --- a/docs/source/zh-Hans/advanced_tutorials/add_your_parallel.md +++ b/docs/source/zh-Hans/advanced_tutorials/add_your_parallel.md @@ -48,7 +48,7 @@ Colossal-AI 为用户提供了一个全局 context，使他们能够轻松地管 world_size: int, config: Config, data_parallel_size: int, - pipeline_parlalel_size: int, + pipeline_parallel_size: int, tensor_parallel_size: int, arg1, arg2): diff --git a/docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md b/docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md index 456878caa147..8ed9a1e43cdd 100644 --- a/docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md +++ b/docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md @@ -9,44 +9,24 @@ - [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) - [Go Wider Instead of Deeper](https://arxiv.org/abs/2107.11817) -（中文版教程将会在近期提供） - ## Introduction -Since the advent of Switch Transformer, the AI community has found Mixture of Experts (MoE) a useful technique to enlarge the capacity of deep learning models. - -Colossal-AI provides an early access version of parallelism specifically designed for MoE models. -The most prominent advantage of MoE in Colossal-AI is convenience. -We aim to help our users to easily combine MoE with model parallelism and data parallelism. - -However, the current implementation has two main drawbacks now. -The first drawback is its poor efficiency in large batch size and long sequence length training. -The second drawback is incompatibility with tensor parallelism. -We are working on system optimization to overcome the training efficiency problem. -The compatibility problem with tensor parallelism requires more adaptation, and we will tackle this issue in the future. - -Here, we will introduce how to use MoE with model parallelism and data parallelism. - -## Table of Content -In this tutorial we will cover: -1. Set up MoE running environment -2. Create MoE layer -3. Train your model +自从`Switch Transformer`出现以来，人工智能社区发现专家混合 (MoE) 是一种扩大深度学习模型容量的有用技术。 +Colossal-AI 提供了专为MoE模型设计的并行性的早期访问版本。Colossal-AI中MoE最突出的优势就是方便。我们的目标是帮助我们的用户轻松地将MoE与模型并行性和数据并行性结合起来。 +但是，当前的实施现在有两个主要缺点。第一个缺点是它在大批量和长序列长度训练中效率低下。第二个缺点是与张量并行性不兼容。我们正在致力于系统优化，以克服训练效率问题。与张量并行的兼容性问题需要更多的适应，我们将在未来解决这个问题。 +在这里，我们将介绍如何使用具有模型并行性和数据并行性的 MoE。 -We provided the [example code](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/widenet) for this tutorial in [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples). -This example uses [WideNet](https://arxiv.org/abs/2107.11817) as an example of MoE-based model. +## 目录 +在本教程中，我们将介绍： +1. [搭建MoE运行环境](#搭建moe运行环境) +2. [创建MoE层](#创建moe层) +3. [定义训练模型](#训练模型) +我们提供[示例](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/widenet)， 详细介绍请参考 [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples). +该示例使用 [WideNet](https://arxiv.org/abs/2107.11817) 作为基于 MoE 的模型的示例. -## Set up MoE running environment -In your project folder, create a `config.py`. - -This file is to specify some features you may want to use to train your model. -In order to enable MoE, you need to add a dict called parallel and specify the value of key moe. -You can assign a value for the key size of moe, which represents the model parallel size of experts (i.e. the number of experts in one group to parallelize training). - -For example, if the size is 4, 4 processes will be assigned to 4 consecutive GPUs and these 4 processes form a moe model parallel group. -Each process on the 4 GPUs will only get a portion of experts. Increasing the model parallel size will reduce communication cost, but increase computation cost in each GPU and activation cost in memory. -The total data parallel size is auto-detected and set as the number of GPUs by default. +## 搭建MoE运行环境 +在您的项目文件夹中，创建`config.py`文件。在该文件中，您可以指定希望用于训练模型的一些功能。为了启用 MoE，您需要在`config.py`中定义`parallel`字段，并指定`moe`的值。`moe`表示一组moe并行化训练组的并行大小。例如，`moe`设置为4，则4个进程将分配给4个连续的GPU，这4个进程组成一个moe模型并行组。每个进程只会得到一部分专家。增加mo e并行的大小将降低通信成本，但会增加每个GPU的计算成本和内存中activation的存储成本。总的数据并行的大小是自动检测的，默认情况下设置为GPU的数量。 ```python MOE_MODEL_PARALLEL_SIZE = ... @@ -55,37 +35,29 @@ parallel = dict( ) ``` -If `MOE_MODEL_PARALLEL_SIZE = E` and set the number of experts as `E` where `E` is a constant number, the process flow of forward pass of a transformer encoder in a model parallel group is shown below. +如果`MOE_MODEL_PARALLEL_SIZE = E`，即设置专家的总数为`E`（`E`为一个常数）。在模型并行中，transformer编码器中前向部分的处理流程如下图所示。
      MoE Transformer, image source: GShard
      -Since all experts are allocated to all GPUs in a model parallel group and a GPU only owns a portion of experts, -original data parallel groups are no longer correct for the parameters of experts during gradient handling in backward pass anymore. -So we create a new kind of parallel group called moe data parallel group. -The difference among different kinds of parallel group, when the configuration is set as `WORLD_SIZE=4`, -`MOE_MODEL_PARALLEL_SIZE=2`, is shown here. +所有专家都分配给模型并行组中的GPU，每一个GPU只拥有一部分专家，原始数据并行组在反向传递的梯度处理期间不再适用于专家参数。所以我们创建了一个新的并行组，叫做moe数据并行组。当配置设置为`WORLD_SIZE=4`，`MOE_MODEL_PARALLEL_SIZE=2`时，两个并行组的区别如下图所示。
      -
      MoE process group
      +
      MoE并行处理
      +至于梯度处理，我们提供了`MoeGradientHandler`来all-reduce模型的每个参数。如果您使用`colossalai.initialize`函数创建您的训练引擎，MoE梯度处理程序将自动添加到您的引擎中。否则，你应该自己处理梯度。MoE运行环境的所有参数都保存在`colossalai.global_variables.moe_env`中。您可以访问您的配置参数来检查您的设置是否正确。 -As for gradient handling, we provide MoeGradientHandler to all-reduce every parameter of the model. -If you use `colossalai.initialize` function to create your training engine, the MoE gradient handler will be added to your engine automatically. -Otherwise, you should take care of gradient by yourself. -All parameters of MoE running environment are stored in colossalai.global_variables.moe_env. -You can access your configuration parameters to check whether your setup is correct. ```python from colossalai.global_variables import moe_env ``` -## Create MoE layer -You can create a MoE layer from `colossalai.nn.moe`. -But before doing that, you should set up random seeds for all processes like this. +## 创建MoE层 + +您可以从`colossalai.nn.moe`创建MoE层。但在此之前，您应该为所有进程设置随机种子。 ```python from colossalai.context.random import moe_set_seed @@ -95,10 +67,7 @@ moe_set_seed(42) model = Widenet(num_experts=4, capacity_factor=1.2) ``` -`moe_set_seed` will set different seed for different processes in a moe model parallel group. -This helps initialize parameters in experts. -Then create an instance of experts and an instance of router. -Here is the example in model zoo. +`moe_set_seed` 会为一个moe模型并行组中的不同进程设置不同的种子（这有助于在专家中初始化参数），创建一个专家实例和一个路由器实例，示例如下。 ```python from colossalai.nn.layer.moe import Experts, MoeLayer, Top2Router, NormalNoiseGenerator @@ -118,16 +87,11 @@ ffn=MoeLayer(dim_model=d_model, num_experts=num_experts, router=shared_router, experts=shared_experts) ``` -Inside the initialization of Experts, the local expert number of each GPU will be calculated automatically. You just need to specify the class of each expert and its parameters used in its initialization. As for routers, we have provided top1 router and top2 router. You can find them in colossalai.nn.layer.moe. After creating the instance of experts and router, the only thing initialized in Moelayer is gate module. More definitions of each class can be found in our API document and code. - +在Experts的初始化中，会自动计算每个GPU的本地expert数量，您只需指定每个专家的类型及其在初始化时使用的参数。此外，我们提供了`Top1Router`和`Top2Router`，您可以在`colossalai.nn.layer.moe` 找到它们。在创建experts和router的实例时，`Moelayer`只初始化了`gate`模块，类型的更多详细信息您可以参考我们的API文档和代码。 -## Train Your Model -Do not to forget to use `colossalai.initialize` function in `colosalai` to add gradient handler for the engine. -We handle the back-propagation of MoE models for you. -In `colossalai.initialize`, we will automatically create a `MoeGradientHandler` object to process gradients. -You can find more information about the handler `MoeGradientHandler` in colossal directory. +## 定义训练模型 -The loss criterion should be wrapped by `Moeloss` to add auxiliary loss of MoE. Example is like this. +使用colossalai中的`colossalai.initialize`函数为引擎添加梯度处理程序以处理 MoE模型的反向传播。在 `colossalai.initialize` 中，我们会自动创建一个`MoeGradientHandler`对象来处理梯度。您可以在colossal目录中找到有关`MoeGradientHandler`的更多信息。为了添加MoE的相关损失处理，损失函数应使用`Moeloss`封装，示例如下。 ```python criterion = MoeLoss( aux_weight=0.01, @@ -135,6 +99,6 @@ criterion = MoeLoss( label_smoothing=0.1 ) ``` +最后，您只需使用 `colossalai` 中的`trainer`或`engine`进行训练即可。 -Finally, just use trainer or engine in `colossalai` to do your training. -Otherwise, you should take care of gradient by yourself. + diff --git a/docs/source/zh-Hans/advanced_tutorials/meet_gemini.md b/docs/source/zh-Hans/advanced_tutorials/meet_gemini.md index 2bf0a9c98c3f..594823862de1 100644 --- a/docs/source/zh-Hans/advanced_tutorials/meet_gemini.md +++ b/docs/source/zh-Hans/advanced_tutorials/meet_gemini.md @@ -8,21 +8,21 @@ ## 用法 -目前Gemini支持和ZeRO并行方式兼容，它的使用方法很简单，在训练策略的配置文件里设置zero的model_config属性tensor_placement_policy='auto' - -``` -zero = dict( - model_config=dict( - reduce_scatter_bucket_size_mb=25, - fp32_reduce_scatter=False, - gradient_predivide_factor=1.0, - tensor_placement_policy="auto", - shard_strategy=TensorShardStrategy(), - ... - ), - optimizer_config=dict( - ... - ) +目前Gemini支持和ZeRO并行方式兼容，它的使用方法很简单：使用booster将`GeminiPlugin`中的特性注入到训练组件中。更多`booster`介绍请参考[booster使用](../basics/booster_api.md)。 + +```python +from torchvision.models import resnet18 +from colossalai.booster import Booster +from colossalai.zero import ColoInitContext +from colossalai.booster.plugin import GeminiPlugin +plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, max_norm=1.0, initial_scale=2**5) +booster = Booster(plugin=plugin) +ctx = ColoInitContext() +with ctx: + model = resnet18() +optimizer = HybridAdam(model.parameters(), lr=1e-3) +criterion = lambda x: x.mean() +model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) ) ``` @@ -48,7 +48,7 @@ zero = dict(
      -ColossalAI设计了Gemini，就像双子星一样，它管理CPU和GPU二者内存空间。它可以让张量在训练过程中动态分布在CPU-GPU的存储空间内，从而让模型训练突破GPU的内存墙。内存管理器由两部分组成，分别是MemStatsCollector(MSC)和StatefuleTensorMgr(STM)。 +ColossalAI设计了Gemini，就像双子星一样，它管理CPU和GPU二者内存空间。它可以让张量在训练过程中动态分布在CPU-GPU的存储空间内，从而让模型训练突破GPU的内存墙。内存管理器由两部分组成，分别是MemStatsCollector(MSC)和StatefulTensorMgr(STM)。 我们利用了深度学习网络训练过程的迭代特性。我们将迭代分为warmup和non-warmup两个阶段，开始时的一个或若干迭代步属于预热阶段，其余的迭代步属于正式阶段。在warmup阶段我们为MSC收集信息，而在non-warmup阶段STM入去MSC收集的信息来移动tensor，以达到最小化CPU-GPU数据移动volume的目的。 @@ -75,7 +75,7 @@ STM管理所有model data tensor的信息。在模型的构造过程中，Coloss 我们在算子的开始和结束计算时，触发内存采样操作，我们称这个时间点为**采样时刻（sampling moment)**，两个采样时刻之间的时间我们称为**period**。计算过程是一个黑盒，由于可能分配临时buffer，内存使用情况很复杂。但是，我们可以较准确的获取period的系统最大内存使用。非模型数据的使用可以通过两个统计时刻之间系统最大内存使用-模型内存使用获得。 -我们如何设计采样时刻呢。我们选择preOp的model data layout adjust之前。如下图所示。我们采样获得上一个period的system memory used，和下一个period的model data memoy used。并行策略会给MSC的工作造成障碍。如图所示，比如对于ZeRO或者Tensor Parallel，由于Op计算前需要gather模型数据，会带来额外的内存需求。因此，我们要求在模型数据变化前进行采样系统内存，这样在一个period内，MSC会把preOp的模型变化内存捕捉。比如在period 2-3内，我们考虑的tensor gather和shard带来的内存变化。 +我们如何设计采样时刻呢。我们选择preOp的model data layout adjust之前。如下图所示。我们采样获得上一个period的system memory used，和下一个period的model data memory used。并行策略会给MSC的工作造成障碍。如图所示，比如对于ZeRO或者Tensor Parallel，由于Op计算前需要gather模型数据，会带来额外的内存需求。因此，我们要求在模型数据变化前进行采样系统内存，这样在一个period内，MSC会把preOp的模型变化内存捕捉。比如在period 2-3内，我们考虑的tensor gather和shard带来的内存变化。 尽管可以将采样时刻放在其他位置，比如排除gather buffer的变动新信息，但是会给造成麻烦。不同并行方式Op的实现有差异，比如对于Linear Op，Tensor Parallel中gather buffer的分配在Op中。而对于ZeRO，gather buffer的分配是在PreOp中。将放在PreOp开始时采样有利于将两种情况统一。 @@ -94,3 +94,5 @@ MSC的重要职责是在调整tensor layout位置，比如在上图S2时刻， 在non-warmup阶段，我们需要利用预热阶段采集的非模型数据内存信息，预留出下一个Period在计算设备上需要的峰值内存，这需要我们移动出一些模型张量。 为了避免频繁在CPU-GPU换入换出相同的tensor，引起类似[cache thrashing](https://en.wikipedia.org/wiki/Thrashing_(computer_science))的现象。我们利用DNN训练迭代特性，设计了OPT cache换出策略。具体来说，在warmup阶段，我们记录每个tensor被计算设备需要的采样时刻。如果我们需要驱逐一些HOLD tensor，那么我们选择在本设备上最晚被需要的tensor作为受害者。 + + diff --git a/docs/source/zh-Hans/advanced_tutorials/opt_service.md b/docs/source/zh-Hans/advanced_tutorials/opt_service.md index a213584fd41d..1f8324a53ecb 100644 --- a/docs/source/zh-Hans/advanced_tutorials/opt_service.md +++ b/docs/source/zh-Hans/advanced_tutorials/opt_service.md @@ -52,7 +52,7 @@ export CHECKPOINT_DIR="your_opt_checkpoint_path" # the ${CONFIG_DIR} must contain a server.sh file as the entry of service export CONFIG_DIR="config_file_path" -docker run --gpus all --rm -it -p 8020:8020 -v ${CHECKPOINT_DIR}:/model_checkpoint -v ${CONFIG_DIR}:/config --ipc=host energonai:lastest +docker run --gpus all --rm -it -p 8020:8020 -v ${CHECKPOINT_DIR}:/model_checkpoint -v ${CONFIG_DIR}:/config --ipc=host energonai:latest ``` 接下来，您就可以在您的浏览器中打开 `https://[IP-ADDRESS]:8020/docs#` 进行测试。 diff --git a/docs/source/zh-Hans/advanced_tutorials/parallelize_your_training_like_Megatron.md b/docs/source/zh-Hans/advanced_tutorials/parallelize_your_training_like_Megatron.md index f3c6247c38e4..3f85d50454ae 100644 --- a/docs/source/zh-Hans/advanced_tutorials/parallelize_your_training_like_Megatron.md +++ b/docs/source/zh-Hans/advanced_tutorials/parallelize_your_training_like_Megatron.md @@ -126,16 +126,16 @@ for mn, module in model.named_modules(): if 'mlp.c_fc' in mn: if 'weight' in pn or 'bias' in pn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice # keep the shape of the output from c_fc param.compute_spec.set_output_replicate(False) elif 'mlp.c_proj' in mn: if 'weight' in pn: split_param_row_tp1d(param, pg) # row slice elif 'wte' in mn or 'wpe' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif 'c_attn' in mn or 'c_proj' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice ``` 修改后的模型如下图所示。 @@ -159,13 +159,13 @@ for mn, module in model.named_modules(): 在我们最新示例中还定义了一个Gemini + ZeRO DDP 的模型从而减小开销，提升效率。这一部分的详细内容可以参考[ZeRO](../features/zero_with_chunk.md)，你可以将这两部分内容结合起来看从而理解我们整个训练流程： ```python -def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: str = "auto"): +def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placement_policy: str = "auto"): from colossalai.nn.parallel import GeminiDDP model = GeminiDDP(model, device=get_current_device(), - placement_policy=placememt_policy, + placement_policy=placement_policy, pin_memory=True, - search_range_mb=32) + search_range_m=32) return model ``` @@ -174,3 +174,6 @@ def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: 我们做的上述优化让我们可以在单GPU上训练GPT-2模型，只需要将`run.sh`中设置参数`GPUNUM`=1，再运行文件时就可以在单个GPU上完成模型的训练。 GPT-2 示例在[Train GPT with Colossal-AI](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). 获得。 + + + diff --git a/docs/source/zh-Hans/advanced_tutorials/train_vit_with_hybrid_parallelism.md b/docs/source/zh-Hans/advanced_tutorials/train_vit_with_hybrid_parallelism.md index 6dc5eccf4421..5ad08392049e 100644 --- a/docs/source/zh-Hans/advanced_tutorials/train_vit_with_hybrid_parallelism.md +++ b/docs/source/zh-Hans/advanced_tutorials/train_vit_with_hybrid_parallelism.md @@ -150,7 +150,7 @@ Colossal-AI 提供了自己的优化器、损失函数和学习率调度器。Py optimizer = colossalai.nn.Lamb(model.parameters(), lr=1.8e-2, weight_decay=0.1) # build loss criterion = torch.nn.CrossEntropyLoss() -# lr_scheduelr +# lr_scheduler lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=gpc.config.NUM_EPOCHS) ``` @@ -477,7 +477,7 @@ def build_cifar(batch_size): return train_dataloader, test_dataloader -# craete dataloaders +# create dataloaders train_dataloader , test_dataloader = build_cifar() # create loss function criterion = CrossEntropyLoss(label_smoothing=0.1) @@ -492,7 +492,7 @@ lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, #### 启动 Colossal-AI 引擎 ```python -# intiailize +# initialize engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model=model, optimizer=optimizer, criterion=criterion, diff --git a/docs/source/zh-Hans/basics/booster_api.md b/docs/source/zh-Hans/basics/booster_api.md new file mode 100644 index 000000000000..b2235b73bca1 --- /dev/null +++ b/docs/source/zh-Hans/basics/booster_api.md @@ -0,0 +1,83 @@ +# booster 使用 + +作者: [Mingyan Jiang](https://github.com/jiangmingyan) [Jianghai Chen](https://github.com/CjhHa1) + +**预备知识:** + +- [分布式训练](../concepts/distributed_training.md) +- [Colossal-AI 总览](../concepts/colossalai_overview.md) + +**示例代码** + + + +- [使用 booster 训练](https://github.com/hpcaitech/ColossalAI/blob/main/examples/tutorial/new_api/cifar_resnet/README.md) + +## 简介 + +在我们的新设计中， `colossalai.booster` 代替 `colossalai.initialize` 将特征(例如，模型、优化器、数据加载器)无缝注入您的训练组件中。 使用 booster API, 您可以更友好地将我们的并行策略整合到待训练模型中. 调用 `colossalai.booster` 是您进入训练循环前的基本操作。 +在下面的章节中，我们将介绍 `colossalai.booster` 是如何工作的以及使用时我们要注意的细节。 + +### Booster 插件 + +Booster 插件是管理并行配置的重要组件（eg：gemini 插件封装了 gemini 加速方案）。目前支持的插件如下： + +**_GeminiPlugin:_** GeminiPlugin 插件封装了 gemini 加速解决方案，即基于块内存管理的 ZeRO 优化方案。 + +**_TorchDDPPlugin:_** TorchDDPPlugin 插件封装了Pytorch的DDP加速方案，实现了模型级别的数据并行，可以跨多机运行。 + +**_LowLevelZeroPlugin:_** LowLevelZeroPlugin 插件封装了零冗余优化器的 1/2 阶段。阶段 1：切分优化器参数，分发到各并发进程或并发 GPU 上。阶段 2：切分优化器参数及梯度，分发到各并发进程或并发 GPU 上。 + +**_TorchFSDPPlugin:_** TorchFSDPPlugin封装了 Pytorch的FSDP加速方案，可以用于零冗余优化器数据并行（ZeroDP）的训练。 + + +### Booster 接口 + + + +{{ autodoc:colossalai.booster.Booster }} + +## 使用方法及示例 + +在使用 colossalai 训练时，首先需要在训练脚本的开头启动分布式环境，并创建需要使用的模型、优化器、损失函数、数据加载器等对象。之后，调用`colossalai.booster` 将特征注入到这些对象中，您就可以使用我们的 booster API 去进行您接下来的训练流程。 + +以下是一个伪代码示例，将展示如何使用我们的 booster API 进行模型训练: + +```python +import torch +from torch.optim import SGD +from torchvision.models import resnet18 + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin + +def train(): + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + model = resnet18() + criterion = lambda x: x.mean() + optimizer = SGD((model.parameters()), lr=0.001) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1) + model, optimizer, criterion, _, scheduler = booster.boost(model, optimizer, criterion, lr_scheduler=scheduler) + + x = torch.randn(4, 3, 224, 224) + x = x.to('cuda') + output = model(x) + loss = criterion(output) + booster.backward(loss, optimizer) + optimizer.clip_grad_by_norm(1.0) + optimizer.step() + scheduler.step() + + save_path = "./model" + booster.save_model(model, save_path, True, True, "", 10, use_safetensors=use_safetensors) + + new_model = resnet18() + booster.load_model(new_model, save_path) +``` + +[更多的设计细节请参考](https://github.com/hpcaitech/ColossalAI/discussions/3046) + + diff --git a/docs/source/zh-Hans/basics/booster_checkpoint.md b/docs/source/zh-Hans/basics/booster_checkpoint.md new file mode 100644 index 000000000000..4ed049dcf44f --- /dev/null +++ b/docs/source/zh-Hans/basics/booster_checkpoint.md @@ -0,0 +1,47 @@ +# Booster Checkpoint + +作者: [Hongxin Liu](https://github.com/ver217) + +**前置教程:** +- [Booster API](./booster_api.md) + +## 引言 + +我们在之前的教程中介绍了 [Booster API](./booster_api.md)。在本教程中，我们将介绍如何使用 booster 保存和加载 checkpoint。 + +## 模型 Checkpoint + +{{ autodoc:colossalai.booster.Booster.save_model }} + +模型在保存前必须被 `colossalai.booster.Booster` 加速。 `checkpoint` 是要保存的 checkpoint 的路径。 如果 `shard=False`，它就是文件。 否则, 它就是文件夹。如果 `shard=True`，checkpoint 将以分片方式保存。当 checkpoint 太大而无法保存在单个文件中时，这很有用。我们的分片 checkpoint 格式与 [huggingface/transformers](https://github.com/huggingface/transformers) 兼容。 + +{{ autodoc:colossalai.booster.Booster.load_model }} + +模型在加载前必须被 `colossalai.booster.Booster` 加速。它会自动检测 checkpoint 格式，并以相应的方式加载。 + +## 优化器 Checkpoint + + +{{ autodoc:colossalai.booster.Booster.save_optimizer }} + +优化器在保存前必须被 `colossalai.booster.Booster` 加速。 + +{{ autodoc:colossalai.booster.Booster.load_optimizer }} + +优化器在加载前必须被 `colossalai.booster.Booster` 加速。 + +## 学习率调度器 Checkpoint + +{{ autodoc:colossalai.booster.Booster.save_lr_scheduler }} + +学习率调度器在保存前必须被 `colossalai.booster.Booster` 加速。 `checkpoint` 是 checkpoint 文件的本地路径. + +{{ autodoc:colossalai.booster.Booster.load_lr_scheduler }} + +学习率调度器在加载前必须被 `colossalai.booster.Booster` 加速。 `checkpoint` 是 checkpoint 文件的本地路径. + +## Checkpoint 设计 + +有关 Checkpoint 设计的更多详细信息，请参见我们的讨论 [A Unified Checkpoint System Design](https://github.com/hpcaitech/ColossalAI/discussions/3339). + + diff --git a/docs/source/zh-Hans/basics/booster_plugins.md b/docs/source/zh-Hans/basics/booster_plugins.md new file mode 100644 index 000000000000..0f355c43901c --- /dev/null +++ b/docs/source/zh-Hans/basics/booster_plugins.md @@ -0,0 +1,73 @@ +# Booster 插件 + +作者: [Hongxin Liu](https://github.com/ver217) + +**前置教程:** +- [Booster API](./booster_api.md) + +## 引言 + +正如 [Booster API](./booster_api.md) 中提到的，我们可以使用 booster 插件来自定义并行训练。在本教程中，我们将介绍如何使用 booster 插件。 + +我们现在提供以下插件: + +- [Low Level Zero 插件](#low-level-zero-plugin): 它包装了 `colossalai.zero.low_level.LowLevelZeroOptimizer`，可用于使用 Zero-dp 训练模型。它仅支持 Zero 阶段1和阶段2。 +- [Gemini 插件](#gemini-plugin): 它包装了 [Gemini](../features/zero_with_chunk.md)，Gemini 实现了基于Chunk内存管理和异构内存管理的 Zero-3。 +- [Torch DDP 插件](#torch-ddp-plugin): 它包装了 `torch.nn.parallel.DistributedDataParallel` 并且可用于使用数据并行训练模型。 +- [Torch FSDP 插件](#torch-fsdp-plugin): 它包装了 `torch.distributed.fsdp.FullyShardedDataParallel` 并且可用于使用 Zero-dp 训练模型。 + +更多插件即将推出。 + +## 插件 + +### Low Level Zero 插件 + +该插件实现了 Zero-1 和 Zero-2（使用/不使用 CPU 卸载），使用`reduce`和`gather`来同步梯度和权重。 + +Zero-1 可以看作是 Torch DDP 更好的替代品，内存效率更高，速度更快。它可以很容易地用于混合并行。 + +Zero-2 不支持局部梯度累积。如果您坚持使用，虽然可以积累梯度，但不能降低通信成本。也就是说，同时使用流水线并行和 Zero-2 并不是一个好主意。 + +{{ autodoc:colossalai.booster.plugin.LowLevelZeroPlugin }} + +我们已经测试了一些主流模型的兼容性，可能不支持以下模型： + +- `timm.models.convit_base` +- dlrm and deepfm models in `torchrec` +- `diffusers.VQModel` +- `transformers.AlbertModel` +- `transformers.AlbertForPreTraining` +- `transformers.BertModel` +- `transformers.BertForPreTraining` +- `transformers.GPT2DoubleHeadsModel` + +兼容性问题将在未来修复。 + +> ⚠ 该插件现在只能加载自己保存的且具有相同进程数的优化器 Checkpoint。这将在未来得到解决。 + +### Gemini 插件 + +这个插件实现了基于Chunk内存管理和异构内存管理的 Zero-3。它可以训练大型模型而不会损失太多速度。它也不支持局部梯度累积。更多详细信息，请参阅 [Gemini 文档](../features/zero_with_chunk.md). + +{{ autodoc:colossalai.booster.plugin.GeminiPlugin }} + + +### Torch DDP 插件 + +更多详细信息，请参阅 [Pytorch 文档](https://pytorch.org/docs/main/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel). + +{{ autodoc:colossalai.booster.plugin.TorchDDPPlugin }} + +### Torch FSDP 插件 + +> ⚠ 如果 torch 版本低于 1.12.0，此插件将不可用。 + +> ⚠ 该插件现在还不支持保存/加载分片的模型 checkpoint。 + +> ⚠ 该插件现在还不支持使用了multi params group的optimizer。 + +更多详细信息，请参阅 [Pytorch 文档](https://pytorch.org/docs/main/fsdp.html). + +{{ autodoc:colossalai.booster.plugin.TorchFSDPPlugin }} + + diff --git a/docs/source/zh-Hans/basics/colotensor_concept.md b/docs/source/zh-Hans/basics/colotensor_concept.md index cac5b9a4b40d..ab2413e990f7 100644 --- a/docs/source/zh-Hans/basics/colotensor_concept.md +++ b/docs/source/zh-Hans/basics/colotensor_concept.md @@ -2,6 +2,8 @@ Author: [Jiarui Fang](https://github.com/feifeibear), [Hongxin Liu](https://github.com/ver217) and [Haichen Huang](https://github.com/1SAA) +> ⚠️ 此页面上的信息已经过时并将被废弃。 + **Prerequisite:** - [Colossal-AI Overview](../concepts/colossalai_overview.md) - [Distributed Training](../concepts/distributed_training.md) @@ -51,18 +53,18 @@ ColoTensor 包含额外的属性[ColoTensorSpec](https://colossalai.readthedocs. ## Example -让我们看一个例子。 使用 tp_degree=4, dp_dgree=2 在 8 个 GPU 上初始化并Shard一个ColoTensor。 然后tensor被沿着 TP 进程组中的最后一个维度进行分片。 最后，我们沿着 TP 进程组中的第一个维度（dim 0）对其进行重新Shard。 我们鼓励用户运行代码并观察每个张量的形状。 +让我们看一个例子。 使用 tp_degree=4, dp_degree=2 在 8 个 GPU 上初始化并Shard一个ColoTensor。 然后tensor被沿着 TP 进程组中的最后一个维度进行分片。 最后，我们沿着 TP 进程组中的第一个维度（dim 0）对其进行重新Shard。 我们鼓励用户运行代码并观察每个张量的形状。 ```python import torch import torch.multiprocessing as mp -from colossalai.utils import free_port, print_rank_0 +from colossalai.utils import print_rank_0 from functools import partial import colossalai from colossalai.tensor import ProcessGroup, ColoTensor, ColoTensorSpec, ShardSpec, ComputeSpec, ComputePattern -from colossalai.utils import free_port +from colossalai.testing import spawn import torch @@ -84,8 +86,7 @@ def run_dist_tests(rank, world_size, port): print_rank_0(f"shape {t1.shape}, {t1.data}") def test_dist_cases(world_size): - run_func = partial(run_dist_tests, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist_tests, world_size) if __name__ == '__main__': test_dist_cases(4) diff --git a/docs/source/zh-Hans/basics/configure_parallelization.md b/docs/source/zh-Hans/basics/configure_parallelization.md index eb4b38f48ddb..0c2a66572d60 100644 --- a/docs/source/zh-Hans/basics/configure_parallelization.md +++ b/docs/source/zh-Hans/basics/configure_parallelization.md @@ -2,6 +2,8 @@ 作者: Shenggui Li, Siqi Mai +> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Booster插件](../basics/booster_plugins.md)页面查阅更新。 + **预备知识:** - [分布式训练](../concepts/distributed_training.md) - [并行技术](../concepts/paradigms_of_parallelism.md) diff --git a/docs/source/zh-Hans/basics/define_your_config.md b/docs/source/zh-Hans/basics/define_your_config.md index d7e49cbf23de..720e75805e8d 100644 --- a/docs/source/zh-Hans/basics/define_your_config.md +++ b/docs/source/zh-Hans/basics/define_your_config.md @@ -2,6 +2,8 @@ 作者: Guangyang Lu, Shenggui Li, Siqi Mai +> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Booster API](../basics/booster_api.md)页面查阅更新。 + **预备知识:** - [分布式训练](../concepts/distributed_training.md) - [Colossal-AI 总览](../concepts/colossalai_overview.md) diff --git a/docs/source/zh-Hans/basics/engine_trainer.md b/docs/source/zh-Hans/basics/engine_trainer.md index a7519bfca14f..a35bd87c44e1 100644 --- a/docs/source/zh-Hans/basics/engine_trainer.md +++ b/docs/source/zh-Hans/basics/engine_trainer.md @@ -2,6 +2,8 @@ 作者: Shenggui Li, Siqi Mai +> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Booster API](../basics/booster_api.md)页面查阅更新。 + **预备知识:** - [初始化功能](./initialize_features.md) diff --git a/docs/source/zh-Hans/basics/initialize_features.md b/docs/source/zh-Hans/basics/initialize_features.md index 67ea114b42b2..1c28d658e1bc 100644 --- a/docs/source/zh-Hans/basics/initialize_features.md +++ b/docs/source/zh-Hans/basics/initialize_features.md @@ -2,6 +2,8 @@ 作者: Shenggui Li, Siqi Mai +> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Booster API](../basics/booster_api.md)页面查阅更新。 + **预备知识:** - [分布式训练](../concepts/distributed_training.md) - [Colossal-AI 总览](../concepts/colossalai_overview.md) diff --git a/docs/source/zh-Hans/basics/launch_colossalai.md b/docs/source/zh-Hans/basics/launch_colossalai.md index ca927de578d5..39b09deae085 100644 --- a/docs/source/zh-Hans/basics/launch_colossalai.md +++ b/docs/source/zh-Hans/basics/launch_colossalai.md @@ -74,7 +74,7 @@ import colossalai args = colossalai.get_default_parser().parse_args() # launch distributed environment -colossalai.launch(config=, +colossalai.launch(config=args.config, rank=args.rank, world_size=args.world_size, host=args.host, @@ -93,12 +93,21 @@ PyTorch自带的启动器需要在每个节点上都启动命令才能启动多 首先，我们需要在代码里指定我们的启动方式。由于这个启动器是PyTorch启动器的封装，那么我们自然而然应该使用`colossalai.launch_from_torch`。 分布式环境所需的参数，如 rank, world size, host 和 port 都是由 PyTorch 启动器设置的，可以直接从环境变量中读取。 +config.py +```python +BATCH_SIZE = 512 +LEARNING_RATE = 3e-3 +WEIGHT_DECAY = 0.3 +NUM_EPOCHS = 2 +``` +train.py ```python import colossalai colossalai.launch_from_torch( - config=, + config="./config.py", ) +... ``` 接下来，我们可以轻松地在终端使用`colossalai run`来启动训练。下面的命令可以在当前机器上启动一个4卡的训练任务。 diff --git a/docs/source/zh-Hans/basics/model_checkpoint.md b/docs/source/zh-Hans/basics/model_checkpoint.md index cec12d451989..a5374b7509c9 100644 --- a/docs/source/zh-Hans/basics/model_checkpoint.md +++ b/docs/source/zh-Hans/basics/model_checkpoint.md @@ -1,7 +1,9 @@ -# 模型检查点 +# 模型Checkpoint 作者 : Guangyang Lu +> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Booster Checkpoint](../basics/booster_checkpoint.md)页面查阅更新。 + **预备知识:** - [Launch Colossal-AI](./launch_colossalai.md) - [Initialize Colossal-AI](./initialize_features.md) @@ -13,9 +15,9 @@ ## 简介 -本教程将介绍如何保存和加载模型检查点。 +本教程将介绍如何保存和加载模型Checkpoint。 -为了充分利用Colossal-AI的强大并行策略，我们需要修改模型和张量，可以直接使用 `torch.save` 或者 `torch.load` 保存或加载模型检查点。在Colossal-AI中，我们提供了应用程序接口实现上述同样的效果。 +为了充分利用Colossal-AI的强大并行策略，我们需要修改模型和张量，可以直接使用 `torch.save` 或者 `torch.load` 保存或加载模型Checkpoint。在Colossal-AI中，我们提供了应用程序接口实现上述同样的效果。 但是，在加载时，你不需要使用与存储相同的保存策略。 @@ -24,7 +26,7 @@ ### 保存 有两种方法可以使用Colossal-AI训练模型，即使用engine或使用trainer。 -**注意我们只保存 `state_dict`.** 因此，在加载检查点时，需要首先定义模型。 +**注意我们只保存 `state_dict`.** 因此，在加载Checkpoint时，需要首先定义模型。 #### 同 engine 保存 diff --git a/docs/source/zh-Hans/concepts/colossalai_overview.md b/docs/source/zh-Hans/concepts/colossalai_overview.md index cfb35e59e64a..8b28baf8e3d5 100755 --- a/docs/source/zh-Hans/concepts/colossalai_overview.md +++ b/docs/source/zh-Hans/concepts/colossalai_overview.md @@ -19,7 +19,7 @@ Colossal-AI 是一个集成的系统，为用户提供一套综合的训练方 1. 准备一个配置文件，指定您要使用的功能和参数。 2. 用 `colossalai.launch` 初始化分布式后端。 -3. 用 `colossalai.initialize` 将训练特征注入您的训练组件（如模型、优化器）中。 +3. 用 `colossalai.booster` 将训练特征注入您的训练组件（如模型、优化器）中。 4. 进行训练和测试. 我们将在`基本教程`部分介绍整个工作流程。 @@ -34,3 +34,5 @@ Colossal-AI 系统将会进一步拓展和优化，包括但不限于: 4. 拓展现有的并行方法 **我们始终欢迎社区的建议和讨论，如果您遇到任何问题，我们将非常愿意帮助您。您可以在GitHub 提 [issue](https://github.com/hpcaitech/ColossalAI/issues) ，或在[论坛](https://github.com/hpcaitech/ColossalAI/discussions)上创建一个讨论主题。** + + diff --git a/docs/source/zh-Hans/features/1D_tensor_parallel.md b/docs/source/zh-Hans/features/1D_tensor_parallel.md index 8f3a3c6209da..4dd45e8783c3 100644 --- a/docs/source/zh-Hans/features/1D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/1D_tensor_parallel.md @@ -7,7 +7,7 @@ - [并行配置](../basics/configure_parallelization.md) **示例代码** -- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_1d.py) +- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **相关论文** - [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://deepakn94.github.io/assets/papers/megatron-sc21.pdf) @@ -17,11 +17,20 @@ 张量并行将模型参数划分到多个设备上，以减少内存负荷。 [Megatron-LM](https://deepakn94.github.io/assets/papers/megatron-sc21.pdf) 介绍了一种高效的一维张量并行化实现。 -让我们以一个线性层为例，它包括一个 GEMM $Y = XA$。 给定2个处理器，我们把列 $A$ 划分为 $[A_1 ~ A_2]$, 并在每个处理器上计算 $Y_i = XA_i$ , which then forms $[Y_1 ~ Y_2] = [XA_1 ~ XA_2]$. This is called a column-parallel fashion. +让我们以一个线性层为例，它包括一个 GEMM $Y = XA$。 给定2个处理器，我们把列 $A$ 划分为 $[A_1 ~ A_2]$, 并在每个处理器上计算 $Y_i = XA_i$ , 然后形成 $[Y_1 ~ Y_2] = [XA_1 ~ XA_2]$. 这被称为列并行方式。 -当第二个线性层 $Z=YB$ 跟随上述列并行层的时候, 我们把 $B$ 划分为 $\left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right]$, +当第二个线性层 $Z=YB$ 跟随上述列并行层的时候, 我们把 $B$ 划分为 +$$ +\left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right] +``` 这就是所谓的行并行方式. -为了计算 $Z = [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right]$, 我们首先在每个处理器上计算 $Y_iB_i$ 然后使用一个all-reduce操作将结果汇总为 $Z=Y_1B_1+Y_2B_2$。 +$$ + +为了计算 +$$ +Z = [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \\ B_2 \end{matrix} \right] +$$ +我们首先在每个处理器上计算 $Y_iB_i$ 然后使用一个all-reduce操作将结果汇总为 $Z=Y_1B_1+Y_2B_2$。 我们还需要注意，在后向计算中，列并行线性层需要聚合输入张量 $X$, 因为在每个处理器 $i$ 上，我们只有 $\dot{X_i}=\dot{Y_i}A_i^T$，因此，我们在各处理器之间进行all-reduce，得到 $\dot{X}=\dot{Y}A^T=\dot{Y_1}A_1^T+\dot{Y_2}A_2^T$。 diff --git a/docs/source/zh-Hans/features/2D_tensor_parallel.md b/docs/source/zh-Hans/features/2D_tensor_parallel.md index c942f82bf9d2..f163432ecceb 100644 --- a/docs/source/zh-Hans/features/2D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/2D_tensor_parallel.md @@ -8,7 +8,7 @@ - [1D 张量并行](./1D_tensor_parallel.md) **示例代码** -- [ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_2d.py) +- [ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **相关论文** - [An Efficient 2D Method for Training Super-Large Deep Learning Models](https://arxiv.org/pdf/2104.05343.pdf) @@ -22,33 +22,33 @@ 给定 $P=q\times q$ 个处理器（必要条件）, 如 $q=2$, 我们把输入 $X$ 和权重A $A$ 都划分为 $$ -\left[\begin{matrix} X_{10} & X_{11} \\ X_{00} & X_{01} \end{matrix} \right] +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \end{matrix} \right] \text{~and~} -\left[\begin{matrix} A_{10} & A_{11} \\ A_{00} & A_{01} \end{matrix} \right]。 +\left[\begin{matrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{matrix} \right]. $$ 该计算包括 $q$ 步。 当 $t=1$ 时, $X_{i0}$ 在其行中被广播, 而 $A_{0j}$ 在其列中被广播。因此，我们有 $$ -\left[\begin{matrix} X_{10},A_{00} & X_{10},A_{01} \\ X_{00},A_{00} & X_{00},A_{01} \end{matrix} \right]。 +\left[\begin{matrix} X_{00},A_{00} & X_{00},A_{01} \\ X_{10},A_{00} & X_{10},A_{01} \end{matrix} \right]. $$ 然后我们在每个处理器 $(i, j)$ 上将 $X_{i0}$ 和 $A_{0j}$ 相乘为 $$ -\left[\begin{matrix} X_{10}A_{00} & X_{10}A_{01} \\ X_{00}A_{00} & X_{00}A_{01} \end{matrix} \right] (1)。 +\left[\begin{matrix} X_{00}A_{00} & X_{00}A_{01} \\ X_{10}A_{00} & X_{10}A_{01} \end{matrix} \right] (1). $$ 同样，当 $t=2$ 时, $X_{i1}$ 在其行中被广播, $A_{1j}$ 在其列中被广播, 我们将它们相乘为 $$ -\left[\begin{matrix} X_{11}A_{10} & X_{11}A_{11} \\ X_{01}A_{10} & X_{01}A_{11} \end{matrix} \right] (2)。 +\left[\begin{matrix} X_{01}A_{10} & X_{01}A_{11} \\ X_{11}A_{10} & X_{11}A_{11} \end{matrix} \right] (2). $$ 通过将 $(1)$ 和 $(2)$ 相加，我们有 $$ -Y = XA = \left[\begin{matrix} X_{10}A_{00}+X_{11}A_{10} & X_{10}A_{01}+X_{11}A_{11} \\ X_{00}A_{00}+X_{01}A_{10} & X_{00}A_{01}+X_{01}A_{11} \end{matrix} \right]。 +Y = XA = \left[\begin{matrix} X_{00}A_{00}+X_{01}A_{10} & X_{00}A_{01}+X_{01}A_{11} \\ X_{10}A_{00}+X_{11}A_{10} & X_{10}A_{01}+X_{11}A_{11} \end{matrix} \right]. $$ ## 效率 diff --git a/docs/source/zh-Hans/features/2p5D_tensor_parallel.md b/docs/source/zh-Hans/features/2p5D_tensor_parallel.md index 59a4be02ce47..5f15202729a7 100644 --- a/docs/source/zh-Hans/features/2p5D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/2p5D_tensor_parallel.md @@ -9,7 +9,7 @@ - [2D 张量并行](./2D_tensor_parallel.md) **示例代码** -- [ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_2p5d.py) +- [ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **相关论文** - [2.5-dimensional distributed model training](https://arxiv.org/pdf/2105.14500.pdf) @@ -22,29 +22,29 @@ 给定 $P=q \times q \times d$ 个处理器（必要条件）, 如 $q=d=2$, 我们把输入 $X$ 划分为 $d\times q$ 行和 $q$ 列 $$ -\left[\begin{matrix} X_{30} & X_{31} \\ X_{20} & X_{21} \\ X_{10} & X_{11} \\ X_{00} & X_{01}\end{matrix} \right], +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \\ X_{20} & X_{21} \\ X_{30} & X_{31}\end{matrix} \right], $$ 它可以被重塑为 $d$ 层 $$ -\left[\begin{matrix} X_{10} & X_{11} \\ X_{00} & X_{01} \end{matrix} \right] \text{~and~}\left[\begin{matrix} X_{30} & X_{31} \\ X_{20} & X_{21} \end{matrix} \right]. +\left[\begin{matrix} X_{00} & X_{01} \\ X_{10} & X_{11} \end{matrix} \right] \text{~and~}\left[\begin{matrix} X_{20} & X_{21} \\ X_{30} & X_{31} \end{matrix} \right]. $$ 另外，权重 $A$ 被分割为 $$ -\left[\begin{matrix} A_{10} & A_{11} \\ A_{00} & A_{01} \end{matrix} \right]. +\left[\begin{matrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{matrix} \right]. $$ 对于 $X$ 相关的每一层, 我们使用SUMMA算法将 $X$ 与 $A$ 相乘。 然后，我们得到输出 $$ -\left[\begin{matrix} Y_{10}=X_{10}A_{00}+X_{11}A_{10} & Y_{11}=X_{10}A_{01}+X_{11}A_{11} \\ Y_{00}=X_{00}A_{00}+X_{01}A_{10} & Y_{01}=X_{00}A_{01}+X_{01}A_{11} \end{matrix} \right] +\left[\begin{matrix} Y_{00}=X_{00}A_{00}+X_{01}A_{10} & Y_{01}=X_{00}A_{01}+X_{01}A_{11} \\ Y_{10}=X_{10}A_{00}+X_{11}A_{10} & Y_{11}=X_{10}A_{01}+X_{11}A_{11} \end{matrix} \right] \text{~and~} $$ $$ -\left[\begin{matrix} Y_{30}=X_{30}A_{00}+X_{31}A_{10} & Y_{31}=X_{30}A_{01}+X_{31}A_{11} \\ Y_{20}=X_{20}A_{00}+X_{21}A_{10} & Y_{21}=X_{20}A_{01}+X_{21}A_{11} \end{matrix} \right]. +\left[\begin{matrix} Y_{20}=X_{20}A_{00}+X_{21}A_{10} & Y_{21}=X_{20}A_{01}+X_{21}A_{11} \\ Y_{30}=X_{30}A_{00}+X_{31}A_{10} & Y_{31}=X_{30}A_{01}+X_{31}A_{11} \end{matrix} \right]. $$ ## 效率 diff --git a/docs/source/zh-Hans/features/3D_tensor_parallel.md b/docs/source/zh-Hans/features/3D_tensor_parallel.md index 440121c94243..5ce0cdf6c068 100644 --- a/docs/source/zh-Hans/features/3D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/3D_tensor_parallel.md @@ -9,7 +9,7 @@ - [2D 张量并行](./2D_tensor_parallel.md) **示例代码** -- [ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/features/tensor_parallel/tensor_parallel_3d.py) +- [ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) **相关论文** - [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://arxiv.org/pdf/2105.14450.pdf) diff --git a/docs/source/zh-Hans/features/cluster_utils.md b/docs/source/zh-Hans/features/cluster_utils.md new file mode 100644 index 000000000000..f54a72c63a66 --- /dev/null +++ b/docs/source/zh-Hans/features/cluster_utils.md @@ -0,0 +1,16 @@ +# 集群实用程序 + +作者: [Hongxin Liu](https://github.com/ver217) + +**前置教程:** +- [分布式训练](../concepts/distributed_training.md) + +## 引言 + +我们提供了一个实用程序类 `colossalai.cluster.DistCoordinator` 来协调分布式训练。它对于获取有关集群的各种信息很有用，例如节点数、每个节点的进程数等。 + +## API 参考 + +{{ autodoc:colossalai.cluster.DistCoordinator }} + + diff --git a/docs/source/zh-Hans/features/gradient_accumulation.md b/docs/source/zh-Hans/features/gradient_accumulation.md index e21e5fcd43d8..fc8b29bbe8f1 100644 --- a/docs/source/zh-Hans/features/gradient_accumulation.md +++ b/docs/source/zh-Hans/features/gradient_accumulation.md @@ -1,4 +1,4 @@ -# 梯度累积 +# 梯度累积 (旧版本) 作者: Shenggui Li, Yongbin Li @@ -38,3 +38,4 @@ iteration 1, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0 iteration 2, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) iteration 3, first 10 elements of param: tensor([-0.0141, 0.0464, 0.0507, 0.0321, 0.0356, -0.0150, 0.0172, -0.0118, 0.0222, 0.0473], device='cuda:0', grad_fn=) ``` + diff --git a/docs/source/zh-Hans/features/gradient_accumulation_with_booster.md b/docs/source/zh-Hans/features/gradient_accumulation_with_booster.md new file mode 100644 index 000000000000..a8422060f0ea --- /dev/null +++ b/docs/source/zh-Hans/features/gradient_accumulation_with_booster.md @@ -0,0 +1,146 @@ +# 梯度累积 (新版本) + +作者: [Mingyan Jiang](https://github.com/jiangmingyan) + +**前置教程** +- [定义配置文件](../basics/define_your_config.md) +- [训练中使用Booster](../basics/booster_api.md) + +## 引言 + +梯度累积是一种常见的增大训练 batch size 的方式。 在训练大模型时，内存经常会成为瓶颈，并且 batch size 通常会很小（如2），这导致收敛性无法保证。梯度累积将多次迭代的梯度累加，并仅在达到预设迭代次数时更新参数。 + +## 使用 + +在 Colossal-AI 中使用梯度累积非常简单，booster提供no_sync返回一个上下文管理器，在该上下文管理器下取消同步并且累积梯度。 + +## 实例 + +我们将介绍如何使用梯度累积。在这个例子中，梯度累积次数被设置为4。 + +### 步骤 1. 在 train.py 导入相关库 +创建train.py并导入必要依赖。 `torch` 的版本应不低于1.8.1。 + +```python +import os +from pathlib import Path + +import torch +from torchvision import transforms +from torchvision.datasets import CIFAR10 +from torchvision.models import resnet18 + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.cluster.dist_coordinator import priority_execution +``` + +### 步骤 2. 初始化分布式环境 + +我们需要初始化分布式环境。为了快速演示，我们使用`launch_from_torch`。你可以参考 [Launch Colossal-AI](../basics/launch_colossalai.md)使用其他初始化方法。 + +```python +# initialize distributed setting +parser = colossalai.get_default_parser() +args = parser.parse_args() + +# launch from torch +colossalai.launch_from_torch(config=dict()) + +``` + +### 步骤 3. 创建训练组件 + +构建你的模型、优化器、损失函数、学习率调整器和数据加载器。注意数据集的路径从环境变量`DATA`获得。你可以通过 `export DATA=/path/to/data` 或 `Path(os.environ['DATA'])`，在你的机器上设置路径。数据将会被自动下载到该路径。 + +```python +# define the training hyperparameters +BATCH_SIZE = 128 +GRADIENT_ACCUMULATION = 4 + +# build resnet +model = resnet18(num_classes=10) + +# build dataloaders +with priority_execution(): + train_dataset = CIFAR10(root=Path(os.environ.get('DATA', './data')), + download=True, + transform=transforms.Compose([ + transforms.RandomCrop(size=32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ])) + +# build criterion +criterion = torch.nn.CrossEntropyLoss() + +# optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) +``` + +### 步骤 4. 注入特性 +创建一个`TorchDDPPlugin`对象，并作为参实例化`Booster`， 调用`booster.boost`注入特性。 + +```python +plugin = TorchDDPPlugin() +booster = Booster(plugin=plugin) +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +model, optimizer, criterion, train_dataloader, _ = booster.boost(model=model, + optimizer=optimizer, + criterion=criterion, + dataloader=train_dataloader) +``` + +### 步骤 5. 使用booster训练 +使用booster构建一个普通的训练循环，验证梯度累积。 `param_by_iter` 记录分布训练的信息。 +```python +optimizer.zero_grad() +for idx, (img, label) in enumerate(train_dataloader): + sync_context = booster.no_sync(model) + img = img.cuda() + label = label.cuda() + if idx % (GRADIENT_ACCUMULATION - 1) != 0: + with sync_context: + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + else: + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + ele_1st = next(model.parameters()).flatten()[0] + param_by_iter.append(str(ele_1st.item())) + + if idx != 0 and idx % (GRADIENT_ACCUMULATION - 1) == 0: + break + + for iteration, val in enumerate(param_by_iter): + print(f'iteration {iteration} - value: {val}') + + if param_by_iter[-1] != param_by_iter[0]: + print('The parameter is only updated in the last iteration') + +``` + +### 步骤 6. 启动训练脚本 +为了验证梯度累积，我们可以只检查参数值的变化。当设置梯度累加时，仅在最后一步更新参数。您可以使用以下命令运行脚本： +```shell +colossalai run --nproc_per_node 1 train.py +``` + +你将会看到类似下方的文本输出。这展现了梯度虽然在前3个迭代中被计算，但直到最后一次迭代，参数才被更新。 + +```text +iteration 0, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 1, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 2, first 10 elements of param: tensor([-0.0208, 0.0189, 0.0234, 0.0047, 0.0116, -0.0283, 0.0071, -0.0359, -0.0267, -0.0006], device='cuda:0', grad_fn=) +iteration 3, first 10 elements of param: tensor([-0.0141, 0.0464, 0.0507, 0.0321, 0.0356, -0.0150, 0.0172, -0.0118, 0.0222, 0.0473], device='cuda:0', grad_fn=) +``` + + diff --git a/docs/source/zh-Hans/features/gradient_clipping.md b/docs/source/zh-Hans/features/gradient_clipping.md index 203f66a3fea2..2f62c31766a6 100644 --- a/docs/source/zh-Hans/features/gradient_clipping.md +++ b/docs/source/zh-Hans/features/gradient_clipping.md @@ -1,4 +1,4 @@ -# 梯度裁剪 +# 梯度裁剪（旧版本） 作者: Boxiang Wang, Haichen Huang, Yongbin Li @@ -49,3 +49,5 @@ clip_grad_norm = 1.0 ```shell python -m torch.distributed.launch --nproc_per_node 1 --master_addr localhost --master_port 29500 train_with_engine.py ``` + + diff --git a/docs/source/zh-Hans/features/gradient_clipping_with_booster.md b/docs/source/zh-Hans/features/gradient_clipping_with_booster.md new file mode 100644 index 000000000000..3c61356dd0d5 --- /dev/null +++ b/docs/source/zh-Hans/features/gradient_clipping_with_booster.md @@ -0,0 +1,140 @@ +# 梯度裁剪 (新版本) + +作者: [Mingyan Jiang](https://github.com/jiangmingyan) + +**前置教程** +- [定义配置文件](../basics/define_your_config.md) +- [booster使用](../basics/booster_api.md) + +**相关论文** +- [On the difficulty of training Recurrent Neural Networks](https://arxiv.org/abs/1211.5063) + +## 引言 + +为了加快训练过程和寻求全局最优以获得更好的性能，越来越多的学习率调度器被提出。人们通过控制学习率来调整训练中的下降速度。这使得梯度向量在每一步都能更好地统一。在这种情况下，下降速度可以按预期被控制。 +因此，梯度裁剪，一种可以将梯度向量归一化，以将其限制在统一长度的技术，对于那些希望模型性能更好的人来说是不可或缺的。 + +在使用 Colossal-AI 时，你不必担心实现梯度剪裁，我们以一种有效而方便的方式支持梯度剪裁。你所需要的只是在你的配置文件中增加一个命令。 + +## 为什么应该使用 Colossal-AI 中的梯度裁剪 + +我们不建议用户自己编写梯度剪裁，因为朴素的梯度剪裁在应用张量并行、流水线并行、MoE 等功能时可能会失败。 + +根据下图，每个 GPU 只拥有线性层中权重的一部分参数。为了得到线性层权重的梯度向量的正确范数，每个 GPU 中的每个梯度向量的范数应该相加。更复杂的是，偏置的分布不同于权重的分布。通信组在求和运算中有所不同。 + +(注: 这种情况是旧版本的 2D 并行，在代码中的实现是不一样的。但这是一个很好的例子，能够说明在梯度剪裁中统一所有通信的困难。) + +
      + +
      参数分布
      +
      + +不用担心它，因为 Colossal-AI 已经为你处理好。 + +### 使用 +要使用梯度裁剪，只需在使用booster注入特性之后，调用optimizer的`clip_grad_by_norm`或者`clip_grad_by_value`函数即可进行梯度裁剪。 + +### 实例 + +下面我们将介绍如何使用梯度裁剪，在本例中，我们将梯度裁剪范数设置为1.0。 + +### 步骤 1. 在训练中导入相关库 +创建`train.py`并导入相关库。 + +```python +import os +from pathlib import Path + +import torch +from torchvision import transforms +from torchvision.datasets import CIFAR10 +from torchvision.models import resnet34 +from tqdm import tqdm + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.nn.lr_scheduler import CosineAnnealingLR +``` + +### 步骤 2. 初始化分布式环境 +我们需要初始化分布式环境. 为了快速演示，我们使用`launch_from_torch`. 您可以参考 [Launch Colossal-AI](../basics/launch_colossalai.md) + +```python +colossalai.launch_from_torch(config=dict()) +logger = get_dist_logger() +``` + +### 步骤 3. 创建训练组件 + +构建你的模型、优化器、损失函数、学习率调整器和数据加载器。注意数据集的路径从环境变量`DATA`获得。你可以通过 `export DATA=/path/to/data` 或 `Path(os.environ['DATA'])`在你的机器上设置路径。数据将会被自动下载到该路径。 +```python +# define training hyperparameters +NUM_EPOCHS = 200 +BATCH_SIZE = 128 +GRADIENT_CLIPPING = 0.1 +# build resnet +model = resnet34(num_classes=10) +# build dataloaders +train_dataset = CIFAR10(root=Path(os.environ.get('DATA', './data')), + download=True, + transform=transforms.Compose([ + transforms.RandomCrop(size=32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]), + ])) +# build criterion +criterion = torch.nn.CrossEntropyLoss() + +# optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) + +# lr_scheduler +lr_scheduler = CosineAnnealingLR(optimizer, total_steps=NUM_EPOCHS) + +``` +### 步骤 4. 注入梯度裁剪特性 + +创建`TorchDDPPlugin`对象并初始化`Booster`, 使用booster注入相关特性。 +```python +plugin = TorchDDPPlugin() +booster = Booster(plugin=plugin) +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +model, optimizer, criterion, train_dataloader, lr_scheduler = booster.boost(model,optimizer, criterion,train_dataloader, lr_scheduler) + +``` + +### 步骤 5. 使用booster训练 +使用booster进行训练。 +```python +# verify gradient clipping +model.train() +for idx, (img, label) in enumerate(train_dataloader): + img = img.cuda() + label = label.cuda() + + model.zero_grad() + output = model(img) + train_loss = criterion(output, label) + booster.backward(train_loss, optimizer) + optimizer.clip_grad_by_norm(max_norm=GRADIENT_CLIPPING) + optimizer.step() + lr_scheduler.step() + + ele_1st = next(model.parameters()).flatten()[0] + logger.info(f'iteration {idx}, loss: {train_loss}, 1st element of parameters: {ele_1st.item()}') + + # only run for 4 iterations + if idx == 3: + break +``` + +### 步骤 6. 启动训练脚本 +你可以使用以下命令运行脚本： + +```shell +colossalai run --nproc_per_node 1 train.py +``` + diff --git a/docs/source/zh-Hans/features/mixed_precision_training.md b/docs/source/zh-Hans/features/mixed_precision_training.md index c9db3a59c1c3..a92e7e093015 100644 --- a/docs/source/zh-Hans/features/mixed_precision_training.md +++ b/docs/source/zh-Hans/features/mixed_precision_training.md @@ -1,4 +1,4 @@ -# 自动混合精度训练 (AMP) +# 自动混合精度训练 (旧版本) 作者: Chuanrui Wang, Shenggui Li, Yongbin Li @@ -203,7 +203,7 @@ Naive AMP 的默认参数: - initial_scale(int): gradient scaler 的初始值 - growth_factor(int): loss scale 的增长率 - backoff_factor(float): loss scale 的下降率 -- hysterisis(int): 动态 loss scaling 的延迟偏移 +- hysteresis(int): 动态 loss scaling 的延迟偏移 - max_scale(int): loss scale 的最大允许值 - verbose(bool): 如果被设为`True`,将打印调试信息 @@ -303,7 +303,7 @@ colossalai.launch_from_torch(config=args.config) # build loss criterion = torch.nn.CrossEntropyLoss() - # lr_scheduelr + # lr_scheduler lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=gpc.config.NUM_EPOCHS) ``` @@ -339,6 +339,7 @@ for epoch in range(gpc.config.NUM_EPOCHS): 使用下列命令启动训练脚本，你可以改变 `--nproc_per_node` 以使用不同数量的 GPU。 -```python +```shell python -m torch.distributed.launch --nproc_per_node 4 --master_addr localhost --master_port 29500 train_with_engine.py --config config/config_AMP_torch.py ``` + diff --git a/docs/source/zh-Hans/features/mixed_precision_training_with_booster.md b/docs/source/zh-Hans/features/mixed_precision_training_with_booster.md new file mode 100644 index 000000000000..0354f92ee7ce --- /dev/null +++ b/docs/source/zh-Hans/features/mixed_precision_training_with_booster.md @@ -0,0 +1,246 @@ +# 自动混合精度训练 (新版本) + +作者: [Mingyan Jiang](https://github.com/jiangmingyan) + +**前置教程** + +- [定义配置文件](../basics/define_your_config.md) +- [booster 使用](../basics/booster_api.md) + +**相关论文** + +- [Accelerating Scientific Computations with Mixed Precision Algorithms](https://arxiv.org/abs/0808.2794) + +## 引言 + +AMP 代表自动混合精度训练。 +在 Colossal-AI 中, 我们结合了混合精度训练的不同实现: + +1. torch.cuda.amp +2. apex.amp +3. naive amp + +| Colossal-AI | 支持张量并行 | 支持流水并行 | fp16 范围 | +| -------------- | ------------ | ------------ | --------------------------------------------------------- | +| AMP_TYPE.TORCH | ✅ | ❌ | 在前向和反向传播期间，模型参数、激活和梯度向下转换至 fp16 | +| AMP_TYPE.APEX | ❌ | ❌ | 更细粒度，我们可以选择 opt_level O0, O1, O2, O3 | +| AMP_TYPE.NAIVE | ✅ | ✅ | 模型参数、前向和反向操作，全都向下转换至 fp16 | + +前两个依赖于 PyTorch (1.6 及以上) 和 NVIDIA Apex 的原始实现。最后一种方法类似 Apex O2。在这些方法中，Apex-AMP 与张量并行不兼容。这是因为张量是以张量并行的方式在设备之间拆分的，因此，需要在不同的进程之间进行通信，以检查整个模型权重中是否出现 inf 或 nan。我们修改了 torch amp 实现，使其现在与张量并行兼容。 + +> ❌️ fp16 与 ZeRO 不兼容 +> +> ⚠️ 流水并行目前仅支持 naive amp + +我们建议使用 torch AMP，因为在不使用流水并行时，它通常比 NVIDIA AMP 提供更好的准确性。 + +## 目录 + +在本教程中，我们将介绍: + +1. [AMP 介绍](#amp-介绍) +2. [Colossal-AI 中的 AMP](#colossal-ai-中的-amp) +3. [练习实例](#实例) + +## AMP 介绍 + +自动混合精度训练是混合 FP16 和 FP32 训练。 + +半精度浮点格式（FP16）具有较低的算法复杂度和较高的计算效率。此外，FP16 仅需要 FP32 所需的一半存储空间，并节省了内存和网络带宽，从而为大 batch size 和大模型提供了更多内存。 + +然而，还有其他操作，如缩减，需要 FP32 的动态范围，以避免数值溢出/下溢。因此，我们引入自动混合精度，尝试将每个操作与其相应的数据类型相匹配，这可以减少内存占用并提高训练效率。 + +
      + +
      AMP 示意图 (图片来自 PatrickStar 论文)
      +
      + +## Colossal-AI 中的 AMP + +我们支持三种 AMP 训练方法，并允许用户在没有改变代码的情况下使用 AMP 进行训练。booster 支持 amp 特性注入，如果您要使用混合精度训练，则在创建 booster 实例时指定`mixed_precision`参数，我们现已支持 torch amp，apex amp, naive amp（现已移植 torch amp 至 booster，apex amp, naive amp 仍由`colossalai.initialize`方式启动，如您需使用，请[参考](./mixed_precision_training.md);后续将会拓展`bf16`,`pf8`的混合精度训练. + +#### booster 启动方式 + +您可以在创建 booster 实例时，指定`mixed_precision="fp16"`即使用 torch amp。 + + + +```python +""" + 初始化映射关系如下： + 'fp16': torch amp + 'fp16_apex': apex amp, + 'bf16': bf16, + 'fp8': fp8, + 'fp16_naive': naive amp +""" +from colossalai import Booster +booster = Booster(mixed_precision='fp16',...) +``` + + + +或者您可以自定义一个`FP16TorchMixedPrecision`对象，如 + + + +```python +from colossalai.mixed_precision import FP16TorchMixedPrecision +mixed_precision = FP16TorchMixedPrecision( + init_scale=2.**16, + growth_factor=2.0, + backoff_factor=0.5, + growth_interval=2000) +booster = Booster(mixed_precision=mixed_precision,...) +``` + + + +其他类型的 amp 使用方式也是一样的。 + +### Torch AMP 配置 + +{{ autodoc:colossalai.booster.mixed_precision.FP16TorchMixedPrecision }} + +### Apex AMP 配置 + +对于这种模式，我们依靠 Apex 实现混合精度训练。我们支持这个插件，因为它允许对混合精度的粒度进行更精细的控制。 +例如, O2 水平 (优化器水平 2) 将保持 batch normalization 为 FP32。 + +如果你想了解更多细节，请参考 [Apex Documentation](https://nvidia.github.io/apex/)。 + +{{ autodoc:colossalai.booster.mixed_precision.FP16ApexMixedPrecision }} + +### Naive AMP 配置 + +在 Naive AMP 模式中, 我们实现了混合精度训练，同时保持了与复杂张量和流水并行的兼容性。该 AMP 模式将所有操作转为 FP16 。下列代码块展示了该模式的 booster 启动方式。 + +{{ autodoc:colossalai.booster.mixed_precision.FP16NaiveMixedPrecision }} + +当使用`colossalai.booster`时, 首先需要实例化一个模型、一个优化器和一个标准。将输出模型转换为内存消耗较小的 AMP 模型。如果您的输入模型已经太大，无法放置在 GPU 中，请使用`dtype=torch.float16`实例化你的模型。或者请尝试更小的模型，或尝试更多的并行化训练技术！ + +## 实例 + +下面我们将展现如何在 Colossal-AI 使用 AMP。在该例程中，我们使用 Torch AMP. + +### 步骤 1. 在 train.py 导入相关库 + +创建`train.py`并导入必要依赖. 请记得通过命令`pip install timm scipy`安装`scipy`和`timm`。 + +```python +import os +from pathlib import Path + +import torch +from timm.models import vit_base_patch16_224 +from titans.utils import barrier_context +from torchvision import datasets, transforms + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.logging import get_dist_logger +from colossalai.nn.lr_scheduler import LinearWarmupLR +``` + +### 步骤 2. 初始化分布式环境 + +我们需要初始化分布式环境。为了快速演示，我们使用`launch_from_torch`。你可以参考 [Launch Colossal-AI](../basics/launch_colossalai.md) +使用其他初始化方法。 + +```python +# 初始化分布式设置 +parser = colossalai.get_default_parser() +args = parser.parse_args() + +# launch from torch +colossalai.launch_from_torch(config=dict()) + +``` + +### 步骤 3. 创建训练组件 + +构建你的模型、优化器、损失函数、学习率调整器和数据加载器。注意数据集的路径从环境变量`DATA`获得。你可以通过 `export DATA=/path/to/data` 或 `Path(os.environ['DATA'])` +在你的机器上设置路径。数据将会被自动下载到该路径。 + +```python +# define the constants +NUM_EPOCHS = 2 +BATCH_SIZE = 128 +# build model +model = vit_base_patch16_224(drop_rate=0.1) + +# build dataloader +train_dataset = datasets.Caltech101( + root=Path(os.environ['DATA']), + download=True, + transform=transforms.Compose([ + transforms.Resize(256), + transforms.RandomResizedCrop(224), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + Gray2RGB(), + transforms.Normalize([0.5, 0.5, 0.5], + [0.5, 0.5, 0.5]) + ])) + +# build optimizer +optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, weight_decay=0.1) + +# build loss +criterion = torch.nn.CrossEntropyLoss() + +# lr_scheduler +lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=50, total_steps=NUM_EPOCHS) +``` + +### 步骤 4. 插入 AMP + +创建一个 MixedPrecision 对象（如果需要）及 torchDDPPlugin 对象，调用 `colossalai.boost` 将所有训练组件转为为 FP16 模式. + +```python +plugin = TorchDDPPlugin() +train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) +booster = Booster(mixed_precision='fp16', plugin=plugin) + +# if you need to customize the config, do like this +# >>> from colossalai.mixed_precision import FP16TorchMixedPrecision +# >>> mixed_precision = FP16TorchMixedPrecision( +# >>> init_scale=2.**16, +# >>> growth_factor=2.0, +# >>> backoff_factor=0.5, +# >>> growth_interval=2000) +# >>> plugin = TorchDDPPlugin() +# >>> booster = Booster(mixed_precision=mixed_precision, plugin=plugin) + +# boost model, optimizer, criterion, dataloader, lr_scheduler +model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model, optimizer, criterion, dataloader, lr_scheduler) +``` + +### 步骤 5. 使用 booster 训练 + +使用 booster 构建一个普通的训练循环。 + +```python +model.train() +for epoch in range(NUM_EPOCHS): + for img, label in enumerate(train_dataloader): + img = img.cuda() + label = label.cuda() + optimizer.zero_grad() + output = model(img) + loss = criterion(output, label) + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() +``` + +### 步骤 6. 启动训练脚本 + +使用下列命令启动训练脚本，你可以改变 `--nproc_per_node` 以使用不同数量的 GPU。 + +```shell +colossalai run --nproc_per_node 1 train.py +``` + + diff --git a/docs/source/zh-Hans/features/nvme_offload.md b/docs/source/zh-Hans/features/nvme_offload.md index f33474efaa78..1feb9dde5725 100644 --- a/docs/source/zh-Hans/features/nvme_offload.md +++ b/docs/source/zh-Hans/features/nvme_offload.md @@ -53,9 +53,8 @@ optimizer = HybridAdam(model.parameters(), lr=1e-3, nvme_offload_fraction=1.0, n > ⚠ 它只会卸载在 CPU 上的优化器状态。这意味着它只会影响 CPU 训练或者使用卸载的 Zero/Gemini。 -## Exampls +## Examples -Let's start from two simple examples -- training GPT with different methods. These examples relies on `transformers`. 首先让我们从两个简单的例子开始 -- 用不同的方法训练 GPT。这些例子依赖`transformers`。 我们首先应该安装依赖： @@ -77,8 +76,9 @@ from transformers.models.gpt2.configuration_gpt2 import GPT2Config from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel import colossalai from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin ``` 然后我们定义一个损失函数： @@ -182,16 +182,24 @@ def train_gemini_cpu(nvme_offload_fraction: float = 0.0): criterion = GPTLMLoss() optimizer = HybridAdam(model.parameters(), nvme_offload_fraction=nvme_offload_fraction) print(f'Model numel: {get_model_numel(model) / 1024**3:.3f} B') - gemini_config = dict(strict_ddp_mode=True, device=torch.cuda.current_device(), - placement_policy='cpu', pin_memory=True, hidden_dim=config.n_embd) - model = zero_model_wrapper(model, zero_stage=3, gemini_config=gemini_config) - optimizer = zero_optim_wrapper(model, optimizer, initial_scale=2**5) + + plugin = GeminiPlugin( + strict_ddp_mode=True, + device=torch.cuda.current_device(), + placement_policy='cpu', + pin_memory=True, + hidden_dim=config.n_embd, + initial_scale=2**5 + ) + booster = Booster(plugin) + model, optimizer, criterion, _* = booster.boost(model, optimizer, criterion) + start = time.time() for step in range(3): data = get_data(4, 128, config.vocab_size) outputs = model(**data) loss = criterion(outputs.logits, data['input_ids']) - optimizer.backward(loss) + booster.backward(loss, optimizer) optimizer.step() optimizer.zero_grad() print(f'[{step}] loss: {loss.item():.3f}') diff --git a/docs/source/zh-Hans/features/zero_with_chunk.md b/docs/source/zh-Hans/features/zero_with_chunk.md index 72403bf610a4..513850f5cab7 100644 --- a/docs/source/zh-Hans/features/zero_with_chunk.md +++ b/docs/source/zh-Hans/features/zero_with_chunk.md @@ -4,7 +4,7 @@ **前置教程:** -- [定义配置文件](../basics/define_your_config.md) +- [booster使用](../basics/booster_api.md) **示例代码** @@ -66,13 +66,13 @@ with ColoInitContext(device='cpu', default_dist_spec=default_dist_spec, default_ chunk_manager = init_chunk_manager(model=module, init_device=device, hidden_dim=hidden_dim, - search_range_mb=search_range_mb, - min_chunk_size_mb=min_chunk_size_mb) + search_range_m=search_range_m, + min_chunk_size_m=min_chunk_size_m) gemini_manager = GeminiManager(placement_policy, chunk_manager) model = ZeroDDP(model, gemini_manager) ``` -`hidden dim`是DNN的隐藏维度。用户可以提供这个参数来加快搜索速度。如果用户在训练前不知道这个参数也可以。 我们将使用默认值 1024。`min_chunk_size_mb`是以兆字节为单位的最小块大小。如果参数的总大小仍然小于最小块大小，则所有参数将被压缩为一个小块。 +`hidden dim`是DNN的隐藏维度。用户可以提供这个参数来加快搜索速度。如果用户在训练前不知道这个参数也可以。 我们将使用默认值 1024。`min_chunk_size_m`是以兆（2^20）为单位的最小块大小。如果参数的总大小仍然小于最小块大小，则所有参数将被压缩为一个小块。 初始化优化器。 ```python @@ -97,6 +97,8 @@ optimizer.step() 首先我们只需要引入`Huggingface transformers` 的 `GPT2LMHeadModel`来定义我们的模型，不需要用户进行模型的定义与修改，方便用户使用。 +定义GPT模型： + ```python class GPTLMModel(nn.Module): @@ -182,34 +184,6 @@ def split_param_col_tp1d(param: ColoParameter, pg: ProcessGroup): split_param_single_dim_tp1d(-1, param, pg) ``` -定义一个使用 Gemini + ZeRO DDP 的模型： - -```python -def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: str = "auto"): - cai_version = colossalai.__version__ - if version.parse(cai_version) > version.parse("0.1.10"): - from colossalai.nn.parallel import GeminiDDP - model = GeminiDDP(model, - device=get_current_device(), - placement_policy=placememt_policy, - pin_memory=True, - search_range_mb=32) - elif version.parse(cai_version) <= version.parse("0.1.10") and version.parse(cai_version) >= version.parse("0.1.9"): - from colossalai.gemini import ChunkManager, GeminiManager - chunk_size = ChunkManager.search_chunk_size(model, 64 * 1024**2, 32) - gemini_manager = GeminiManager(placememt_policy, chunk_manager) - chunk_manager = ChunkManager(chunk_size, - pg, - enable_distributed_storage=True, - init_device=GeminiManager.get_default_device(placememt_policy)) - model = ZeroDDP(model, gemini_manager) - else: - raise NotImplemented(f"CAI version {cai_version} is not supported") - return model -``` - -由于我们在这个例子中对GPT进行预训练，因此只使用了一个简单的语言模型损失函数。 - 写一个获得随机输入的函数: ```python @@ -219,9 +193,16 @@ def get_data(batch_size, seq_len, vocab_size): return input_ids, attention_mask ``` -最后，我们可以定义我们的训练循环: + +最后，使用booster注入 Gemini + ZeRO DDP 特性, 并定义训练循环。由于我们在这个例子中对GPT进行预训练，因此只使用了一个简单的语言模型损失函数： ```python +from colossalai.nn.optimizer import HybridAdam + +from colossalai.booster import Booster +from colossalai.zero import ColoInitContext +from colossalai.booster.plugin import GeminiPlugin + def main(): args = parse_args() BATCH_SIZE = 8 @@ -232,22 +213,23 @@ def main(): # build criterion criterion = GPTLMLoss() + optimizer = HybridAdam(model.parameters(), lr=0.001) torch.manual_seed(123) default_pg = ProcessGroup(tp_degree=args.tp_degree) - default_dist_spec = ShardSpec([-1], [args.tp_degree]) if args.shardinit else None + default_dist_spec = ShardSpec([-1], [args.tp_degree]) # build GPT model with ColoInitContext(device='cpu', default_dist_spec=default_dist_spec, default_pg=default_pg): model = gpt2_medium(checkpoint=True) pg = default_pg # Tensor Parallelism (TP) tensor_parallelize(model, pg) + # Gemini + ZeRO DP, Note it must be used after TP - model = gemini_zero_dpp(model, pg, args.placement) - # build optimizer - optimizer = GeminiAdamOptimizer(model, lr=1e-3, initial_scale=2**5) - numel = sum([p.numel() for p in model.parameters()]) - get_tflops_func = partial(get_tflops, numel, BATCH_SIZE, SEQ_LEN) + plugin = GeminiPlugin(placement_policy='cuda', max_norm=1.0, initial_scale=2**5) + booster = Booster(plugin=plugin) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + torch.cuda.synchronize() model.train() for n in range(NUM_STEPS): @@ -256,10 +238,12 @@ def main(): optimizer.zero_grad() outputs = model(input_ids, attn_mask) loss = criterion(outputs, input_ids) - optimizer.backward(loss) + booster.backward(loss, optimizer) optimizer.step() torch.cuda.synchronize() ``` > ⚠️ 注意：如果你使用Gemini模块的话，请不要使用我们之前提到过的[梯度累加](../features/gradient_accumulation.md)。 完整的例子代码可以在 [Train GPT with Colossal-AI](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). 获得。 + + diff --git a/docs/source/zh-Hans/get_started/installation.md b/docs/source/zh-Hans/get_started/installation.md index 7a9b20255e77..a6c88672b907 100755 --- a/docs/source/zh-Hans/get_started/installation.md +++ b/docs/source/zh-Hans/get_started/installation.md @@ -5,6 +5,8 @@ - PyTorch >= 1.11 (PyTorch 2.x 正在适配中) - Python >= 3.7 - CUDA >= 11.0 +- [NVIDIA GPU Compute Capability](https://developer.nvidia.com/cuda-gpus) >= 7.0 (V100/RTX20 and higher) +- Linux OS 如果你遇到安装问题，可以向本项目 [反馈](https://github.com/hpcaitech/ColossalAI/issues/new/choose)。 @@ -26,7 +28,7 @@ CUDA_EXT=1 pip install colossalai ## 从源安装 -> 此文档将与版本库的主分支保持一致。如果您遇到任何问题，欢迎给我们提 issue :) +> 此文档将与版本库的主分支保持一致。如果您遇到任何问题，欢迎给我们提 issue。 ```shell git clone https://github.com/hpcaitech/ColossalAI.git @@ -36,13 +38,29 @@ cd ColossalAI pip install -r requirements/requirements.txt # install colossalai -pip install . +CUDA_EXT=1 pip install . ``` -如果您不想安装和启用 CUDA 内核融合（使用融合优化器时强制安装）： +如果您不想安装和启用 CUDA 内核融合（使用融合优化器时强制安装），您可以不添加`CUDA_EXT=1`： ```shell -NO_CUDA_EXT=1 pip install . +pip install . +``` + +如果您在使用CUDA 10.2，您仍然可以从源码安装ColossalAI。但是您需要手动下载cub库并将其复制到相应的目录。 + +```bash +# clone the repository +git clone https://github.com/hpcaitech/ColossalAI.git +cd ColossalAI + +# download the cub library +wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip +unzip 1.8.0.zip +cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/ + +# install +CUDA_EXT=1 pip install . ``` diff --git a/docs/source/zh-Hans/get_started/run_demo.md b/docs/source/zh-Hans/get_started/run_demo.md index edfc246c22d5..70ed5ebe251b 100755 --- a/docs/source/zh-Hans/get_started/run_demo.md +++ b/docs/source/zh-Hans/get_started/run_demo.md @@ -4,8 +4,8 @@ Colossal-AI 是一个集成的大规模深度学习系统，具有高效的并 ## 单 GPU -Colossal-AI 可以用在只有一个 GPU 的系统上训练深度学习模型，并达到 baseline 的性能。 我们提供了一个 [在CIFAR10数据集上训练ResNet](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/resnet) 的例子，该例子只需要一个 GPU。 -您可以在 [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples) 中获取该例子。详细说明可以在其 `README.md` 中获取。 +Colossal-AI 可以用在只有一个 GPU 的系统上训练深度学习模型，并达到 baseline 的性能。 我们提供了一个 [在 CIFAR10 数据集上训练 ResNet](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/resnet) 的例子，该例子只需要一个 GPU。 +您可以在 [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples) 中获取该例子。详细说明可以在其 `README.md` 中获取。 ## 多 GPU @@ -13,16 +13,20 @@ Colossal-AI 可用于在具有多个 GPU 的分布式系统上训练深度学习 #### 1. 数据并行 -您可以使用与上述单 GPU 演示相同的 [ResNet例子](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/resnet)。 通过设置 `--nproc_per_node` 为您机器上的 GPU 数量，您就能把数据并行应用在您的例子上了。 +您可以使用与上述单 GPU 演示相同的 [ResNet 例子](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/resnet)。 通过设置 `--nproc_per_node` 为您机器上的 GPU 数量，您就能把数据并行应用在您的例子上了。 #### 2. 混合并行 -混合并行包括数据、张量和流水线并行。在 Colossal-AI 中，我们支持不同类型的张量并行（即 1D、2D、2.5D 和 3D）。您可以通过简单地改变 `config.py` 中的配置在不同的张量并行之间切换。您可以参考 [GPT example](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/language/gpt), 更多细节能在它的 `README.md` 中被找到。 +混合并行包括数据、张量和流水线并行。在 Colossal-AI 中，我们支持不同类型的张量并行（即 1D、2D、2.5D 和 3D）。您可以通过简单地改变 `config.py` 中的配置在不同的张量并行之间切换。您可以参考 [GPT example](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt), 更多细节能在它的 `README.md` 中被找到。 -#### 3. MoE并行 +#### 3. MoE 并行 -我们提供了一个 [WideNet例子](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/widenet) 来验证 MoE 的并行性。 WideNet 使用 Mixture of Experts（MoE）来实现更好的性能。更多的细节可以在我们的教程中获取：[教会您如何把Mixture of Experts整合到模型中](../advanced_tutorials/integrate_mixture_of_experts_into_your_model.md)。 + + +我们提供了一个 [ViT-MoE 例子](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/image/moe) 来验证 MoE 的并行性。 WideNet 使用 Mixture of Experts（MoE）来实现更好的性能。更多的细节可以在我们的教程中获取：[教会您如何把 Mixture of Experts 整合到模型中](../advanced_tutorials/integrate_mixture_of_experts_into_your_model.md)。 #### 4. 序列并行 -序列并行是为了解决NLP任务中的内存效率和序列长度限制问题。 我们在 [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI-Examples) 中提供了一个 [BERT例子](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/language/bert/sequene_parallel)。您可以按照 `README.md` 来执行代码。 +序列并行是为了解决 NLP 任务中的内存效率和序列长度限制问题。 我们在 [ColossalAI-Examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples) 中提供了一个 [Sequence Parallelism 例子](https://github.com/hpcaitech/ColossalAI/tree/main/examples/tutorial/sequence_parallel)。您可以按照 `README.md` 来执行代码。 + + diff --git a/examples/README.md b/examples/README.md index 710ced101768..142a735c6819 100644 --- a/examples/README.md +++ b/examples/README.md @@ -10,7 +10,12 @@ ## Overview -This folder provides several examples accelerated by Colossal-AI. The `tutorial` folder is for everyone to quickly try out the different features in Colossal-AI. Other folders such as `images` and `language` include a wide range of deep learning tasks and applications. +This folder provides several examples accelerated by Colossal-AI. +Folders such as `images` and `language` include a wide range of deep learning tasks and applications. +The `community` folder aim to create a collaborative platform for developers to contribute exotic features built on top of Colossal-AI. +The `tutorial` folder is for everyone to quickly try out the different features in Colossal-AI. + +You can find applications such as Chatbot, AIGC and Biomedicine in the [Applications](https://github.com/hpcaitech/ColossalAI/tree/main/applications) directory. ## Folder Structure @@ -50,3 +55,10 @@ Therefore, it is essential for the example contributors to know how to integrate 2. Configure your testing parameters such as number steps, batch size in `test_ci.sh`, e.t.c. Keep these parameters small such that each example only takes several minutes. 3. Export your dataset path with the prefix `/data` and make sure you have a copy of the dataset in the `/data/scratch/examples-data` directory on the CI machine. Community contributors can contact us via slack to request for downloading the dataset on the CI machine. 4. Implement the logic such as dependency setup and example execution + +## Community Dependency +We are happy to introduce the following nice community dependency repos that are powered by Colossal-AI: +- [lightning-ColossalAI](https://github.com/Lightning-AI/lightning) +- [HCP-Diffusion](https://github.com/7eu7d7/HCP-Diffusion) +- [KoChatGPT](https://github.com/airobotlab/KoChatGPT) +- [minichatgpt](https://github.com/juncongmoo/minichatgpt) diff --git a/examples/community/README.md b/examples/community/README.md new file mode 100644 index 000000000000..fb2ca37ed988 --- /dev/null +++ b/examples/community/README.md @@ -0,0 +1,28 @@ +## Community Examples + +Community-driven Examples is an initiative that allows users to share their own examples to the Colossal-AI community, fostering a sense of community and making it easy for others to access and benefit from shared work. The primary goal with community-driven examples is to have a community-maintained collection of diverse and exotic functionalities built on top of the Colossal-AI package. + +If a community example doesn't work as expected, you can [open an issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose) and @ the author to report it. + + +| Example | Description | Code Example | Colab |Author | +|:------------------|:---------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------|:-----------------------------------------|-----------------------------------------------------:| +| RoBERTa | Adding RoBERTa for SFT and Prompts model training | [RoBERTa](./roberta) | - | [YY Lin](https://github.com/yynil) (Moore Threads) | +| TransformerEngine FP8 | Adding TransformerEngine with FP8 training | [TransformerEngine FP8](./fp8) | - | [Kirthi Shankar Sivamani](https://github.com/ksivaman) (NVIDIA) | +|...|...|...|...|...| + +## Looking for Examples +* [Swin-Transformer](https://github.com/microsoft/Swin-Transformer) +* [T-5](https://github.com/google-research/text-to-text-transfer-transformer) +* [Segment Anything (SAM)](https://github.com/facebookresearch/segment-anything) +* [ControlNet](https://github.com/lllyasviel/ControlNet) +* [Consistency Models](https://github.com/openai/consistency_models) +* [MAE](https://github.com/facebookresearch/mae) +* [CLIP](https://github.com/openai/CLIP) + +Welcome to [open an issue](https://github.com/hpcaitech/ColossalAI/issues/new/choose) to share your insights and needs. + +## How to get involved +To join our community-driven initiative, please visit the [Colossal-AI examples](https://github.com/hpcaitech/ColossalAI/tree/main/examples), review the provided information, and explore the codebase. + +To contribute, create a new issue outlining your proposed feature or enhancement, and our team will review and provide feedback. If you are confident enough you can also submit a PR directly. We look forward to collaborating with you on this exciting project! diff --git a/examples/tutorial/fp8/mnist/README.md b/examples/community/fp8/mnist/README.md similarity index 89% rename from examples/tutorial/fp8/mnist/README.md rename to examples/community/fp8/mnist/README.md index 46711f9ebdd8..e1128c1054b7 100644 --- a/examples/tutorial/fp8/mnist/README.md +++ b/examples/community/fp8/mnist/README.md @@ -1,13 +1,13 @@ -# Basic MNIST Example with optional FP8 of TransformerEngine - -[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) is a library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper GPUs, to provide better performance with lower memory utilization in both training and inference. - -Thanks for the contribution to this tutorial from NVIDIA. - -```bash -python main.py -python main.py --use-te # Linear layers from TransformerEngine -python main.py --use-fp8 # FP8 + TransformerEngine for Linear layers -``` - -> We are working to integrate it with Colossal-AI and will finish it soon. +# Basic MNIST Example with optional FP8 of TransformerEngine + +[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) is a library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper GPUs, to provide better performance with lower memory utilization in both training and inference. + +Thanks for the contribution to this tutorial from NVIDIA. + +```bash +python main.py +python main.py --use-te # Linear layers from TransformerEngine +python main.py --use-fp8 # FP8 + TransformerEngine for Linear layers +``` + +> We are working to integrate it with Colossal-AI and will finish it soon. diff --git a/examples/tutorial/fp8/mnist/main.py b/examples/community/fp8/mnist/main.py similarity index 81% rename from examples/tutorial/fp8/mnist/main.py rename to examples/community/fp8/mnist/main.py index 000ded2f111f..a534663d380f 100644 --- a/examples/tutorial/fp8/mnist/main.py +++ b/examples/community/fp8/mnist/main.py @@ -3,12 +3,13 @@ # See LICENSE for license information. import argparse + import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim -from torchvision import datasets, transforms from torch.optim.lr_scheduler import StepLR +from torchvision import datasets, transforms try: from transformer_engine import pytorch as te @@ -18,6 +19,7 @@ class Net(nn.Module): + def __init__(self, use_te=False): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) @@ -62,12 +64,10 @@ def train(args, model, device, train_loader, optimizer, epoch, use_fp8): loss.backward() optimizer.step() if batch_idx % args.log_interval == 0: - print( - f"Train Epoch: {epoch} " - f"[{batch_idx * len(data)}/{len(train_loader.dataset)} " - f"({100. * batch_idx / len(train_loader):.0f}%)]\t" - f"Loss: {loss.item():.6f}" - ) + print(f"Train Epoch: {epoch} " + f"[{batch_idx * len(data)}/{len(train_loader.dataset)} " + f"({100. * batch_idx / len(train_loader):.0f}%)]\t" + f"Loss: {loss.item():.6f}") if args.dry_run: break @@ -83,6 +83,7 @@ def calibrate(model, device, test_loader): with te.fp8_autocast(enabled=False, calibrating=True): output = model(data) + def test(model, device, test_loader, use_fp8): """Testing function.""" model.eval() @@ -93,21 +94,15 @@ def test(model, device, test_loader, use_fp8): data, target = data.to(device), target.to(device) with te.fp8_autocast(enabled=use_fp8): output = model(data) - test_loss += F.nll_loss( - output, target, reduction="sum" - ).item() # sum up batch loss - pred = output.argmax( - dim=1, keepdim=True - ) # get the index of the max log-probability + test_loss += F.nll_loss(output, target, reduction="sum").item() # sum up batch loss + pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) - print( - f"\nTest set: Average loss: {test_loss:.4f}, " - f"Accuracy: {correct}/{len(test_loader.dataset)} " - f"({100. * correct / len(test_loader.dataset):.0f}%)\n" - ) + print(f"\nTest set: Average loss: {test_loss:.4f}, " + f"Accuracy: {correct}/{len(test_loader.dataset)} " + f"({100. * correct / len(test_loader.dataset):.0f}%)\n") def main(): @@ -154,9 +149,7 @@ def main(): default=False, help="quickly check a single pass", ) - parser.add_argument( - "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)" - ) + parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)") parser.add_argument( "--log-interval", type=int, @@ -170,15 +163,12 @@ def main(): default=False, help="For Saving the current Model", ) - parser.add_argument( - "--use-fp8", action="store_true", default=False, help="Use FP8 for inference and training without recalibration" - ) - parser.add_argument( - "--use-fp8-infer", action="store_true", default=False, help="Use FP8 inference only" - ) - parser.add_argument( - "--use-te", action="store_true", default=False, help="Use Transformer Engine" - ) + parser.add_argument("--use-fp8", + action="store_true", + default=False, + help="Use FP8 for inference and training without recalibration") + parser.add_argument("--use-fp8-infer", action="store_true", default=False, help="Use FP8 inference only") + parser.add_argument("--use-te", action="store_true", default=False, help="Use Transformer Engine") args = parser.parse_args() use_cuda = torch.cuda.is_available() @@ -205,9 +195,7 @@ def main(): train_kwargs.update(cuda_kwargs) test_kwargs.update(cuda_kwargs) - transform = transforms.Compose( - [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] - ) + transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) dataset1 = datasets.MNIST("../data", train=True, download=True, transform=transform) dataset2 = datasets.MNIST("../data", train=False, transform=transform) train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs) @@ -227,7 +215,7 @@ def main(): if args.save_model or args.use_fp8_infer: torch.save(model.state_dict(), "mnist_cnn.pt") - print('Eval with reloaded checkpoint : fp8='+str(args.use_fp8_infer)) + print('Eval with reloaded checkpoint : fp8=' + str(args.use_fp8_infer)) weights = torch.load("mnist_cnn.pt") model.load_state_dict(weights) test(model, device, test_loader, args.use_fp8_infer) diff --git a/examples/community/roberta/README.md b/examples/community/roberta/README.md new file mode 100644 index 000000000000..000fce63f35f --- /dev/null +++ b/examples/community/roberta/README.md @@ -0,0 +1,50 @@ +# Introduction +This example introduce how to pretrain roberta from scratch, including preprocessing, pretraining, finetune. The example can help you quickly train a high-quality roberta. + +## 0. Prerequisite +- Install Colossal-AI +- Editing the port from `/etc/ssh/sshd_config` and `/etc/ssh/ssh_config`, every host expose the same ssh port of server and client. If you are a root user, you also set the **PermitRootLogin** from `/etc/ssh/sshd_config` to "yes" +- Ensure that each host can log in to each other without password. If you have n hosts, need to execute n2 times + +``` +ssh-keygen +ssh-copy-id -i ~/.ssh/id_rsa.pub ip_destination +``` + +- In all hosts, edit /etc/hosts to record all hosts' name and ip.The example is shown below. + +```bash +192.168.2.1 GPU001 +192.168.2.2 GPU002 +192.168.2.3 GPU003 +192.168.2.4 GPU004 +192.168.2.5 GPU005 +192.168.2.6 GPU006 +192.168.2.7 GPU007 +... +``` + +- restart ssh +``` +service ssh restart +``` + +## 1. Corpus Preprocessing +```bash +cd preprocessing +``` +following the `README.md`, preprocess original corpus to h5py plus numpy + +## 2. Pretrain + +```bash +cd pretraining +``` +following the `README.md`, load the h5py generated by preprocess of step 1 to pretrain the model + +## 3. Finetune + +The checkpoint produced by this repo can replace `pytorch_model.bin` from [hfl/chinese-roberta-wwm-ext-large](https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/tree/main) directly. Then use transformers from Hugging Face to finetune downstream application. + +## Contributors +The example is contributed by AI team from [Moore Threads](https://www.mthreads.com/). If you find any problems for pretraining, please file an issue or send an email to yehua.zhang@mthreads.com. At last, welcome any form of contribution! diff --git a/examples/language/roberta/preprocessing/Makefile b/examples/community/roberta/preprocessing/Makefile similarity index 100% rename from examples/language/roberta/preprocessing/Makefile rename to examples/community/roberta/preprocessing/Makefile diff --git a/examples/language/roberta/preprocessing/README.md b/examples/community/roberta/preprocessing/README.md similarity index 90% rename from examples/language/roberta/preprocessing/README.md rename to examples/community/roberta/preprocessing/README.md index 1dbd745ab9bd..2ed747541280 100644 --- a/examples/language/roberta/preprocessing/README.md +++ b/examples/community/roberta/preprocessing/README.md @@ -21,14 +21,14 @@ This folder is used to preprocess chinese corpus with Whole Word Masked. You can ### 2.1. Split Sentence & Split data into multiple shard: -Firstly, each file has multiple documents, and each document contains multiple sentences. Split sentence through punctuation, such as `。！`. **Secondly, split data into multiple shard based on server hardware (cpu, cpu memory, hard disk) and corpus size.** Each shard contains a part of corpus, and the model needs to train all the shards as one epoch. +Firstly, each file has multiple documents, and each document contains multiple sentences. Split sentence through punctuation, such as `。！`. **Secondly, split data into multiple shard based on server hardware (cpu, cpu memory, hard disk) and corpus size.** Each shard contains a part of corpus, and the model needs to train all the shards as one epoch. In this example, split 200G Corpus into 100 shard, and each shard is about 2G. The size of the shard is memory-dependent, taking into account the number of servers, the memory used by the tokenizer, and the memory used by the multi-process training to read the shard (n data parallel requires n\*shard_size memory). **To sum up, data preprocessing and model pretraining requires fighting with hardware, not just GPU.** ```python -python sentence_split.py --input_path /orginal_corpus --output_path /shard --shard 100 +python sentence_split.py --input_path /original_corpus --output_path /shard --shard 100 # This step takes a short time ``` -* `--input_path`: all original corpus, e.g., /orginal_corpus/0.json /orginal_corpus/1.json ... +* `--input_path`: all original corpus, e.g., /original_corpus/0.json /original_corpus/1.json ... * `--output_path`: all shard with split sentences, e.g., /shard/0.txt, /shard/1.txt ... * `--shard`: Number of shard, e.g., 10, 50, or 100 @@ -49,7 +49,7 @@ python sentence_split.py --input_path /orginal_corpus --output_path /shard --sha ] ``` -Output txt: +Output txt: ``` 我今天去打篮球。 @@ -76,7 +76,7 @@ make * `--input_path`: location of all shard with split sentences, e.g., /shard/0.txt, /shard/1.txt ... * `--output_path`: location of all h5 with token_id, input_mask, segment_ids and masked_lm_positions, e.g., /h5/0.h5, /h5/1.h5 ... -* `--tokenizer_path`: tokenizer path contains huggingface tokenizer.json. Download config.json, special_tokens_map.json, vocab.txt and tokenzier.json from [hfl/chinese-roberta-wwm-ext-large](https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/tree/main) +* `--tokenizer_path`: tokenizer path contains huggingface tokenizer.json. Download config.json, special_tokens_map.json, vocab.txt and tokenizer.json from [hfl/chinese-roberta-wwm-ext-large](https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/tree/main) * `--backend`: python or c++, **specifies c++ can obtain faster preprocess speed** * `--dupe_factor`: specifies how many times the preprocessor repeats to create the input from the same article/document * `--worker`: number of process @@ -91,7 +91,7 @@ make 下周请假。 ``` -Output h5+numpy: +Output h5+numpy: ``` 'input_ids': [[id0,id1,id2,id3,id4,id5,id6,0,0..], @@ -102,4 +102,4 @@ make ...] 'masked_lm_positions': [[label1,-1,-1,label2,-1...], ...] -``` \ No newline at end of file +``` diff --git a/examples/language/roberta/preprocessing/get_mask.py b/examples/community/roberta/preprocessing/get_mask.py similarity index 81% rename from examples/language/roberta/preprocessing/get_mask.py rename to examples/community/roberta/preprocessing/get_mask.py index da297f98e6c9..74c97a63a9f3 100644 --- a/examples/language/roberta/preprocessing/get_mask.py +++ b/examples/community/roberta/preprocessing/get_mask.py @@ -1,20 +1,22 @@ -import torch +import collections +import logging import os -from enum import IntEnum -from random import choice import random -import collections import time -import logging +from enum import IntEnum +from random import choice + import jieba +import torch + jieba.setLogLevel(logging.CRITICAL) import re -import numpy as np + import mask +import numpy as np PAD = 0 -MaskedLMInstance = collections.namedtuple("MaskedLMInstance", - ["index", "label"]) +MaskedLMInstance = collections.namedtuple("MaskedLMInstance", ["index", "label"]) def map_to_numpy(data): @@ -22,6 +24,7 @@ def map_to_numpy(data): class PreTrainingDataset(): + def __init__(self, tokenizer, max_seq_length, @@ -43,17 +46,15 @@ def __init__(self, self.mlm_tamper_p = 0.05 self.mlm_maintain_p = 0.1 - def tokenize(self, doc): temp = [] for d in doc: temp.append(self.tokenizer.tokenize(d)) return temp - def create_training_instance(self, instance): is_next = 1 - raw_text_list = self.get_new_segment(instance) + raw_text_list = self.get_new_segment(instance) tokens_a = raw_text_list assert len(tokens_a) == len(instance) # tokens_a, tokens_b, is_next = instance.get_values() @@ -83,8 +84,9 @@ def create_training_instance(self, instance): # Get Masked LM predictions if self.backend == 'c++': - output_tokens, masked_lm_output = mask.create_whole_masked_lm_predictions(tokens, original_tokens, self.vocab_words, - self.tokenizer.vocab, self.max_predictions_per_seq, self.masked_lm_prob) + output_tokens, masked_lm_output = mask.create_whole_masked_lm_predictions( + tokens, original_tokens, self.vocab_words, self.tokenizer.vocab, self.max_predictions_per_seq, + self.masked_lm_prob) elif self.backend == 'python': output_tokens, masked_lm_output = self.create_whole_masked_lm_predictions(tokens) @@ -102,29 +104,25 @@ def create_training_instance(self, instance): map_to_numpy(input_mask), map_to_numpy(segment_ids), map_to_numpy(masked_lm_output), - map_to_numpy([is_next]) + map_to_numpy([is_next]) ]) - def create_masked_lm_predictions(self, tokens): cand_indexes = [] for i, token in enumerate(tokens): if token == "[CLS]" or token == "[SEP]": continue - if (self.do_whole_word_mask and len(cand_indexes) >= 1 and - token.startswith("##")): + if (self.do_whole_word_mask and len(cand_indexes) >= 1 and token.startswith("##")): cand_indexes[-1].append(i) else: cand_indexes.append([i]) - + # cand_indexes.append(i) random.shuffle(cand_indexes) output_tokens = list(tokens) - num_to_predict = min( - self.max_predictions_per_seq, - max(1, int(round(len(tokens) * self.masked_lm_prob)))) + num_to_predict = min(self.max_predictions_per_seq, max(1, int(round(len(tokens) * self.masked_lm_prob)))) masked_lms = [] covered_indexes = set() @@ -145,13 +143,10 @@ def create_masked_lm_predictions(self, tokens): masked_token = tokens[index] # 10% replace w/ random word else: - masked_token = self.vocab_words[random.randint( - 0, - len(self.vocab_words) - 1)] + masked_token = self.vocab_words[random.randint(0, len(self.vocab_words) - 1)] output_tokens[index] = masked_token - masked_lms.append( - MaskedLMInstance(index=index, label=tokens[index])) + masked_lms.append(MaskedLMInstance(index=index, label=tokens[index])) masked_lms = sorted(masked_lms, key=lambda x: x.index) masked_lm_output = [-1] * len(output_tokens) @@ -160,19 +155,17 @@ def create_masked_lm_predictions(self, tokens): return (output_tokens, masked_lm_output) - def get_new_segment(self, segment): """ - 输入一句话，返回一句经过处理的话: 为了支持中文全称mask，将被分开的词，将上特殊标记("#")，使得后续处理模块，能够知道哪些字是属于同一个词的。 - :param segment: 一句话 - :return: 一句处理过的话 + Input a sentence, return a processed sentence: In order to support the Chinese whole word mask, the words that are separated will be marked with a special mark ("#"), so that the subsequent processing module can know which words belong to the same word. + :param segment: a sentence """ seq_cws = jieba.lcut(''.join(segment)) seq_cws_dict = {x: 1 for x in seq_cws} new_segment = [] i = 0 while i < len(segment): - if len(self.rec.findall(segment[i])) == 0: # 不是中文的，原文加进去。 + if len(self.rec.findall(segment[i])) == 0: new_segment.append(segment[i]) i += 1 continue @@ -181,10 +174,10 @@ def get_new_segment(self, segment): for length in range(3, 0, -1): if i + length > len(segment): continue - if ''.join(segment[i: i+length]) in seq_cws_dict: + if ''.join(segment[i:i + length]) in seq_cws_dict: new_segment.append(segment[i]) for l in range(1, length): - new_segment.append('##' + segment[i+l]) + new_segment.append('##' + segment[i + l]) i += length has_add = True break @@ -193,7 +186,6 @@ def get_new_segment(self, segment): i += 1 return new_segment - def create_whole_masked_lm_predictions(self, tokens): """Creates the predictions for the masked LM objective.""" @@ -210,18 +202,16 @@ def create_whole_masked_lm_predictions(self, tokens): # Note that Whole Word Masking does *not* change the training code # at all -- we still predict each WordPiece independently, softmaxed # over the entire vocabulary. - if (self.do_whole_word_mask and len(cand_indexes) >= 1 and - token.startswith("##")): + if (self.do_whole_word_mask and len(cand_indexes) >= 1 and token.startswith("##")): cand_indexes[-1].append(i) else: cand_indexes.append([i]) random.shuffle(cand_indexes) - output_tokens = [t[2:] if len(self.whole_rec.findall(t))>0 else t for t in tokens] # 去掉"##" + output_tokens = [t[2:] if len(self.whole_rec.findall(t)) > 0 else t for t in tokens] # 去掉"##" - num_to_predict = min(self.max_predictions_per_seq, - max(1, int(round(len(tokens) * self.masked_lm_prob)))) + num_to_predict = min(self.max_predictions_per_seq, max(1, int(round(len(tokens) * self.masked_lm_prob)))) masked_lms = [] covered_indexes = set() @@ -249,14 +239,18 @@ def create_whole_masked_lm_predictions(self, tokens): else: # 10% of the time, keep original if random.random() < 0.5: - masked_token = tokens[index][2:] if len(self.whole_rec.findall(tokens[index]))>0 else tokens[index] # 去掉"##" + masked_token = tokens[index][2:] if len(self.whole_rec.findall( + tokens[index])) > 0 else tokens[index] # 去掉"##" # 10% of the time, replace with random word else: masked_token = self.vocab_words[random.randint(0, len(self.vocab_words) - 1)] output_tokens[index] = masked_token - masked_lms.append(MaskedLMInstance(index=index, label=tokens[index][2:] if len(self.whole_rec.findall(tokens[index]))>0 else tokens[index])) + masked_lms.append( + MaskedLMInstance( + index=index, + label=tokens[index][2:] if len(self.whole_rec.findall(tokens[index])) > 0 else tokens[index])) assert len(masked_lms) <= num_to_predict masked_lms = sorted(masked_lms, key=lambda x: x.index) masked_lm_output = [-1] * len(output_tokens) diff --git a/examples/community/roberta/preprocessing/mask.cpp b/examples/community/roberta/preprocessing/mask.cpp new file mode 100644 index 000000000000..d44f58eccfc2 --- /dev/null +++ b/examples/community/roberta/preprocessing/mask.cpp @@ -0,0 +1,190 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +const int32_t LONG_SENTENCE_LEN = 512; + +struct MaskedLMInstance { + int index; + std::string label; + MaskedLMInstance(int index, std::string label) { + this->index = index; + this->label = label; + } +}; + +auto get_new_segment( + std::vector segment, std::vector segment_jieba, + const std::vector chinese_vocab) { // const + // std::unordered_set + // &chinese_vocab + std::unordered_set seq_cws_dict; + for (auto word : segment_jieba) { + seq_cws_dict.insert(word); + } + int i = 0; + std::vector new_segment; + int segment_size = segment.size(); + while (i < segment_size) { + if (!chinese_vocab[i]) { // chinese_vocab.find(segment[i]) == + // chinese_vocab.end() + new_segment.emplace_back(segment[i]); + i += 1; + continue; + } + bool has_add = false; + for (int length = 3; length >= 1; length--) { + if (i + length > segment_size) { + continue; + } + std::string chinese_word = ""; + for (int j = i; j < i + length; j++) { + chinese_word += segment[j]; + } + if (seq_cws_dict.find(chinese_word) != seq_cws_dict.end()) { + new_segment.emplace_back(segment[i]); + for (int j = i + 1; j < i + length; j++) { + new_segment.emplace_back("##" + segment[j]); + } + i += length; + has_add = true; + break; + } + } + if (!has_add) { + new_segment.emplace_back(segment[i]); + i += 1; + } + } + + return new_segment; +} + +bool startsWith(const std::string &s, const std::string &sub) { + return s.find(sub) == 0 ? true : false; +} + +auto create_whole_masked_lm_predictions( + std::vector &tokens, + const std::vector &original_tokens, + const std::vector &vocab_words, + std::map &vocab, const int max_predictions_per_seq, + const double masked_lm_prob) { + // for (auto item : vocab) { + // std::cout << "key=" << std::string(py::str(item.first)) << ", " + // << "value=" << std::string(py::str(item.second)) << + // std::endl; + // } + std::vector > cand_indexes; + std::vector cand_temp; + int tokens_size = tokens.size(); + std::string prefix = "##"; + bool do_whole_masked = true; + + for (int i = 0; i < tokens_size; i++) { + if (tokens[i] == "[CLS]" || tokens[i] == "[SEP]") { + continue; + } + if (do_whole_masked && (cand_indexes.size() > 0) && + (tokens[i].rfind(prefix, 0) == 0)) { + cand_temp.emplace_back(i); + } else { + if (cand_temp.size() > 0) { + cand_indexes.emplace_back(cand_temp); + } + cand_temp.clear(); + cand_temp.emplace_back(i); + } + } + auto seed = std::chrono::system_clock::now().time_since_epoch().count(); + std::shuffle(cand_indexes.begin(), cand_indexes.end(), + std::default_random_engine(seed)); + // for (auto i : cand_indexes) { + // for (auto j : i) { + // std::cout << tokens[j] << " "; + // } + // std::cout << std::endl; + // } + // for (auto i : output_tokens) { + // std::cout << i; + // } + // std::cout << std::endl; + + int num_to_predict = std::min(max_predictions_per_seq, + std::max(1, int(tokens_size * masked_lm_prob))); + // std::cout << num_to_predict << std::endl; + + std::set covered_indexes; + std::vector masked_lm_output(tokens_size, -1); + int vocab_words_len = vocab_words.size(); + std::default_random_engine e(seed); + std::uniform_real_distribution u1(0.0, 1.0); + std::uniform_int_distribution u2(0, vocab_words_len - 1); + int mask_cnt = 0; + std::vector output_tokens; + output_tokens = original_tokens; + + for (auto index_set : cand_indexes) { + if (mask_cnt > num_to_predict) { + break; + } + int index_set_size = index_set.size(); + if (mask_cnt + index_set_size > num_to_predict) { + continue; + } + bool is_any_index_covered = false; + for (auto index : index_set) { + if (covered_indexes.find(index) != covered_indexes.end()) { + is_any_index_covered = true; + break; + } + } + if (is_any_index_covered) { + continue; + } + for (auto index : index_set) { + covered_indexes.insert(index); + std::string masked_token; + if (u1(e) < 0.8) { + masked_token = "[MASK]"; + } else { + if (u1(e) < 0.5) { + masked_token = output_tokens[index]; + } else { + int random_index = u2(e); + masked_token = vocab_words[random_index]; + } + } + // masked_lms.emplace_back(MaskedLMInstance(index, output_tokens[index])); + masked_lm_output[index] = vocab[output_tokens[index]]; + output_tokens[index] = masked_token; + mask_cnt++; + } + } + + // for (auto p : masked_lms) { + // masked_lm_output[p.index] = vocab[p.label]; + // } + return std::make_tuple(output_tokens, masked_lm_output); +} + +PYBIND11_MODULE(mask, m) { + m.def("create_whole_masked_lm_predictions", + &create_whole_masked_lm_predictions); + m.def("get_new_segment", &get_new_segment); +} diff --git a/examples/language/roberta/preprocessing/sentence_split.py b/examples/community/roberta/preprocessing/sentence_split.py similarity index 81% rename from examples/language/roberta/preprocessing/sentence_split.py rename to examples/community/roberta/preprocessing/sentence_split.py index 231be152b067..76e8bd428723 100644 --- a/examples/language/roberta/preprocessing/sentence_split.py +++ b/examples/community/roberta/preprocessing/sentence_split.py @@ -1,35 +1,30 @@ - +import argparse +import functools +import json import multiprocessing import os import re -from tqdm import tqdm -from typing import List -import json import time -import argparse -import functools +from typing import List + +from tqdm import tqdm + def split_sentence(document: str, flag: str = "all", limit: int = 510) -> List[str]: - """ - Args: - document: - flag: Type:str, "all" 中英文标点分句，"zh" 中文标点分句，"en" 英文标点分句 - limit: 默认单句最大长度为510个字符 - Returns: Type:list - """ sent_list = [] try: if flag == "zh": - document = re.sub('(?P([。？！…](?![”’"\'])))', r'\g\n', document) # 单字符断句符 - document = re.sub('(?P([。？！]|…{1,2})[”’"\'])', r'\g\n', document) # 特殊引号 + document = re.sub('(?P([。？！…](?![”’"\'])))', r'\g\n', document) + document = re.sub('(?P([。？！]|…{1,2})[”’"\'])', r'\g\n', document) elif flag == "en": - document = re.sub('(?P([.?!](?![”’"\'])))', r'\g\n', document) # 英文单字符断句符 - document = re.sub('(?P([?!.]["\']))', r'\g\n', document) # 特殊引号 + document = re.sub('(?P([.?!](?![”’"\'])))', r'\g\n', document) + document = re.sub('(?P([?!.]["\']))', r'\g\n', + document) # Special quotation marks else: - document = re.sub('(?P([。？！….?!](?![”’"\'])))', r'\g\n', document) # 单字符断句符 - + document = re.sub('(?P([。？！….?!](?![”’"\'])))', r'\g\n', document) + document = re.sub('(?P(([。？！.!?]|…{1,2})[”’"\']))', r'\g\n', - document) # 特殊引号 + document) # Special quotation marks sent_list_ori = document.splitlines() for sent in sent_list_ori: @@ -50,17 +45,15 @@ def split_sentence(document: str, flag: str = "all", limit: int = 510) -> List[s return sent_list -def get_sent(output_path, - input_path, - fin_list=[], host=-1, seq_len=512) -> None: +def get_sent(output_path, input_path, fin_list=[], host=-1, seq_len=512) -> None: workers = 32 if input_path[-1] == '/': input_path = input_path[:-1] - + cur_path = os.path.join(output_path, str(host) + '.txt') - new_split_sentence = functools.partial(split_sentence, limit=seq_len-2) + new_split_sentence = functools.partial(split_sentence, limit=seq_len - 2) with open(cur_path, 'w', encoding='utf-8') as f: for fi, fin_path in enumerate(fin_list): if not os.path.exists(os.path.join(input_path, fin_path[0])): @@ -69,7 +62,7 @@ def get_sent(output_path, continue print("Processing ", fin_path[0], " ", fi) - + with open(os.path.join(input_path, fin_path[0]), 'r') as fin: f_data = [l['content'] for l in json.load(fin)] @@ -106,17 +99,17 @@ def getFileSize(filepath, shard): real_shard.append(temp) accu_size = 0 temp = [] - + if len(temp) > 0: real_shard.append(temp) - + return real_shard def get_start_end(real_shard, base=0, server_num=10, server_name='GPU'): import socket host = int(socket.gethostname().split(server_name)[-1]) - + fin_list = real_shard[server_num * base + host - 1] print(fin_list) print(f'I am server {host}, process {server_num * base + host - 1}, len {len(fin_list)}') @@ -133,28 +126,24 @@ def get_start_end(real_shard, base=0, server_num=10, server_name='GPU'): parser.add_argument('--output_path', type=str, required=True, help='output path of shard which has split sentence') args = parser.parse_args() - server_num = args.server_num + server_num = args.server_num seq_len = args.seq_len - shard = args.shard + shard = args.shard input_path = args.input_path - output_path = args.output_path + output_path = args.output_path real_shard = getFileSize(input_path, shard) start = time.time() for index, shard in enumerate(real_shard): - get_sent(output_path, - input_path, - fin_list=shard, - host=index, - seq_len=seq_len) + get_sent(output_path, input_path, fin_list=shard, host=index, seq_len=seq_len) print(f'cost {str(time.time() - start)}') # if you have multiple server, you can use code below or modify code to openmpi - + # for i in range(len(real_shard) // server_num + 1): # fin_list, host = get_start_end(real_shard, i) - + # start = time.time() # get_sent(output_path, # input_path, diff --git a/examples/language/roberta/preprocessing/tokenize_mask.py b/examples/community/roberta/preprocessing/tokenize_mask.py similarity index 72% rename from examples/language/roberta/preprocessing/tokenize_mask.py rename to examples/community/roberta/preprocessing/tokenize_mask.py index b33871d5d037..f3d49c3d965f 100644 --- a/examples/language/roberta/preprocessing/tokenize_mask.py +++ b/examples/community/roberta/preprocessing/tokenize_mask.py @@ -1,22 +1,22 @@ -import time +import argparse +import multiprocessing import os -import psutil -import h5py import socket -import argparse +import time +from random import shuffle + +import h5py import numpy as np -import multiprocessing +import psutil +from get_mask import PreTrainingDataset from tqdm import tqdm -from random import shuffle from transformers import AutoTokenizer -from get_mask import PreTrainingDataset def get_raw_instance(document, max_sequence_length=512): - """ - 获取初步的训练实例，将整段按照max_sequence_length切分成多个部分,并以多个处理好的实例的形式返回。 - :param document: 一整段 + Get the initial training instances, split the whole segment into multiple parts according to the max_sequence_length, and return as multiple processed instances. + :param document: document :param max_sequence_length: :return: a list. each element is a sequence of text """ @@ -26,31 +26,29 @@ def get_raw_instance(document, max_sequence_length=512): sizes = [len(seq) for seq in document] result_list = [] - curr_seq = [] # 当前处理的序列 + curr_seq = [] sz_idx = 0 while sz_idx < len(sizes): - # 当前句子加上新的句子，如果长度小于最大限制，则合并当前句子和新句子；否则即超过了最大限制，那么做为一个新的序列加到目标列表中 - - if len(curr_seq) + sizes[sz_idx] <= max_sequence_length_allowed: # or len(curr_seq)==0: + + if len(curr_seq) + sizes[sz_idx] <= max_sequence_length_allowed: # or len(curr_seq)==0: curr_seq += document[sz_idx] sz_idx += 1 elif sizes[sz_idx] >= max_sequence_length_allowed: if len(curr_seq) > 0: result_list.append(curr_seq) curr_seq = [] - result_list.append(document[sz_idx][ : max_sequence_length_allowed]) + result_list.append(document[sz_idx][:max_sequence_length_allowed]) sz_idx += 1 else: result_list.append(curr_seq) curr_seq = [] - # 对最后一个序列进行处理，如果太短的话，丢弃掉。 - if len(curr_seq) > max_sequence_length_allowed / 2: # /2 + + if len(curr_seq) > max_sequence_length_allowed / 2: # /2 result_list.append(curr_seq) - # # 计算总共可以得到多少份 # num_instance=int(len(big_list)/max_sequence_length_allowed)+1 # print("num_instance:",num_instance) - # # 切分成多份，添加到列表中 + # result_list=[] # for j in range(num_instance): # index=j*max_sequence_length_allowed @@ -72,8 +70,7 @@ def split_numpy_chunk(path, tokenizer, pretrain_data, host): # document = line # if len(document.split("")) <= 3: # continue - if len(line - ) > 0 and line[:2] == "]]": # This is end of document + if len(line) > 0 and line[:2] == "]]": # This is end of document documents.append(document) document = [] elif len(line) >= 2: @@ -86,8 +83,8 @@ def split_numpy_chunk(path, tokenizer, pretrain_data, host): # print(len(documents)) # print(len(documents[0])) # print(documents[0][0:10]) - from typing import List import multiprocessing + from typing import List ans = [] for docs in tqdm(documents): @@ -100,7 +97,7 @@ def split_numpy_chunk(path, tokenizer, pretrain_data, host): raw_ins = get_raw_instance(a) instances.extend(raw_ins) del ans - + print('len instance', len(instances)) sen_num = len(instances) @@ -118,21 +115,15 @@ def split_numpy_chunk(path, tokenizer, pretrain_data, host): masked_lm_output[index] = mask_dict[3] with h5py.File(f'/output/{host}.h5', 'w') as hf: - hf.create_dataset("input_ids", data=input_ids) - hf.create_dataset("input_mask", data=input_ids) - hf.create_dataset("segment_ids", data=segment_ids) - hf.create_dataset("masked_lm_positions", data=masked_lm_output) + hf.create_dataset("input_ids", data=input_ids) + hf.create_dataset("input_mask", data=input_ids) + hf.create_dataset("segment_ids", data=segment_ids) + hf.create_dataset("masked_lm_positions", data=masked_lm_output) del instances -def split_numpy_chunk_pool(input_path, - output_path, - pretrain_data, - worker, - dupe_factor, - seq_len, - file_name): +def split_numpy_chunk_pool(input_path, output_path, pretrain_data, worker, dupe_factor, seq_len, file_name): if os.path.exists(os.path.join(output_path, f'{file_name}.h5')): print(f'{file_name}.h5 exists') @@ -146,8 +137,7 @@ def split_numpy_chunk_pool(input_path, document = [] for i, line in enumerate(tqdm(fd)): line = line.strip() - if len(line - ) > 0 and line[:2] == "]]": # This is end of document + if len(line) > 0 and line[:2] == "]]": # This is end of document documents.append(document) document = [] elif len(line) >= 2: @@ -155,7 +145,7 @@ def split_numpy_chunk_pool(input_path, if len(document) > 0: documents.append(document) print(f'read_file cost {time.time() - s}, length is {len(documents)}') - + ans = [] s = time.time() pool = multiprocessing.Pool(worker) @@ -171,7 +161,7 @@ def split_numpy_chunk_pool(input_path, raw_ins = get_raw_instance(a, max_sequence_length=seq_len) instances.extend(raw_ins) del ans - + print('len instance', len(instances)) new_instances = [] @@ -201,10 +191,10 @@ def split_numpy_chunk_pool(input_path, print((time.time() - s) / 60) with h5py.File(os.path.join(output_path, f'{file_name}.h5'), 'w') as hf: - hf.create_dataset("input_ids", data=input_ids) - hf.create_dataset("input_mask", data=input_mask) - hf.create_dataset("segment_ids", data=segment_ids) - hf.create_dataset("masked_lm_positions", data=masked_lm_output) + hf.create_dataset("input_ids", data=input_ids) + hf.create_dataset("input_mask", data=input_mask) + hf.create_dataset("segment_ids", data=segment_ids) + hf.create_dataset("masked_lm_positions", data=masked_lm_output) del instances @@ -214,22 +204,31 @@ def split_numpy_chunk_pool(input_path, parser = argparse.ArgumentParser() parser.add_argument('--tokenizer_path', type=str, required=True, default=10, help='path of tokenizer') parser.add_argument('--seq_len', type=int, default=512, help='sequence length') - parser.add_argument('--max_predictions_per_seq', type=int, default=80, help='number of shards, e.g., 10, 50, or 100') + parser.add_argument('--max_predictions_per_seq', + type=int, + default=80, + help='number of shards, e.g., 10, 50, or 100') parser.add_argument('--input_path', type=str, required=True, help='input path of shard which has split sentence') parser.add_argument('--output_path', type=str, required=True, help='output path of h5 contains token id') - parser.add_argument('--backend', type=str, default='python', help='backend of mask token, python, c++, numpy respectively') - parser.add_argument('--dupe_factor', type=int, default=1, help='specifies how many times the preprocessor repeats to create the input from the same article/document') + parser.add_argument('--backend', + type=str, + default='python', + help='backend of mask token, python, c++, numpy respectively') + parser.add_argument( + '--dupe_factor', + type=int, + default=1, + help='specifies how many times the preprocessor repeats to create the input from the same article/document') parser.add_argument('--worker', type=int, default=32, help='number of process') parser.add_argument('--server_num', type=int, default=10, help='number of servers') args = parser.parse_args() tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) - pretrain_data = PreTrainingDataset(tokenizer, - args.seq_len, - args.backend, - max_predictions_per_seq=args.max_predictions_per_seq) - - + pretrain_data = PreTrainingDataset(tokenizer, + args.seq_len, + args.backend, + max_predictions_per_seq=args.max_predictions_per_seq) + data_len = len(os.listdir(args.input_path)) for i in range(data_len): @@ -237,15 +236,10 @@ def split_numpy_chunk_pool(input_path, if os.path.exists(input_path): start = time.time() print(f'process {input_path}') - split_numpy_chunk_pool(input_path, - args.output_path, - pretrain_data, - args.worker, - args.dupe_factor, - args.seq_len, - i) + split_numpy_chunk_pool(input_path, args.output_path, pretrain_data, args.worker, args.dupe_factor, + args.seq_len, i) end_ = time.time() - print(u'memory：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ) + print(u'memory：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024)) print(f'has cost {(end_ - start) / 60}') print('-' * 100) print('') @@ -259,9 +253,9 @@ def split_numpy_chunk_pool(input_path, # if os.path.exists(input_path): # start = time.time() # print(f'I am server {host}, process {input_path}') - # split_numpy_chunk_pool(input_path, - # args.output_path, - # pretrain_data, + # split_numpy_chunk_pool(input_path, + # args.output_path, + # pretrain_data, # args.worker, # args.dupe_factor, # args.seq_len, @@ -271,5 +265,3 @@ def split_numpy_chunk_pool(input_path, # print(f'has cost {(end_ - start) / 60}') # print('-' * 100) # print('') - - diff --git a/examples/language/roberta/pretraining/README.md b/examples/community/roberta/pretraining/README.md similarity index 91% rename from examples/language/roberta/pretraining/README.md rename to examples/community/roberta/pretraining/README.md index 055d6969654d..8abe48aa6c0e 100644 --- a/examples/language/roberta/pretraining/README.md +++ b/examples/community/roberta/pretraining/README.md @@ -13,12 +13,11 @@ bash run_pretrain.sh * `--bert_config`: config.json which represent model * `--mlm`: model type of backbone, bert or deberta_v2 -2. if resume training from earylier checkpoint, run the script below. +2. if resume training from earlier checkpoint, run the script below. ```shell bash run_pretrain_resume.sh ``` * `--resume_train`: whether to resume training -* `--load_pretrain_model`: absolute path which contains model checkpoint -* `--load_optimizer_lr`: absolute path which contains optimizer checkpoint - +* `--load_pretrain_model`: absolute path which contains model checkpoint +* `--load_optimizer_lr`: absolute path which contains optimizer checkpoint diff --git a/examples/community/roberta/pretraining/arguments.py b/examples/community/roberta/pretraining/arguments.py new file mode 100644 index 000000000000..e0702ceb59b0 --- /dev/null +++ b/examples/community/roberta/pretraining/arguments.py @@ -0,0 +1,87 @@ +from numpy import require + +import colossalai + +__all__ = ['parse_args'] + + +def parse_args(): + parser = colossalai.get_default_parser() + + parser.add_argument( + "--distplan", + type=str, + default='CAI_Gemini', + help="The distributed plan [colossalai, zero1, zero2, torch_ddp, torch_zero].", + ) + parser.add_argument( + "--tp_degree", + type=int, + default=1, + help="Tensor Parallelism Degree. Valid when using colossalai as dist plan.", + ) + parser.add_argument( + "--placement", + type=str, + default='cpu', + help="Placement Policy for Gemini. Valid when using colossalai as dist plan.", + ) + parser.add_argument( + "--shardinit", + action='store_true', + help= + "Shard the tensors when init the model to shrink peak memory size on the assigned device. Valid when using colossalai as dist plan.", + ) + + parser.add_argument('--lr', type=float, required=True, help='initial learning rate') + parser.add_argument('--epoch', type=int, required=True, help='number of epoch') + parser.add_argument('--data_path_prefix', type=str, required=True, help="location of the train data corpus") + parser.add_argument('--eval_data_path_prefix', + type=str, + required=True, + help='location of the evaluation data corpus') + parser.add_argument('--tokenizer_path', type=str, required=True, help='location of the tokenizer') + parser.add_argument('--max_seq_length', type=int, default=512, help='sequence length') + parser.add_argument('--refresh_bucket_size', + type=int, + default=1, + help="This param makes sure that a certain task is repeated for this time steps to \ + optimize on the back propagation speed with APEX's DistributedDataParallel") + parser.add_argument("--max_predictions_per_seq", + "--max_pred", + default=80, + type=int, + help="The maximum number of masked tokens in a sequence to be predicted.") + parser.add_argument("--gradient_accumulation_steps", default=1, type=int, help="accumulation_steps") + parser.add_argument("--train_micro_batch_size_per_gpu", default=2, type=int, required=True, help="train batch size") + parser.add_argument("--eval_micro_batch_size_per_gpu", default=2, type=int, required=True, help="eval batch size") + parser.add_argument("--num_workers", default=8, type=int, help="") + parser.add_argument("--async_worker", action='store_true', help="") + parser.add_argument("--bert_config", required=True, type=str, help="location of config.json") + parser.add_argument("--wandb", action='store_true', help="use wandb to watch model") + parser.add_argument("--wandb_project_name", default='roberta', help="wandb project name") + parser.add_argument("--log_interval", default=100, type=int, help="report interval") + parser.add_argument("--log_path", type=str, required=True, help="log file which records train step") + parser.add_argument("--tensorboard_path", type=str, required=True, help="location of tensorboard file") + parser.add_argument("--colossal_config", + type=str, + required=True, + help="colossal config, which contains zero config and so on") + parser.add_argument("--ckpt_path", + type=str, + required=True, + help="location of saving checkpoint, which contains model and optimizer") + parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") + parser.add_argument('--vscode_debug', action='store_true', help="use vscode to debug") + parser.add_argument('--load_pretrain_model', default='', type=str, help="location of model's checkpoint") + parser.add_argument( + '--load_optimizer_lr', + default='', + type=str, + help="location of checkpoint, which contains optimizer, learning rate, epoch, shard and global_step") + parser.add_argument('--resume_train', action='store_true', help="whether resume training from a early checkpoint") + parser.add_argument('--mlm', default='bert', type=str, help="model type, bert or deberta") + parser.add_argument('--checkpoint_activations', action='store_true', help="whether to use gradient checkpointing") + + args = parser.parse_args() + return args diff --git a/examples/language/roberta/pretraining/bert_dataset_provider.py b/examples/community/roberta/pretraining/bert_dataset_provider.py similarity index 99% rename from examples/language/roberta/pretraining/bert_dataset_provider.py rename to examples/community/roberta/pretraining/bert_dataset_provider.py index 1d8cf2a910e9..eaf165ed18f4 100644 --- a/examples/language/roberta/pretraining/bert_dataset_provider.py +++ b/examples/community/roberta/pretraining/bert_dataset_provider.py @@ -1,4 +1,5 @@ class BertDatasetProviderInterface: + def get_shard(self, index, shuffle=True): raise NotImplementedError diff --git a/examples/language/roberta/pretraining/evaluation.py b/examples/community/roberta/pretraining/evaluation.py similarity index 68% rename from examples/language/roberta/pretraining/evaluation.py rename to examples/community/roberta/pretraining/evaluation.py index 83f94082f6c0..009242cd1cf5 100644 --- a/examples/language/roberta/pretraining/evaluation.py +++ b/examples/community/roberta/pretraining/evaluation.py @@ -1,15 +1,17 @@ -import os import math +import os + import torch +from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider from tqdm import tqdm -from utils.global_vars import get_timers, get_tensorboard_writer -from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider +from utils.global_vars import get_tensorboard_writer, get_timers -def evaluate(engine, args, logger, global_step): + +def evaluate(model, args, logger, global_step, criterion): evaluate_dataset_provider = NvidiaBertDatasetProvider(args, evaluate=True) start_shard = 0 - engine.eval() + model.eval() timers = get_timers() eval_step = 0 eval_loss = 0 @@ -20,16 +22,19 @@ def evaluate(engine, args, logger, global_step): for shard in range(start_shard, len(os.listdir(args.eval_data_path_prefix))): - timers('eval_shard_time').start() + timers('eval_shard_time').start() dataset_iterator, total_length = evaluate_dataset_provider.get_shard(shard) # evaluate_dataset_provider.prefetch_shard(shard + 1) if torch.distributed.get_rank() == 0: - iterator_data = tqdm(enumerate(dataset_iterator), total=(total_length // args.eval_micro_batch_size_per_gpu // world_size), colour='MAGENTA', smoothing=1) + iterator_data = tqdm(enumerate(dataset_iterator), + total=(total_length // args.eval_micro_batch_size_per_gpu // world_size), + colour='MAGENTA', + smoothing=1) else: iterator_data = enumerate(dataset_iterator) - - for step, batch_data in iterator_data: #tqdm(enumerate(dataset_iterator), total=(total_length // args.train_micro_batch_size_per_gpu // world_size), colour='cyan', smoothing=1): + + for step, batch_data in iterator_data: #tqdm(enumerate(dataset_iterator), total=(total_length // args.train_micro_batch_size_per_gpu // world_size), colour='cyan', smoothing=1): # batch_data = pretrain_dataset_provider.get_batch(batch_index) eval_step += 1 @@ -39,9 +44,9 @@ def evaluate(engine, args, logger, global_step): mlm_label = batch_data[3].cuda() # nsp_label = batch_data[5].cuda() - output = engine(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) - - loss = engine.criterion(output.logits, mlm_label)#prediction_scores + output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) + + loss = criterion(output.logits, mlm_label) #prediction_scores evaluate_dataset_provider.prefetch_batch() eval_loss += loss.float().item() @@ -54,10 +59,10 @@ def evaluate(engine, args, logger, global_step): if args.wandb and torch.distributed.get_rank() == 0: tensorboard_log = get_tensorboard_writer() tensorboard_log.log_eval({ - 'loss': cur_loss, - 'ppl': ppl, - 'mins_batch': elapsed_time_per_iteration - }, global_step) + 'loss': cur_loss, + 'ppl': ppl, + 'mins_batch': elapsed_time_per_iteration + }, global_step) eval_log_str = f'evaluation shard: {shard} | step: {eval_step} | elapsed_time: {elapsed_time / 60 :.3f} minutes ' + \ f'| mins/batch: {elapsed_time_per_iteration :.3f} seconds | loss: {cur_loss:.7f} | ppl: {ppl:.7f}' @@ -67,5 +72,5 @@ def evaluate(engine, args, logger, global_step): logger.info('') evaluate_dataset_provider.release_shard() - engine.train() + model.train() return cur_loss diff --git a/examples/language/roberta/pretraining/hostfile b/examples/community/roberta/pretraining/hostfile similarity index 100% rename from examples/language/roberta/pretraining/hostfile rename to examples/community/roberta/pretraining/hostfile diff --git a/examples/language/roberta/pretraining/loss.py b/examples/community/roberta/pretraining/loss.py similarity index 91% rename from examples/language/roberta/pretraining/loss.py rename to examples/community/roberta/pretraining/loss.py index dc4f872a755d..989c2bd5c450 100644 --- a/examples/language/roberta/pretraining/loss.py +++ b/examples/community/roberta/pretraining/loss.py @@ -13,5 +13,5 @@ def __init__(self, vocab_size): def forward(self, prediction_scores, masked_lm_labels, next_sentence_labels=None): masked_lm_loss = self.loss_fn(prediction_scores.view(-1, self.vocab_size), masked_lm_labels.view(-1)) # next_sentence_loss = self.loss_fn(seq_relationship_score.view(-1, 2), next_sentence_labels.view(-1)) - total_loss = masked_lm_loss #+ next_sentence_loss + total_loss = masked_lm_loss #+ next_sentence_loss return total_loss diff --git a/examples/language/roberta/pretraining/model/bert.py b/examples/community/roberta/pretraining/model/bert.py similarity index 96% rename from examples/language/roberta/pretraining/model/bert.py rename to examples/community/roberta/pretraining/model/bert.py index 67c85f760776..abdf925d0540 100644 --- a/examples/language/roberta/pretraining/model/bert.py +++ b/examples/community/roberta/pretraining/model/bert.py @@ -15,7 +15,6 @@ # limitations under the License. """PyTorch BERT model.""" - import math import os import warnings @@ -27,7 +26,6 @@ from packaging import version from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss - from transformers.activations import ACT2FN from transformers.modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, @@ -41,8 +39,9 @@ TokenClassifierOutput, ) from transformers.modeling_utils import PreTrainedModel +from transformers.models.bert.configuration_bert import BertConfig from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer -from transformers.utils import ( +from transformers.utils import ( ModelOutput, add_code_sample_docstrings, add_start_docstrings, @@ -50,8 +49,6 @@ logging, replace_return_docstrings, ) -from transformers.models.bert.configuration_bert import BertConfig - logger = logging.get_logger(__name__) @@ -62,8 +59,7 @@ # TokenClassification docstring _CHECKPOINT_FOR_TOKEN_CLASSIFICATION = "dbmdz/bert-large-cased-finetuned-conll03-english" _TOKEN_CLASS_EXPECTED_OUTPUT = ( - "['O', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 'O', 'O', 'O', 'O', 'I-LOC', 'O', 'I-LOC', 'I-LOC'] " -) + "['O', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 'O', 'O', 'O', 'O', 'I-LOC', 'O', 'I-LOC', 'I-LOC'] ") _TOKEN_CLASS_EXPECTED_LOSS = 0.01 # QuestionAnswering docstring @@ -78,7 +74,6 @@ _SEQ_CLASS_EXPECTED_OUTPUT = "'LABEL_1'" _SEQ_CLASS_EXPECTED_LOSS = 0.01 - BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [ "bert-base-uncased", "bert-large-uncased", @@ -114,10 +109,8 @@ def load_tf_weights_in_bert(model, config, tf_checkpoint_path): import numpy as np import tensorflow as tf except ImportError: - logger.error( - "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " - "https://www.tensorflow.org/install/ for installation instructions." - ) + logger.error("Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions.") raise tf_path = os.path.abspath(tf_checkpoint_path) logger.info(f"Converting TensorFlow checkpoint from {tf_path}") @@ -135,10 +128,8 @@ def load_tf_weights_in_bert(model, config, tf_checkpoint_path): name = name.split("/") # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v # which are not required for using pretrained model - if any( - n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] - for n in name - ): + if any(n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] + for n in name): logger.info(f"Skipping {'/'.join(name)}") continue pointer = model @@ -218,7 +209,7 @@ def forward( seq_length = input_shape[1] if position_ids is None: - position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length] + position_ids = self.position_ids[:, past_key_values_length:seq_length + past_key_values_length] # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves @@ -245,13 +236,12 @@ def forward( class BertSelfAttention(nn.Module): + def __init__(self, config, position_embedding_type=None): super().__init__() if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): - raise ValueError( - f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " - f"heads ({config.num_attention_heads})" - ) + raise ValueError(f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})") self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) @@ -262,9 +252,7 @@ def __init__(self, config, position_embedding_type=None): self.value = nn.Linear(config.hidden_size, self.all_head_size) self.dropout = nn.Dropout(config.attention_probs_dropout_prob) - self.position_embedding_type = position_embedding_type or getattr( - config, "position_embedding_type", "absolute" - ) + self.position_embedding_type = position_embedding_type or getattr(config, "position_embedding_type", "absolute") if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query": self.max_position_embeddings = config.max_position_embeddings self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size) @@ -332,14 +320,14 @@ def forward( position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1) distance = position_ids_l - position_ids_r positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1) - positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility + positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility if self.position_embedding_type == "relative_key": relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) attention_scores = attention_scores + relative_position_scores elif self.position_embedding_type == "relative_key_query": relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding) - relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding) + relative_position_scores_key = torch.einsum("bhld,lrd->bhlr", key_layer, positional_embedding) attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key attention_scores = attention_scores / math.sqrt(self.attention_head_size) @@ -372,6 +360,7 @@ def forward( class BertSelfOutput(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -386,6 +375,7 @@ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> to class BertAttention(nn.Module): + def __init__(self, config, position_embedding_type=None): super().__init__() self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type) @@ -395,9 +385,8 @@ def __init__(self, config, position_embedding_type=None): def prune_heads(self, heads): if len(heads) == 0: return - heads, index = find_pruneable_heads_and_indices( - heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads - ) + heads, index = find_pruneable_heads_and_indices(heads, self.self.num_attention_heads, + self.self.attention_head_size, self.pruned_heads) # Prune linear layers self.self.query = prune_linear_layer(self.self.query, index) @@ -430,11 +419,12 @@ def forward( output_attentions, ) attention_output = self.output(self_outputs[0], hidden_states) - outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them return outputs class BertIntermediate(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.intermediate_size) @@ -450,6 +440,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class BertOutput(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.intermediate_size, config.hidden_size) @@ -464,6 +455,7 @@ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> to class BertLayer(nn.Module): + def __init__(self, config): super().__init__() self.chunk_size_feed_forward = config.chunk_size_feed_forward @@ -504,15 +496,14 @@ def forward( outputs = self_attention_outputs[1:-1] present_key_value = self_attention_outputs[-1] else: - outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights cross_attn_present_key_value = None if self.is_decoder and encoder_hidden_states is not None: if not hasattr(self, "crossattention"): raise ValueError( f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers" - " by setting `config.add_cross_attention=True`" - ) + " by setting `config.add_cross_attention=True`") # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None @@ -526,15 +517,14 @@ def forward( output_attentions, ) attention_output = cross_attention_outputs[0] - outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights + outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights # add cross-attn cache to positions 3,4 of present_key_value tuple cross_attn_present_key_value = cross_attention_outputs[-1] present_key_value = present_key_value + cross_attn_present_key_value - layer_output = apply_chunking_to_forward( - self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output - ) + layer_output = apply_chunking_to_forward(self.feed_forward_chunk, self.chunk_size_feed_forward, + self.seq_len_dim, attention_output) outputs = (layer_output,) + outputs # if decoder, return the attn key/values as the last output @@ -550,6 +540,7 @@ def feed_forward_chunk(self, attention_output): class BertEncoder(nn.Module): + def __init__(self, config): super().__init__() self.config = config @@ -585,11 +576,11 @@ def forward( if use_cache: logger.warning( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...") use_cache = False def create_custom_forward(module): + def custom_forward(*inputs): return module(*inputs, past_key_value, output_attentions) @@ -626,17 +617,13 @@ def custom_forward(*inputs): all_hidden_states = all_hidden_states + (hidden_states,) if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_decoder_cache, - all_hidden_states, - all_self_attentions, - all_cross_attentions, - ] - if v is not None - ) + return tuple(v for v in [ + hidden_states, + next_decoder_cache, + all_hidden_states, + all_self_attentions, + all_cross_attentions, + ] if v is not None) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=next_decoder_cache, @@ -647,6 +634,7 @@ def custom_forward(*inputs): class BertPooler(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -662,6 +650,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class BertPredictionHeadTransform(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -679,6 +668,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class BertLMPredictionHead(nn.Module): + def __init__(self, config): super().__init__() self.transform = BertPredictionHeadTransform(config) @@ -699,6 +689,7 @@ def forward(self, hidden_states): class BertOnlyMLMHead(nn.Module): + def __init__(self, config): super().__init__() self.predictions = BertLMPredictionHead(config) @@ -709,6 +700,7 @@ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: class BertOnlyNSPHead(nn.Module): + def __init__(self, config): super().__init__() self.seq_relationship = nn.Linear(config.hidden_size, 2) @@ -719,6 +711,7 @@ def forward(self, pooled_output): class BertPreTrainingHeads(nn.Module): + def __init__(self, config): super().__init__() self.predictions = BertLMPredictionHead(config) @@ -950,9 +943,8 @@ def forward( `past_key_values`). """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if self.config.is_decoder: @@ -1051,6 +1043,7 @@ def forward( BERT_START_DOCSTRING, ) class BertForPreTraining(BertPreTrainedModel): + def __init__(self, config): super().__init__(config) @@ -1151,9 +1144,8 @@ def forward( ) -@add_start_docstrings( - """Bert Model with a `language modeling` head on top for CLM fine-tuning.""", BERT_START_DOCSTRING -) +@add_start_docstrings("""Bert Model with a `language modeling` head on top for CLM fine-tuning.""", + BERT_START_DOCSTRING) class BertLMHeadModel(BertPreTrainedModel): _keys_to_ignore_on_load_unexpected = [r"pooler"] @@ -1298,10 +1290,8 @@ def __init__(self, config): super().__init__(config) if config.is_decoder: - logger.warning( - "If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for " - "bi-directional self-attention." - ) + logger.warning("If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for " + "bi-directional self-attention.") self.bert = BertModel(config, add_pooling_layer=False) self.cls = BertOnlyMLMHead(config) @@ -1367,7 +1357,7 @@ def forward( masked_lm_loss = None if labels is not None: - loss_fct = CrossEntropyLoss() # -100 index = padding token + loss_fct = CrossEntropyLoss() # -100 index = padding token masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) if not return_dict: @@ -1390,9 +1380,10 @@ def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_ raise ValueError("The PAD token should be defined for generation") attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1) - dummy_token = torch.full( - (effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device - ) + dummy_token = torch.full((effective_batch_size, 1), + self.config.pad_token_id, + dtype=torch.long, + device=input_ids.device) input_ids = torch.cat([input_ids, dummy_token], dim=1) return {"input_ids": input_ids, "attention_mask": attention_mask} @@ -1403,6 +1394,7 @@ def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_ BERT_START_DOCSTRING, ) class BertForNextSentencePrediction(BertPreTrainedModel): + def __init__(self, config): super().__init__(config) @@ -1508,15 +1500,15 @@ def forward( BERT_START_DOCSTRING, ) class BertForSequenceClassification(BertPreTrainedModel): + def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.config = config self.bert = BertModel(config) - classifier_dropout = ( - config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob - ) + classifier_dropout = (config.classifier_dropout + if config.classifier_dropout is not None else config.hidden_dropout_prob) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, config.num_labels) @@ -1612,13 +1604,13 @@ def forward( BERT_START_DOCSTRING, ) class BertForMultipleChoice(BertPreTrainedModel): + def __init__(self, config): super().__init__(config) self.bert = BertModel(config) - classifier_dropout = ( - config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob - ) + classifier_dropout = (config.classifier_dropout + if config.classifier_dropout is not None else config.hidden_dropout_prob) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, 1) @@ -1658,11 +1650,8 @@ def forward( attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None - inputs_embeds = ( - inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) - if inputs_embeds is not None - else None - ) + inputs_embeds = (inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None else None) outputs = self.bert( input_ids, @@ -1715,9 +1704,8 @@ def __init__(self, config): self.num_labels = config.num_labels self.bert = BertModel(config, add_pooling_layer=False) - classifier_dropout = ( - config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob - ) + classifier_dropout = (config.classifier_dropout + if config.classifier_dropout is not None else config.hidden_dropout_prob) self.dropout = nn.Dropout(classifier_dropout) self.classifier = nn.Linear(config.hidden_size, config.num_labels) diff --git a/examples/language/roberta/pretraining/model/deberta_v2.py b/examples/community/roberta/pretraining/model/deberta_v2.py similarity index 92% rename from examples/language/roberta/pretraining/model/deberta_v2.py rename to examples/community/roberta/pretraining/model/deberta_v2.py index c6ce82847f75..5fc284911e38 100644 --- a/examples/language/roberta/pretraining/model/deberta_v2.py +++ b/examples/community/roberta/pretraining/model/deberta_v2.py @@ -23,7 +23,7 @@ import torch.utils.checkpoint from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss - +from transformers import FillMaskPipeline, T5ForConditionalGeneration, T5Tokenizer from transformers.activations import ACT2FN from transformers.modeling_outputs import ( BaseModelOutput, @@ -34,10 +34,14 @@ TokenClassifierOutput, ) from transformers.modeling_utils import PreTrainedModel -from transformers.pytorch_utils import softmax_backward_data -from transformers.utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging from transformers.models.deberta_v2.configuration_deberta_v2 import DebertaV2Config -from transformers import T5Tokenizer, T5ForConditionalGeneration, FillMaskPipeline +from transformers.pytorch_utils import softmax_backward_data +from transformers.utils import ( + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, +) logger = logging.get_logger(__name__) @@ -55,6 +59,7 @@ # Copied from transformers.models.deberta.modeling_deberta.ContextPooler class ContextPooler(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size) @@ -133,15 +138,15 @@ def symbolic(g, self, mask, dim): g.op("Sub", g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)), mask_cast_value), to_i=sym_help.cast_pytorch_to_onnx["Byte"], ) - output = masked_fill( - g, self, r_mask, g.op("Constant", value_t=torch.tensor(torch.finfo(self.type().dtype()).min)) - ) + output = masked_fill(g, self, r_mask, + g.op("Constant", value_t=torch.tensor(torch.finfo(self.type().dtype()).min))) output = softmax(g, output, dim) return masked_fill(g, output, r_mask, g.op("Constant", value_t=torch.tensor(0, dtype=torch.uint8))) # Copied from transformers.models.deberta.modeling_deberta.DropoutContext class DropoutContext(object): + def __init__(self): self.dropout = 0 self.mask = None @@ -244,6 +249,7 @@ def get_context(self): # Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm class DebertaV2SelfOutput(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -259,6 +265,7 @@ def forward(self, hidden_states, input_tensor): # Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2 class DebertaV2Attention(nn.Module): + def __init__(self, config): super().__init__() self.self = DisentangledSelfAttention(config) @@ -296,6 +303,7 @@ def forward( # Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2 class DebertaV2Intermediate(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.intermediate_size) @@ -312,6 +320,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Copied from transformers.models.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm class DebertaV2Output(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.intermediate_size, config.hidden_size) @@ -328,6 +337,7 @@ def forward(self, hidden_states, input_tensor): # Copied from transformers.models.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2 class DebertaV2Layer(nn.Module): + def __init__(self, config): super().__init__() self.attention = DebertaV2Attention(config) @@ -362,14 +372,17 @@ def forward( class ConvLayer(nn.Module): + def __init__(self, config): super().__init__() kernel_size = getattr(config, "conv_kernel_size", 3) groups = getattr(config, "conv_groups", 1) self.conv_act = getattr(config, "conv_act", "tanh") - self.conv = nn.Conv1d( - config.hidden_size, config.hidden_size, kernel_size, padding=(kernel_size - 1) // 2, groups=groups - ) + self.conv = nn.Conv1d(config.hidden_size, + config.hidden_size, + kernel_size, + padding=(kernel_size - 1) // 2, + groups=groups) self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps) self.dropout = StableDropout(config.hidden_dropout_prob) self.config = config @@ -452,9 +465,10 @@ def get_attention_mask(self, attention_mask): def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None): if self.relative_attention and relative_pos is None: q = query_states.size(-2) if query_states is not None else hidden_states.size(-2) - relative_pos = build_relative_position( - q, hidden_states.size(-2), bucket_size=self.position_buckets, max_position=self.max_relative_positions - ) + relative_pos = build_relative_position(q, + hidden_states.size(-2), + bucket_size=self.position_buckets, + max_position=self.max_relative_positions) return relative_pos def forward( @@ -491,6 +505,7 @@ def forward( if self.gradient_checkpointing and self.training: def create_custom_forward(module): + def custom_forward(*inputs): return module(*inputs, output_attentions) @@ -535,9 +550,9 @@ def custom_forward(*inputs): if not return_dict: return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None) - return BaseModelOutput( - last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions - ) + return BaseModelOutput(last_hidden_state=output_states, + hidden_states=all_hidden_states, + attentions=all_attentions) def make_log_bucket_position(relative_pos, bucket_size, max_position): @@ -610,10 +625,8 @@ class DisentangledSelfAttention(nn.Module): def __init__(self, config): super().__init__() if config.hidden_size % config.num_attention_heads != 0: - raise ValueError( - f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " - f"heads ({config.num_attention_heads})" - ) + raise ValueError(f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})") self.num_attention_heads = config.num_attention_heads _attention_head_size = config.hidden_size // config.num_attention_heads self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size) @@ -706,28 +719,22 @@ def forward( attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / scale if self.relative_attention: rel_embeddings = self.pos_dropout(rel_embeddings) - rel_att = self.disentangled_attention_bias( - query_layer, key_layer, relative_pos, rel_embeddings, scale_factor - ) + rel_att = self.disentangled_attention_bias(query_layer, key_layer, relative_pos, rel_embeddings, + scale_factor) if rel_att is not None: attention_scores = attention_scores + rel_att attention_scores = attention_scores - attention_scores = attention_scores.view( - -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1) - ) + attention_scores = attention_scores.view(-1, self.num_attention_heads, attention_scores.size(-2), + attention_scores.size(-1)) # bsz x height x length x dimension attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1) attention_probs = self.dropout(attention_probs) - context_layer = torch.bmm( - attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer - ) - context_layer = ( - context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1)) - .permute(0, 2, 1, 3) - .contiguous() - ) + context_layer = torch.bmm(attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), + value_layer) + context_layer = (context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), + context_layer.size(-1)).permute(0, 2, 1, 3).contiguous()) new_context_layer_shape = context_layer.size()[:-2] + (-1,) context_layer = context_layer.view(new_context_layer_shape) if output_attentions: @@ -738,9 +745,10 @@ def forward( def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor): if relative_pos is None: q = query_layer.size(-2) - relative_pos = build_relative_position( - q, key_layer.size(-2), bucket_size=self.position_buckets, max_position=self.max_relative_positions - ) + relative_pos = build_relative_position(q, + key_layer.size(-2), + bucket_size=self.position_buckets, + max_position=self.max_relative_positions) if relative_pos.dim() == 2: relative_pos = relative_pos.unsqueeze(0).unsqueeze(0) elif relative_pos.dim() == 3: @@ -758,25 +766,22 @@ def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_ # rel_embeddings = rel_embeddings.unsqueeze(0) # rel_embeddings = rel_embeddings[0 : att_span * 2, :].unsqueeze(0) if self.share_att_key: - pos_query_layer = self.transpose_for_scores( - self.query_proj(rel_embeddings), self.num_attention_heads - ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) + pos_query_layer = self.transpose_for_scores(self.query_proj(rel_embeddings), + self.num_attention_heads).repeat( + query_layer.size(0) // self.num_attention_heads, 1, 1) pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat( - query_layer.size(0) // self.num_attention_heads, 1, 1 - ) + query_layer.size(0) // self.num_attention_heads, 1, 1) else: if "c2p" in self.pos_att_type: - pos_key_layer = self.transpose_for_scores( - self.pos_key_proj(rel_embeddings), self.num_attention_heads - ).repeat( - query_layer.size(0) // self.num_attention_heads, 1, 1 - ) # .split(self.all_head_size, dim=-1) + pos_key_layer = self.transpose_for_scores(self.pos_key_proj(rel_embeddings), + self.num_attention_heads).repeat( + query_layer.size(0) // self.num_attention_heads, 1, + 1) # .split(self.all_head_size, dim=-1) if "p2c" in self.pos_att_type: - pos_query_layer = self.transpose_for_scores( - self.pos_query_proj(rel_embeddings), self.num_attention_heads - ).repeat( - query_layer.size(0) // self.num_attention_heads, 1, 1 - ) # .split(self.all_head_size, dim=-1) + pos_query_layer = self.transpose_for_scores(self.pos_query_proj(rel_embeddings), + self.num_attention_heads).repeat( + query_layer.size(0) // self.num_attention_heads, 1, + 1) # .split(self.all_head_size, dim=-1) score = 0 # content->position @@ -787,7 +792,9 @@ def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_ c2p_att = torch.gather( c2p_att, dim=-1, - index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]), + index=c2p_pos.squeeze(0).expand([query_layer.size(0), + query_layer.size(1), + relative_pos.size(-1)]), ) score += c2p_att / scale @@ -810,7 +817,9 @@ def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_ p2c_att = torch.gather( p2c_att, dim=-1, - index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]), + index=p2c_pos.squeeze(0).expand([query_layer.size(0), + key_layer.size(-2), + key_layer.size(-2)]), ).transpose(-1, -2) score += p2c_att / scale @@ -990,6 +999,7 @@ def _set_gradient_checkpointing(self, module, value=False): ) # Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2 class DebertaV2Model(DebertaV2PreTrainedModel): + def __init__(self, config): super().__init__(config) @@ -1032,9 +1042,8 @@ def forward( return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutput]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_ids is not None and inputs_embeds is not None: @@ -1091,7 +1100,7 @@ def forward( sequence_output = encoded_layers[-1] if not return_dict: - return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :] + return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2):] return BaseModelOutput( last_hidden_state=sequence_output, @@ -1165,7 +1174,7 @@ def forward( masked_lm_loss = None if labels is not None: - loss_fct = CrossEntropyLoss() # -100 index = padding token + loss_fct = CrossEntropyLoss() # -100 index = padding token masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) if not return_dict: @@ -1182,6 +1191,7 @@ def forward( # copied from transformers.models.bert.BertPredictionHeadTransform with bert -> deberta class DebertaV2PredictionHeadTransform(nn.Module): + def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -1200,6 +1210,7 @@ def forward(self, hidden_states): # copied from transformers.models.bert.BertLMPredictionHead with bert -> deberta class DebertaV2LMPredictionHead(nn.Module): + def __init__(self, config): super().__init__() self.transform = DebertaV2PredictionHeadTransform(config) @@ -1221,6 +1232,7 @@ def forward(self, hidden_states): # copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta class DebertaV2OnlyMLMHead(nn.Module): + def __init__(self, config): super().__init__() self.predictions = DebertaV2LMPredictionHead(config) @@ -1239,6 +1251,7 @@ def forward(self, sequence_output): ) # Copied from transformers.models.deberta.modeling_deberta.DebertaForSequenceClassification with Deberta->DebertaV2 class DebertaV2ForSequenceClassification(DebertaV2PreTrainedModel): + def __init__(self, config): super().__init__(config) @@ -1318,9 +1331,8 @@ def forward( label_index = (labels >= 0).nonzero() labels = labels.long() if label_index.size(0) > 0: - labeled_logits = torch.gather( - logits, 0, label_index.expand(label_index.size(0), logits.size(1)) - ) + labeled_logits = torch.gather(logits, 0, label_index.expand(label_index.size(0), + logits.size(1))) labels = torch.gather(labels, 0, label_index.view(-1)) loss_fct = CrossEntropyLoss() loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1)) @@ -1345,9 +1357,10 @@ def forward( output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output - return SequenceClassifierOutput( - loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions - ) + return SequenceClassifierOutput(loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions) @add_start_docstrings( @@ -1422,9 +1435,10 @@ def forward( output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output - return TokenClassifierOutput( - loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions - ) + return TokenClassifierOutput(loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions) @add_start_docstrings( @@ -1536,6 +1550,7 @@ def forward( DEBERTA_START_DOCSTRING, ) class DebertaV2ForMultipleChoice(DebertaV2PreTrainedModel): + def __init__(self, config): super().__init__(config) @@ -1591,11 +1606,8 @@ def forward( flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None - flat_inputs_embeds = ( - inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) - if inputs_embeds is not None - else None - ) + flat_inputs_embeds = (inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None else None) outputs = self.deberta( flat_input_ids, diff --git a/examples/language/roberta/pretraining/nvidia_bert_dataset_provider.py b/examples/community/roberta/pretraining/nvidia_bert_dataset_provider.py similarity index 76% rename from examples/language/roberta/pretraining/nvidia_bert_dataset_provider.py rename to examples/community/roberta/pretraining/nvidia_bert_dataset_provider.py index cce836913505..72c7bd852a40 100644 --- a/examples/language/roberta/pretraining/nvidia_bert_dataset_provider.py +++ b/examples/community/roberta/pretraining/nvidia_bert_dataset_provider.py @@ -1,24 +1,25 @@ +import json +import logging import os import random -import h5py -import logging -import json import time from concurrent.futures import ProcessPoolExecutor +import h5py import numpy as np - import torch import torch.distributed as dist +from bert_dataset_provider import BertDatasetProviderInterface from torch.utils.data import DataLoader, Dataset -from torch.utils.data.sampler import RandomSampler from torch.utils.data.distributed import DistributedSampler +from torch.utils.data.sampler import RandomSampler -from bert_dataset_provider import BertDatasetProviderInterface import colossalai.utils as utils + # Workaround because python functions are not picklable class WorkerInitObj(object): + def __init__(self, seed): self.seed = seed @@ -27,29 +28,25 @@ def __call__(self, id): random.seed(self.seed + id) -def create_pretraining_dataset(input_file, max_predictions_per_seq, - num_workers, train_batch_size, worker_init, +def create_pretraining_dataset(input_file, max_predictions_per_seq, num_workers, train_batch_size, worker_init, data_sampler): - train_data = pretraining_dataset( - input_file=input_file, max_predictions_per_seq=max_predictions_per_seq) + train_data = pretraining_dataset(input_file=input_file, max_predictions_per_seq=max_predictions_per_seq) train_dataloader = DataLoader(train_data, sampler=data_sampler(train_data), batch_size=train_batch_size, num_workers=num_workers, worker_init_fn=worker_init, - pin_memory=True - ) + pin_memory=True) return train_dataloader, len(train_data) class pretraining_dataset(Dataset): + def __init__(self, input_file, max_predictions_per_seq): self.input_file = input_file self.max_predictions_per_seq = max_predictions_per_seq f = h5py.File(input_file, "r") - keys = [ - 'input_ids', 'input_mask', 'segment_ids', 'masked_lm_positions' - ] + keys = ['input_ids', 'input_mask', 'segment_ids', 'masked_lm_positions'] self.inputs = [np.asarray(f[key][:]) for key in keys] f.close() @@ -59,21 +56,16 @@ def __len__(self): def __getitem__(self, index): - [ - input_ids, input_mask, segment_ids, masked_lm_labels - ] = [ - torch.from_numpy(input[index].astype(np.int64)) if indice < 5 else - torch.from_numpy(np.asarray(input[index].astype(np.int64))) - for indice, input in enumerate(self.inputs) + [input_ids, input_mask, segment_ids, masked_lm_labels] = [ + torch.from_numpy(input[index].astype(np.int64)) if indice < 5 else torch.from_numpy( + np.asarray(input[index].astype(np.int64))) for indice, input in enumerate(self.inputs) ] - return [ - input_ids, input_mask, - segment_ids, masked_lm_labels - ] + return [input_ids, input_mask, segment_ids, masked_lm_labels] class NvidiaBertDatasetProvider(BertDatasetProviderInterface): + def __init__(self, args, evaluate=False): self.num_workers = args.num_workers self.max_seq_length = args.max_seq_length @@ -85,22 +77,24 @@ def __init__(self, args, evaluate=False): else: self.train_micro_batch_size_per_gpu = args.eval_micro_batch_size_per_gpu self.logger = args.logger - + self.global_rank = dist.get_rank() self.world_size = dist.get_world_size() # Initialize dataset files if not evaluate: self.dataset_files = [ - os.path.join(args.data_path_prefix, f) for f in os.listdir(args.data_path_prefix) if - os.path.isfile(os.path.join(args.data_path_prefix, f)) and 'h5' in f + os.path.join(args.data_path_prefix, f) + for f in os.listdir(args.data_path_prefix) + if os.path.isfile(os.path.join(args.data_path_prefix, f)) and 'h5' in f ] else: self.dataset_files = [ - os.path.join(args.eval_data_path_prefix, f) for f in os.listdir(args.eval_data_path_prefix) if - os.path.isfile(os.path.join(args.eval_data_path_prefix, f)) and 'h5' in f + os.path.join(args.eval_data_path_prefix, f) + for f in os.listdir(args.eval_data_path_prefix) + if os.path.isfile(os.path.join(args.eval_data_path_prefix, f)) and 'h5' in f ] - + self.dataset_files.sort() # random.shuffle(self.dataset_files) self.num_files = len(self.dataset_files) @@ -114,9 +108,7 @@ def __init__(self, args, evaluate=False): self.shuffle = True if self.global_rank == 0: - self.logger.info( - f"NvidiaBertDatasetProvider - Initialization: num_files = {self.num_files}" - ) + self.logger.info(f"NvidiaBertDatasetProvider - Initialization: num_files = {self.num_files}") def get_shard(self, index): start = time.time() @@ -130,9 +122,8 @@ def get_shard(self, index): worker_init=self.worker_init, data_sampler=self.data_sampler) else: - self.train_dataloader, sample_count = self.dataset_future.result( - timeout=None) - + self.train_dataloader, sample_count = self.dataset_future.result(timeout=None) + self.logger.info( f"Data Loading Completed for Pretraining Data from {self.data_file} with {sample_count} samples took {time.time()-start:.2f}s." ) @@ -145,11 +136,9 @@ def release_shard(self): def prefetch_shard(self, index): self.data_file = self._get_shard_file(index) - self.dataset_future = self.pool.submit( - create_pretraining_dataset, self.data_file, - self.max_predictions_per_seq, self.num_workers, - self.train_micro_batch_size_per_gpu, self.worker_init, - self.data_sampler) + self.dataset_future = self.pool.submit(create_pretraining_dataset, self.data_file, self.max_predictions_per_seq, + self.num_workers, self.train_micro_batch_size_per_gpu, self.worker_init, + self.data_sampler) def get_batch(self, batch_iter): return batch_iter @@ -179,4 +168,3 @@ def shuffle_dataset(self, epoch): indices = torch.randperm(self.num_files, generator=g).tolist() new_dataset = [self.dataset_files[i] for i in indices] self.dataset_files = new_dataset - \ No newline at end of file diff --git a/examples/language/roberta/pretraining/pretrain_utils.py b/examples/community/roberta/pretraining/pretrain_utils.py similarity index 74% rename from examples/language/roberta/pretraining/pretrain_utils.py rename to examples/community/roberta/pretraining/pretrain_utils.py index ba17b0f5ee09..cea6ac2c36e5 100644 --- a/examples/language/roberta/pretraining/pretrain_utils.py +++ b/examples/community/roberta/pretraining/pretrain_utils.py @@ -1,35 +1,45 @@ -import transformers import logging -from colossalai.nn.lr_scheduler import LinearWarmupLR -from transformers import get_linear_schedule_with_warmup -from transformers import BertForPreTraining, RobertaForMaskedLM, RobertaConfig -from transformers import GPT2Config, GPT2LMHeadModel -from transformers import AutoTokenizer, AutoModelForMaskedLM -from colossalai.nn.optimizer import FusedAdam -from torch.optim import AdamW -from colossalai.core import global_context as gpc -import torch import os import sys -sys.path.append(os.getcwd()) -from model.deberta_v2 import DebertaV2ForMaskedLM -from model.bert import BertForMaskedLM -import torch.nn as nn +import torch +import transformers +from torch.optim import AdamW +from transformers import ( + AutoModelForMaskedLM, + AutoTokenizer, + BertForPreTraining, + GPT2Config, + GPT2LMHeadModel, + RobertaConfig, + RobertaForMaskedLM, + get_linear_schedule_with_warmup, +) + +from colossalai.core import global_context as gpc +from colossalai.nn.lr_scheduler import LinearWarmupLR +from colossalai.nn.optimizer import FusedAdam, HybridAdam + +sys.path.append(os.getcwd()) from collections import OrderedDict +import torch.nn as nn +from model.bert import BertForMaskedLM +from model.deberta_v2 import DebertaV2ForMaskedLM + __all__ = ['get_model', 'get_optimizer', 'get_lr_scheduler', 'get_dataloader_for_pretraining'] def get_new_state_dict(state_dict, start_index=13): - new_state_dict = OrderedDict() + new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[start_index:] - new_state_dict[name] = v + new_state_dict[name] = v return new_state_dict class LMModel(nn.Module): + def __init__(self, model, config, args): super().__init__() @@ -58,16 +68,18 @@ def get_model(args, logger): if len(args.load_pretrain_model) > 0: assert os.path.exists(args.load_pretrain_model) # load_checkpoint(args.load_pretrain_model, model, strict=False) - m_state_dict = torch.load(args.load_pretrain_model, map_location=torch.device(f"cuda:{torch.cuda.current_device()}")) + m_state_dict = torch.load(args.load_pretrain_model, + map_location=torch.device(f"cuda:{torch.cuda.current_device()}")) # new_state_dict = get_new_state_dict(m_state_dict) - model.load_state_dict(m_state_dict, strict=True) # must insure that every process have identical parameters !!!!!!! + model.load_state_dict(m_state_dict, + strict=True) # must insure that every process have identical parameters !!!!!!! logger.info("load model success") - + numel = sum([p.numel() for p in model.parameters()]) if args.checkpoint_activations: model.gradient_checkpointing_enable() # model = LMModel(model, config, args) - + return config, model, numel @@ -83,13 +95,16 @@ def get_optimizer(model, lr): 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] - optimizer = FusedAdam(optimizer_grouped_parameters, lr=lr, betas=[0.9, 0.95]) + optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, betas=[0.9, 0.95]) return optimizer def get_lr_scheduler(optimizer, total_steps, warmup_steps=2000, last_epoch=-1): # warmup_steps = int(total_steps * warmup_ratio) - lr_scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, last_epoch=last_epoch) + lr_scheduler = get_linear_schedule_with_warmup(optimizer, + num_warmup_steps=warmup_steps, + num_training_steps=total_steps, + last_epoch=last_epoch) # lr_scheduler = LinearWarmupLR(optimizer, total_steps=total_steps, warmup_steps=warmup_steps) return lr_scheduler @@ -103,10 +118,7 @@ def save_ckpt(model, optimizer, lr_scheduler, path, epoch, shard, global_step): checkpoint['epoch'] = epoch checkpoint['shard'] = shard checkpoint['global_step'] = global_step - model_state = model.state_dict() #each process must run model.state_dict() + model_state = model.state_dict() #each process must run model.state_dict() if gpc.get_global_rank() == 0: torch.save(checkpoint, optimizer_lr_path) torch.save(model_state, model_path) - - - diff --git a/examples/language/roberta/pretraining/run_pretrain.sh b/examples/community/roberta/pretraining/run_pretrain.sh similarity index 90% rename from examples/language/roberta/pretraining/run_pretrain.sh rename to examples/community/roberta/pretraining/run_pretrain.sh index 144cd0ab96fd..280dba714de5 100644 --- a/examples/language/roberta/pretraining/run_pretrain.sh +++ b/examples/community/roberta/pretraining/run_pretrain.sh @@ -7,7 +7,6 @@ tensorboard_path="$root_path/tensorboard" log_path="$root_path/exp_log" ckpt_path="$root_path/ckpt" -colossal_config="$root_path/../configs/colossalai_ddp.py" mkdir -p $tensorboard_path mkdir -p $log_path @@ -32,9 +31,7 @@ env OMP_NUM_THREADS=40 colossalai run --hostfile ./hostfile \ --tensorboard_path $tensorboard_path \ --log_path $log_path \ --ckpt_path $ckpt_path \ - --colossal_config $colossal_config \ --log_interval 50 \ --mlm bert \ --wandb \ --checkpoint_activations \ - \ No newline at end of file diff --git a/examples/language/roberta/pretraining/run_pretrain_resume.sh b/examples/community/roberta/pretraining/run_pretrain_resume.sh similarity index 91% rename from examples/language/roberta/pretraining/run_pretrain_resume.sh rename to examples/community/roberta/pretraining/run_pretrain_resume.sh index a0704cf7c517..8f443b454d7d 100644 --- a/examples/language/roberta/pretraining/run_pretrain_resume.sh +++ b/examples/community/roberta/pretraining/run_pretrain_resume.sh @@ -7,7 +7,6 @@ tensorboard_path="$root_path/tensorboard" log_path="$root_path/exp_log" ckpt_path="$root_path/ckpt" -colossal_config="$root_path/../configs/colossalai_ddp.py" mkdir -p $tensorboard_path mkdir -p $log_path @@ -32,7 +31,6 @@ env OMP_NUM_THREADS=40 colossalai run --hostfile ./hostfile \ --tensorboard_path $tensorboard_path \ --log_path $log_path \ --ckpt_path $ckpt_path \ - --colossal_config $colossal_config \ --log_interval 50 \ --mlm bert \ --wandb \ @@ -40,4 +38,3 @@ env OMP_NUM_THREADS=40 colossalai run --hostfile ./hostfile \ --resume_train \ --load_pretrain_model /ckpt/1.pt \ --load_optimizer_lr /ckpt/1.op_lrs \ - \ No newline at end of file diff --git a/examples/language/roberta/pretraining/run_pretraining.py b/examples/community/roberta/pretraining/run_pretraining.py similarity index 52% rename from examples/language/roberta/pretraining/run_pretraining.py rename to examples/community/roberta/pretraining/run_pretraining.py index 9840a122cbc4..9fae4bef227a 100644 --- a/examples/language/roberta/pretraining/run_pretraining.py +++ b/examples/community/roberta/pretraining/run_pretraining.py @@ -1,93 +1,132 @@ -import colossalai import math +import os +import time +from functools import partial + import torch -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -import colossalai.nn as col_nn from arguments import parse_args -from pretrain_utils import get_model, get_optimizer, get_lr_scheduler, save_ckpt -from utils.exp_util import get_tflops, get_mem_info, throughput_calculator, log_args -from utils.global_vars import set_global_variables, get_timers, get_tensorboard_writer -from utils.logger import Logger from evaluation import evaluate from loss import LossForPretraining - -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_optim import ShardedOptimizerV2 from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider +from pretrain_utils import get_lr_scheduler, get_model, get_optimizer, save_ckpt from tqdm import tqdm -import os -import time -from functools import partial - from transformers import AutoTokenizer +from utils.exp_util import get_mem_info, get_tflops, log_args, throughput_calculator +from utils.global_vars import get_tensorboard_writer, get_timers, set_global_variables +from utils.logger import Logger -from colossalai.gemini import ChunkManager, GeminiManager -from colossalai.utils.model.colo_init_context import ColoInitContext +import colossalai +from colossalai.context import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.nn.parallel import GeminiDDP, zero_model_wrapper, zero_optim_wrapper +from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec from colossalai.utils import get_current_device -from colossalai.nn.parallel import ZeroDDP +from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.zero import ZeroOptimizer -from colossalai.tensor import ProcessGroup -from colossalai.nn.optimizer import HybridAdam def main(): args = parse_args() launch_time = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime()) - + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) - os.environ['CUDA_LAUNCH_BLOCKING'] = '1' - + # os.environ['CUDA_LAUNCH_BLOCKING'] = '1' + logger = Logger(os.path.join(args.log_path, launch_time), cuda=torch.cuda.is_available(), debug=args.vscode_debug) - + if args.vscode_debug: colossalai.launch(config={}, - rank=args.rank, - world_size=args.world_size, - host=args.host, - port=args.port, - backend=args.backend) + rank=args.rank, + world_size=args.world_size, + host=args.host, + port=args.port, + backend=args.backend) args.local_rank = -1 args.log_interval = 1 else: - colossalai.launch_from_torch(args.colossal_config) #args.colossal_config + colossalai.launch_from_torch(config={}) #args.colossal_config args.local_rank = int(os.environ["LOCAL_RANK"]) - logger.info(f'launch_from_torch, world size: {torch.distributed.get_world_size()} | ' + - f'ParallelMode.MODEL: {ParallelMode.MODEL} | ParallelMode.DATA: {ParallelMode.DATA} | ParallelMode.TENSOR: {ParallelMode.TENSOR}') + logger.info( + f'launch_from_torch, world size: {torch.distributed.get_world_size()} | ' + + f'ParallelMode.MODEL: {ParallelMode.MODEL} | ParallelMode.DATA: {ParallelMode.DATA} | ParallelMode.TENSOR: {ParallelMode.TENSOR}' + ) log_args(logger, args) args.tokenizer = tokenizer args.logger = logger set_global_variables(launch_time, args.tensorboard_path) - - use_zero = hasattr(gpc.config, 'zero') + world_size = torch.distributed.get_world_size() + init_dev = get_current_device() # build model, optimizer and criterion - if use_zero: - shard_strategy = TensorShardStrategy() - with ZeroInitContext(target_device=torch.cuda.current_device(), shard_strategy=shard_strategy, - shard_param=True): - + if args.distplan.startswith("CAI"): + # all param must use the same process group. + world_size = torch.distributed.get_world_size() + shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None + default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None + + if args.shardinit and args.distplan != "CAI_Gemini": + raise RuntimeError("You can only use shardinit with CAI_Gemini") + + # build GPT model + with ColoInitContext(device=get_current_device(), + dtype=torch.half, + default_dist_spec=default_dist_spec, + default_pg=shard_pg): config, model, numel = get_model(args, logger) - # model = ShardedModelV2(model, shard_strategy, tensor_placement_policy='cpu', reuse_fp16_shard=True) + + # assign running configurations + gemini_config = None + if args.distplan.startswith("CAI_ZeRO"): + optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True) + elif args.distplan == "CAI_Gemini": + gemini_config = dict(strict_ddp_mode=args.tp_degree == 1, + device=get_current_device(), + placement_policy=args.placement, + pin_memory=True, + hidden_dim=model.config.hidden_size, + search_range_m=128) + optim_config = dict(gpu_margin_mem_ratio=0.) + else: + raise RuntimeError + + # build a highly optimized gpu/cpu optimizer + optimizer = get_optimizer(model, lr=args.lr) + + if args.distplan == "CAI_ZeRO1": + zero_stage = 1 + elif args.distplan == "CAI_ZeRO2": + zero_stage = 2 + elif args.distplan == "CAI_Gemini": + zero_stage = 3 + else: + raise RuntimeError + + # wrap your model and optimizer + model = zero_model_wrapper(model, zero_stage, gemini_config) + optimizer = zero_optim_wrapper(model, optimizer, optim_config=optim_config) + + logger.info(get_mem_info(prefix='After init optim, ')) + else: config, model, numel = get_model(args, logger) logger.info("no_zero") + if torch.distributed.get_rank() == 0: os.mkdir(os.path.join(args.ckpt_path, launch_time)) logger.info(f'Model numel: {numel}') - + get_tflops_func = partial(get_tflops, numel, args.train_micro_batch_size_per_gpu, args.max_seq_length) - steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size #len(dataloader) + + # 144003367 is is the length of the entire dataset + steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size #len(dataloader) total_steps = steps_per_epoch * args.epoch - # build optimizer and lr_scheduler + lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=-1) start_epoch = 0 start_shard = 0 @@ -96,40 +135,31 @@ def main(): assert os.path.exists(args.load_optimizer_lr) o_l_state_dict = torch.load(args.load_optimizer_lr, map_location='cpu') o_l_state_dict['lr_scheduler']['last_epoch'] = o_l_state_dict['lr_scheduler']['last_epoch'] - 1 - optimizer = get_optimizer(model, lr=args.lr) optimizer.load_state_dict(o_l_state_dict['optimizer']) - lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=o_l_state_dict['lr_scheduler']['last_epoch']) #o_l_state_dict['lr_scheduler']['last_epoch'] + # o_l_state_dict['lr_scheduler']['last_epoch'] + lr_scheduler = get_lr_scheduler(optimizer, + total_steps=total_steps, + last_epoch=o_l_state_dict['lr_scheduler']['last_epoch']) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.cuda(f"cuda:{torch.cuda.current_device()}") - # if you want delete the above three code, have to move the model to gpu, because in optimizer.step() + # if you want delete the above three code, must move the model to gpu. Because in optimizer.step() lr_scheduler.load_state_dict(o_l_state_dict['lr_scheduler']) - + start_epoch = o_l_state_dict['epoch'] start_shard = o_l_state_dict['shard'] + 1 # global_step = o_l_state_dict['global_step'] + 1 - logger.info(f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}') - else: - optimizer = get_optimizer(model, lr=args.lr) - lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=-1) + logger.info( + f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}' + ) - # optimizer = gpc.config.optimizer.pop('type')( - # model.parameters(), **gpc.config.optimizer) - # optimizer = ShardedOptimizerV2(model, optimizer, initial_scale=2**5) criterion = LossForPretraining(config.vocab_size) # build dataloader pretrain_dataset_provider = NvidiaBertDatasetProvider(args) - # initialize with colossalai - engine, _, _, lr_scheduelr = colossalai.initialize(model=model, - optimizer=optimizer, - criterion=criterion, - lr_scheduler=lr_scheduler) - logger.info(get_mem_info(prefix='After init model, ')) - best_loss = None eval_loss = 0 @@ -146,13 +176,16 @@ def main(): dataset_iterator, total_length = pretrain_dataset_provider.get_shard(shard) # pretrain_dataset_provider.prefetch_shard(shard + 1) # may cause cpu memory overload if torch.distributed.get_rank() == 0: - iterator_data = tqdm(enumerate(dataset_iterator), total=(total_length // args.train_micro_batch_size_per_gpu // world_size), colour='cyan', smoothing=1) + iterator_data = tqdm(enumerate(dataset_iterator), + total=(total_length // args.train_micro_batch_size_per_gpu // world_size), + colour='cyan', + smoothing=1) else: iterator_data = enumerate(dataset_iterator) - engine.train() - - for step, batch_data in iterator_data: + model.train() + + for step, batch_data in iterator_data: # batch_data = pretrain_dataset_provider.get_batch(batch_index) input_ids = batch_data[0].cuda(f"cuda:{torch.cuda.current_device()}") @@ -161,53 +194,57 @@ def main(): mlm_label = batch_data[3].cuda(f"cuda:{torch.cuda.current_device()}") # nsp_label = batch_data[5].cuda() - output = engine(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) - - loss = engine.criterion(output.logits, mlm_label) + output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) + + loss = criterion(output.logits, mlm_label) pretrain_dataset_provider.prefetch_batch() - engine.backward(loss) + optimizer.backward(loss) train_loss += loss.float().item() # if (step + 1) % args.accumulation_step == 0: - engine.step() - lr_scheduelr.step() - engine.zero_grad() - + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + global_step += 1 if global_step % args.log_interval == 0 and global_step != 0 \ - and torch.distributed.get_rank() == 0: + and torch.distributed.get_rank() == 0: elapsed_time = timers('interval_time').elapsed(reset=False) elapsed_time_per_iteration = elapsed_time / global_step - samples_per_sec, tflops, approx_parameters_in_billions = throughput_calculator(numel, args, config, elapsed_time, global_step, world_size) + samples_per_sec, tflops, approx_parameters_in_billions = throughput_calculator( + numel, args, config, elapsed_time, global_step, world_size) cur_loss = train_loss / args.log_interval - current_lr = lr_scheduelr.get_last_lr()[0] + current_lr = lr_scheduler.get_last_lr()[0] log_str = f'| epoch: {epoch} | shard: {shard} | step: {global_step} | lr {current_lr:.7f} | elapsed_time: {elapsed_time / 60 :.3f} minutes ' + \ f'| mins/batch: {elapsed_time_per_iteration :.3f} seconds | loss: {cur_loss:.7f} | ppl: {math.exp(cur_loss):.3f} | TFLOPS: {get_tflops_func(elapsed_time_per_iteration):.3f} or {tflops:.3f}' logger.info(log_str, print_=False) if args.wandb: tensorboard_log = get_tensorboard_writer() - tensorboard_log.log_train({ - 'lr': current_lr, - 'loss': cur_loss, - 'ppl': math.exp(cur_loss), - 'mins_batch': elapsed_time_per_iteration - }, global_step) + tensorboard_log.log_train( + { + 'lr': current_lr, + 'loss': cur_loss, + 'ppl': math.exp(cur_loss), + 'mins_batch': elapsed_time_per_iteration + }, global_step) train_loss = 0 logger.info(f'epoch {epoch} shard {shard} has cost {timers("shard_time").elapsed() / 60 :.3f} mins') logger.info('*' * 100) - eval_loss += evaluate(engine, args, logger, global_step) - save_ckpt(engine.model, optimizer, lr_scheduelr, os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch, shard, global_step) - - + eval_loss += evaluate(model, args, logger, global_step, criterion) + save_ckpt(model, optimizer, lr_scheduler, + os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch, + shard, global_step) + eval_loss /= len(os.listdir(args.data_path_prefix)) - logger.info(f'epoch {epoch} | shard_length {len(os.listdir(args.data_path_prefix))} | elapsed_time: {timers("epoch_time").elapsed() / 60 :.3f} mins' + \ - f'eval_loss: {eval_loss} | ppl: {math.exp(eval_loss)}') + logger.info( + f'epoch {epoch} | shard_length {len(os.listdir(args.data_path_prefix))} | elapsed_time: {timers("epoch_time").elapsed() / 60 :.3f} mins' + + f'eval_loss: {eval_loss} | ppl: {math.exp(eval_loss)}') logger.info('-' * 100) if args.wandb and torch.distributed.get_rank() == 0: tensorboard_log = get_tensorboard_writer() diff --git a/examples/language/roberta/pretraining/utils/WandbLog.py b/examples/community/roberta/pretraining/utils/WandbLog.py similarity index 98% rename from examples/language/roberta/pretraining/utils/WandbLog.py rename to examples/community/roberta/pretraining/utils/WandbLog.py index 9dd28a98186b..b68ba8387dcd 100644 --- a/examples/language/roberta/pretraining/utils/WandbLog.py +++ b/examples/community/roberta/pretraining/utils/WandbLog.py @@ -1,8 +1,10 @@ +import os import time + import wandb -import os from torch.utils.tensorboard import SummaryWriter + class WandbLog: @classmethod @@ -15,7 +17,7 @@ def log(cls, result, model=None, gradient=None): if model: wandb.watch(model) - + if gradient: wandb.watch(gradient) @@ -30,7 +32,7 @@ def __init__(self, location, name=time.strftime("%Y-%m-%d %H:%M:%S", time.localt def log_train(self, result, step): for k, v in result.items(): self.writer.add_scalar(f'{k}/train', v, step) - + def log_eval(self, result, step): for k, v in result.items(): self.writer.add_scalar(f'{k}/eval', v, step) @@ -38,9 +40,3 @@ def log_eval(self, result, step): def log_zeroshot(self, result, step): for k, v in result.items(): self.writer.add_scalar(f'{k}_acc/eval', v, step) - - - - - - diff --git a/examples/language/roberta/pretraining/utils/exp_util.py b/examples/community/roberta/pretraining/utils/exp_util.py similarity index 85% rename from examples/language/roberta/pretraining/utils/exp_util.py rename to examples/community/roberta/pretraining/utils/exp_util.py index a02b0872acbc..4a2c9d8a47ad 100644 --- a/examples/language/roberta/pretraining/utils/exp_util.py +++ b/examples/community/roberta/pretraining/utils/exp_util.py @@ -1,9 +1,13 @@ import functools -import os, shutil -import torch +import os +import shutil + import psutil +import torch + from colossalai.core import global_context as gpc + def logging(s, log_path, print_=True, log_=True): if print_: print(s) @@ -11,9 +15,11 @@ def logging(s, log_path, print_=True, log_=True): with open(log_path, 'a+') as f_log: f_log.write(s + '\n') + def get_logger(log_path, **kwargs): return functools.partial(logging, log_path=log_path, **kwargs) + def create_exp_dir(dir_path, scripts_to_save=None, debug=False): if debug: print('Debug Mode : no experiment dir created') @@ -33,6 +39,7 @@ def create_exp_dir(dir_path, scripts_to_save=None, debug=False): return get_logger(log_path=os.path.join(dir_path, 'log.txt')) + def get_cpu_mem(): return psutil.Process().memory_info().rss / 1024**2 @@ -52,11 +59,15 @@ def get_tflops(model_numel, batch_size, seq_len, step_time): def get_parameters_in_billions(model, world_size=1): gpus_per_model = world_size - approx_parameters_in_billions = sum([sum([p.ds_numel if hasattr(p,'ds_id') else p.nelement() for p in model_module.parameters()]) - for model_module in model]) + approx_parameters_in_billions = sum([ + sum([p.ds_numel if hasattr(p, 'ds_id') else p.nelement() + for p in model_module.parameters()]) + for model_module in model + ]) return approx_parameters_in_billions * gpus_per_model / (1e9) + def throughput_calculator(numel, args, config, iteration_time, total_iterations, world_size=1): gpus_per_model = 1 batch_size = args.train_micro_batch_size_per_gpu @@ -76,24 +87,28 @@ def throughput_calculator(numel, args, config, iteration_time, total_iterations, # The factor of 4 is when used with activation check-pointing, # otherwise it will be 3. checkpoint_activations_factor = 4 if args.checkpoint_activations else 3 - flops_per_iteration = (24 * checkpoint_activations_factor * batch_size * args.max_seq_length * num_layers * (hidden_size**2)) * (1. + (args.max_seq_length / (6. * hidden_size)) + (vocab_size / (16. * num_layers * hidden_size))) + flops_per_iteration = (24 * checkpoint_activations_factor * batch_size * args.max_seq_length * num_layers * + (hidden_size**2)) * (1. + (args.max_seq_length / (6. * hidden_size)) + + (vocab_size / (16. * num_layers * hidden_size))) tflops = flops_per_iteration / (elapsed_time_per_iter * (10**12)) return samples_per_second, tflops, approx_parameters_in_billions + def synchronize(): if not torch.distributed.is_available(): return - if not torch.distributed.is_intialized(): + if not torch.distributed.is_initialized(): return world_size = torch.distributed.get_world_size() if world_size == 1: return torch.distributed.barrier() + def log_args(logger, args): logger.info('--------args----------') message = '\n'.join([f'{k:<30}: {v}' for k, v in vars(args).items()]) message += '\n' message += '\n'.join([f'{k:<30}: {v}' for k, v in gpc.config.items()]) logger.info(message) - logger.info('--------args----------\n') \ No newline at end of file + logger.info('--------args----------\n') diff --git a/examples/language/roberta/pretraining/utils/global_vars.py b/examples/community/roberta/pretraining/utils/global_vars.py similarity index 89% rename from examples/language/roberta/pretraining/utils/global_vars.py rename to examples/community/roberta/pretraining/utils/global_vars.py index 363cbf91c065..9eef19e71614 100644 --- a/examples/language/roberta/pretraining/utils/global_vars.py +++ b/examples/community/roberta/pretraining/utils/global_vars.py @@ -1,5 +1,7 @@ import time + import torch + from .WandbLog import TensorboardLog _GLOBAL_TIMERS = None @@ -10,30 +12,34 @@ def set_global_variables(launch_time, tensorboard_path): _set_timers() _set_tensorboard_writer(launch_time, tensorboard_path) + def _set_timers(): """Initialize timers.""" global _GLOBAL_TIMERS _ensure_var_is_not_initialized(_GLOBAL_TIMERS, 'timers') _GLOBAL_TIMERS = Timers() + def _set_tensorboard_writer(launch_time, tensorboard_path): """Set tensorboard writer.""" global _GLOBAL_TENSORBOARD_WRITER - _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER, - 'tensorboard writer') + _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER, 'tensorboard writer') if torch.distributed.get_rank() == 0: _GLOBAL_TENSORBOARD_WRITER = TensorboardLog(tensorboard_path + f'/{launch_time}', launch_time) - + + def get_timers(): """Return timers.""" _ensure_var_is_initialized(_GLOBAL_TIMERS, 'timers') return _GLOBAL_TIMERS + def get_tensorboard_writer(): """Return tensorboard writer. It can be None so no need to check if it is initialized.""" return _GLOBAL_TENSORBOARD_WRITER + def _ensure_var_is_initialized(var, name): """Make sure the input variable is not None.""" assert var is not None, '{} is not initialized.'.format(name) @@ -104,7 +110,7 @@ def write(self, names, writer, iteration, normalizer=1.0, reset=False): """Write timers to a tensorboard writer""" # currently when using add_scalars, # torch.utils.add_scalars makes each timer its own run, which - # polutes the runs list, so we just add each as a scalar + # pollutes the runs list, so we just add each as a scalar assert normalizer > 0.0 for name in names: value = self.timers[name].elapsed(reset=reset) / normalizer @@ -115,12 +121,10 @@ def log(self, names, normalizer=1.0, reset=True): assert normalizer > 0.0 string = 'time (ms)' for name in names: - elapsed_time = self.timers[name].elapsed( - reset=reset) * 1000.0 / normalizer + elapsed_time = self.timers[name].elapsed(reset=reset) * 1000.0 / normalizer string += ' | {}: {:.2f}'.format(name, elapsed_time) if torch.distributed.is_initialized(): - if torch.distributed.get_rank() == ( - torch.distributed.get_world_size() - 1): + if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): print(string, flush=True) else: print(string, flush=True) diff --git a/examples/language/roberta/pretraining/utils/logger.py b/examples/community/roberta/pretraining/utils/logger.py similarity index 81% rename from examples/language/roberta/pretraining/utils/logger.py rename to examples/community/roberta/pretraining/utils/logger.py index 481c4c6ce61b..75c9bf4bef25 100644 --- a/examples/language/roberta/pretraining/utils/logger.py +++ b/examples/community/roberta/pretraining/utils/logger.py @@ -1,22 +1,22 @@ -import os import logging +import os + import torch.distributed as dist -logging.basicConfig( - format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', - datefmt='%m/%d/%Y %H:%M:%S', - level=logging.INFO) +logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', + datefmt='%m/%d/%Y %H:%M:%S', + level=logging.INFO) logger = logging.getLogger(__name__) class Logger(): + def __init__(self, log_path, cuda=False, debug=False): self.logger = logging.getLogger(__name__) self.cuda = cuda self.log_path = log_path self.debug = debug - def info(self, message, log_=True, print_=True, *args, **kwargs): if (self.cuda and dist.get_rank() == 0) or not self.cuda: if print_: @@ -26,6 +26,5 @@ def info(self, message, log_=True, print_=True, *args, **kwargs): with open(self.log_path, 'a+') as f_log: f_log.write(message + '\n') - def error(self, message, *args, **kwargs): self.logger.error(message, *args, **kwargs) diff --git a/examples/community/roberta/requirements.txt b/examples/community/roberta/requirements.txt new file mode 100644 index 000000000000..de082defb14a --- /dev/null +++ b/examples/community/roberta/requirements.txt @@ -0,0 +1,7 @@ +colossalai >= 0.1.12 +torch >= 1.8.1 +tqdm +tensorboard +numpy +h5py +wandb diff --git a/examples/community/roberta/test_ci.sh b/examples/community/roberta/test_ci.sh new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/examples/images/diffusion/README.md b/examples/images/diffusion/README.md index 2a99094b703a..0c7f42ded318 100644 --- a/examples/images/diffusion/README.md +++ b/examples/images/diffusion/README.md @@ -37,35 +37,40 @@ This project is in rapid development. ## Installation -### Option #1: install from source +### Option #1: Install from source #### Step 1: Requirements -A suitable [conda](https://conda.io/) environment named `ldm` can be created -and activated with: +To begin with, make sure your operating system has the cuda version suitable for this exciting training session, which is cuda11.6/11.8. For your convience, we have set up the rest of packages here. You can create and activate a suitable [conda](https://conda.io/) environment named `ldm` : ``` conda env create -f environment.yaml conda activate ldm ``` -You can also update an existing [latent diffusion](https://github.com/CompVis/latent-diffusion) environment by running +You can also update an existing [latent diffusion](https://github.com/CompVis/latent-diffusion) environment by running: ``` conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch pip install transformers diffusers invisible-watermark ``` -#### Step 2:Install [Colossal-AI](https://colossalai.org/download/) From Our Official Website +#### Step 2: Install [Colossal-AI](https://colossalai.org/download/) From Our Official Website -##### From pip +You can install the latest version (0.2.7) from our official website or from source. Notice that the suitable version for this training is colossalai(0.2.5), which stands for torch(1.12.1). -For example, you can install v0.2.0 from our official website. +##### Download suggested version for this training + +``` +pip install colossalai==0.2.5 +``` + +##### Download the latest version from pip for latest torch version ``` pip install colossalai ``` -##### From source +##### From source: ``` git clone https://github.com/hpcaitech/ColossalAI.git @@ -75,10 +80,12 @@ cd ColossalAI CUDA_EXT=1 pip install . ``` -#### Step 3:Accelerate with flash attention by xformers(Optional) +#### Step 3: Accelerate with flash attention by xformers (Optional) + +Notice that xformers will accelerate the training process at the cost of extra disk space. The suitable version of xformers for this training process is 0.0.12, which can be downloaded directly via pip. For more release versions, feel free to check its official website: [XFormers](https://pypi.org/project/xformers/) ``` -pip install xformers +pip install xformers==0.0.12 ``` ### Option #2: Use Docker @@ -87,21 +94,21 @@ To use the stable diffusion Docker image, you can either build using the provide ``` # 1. build from dockerfile -cd docker +cd ColossalAI/examples/images/diffusion/docker docker build -t hpcaitech/diffusion:0.2.0 . # 2. pull from our docker hub docker pull hpcaitech/diffusion:0.2.0 ``` -Once you have the image ready, you can launch the image with the following command: +Once you have the image ready, you can launch the image with the following command ```bash ######################## # On Your Host Machine # ######################## # make sure you start your image in the repository root directory -cd Colossal-AI +cd ColossalAI # run the docker container docker run --rm \ @@ -113,24 +120,26 @@ docker run --rm \ /bin/bash ######################## -# Insider Container # +# Inside a Container # ######################## # Once you have entered the docker container, go to the stable diffusion directory for training cd examples/images/diffusion/ +# Download the model checkpoint from pretrained (See the following steps) +# Set up your configuration the "train_colossalai.sh" (See the following steps) # start training with colossalai bash train_colossalai.sh ``` It is important for you to configure your volume mapping in order to get the best training experience. -1. **Mandatory**, mount your prepared data to `/data/scratch` via `-v :/data/scratch`, where you need to replace `` with the actual data path on your machine. -2. **Recommended**, store the downloaded model weights to your host machine instead of the container directory via `-v :/root/.cache/huggingface`, where you need to repliace the `` with the actual path. In this way, you don't have to repeatedly download the pretrained weights for every `docker run`. +1. **Mandatory**, mount your prepared data to `/data/scratch` via `-v :/data/scratch`, where you need to replace `` with the actual data path on your machine. Notice that within docker we need to transform the Windows path to a Linux one, e.g. `C:\User\Desktop` into `/mnt/c/User/Desktop`. +2. **Recommended**, store the downloaded model weights to your host machine instead of the container directory via `-v :/root/.cache/huggingface`, where you need to replace the `` with the actual path. In this way, you don't have to repeatedly download the pretrained weights for every `docker run`. 3. **Optional**, if you encounter any problem stating that shared memory is insufficient inside container, please add `-v /dev/shm:/dev/shm` to your `docker run` command. ## Download the model checkpoint from pretrained -### stable-diffusion-v2-base(Recommand) +### stable-diffusion-v2-base (Recommended) ``` wget https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt @@ -157,10 +166,9 @@ you should the change the `data.file_path` in the `config/train_colossalai.yaml` ## Training -We provide the script `train_colossalai.sh` to run the training task with colossalai, -and can also use `train_ddp.sh` to run the training task with ddp to compare. +We provide the script `train_colossalai.sh` to run the training task with colossalai. Meanwhile, we have enlightened other training process such as DDP model in PyTorch. You can also use `train_ddp.sh` to run the training task with ddp to compare the corresponding performance. -In `train_colossalai.sh` the main command is: +In `train_colossalai.sh` the main command is ``` python main.py --logdir /tmp/ --train --base configs/train_colossalai.yaml --ckpt 512-base-ema.ckpt @@ -174,11 +182,12 @@ python main.py --logdir /tmp/ --train --base configs/train_colossalai.yaml --ckp ### Training config -You can change the trainging config in the yaml file +You can change the training config in the yaml file -- devices: device number used for training, default 8 -- max_epochs: max training epochs, default 2 -- precision: the precision type used in training, default 16 (fp16), you must use fp16 if you want to apply colossalai +- devices: device number used for training, default = 8 +- max_epochs: max training epochs, default = 2 +- precision: the precision type used in training, default = 16 (fp16), you must use fp16 if you want to apply colossalai +- placement_policy: the training strategy supported by Colossal AI, default = 'cuda', which refers to loading all the parameters into cuda memory. On the other hand, 'cpu' refers to 'cpu offload' strategy while 'auto' enables 'Gemini', both featured by Colossal AI. - more information about the configuration of ColossalAIStrategy can be found [here](https://pytorch-lightning.readthedocs.io/en/latest/advanced/model_parallel.html#colossal-ai) @@ -193,7 +202,8 @@ python main.py --logdir /tmp/ -t -b configs/Teyvat/train_colossalai_teyvat.yaml ``` ## Inference -you can get yout training last.ckpt and train config.yaml in your `--logdir`, and run by + +You can get your training last.ckpt and train config.yaml in your `--logdir`, and run by ``` python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms --outdir ./output \ diff --git a/examples/images/diffusion/configs/Inference/v2-inference-v.yaml b/examples/images/diffusion/configs/Inference/v2-inference-v.yaml index 8ec8dfbfefe9..b05955d3faf7 100644 --- a/examples/images/diffusion/configs/Inference/v2-inference-v.yaml +++ b/examples/images/diffusion/configs/Inference/v2-inference-v.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: parameterization: "v" linear_start: 0.00085 @@ -19,50 +18,42 @@ model: use_ema: False # we set this to false because this is an inference only config unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" diff --git a/examples/images/diffusion/configs/Inference/v2-inference.yaml b/examples/images/diffusion/configs/Inference/v2-inference.yaml index 152c4f3c2b36..5d8d583d06d1 100644 --- a/examples/images/diffusion/configs/Inference/v2-inference.yaml +++ b/examples/images/diffusion/configs/Inference/v2-inference.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: linear_start: 0.00085 linear_end: 0.0120 @@ -18,50 +17,42 @@ model: use_ema: False # we set this to false because this is an inference only config unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" diff --git a/examples/images/diffusion/configs/Inference/v2-inpainting-inference.yaml b/examples/images/diffusion/configs/Inference/v2-inpainting-inference.yaml index 32a9471d71b8..ffaa5e8da2ad 100644 --- a/examples/images/diffusion/configs/Inference/v2-inpainting-inference.yaml +++ b/examples/images/diffusion/configs/Inference/v2-inpainting-inference.yaml @@ -19,106 +19,97 @@ model: use_ema: False unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - image_size: 32 # unused - in_channels: 9 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + image_size: 32 # unused + in_channels: 9 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [ ] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [ ] + dropout: 0.0 + lossconfig: cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" data: - target: ldm.data.laion.WebDataModuleFromConfig - params: - tar_base: null # for concat as in LAION-A - p_unsafe_threshold: 0.1 - filter_word_list: "data/filters.yaml" - max_pwatermark: 0.45 - batch_size: 8 - num_workers: 6 - multinode: True - min_size: 512 - train: - shards: - - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-0/{00000..18699}.tar -" - - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-1/{00000..18699}.tar -" - - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-2/{00000..18699}.tar -" - - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-3/{00000..18699}.tar -" - - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-4/{00000..18699}.tar -" #{00000-94333}.tar" - shuffle: 10000 - image_key: jpg - image_transforms: - - target: torchvision.transforms.Resize - params: - size: 512 - interpolation: 3 - - target: torchvision.transforms.RandomCrop - params: - size: 512 - postprocess: - target: ldm.data.laion.AddMask - params: - mode: "512train-large" - p_drop: 0.25 - # NOTE use enough shards to avoid empty validation loops in workers - validation: - shards: - - "pipe:aws s3 cp s3://deep-floyd-s3/datasets/laion_cleaned-part5/{93001..94333}.tar - " - shuffle: 0 - image_key: jpg - image_transforms: - - target: torchvision.transforms.Resize - params: - size: 512 - interpolation: 3 - - target: torchvision.transforms.CenterCrop - params: - size: 512 - postprocess: - target: ldm.data.laion.AddMask - params: - mode: "512train-large" - p_drop: 0.25 + tar_base: null # for concat as in LAION-A + p_unsafe_threshold: 0.1 + filter_word_list: "data/filters.yaml" + max_pwatermark: 0.45 + batch_size: 8 + num_workers: 6 + multinode: True + min_size: 512 + train: + shards: + - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-0/{00000..18699}.tar -" + - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-1/{00000..18699}.tar -" + - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-2/{00000..18699}.tar -" + - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-3/{00000..18699}.tar -" + - "pipe:aws s3 cp s3://stability-aws/laion-a-native/part-4/{00000..18699}.tar -" #{00000-94333}.tar" + shuffle: 10000 + image_key: jpg + image_transforms: + - target: torchvision.transforms.Resize + params: + size: 512 + interpolation: 3 + - target: torchvision.transforms.RandomCrop + params: + size: 512 + postprocess: + target: ldm.data.laion.AddMask + params: + mode: "512train-large" + p_drop: 0.25 + # NOTE use enough shards to avoid empty validation loops in workers + validation: + shards: + - "pipe:aws s3 cp s3://deep-floyd-s3/datasets/laion_cleaned-part5/{93001..94333}.tar - " + shuffle: 0 + image_key: jpg + image_transforms: + - target: torchvision.transforms.Resize + params: + size: 512 + interpolation: 3 + - target: torchvision.transforms.CenterCrop + params: + size: 512 + postprocess: + target: ldm.data.laion.AddMask + params: + mode: "512train-large" + p_drop: 0.25 lightning: find_unused_parameters: True @@ -132,8 +123,6 @@ lightning: every_n_train_steps: 10000 image_logger: - target: main.ImageLogger - params: enable_autocast: False disabled: False batch_frequency: 1000 diff --git a/examples/images/diffusion/configs/Inference/v2-midas-inference.yaml b/examples/images/diffusion/configs/Inference/v2-midas-inference.yaml index 531199de4878..01d3729f1590 100644 --- a/examples/images/diffusion/configs/Inference/v2-midas-inference.yaml +++ b/examples/images/diffusion/configs/Inference/v2-midas-inference.yaml @@ -19,54 +19,45 @@ model: use_ema: False depth_stage_config: - target: ldm.modules.midas.api.MiDaSInference - params: - model_type: "dpt_hybrid" + model_type: "dpt_hybrid" unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - image_size: 32 # unused - in_channels: 5 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + image_size: 32 # unused + in_channels: 5 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [ ] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [ ] + dropout: 0.0 + lossconfig: cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" diff --git a/examples/images/diffusion/configs/Inference/x4-upscaling.yaml b/examples/images/diffusion/configs/Inference/x4-upscaling.yaml index 45ecbf9ad863..426d387ca611 100644 --- a/examples/images/diffusion/configs/Inference/x4-upscaling.yaml +++ b/examples/images/diffusion/configs/Inference/x4-upscaling.yaml @@ -20,56 +20,47 @@ model: use_ema: False low_scale_config: - target: ldm.modules.diffusionmodules.upscaling.ImageConcatWithNoiseAugmentation - params: - noise_schedule_config: # image space - linear_start: 0.0001 - linear_end: 0.02 - max_noise_level: 350 + noise_schedule_config: # image space + linear_start: 0.0001 + linear_end: 0.02 + max_noise_level: 350 unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - num_classes: 1000 # timesteps for noise conditioning (here constant, just need one) - image_size: 128 - in_channels: 7 - out_channels: 4 - model_channels: 256 - attention_resolutions: [ 2,4,8] - num_res_blocks: 2 - channel_mult: [ 1, 2, 2, 4] - disable_self_attentions: [True, True, True, False] - disable_middle_self_attn: False - num_heads: 8 - use_spatial_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False - use_linear_in_transformer: True + use_checkpoint: True + num_classes: 1000 # timesteps for noise conditioning (here constant, just need one) + image_size: 128 + in_channels: 7 + out_channels: 4 + model_channels: 256 + attention_resolutions: [ 2,4,8] + num_res_blocks: 2 + channel_mult: [ 1, 2, 2, 4] + disable_self_attentions: [True, True, True, False] + disable_middle_self_attn: False + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False + use_linear_in_transformer: True first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - ddconfig: - # attn_type: "vanilla-xformers" this model needs efficient attention to be feasible on HR data, also the decoder seems to break in half precision (UNet is fine though) - double_z: True - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: [ 1,2,4 ] # num_down = len(ch_mult)-1 - num_res_blocks: 2 - attn_resolutions: [ ] - dropout: 0.0 + embed_dim: 4 + ddconfig: + # attn_type: "vanilla-xformers" this model needs efficient attention to be feasible on HR data, also the decoder seems to break in half precision (UNet is fine though) + double_z: True + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: [ 1,2,4 ] # num_down = len(ch_mult)-1 + num_res_blocks: 2 + attn_resolutions: [ ] + dropout: 0.0 + lossconfig: - lossconfig: - target: torch.nn.Identity cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" diff --git a/examples/images/diffusion/configs/Teyvat/train_colossalai_teyvat.yaml b/examples/images/diffusion/configs/Teyvat/train_colossalai_teyvat.yaml index ff0f4c5a0463..9e760124c7a4 100644 --- a/examples/images/diffusion/configs/Teyvat/train_colossalai_teyvat.yaml +++ b/examples/images/diffusion/configs/Teyvat/train_colossalai_teyvat.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: parameterization: "v" linear_start: 0.00085 @@ -20,81 +19,70 @@ model: use_ema: False scheduler_config: # 10000 warmup steps - target: ldm.lr_scheduler.LambdaLinearScheduler - params: - warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch - cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases - f_start: [ 1.e-6 ] - f_max: [ 1.e-4 ] - f_min: [ 1.e-10 ] + warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1.e-4 ] + f_min: [ 1.e-10 ] unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" data: - target: main.DataModuleFromConfig - params: - batch_size: 16 - num_workers: 4 - train: - target: ldm.data.teyvat.hf_dataset - params: - path: Fazzie/Teyvat - image_transforms: - - target: torchvision.transforms.Resize - params: - size: 512 - - target: torchvision.transforms.RandomCrop - params: - size: 512 - - target: torchvision.transforms.RandomHorizontalFlip + batch_size: 16 + num_workers: 4 + train: + target: ldm.data.teyvat.hf_dataset + params: + path: Fazzie/Teyvat + image_transforms: + - target: torchvision.transforms.Resize + params: + size: 512 + - target: torchvision.transforms.RandomCrop + params: + size: 512 + - target: torchvision.transforms.RandomHorizontalFlip lightning: trainer: @@ -105,13 +93,11 @@ lightning: precision: 16 auto_select_gpus: False strategy: - target: strategies.ColossalAIStrategy - params: - use_chunk: True - enable_distributed_storage: True - placement_policy: cuda - force_outputs_fp32: true - min_chunk_size: 64 + use_chunk: True + enable_distributed_storage: True + placement_policy: cuda + force_outputs_fp32: true + min_chunk_size: 64 log_every_n_steps: 2 logger: True @@ -120,9 +106,7 @@ lightning: logger_config: wandb: - target: loggers.WandbLogger - params: - name: nowname - save_dir: "/tmp/diff_log/" - offline: opt.debug - id: nowname + name: nowname + save_dir: "/tmp/diff_log/" + offline: opt.debug + id: nowname diff --git a/examples/images/diffusion/configs/train_colossalai.yaml b/examples/images/diffusion/configs/train_colossalai.yaml index 88432e978a0f..5f745286a719 100644 --- a/examples/images/diffusion/configs/train_colossalai.yaml +++ b/examples/images/diffusion/configs/train_colossalai.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: parameterization: "v" linear_start: 0.00085 @@ -19,95 +18,83 @@ model: use_ema: False # we set this to false because this is an inference only config scheduler_config: # 10000 warmup steps - target: ldm.lr_scheduler.LambdaLinearScheduler - params: - warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch - cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases - f_start: [ 1.e-6 ] - f_max: [ 1.e-4 ] - f_min: [ 1.e-10 ] + warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1.e-4 ] + f_min: [ 1.e-10 ] unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" data: - target: main.DataModuleFromConfig - params: - batch_size: 128 - wrap: False - # num_workwers should be 2 * batch_size, and total num less than 1024 - # e.g. if use 8 devices, no more than 128 - num_workers: 128 - train: - target: ldm.data.base.Txt2ImgIterableBaseDataset - params: - file_path: # YOUR DATASET_PATH - world_size: 1 - rank: 0 + batch_size: 128 + wrap: False + # num_workwers should be 2 * batch_size, and total num less than 1024 + # e.g. if use 8 devices, no more than 128 + num_workers: 128 + train: + target: ldm.data.base.Txt2ImgIterableBaseDataset + params: + file_path: # YOUR DATASET_PATH + world_size: 1 + rank: 0 lightning: trainer: accelerator: 'gpu' - devices: 8 + devices: 2 log_gpu_memory: all max_epochs: 2 precision: 16 auto_select_gpus: False strategy: - target: strategies.ColossalAIStrategy - params: - use_chunk: True - enable_distributed_storage: True - placement_policy: cuda - force_outputs_fp32: true - min_chunk_size: 64 + use_chunk: True + enable_distributed_storage: True + placement_policy: cuda + force_outputs_fp32: true + min_chunk_size: 64 log_every_n_steps: 2 logger: True @@ -116,9 +103,7 @@ lightning: logger_config: wandb: - target: loggers.WandbLogger - params: - name: nowname - save_dir: "/tmp/diff_log/" - offline: opt.debug - id: nowname + name: nowname + save_dir: "/tmp/diff_log/" + offline: opt.debug + id: nowname diff --git a/examples/images/diffusion/configs/train_colossalai_cifar10.yaml b/examples/images/diffusion/configs/train_colossalai_cifar10.yaml index 0ba06f832178..0d0f185426c2 100644 --- a/examples/images/diffusion/configs/train_colossalai_cifar10.yaml +++ b/examples/images/diffusion/configs/train_colossalai_cifar10.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: parameterization: "v" linear_start: 0.00085 @@ -19,82 +18,71 @@ model: use_ema: False # we set this to false because this is an inference only config scheduler_config: # 10000 warmup steps - target: ldm.lr_scheduler.LambdaLinearScheduler - params: - warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch - cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases - f_start: [ 1.e-6 ] - f_max: [ 1.e-4 ] - f_min: [ 1.e-10 ] + warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1.e-4 ] + f_min: [ 1.e-10 ] unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" data: - target: main.DataModuleFromConfig - params: - batch_size: 4 - num_workers: 4 - train: - target: ldm.data.cifar10.hf_dataset - params: - name: cifar10 - image_transforms: - - target: torchvision.transforms.Resize - params: - size: 512 - interpolation: 3 - - target: torchvision.transforms.RandomCrop - params: - size: 512 - - target: torchvision.transforms.RandomHorizontalFlip + batch_size: 4 + num_workers: 4 + train: + target: ldm.data.cifar10.hf_dataset + params: + name: cifar10 + image_transforms: + - target: torchvision.transforms.Resize + params: + size: 512 + interpolation: 3 + - target: torchvision.transforms.RandomCrop + params: + size: 512 + - target: torchvision.transforms.RandomHorizontalFlip lightning: trainer: @@ -105,13 +93,11 @@ lightning: precision: 16 auto_select_gpus: False strategy: - target: strategies.ColossalAIStrategy - params: - use_chunk: True - enable_distributed_storage: True - placement_policy: cuda - force_outputs_fp32: true - min_chunk_size: 64 + use_chunk: True + enable_distributed_storage: True + placement_policy: cuda + force_outputs_fp32: true + min_chunk_size: 64 log_every_n_steps: 2 logger: True @@ -120,9 +106,7 @@ lightning: logger_config: wandb: - target: loggers.WandbLogger - params: - name: nowname - save_dir: "/tmp/diff_log/" - offline: opt.debug - id: nowname + name: nowname + save_dir: "/tmp/diff_log/" + offline: opt.debug + id: nowname diff --git a/examples/images/diffusion/configs/train_ddp.yaml b/examples/images/diffusion/configs/train_ddp.yaml index a63df887e719..f3ae3ddb5ff6 100644 --- a/examples/images/diffusion/configs/train_ddp.yaml +++ b/examples/images/diffusion/configs/train_ddp.yaml @@ -1,6 +1,5 @@ model: base_learning_rate: 1.0e-4 - target: ldm.models.diffusion.ddpm.LatentDiffusion params: parameterization: "v" linear_start: 0.00085 @@ -19,77 +18,65 @@ model: use_ema: False # we set this to false because this is an inference only config scheduler_config: # 10000 warmup steps - target: ldm.lr_scheduler.LambdaLinearScheduler - params: - warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch - cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases - f_start: [ 1.e-6 ] - f_max: [ 1.e-4 ] - f_min: [ 1.e-10 ] + warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1.e-4 ] + f_min: [ 1.e-10 ] unet_config: - target: ldm.modules.diffusionmodules.openaimodel.UNetModel - params: - use_checkpoint: True - use_fp16: True - image_size: 32 # unused - in_channels: 4 - out_channels: 4 - model_channels: 320 - attention_resolutions: [ 4, 2, 1 ] - num_res_blocks: 2 - channel_mult: [ 1, 2, 4, 4 ] - num_head_channels: 64 # need to fix for flash-attn - use_spatial_transformer: True - use_linear_in_transformer: True - transformer_depth: 1 - context_dim: 1024 - legacy: False + use_checkpoint: True + use_fp16: True + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_head_channels: 64 # need to fix for flash-attn + use_spatial_transformer: True + use_linear_in_transformer: True + transformer_depth: 1 + context_dim: 1024 + legacy: False first_stage_config: - target: ldm.models.autoencoder.AutoencoderKL - params: - embed_dim: 4 - monitor: val/rec_loss - ddconfig: - #attn_type: "vanilla-xformers" - double_z: true - z_channels: 4 - resolution: 256 - in_channels: 3 - out_ch: 3 - ch: 128 - ch_mult: - - 1 - - 2 - - 4 - - 4 - num_res_blocks: 2 - attn_resolutions: [] - dropout: 0.0 - lossconfig: - target: torch.nn.Identity + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + #attn_type: "vanilla-xformers" + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 cond_stage_config: - target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder - params: - freeze: True - layer: "penultimate" + freeze: True + layer: "penultimate" data: - target: main.DataModuleFromConfig - params: - batch_size: 128 - # num_workwers should be 2 * batch_size, and the total num less than 1024 - # e.g. if use 8 devices, no more than 128 - num_workers: 128 - train: - target: ldm.data.base.Txt2ImgIterableBaseDataset - params: - file_path: # YOUR DATAPATH - world_size: 1 - rank: 0 + batch_size: 128 + # num_workwers should be 2 * batch_size, and the total num less than 1024 + # e.g. if use 8 devices, no more than 128 + num_workers: 128 + train: + target: ldm.data.base.Txt2ImgIterableBaseDataset + params: + file_path: # YOUR DATAPATH + world_size: 1 + rank: 0 lightning: trainer: @@ -100,9 +87,7 @@ lightning: precision: 16 auto_select_gpus: False strategy: - target: strategies.DDPStrategy - params: - find_unused_parameters: False + find_unused_parameters: False log_every_n_steps: 2 # max_steps: 6o logger: True @@ -111,9 +96,7 @@ lightning: logger_config: wandb: - target: loggers.WandbLogger - params: - name: nowname - save_dir: "/data2/tmp/diff_log/" - offline: opt.debug - id: nowname + name: nowname + save_dir: "/data2/tmp/diff_log/" + offline: opt.debug + id: nowname diff --git a/examples/images/diffusion/ldm/data/lsun.py b/examples/images/diffusion/ldm/data/lsun.py index 6256e45715ff..f5bf26c14254 100644 --- a/examples/images/diffusion/ldm/data/lsun.py +++ b/examples/images/diffusion/ldm/data/lsun.py @@ -5,87 +5,105 @@ from torch.utils.data import Dataset from torchvision import transforms - +# This class is used to create a dataset of images from LSUN dataset for training class LSUNBase(Dataset): def __init__(self, - txt_file, - data_root, - size=None, - interpolation="bicubic", - flip_p=0.5 + txt_file, # path to the text file containing the list of image paths + data_root, # root directory of the LSUN dataset + size=None, # the size of images to resize to + interpolation="bicubic", # interpolation method to be used while resizing + flip_p=0.5 # probability of random horizontal flipping ): - self.data_paths = txt_file - self.data_root = data_root - with open(self.data_paths, "r") as f: - self.image_paths = f.read().splitlines() - self._length = len(self.image_paths) + self.data_paths = txt_file # store path to text file containing list of images + self.data_root = data_root # store path to root directory of the dataset + with open(self.data_paths, "r") as f: # open and read the text file + self.image_paths = f.read().splitlines() # read the lines of the file and store as list + self._length = len(self.image_paths) # store the number of images + + # create dictionary to hold image path information self.labels = { "relative_file_path_": [l for l in self.image_paths], "file_path_": [os.path.join(self.data_root, l) for l in self.image_paths], } - self.size = size + # set the image size to be resized + self.size = size + # set the interpolation method for resizing the image self.interpolation = {"linear": PIL.Image.LINEAR, "bilinear": PIL.Image.BILINEAR, "bicubic": PIL.Image.BICUBIC, "lanczos": PIL.Image.LANCZOS, }[interpolation] + # randomly flip the image horizontally with a given probability self.flip = transforms.RandomHorizontalFlip(p=flip_p) def __len__(self): + # return the length of dataset return self._length + def __getitem__(self, i): + # get the image path for the given index example = dict((k, self.labels[k][i]) for k in self.labels) image = Image.open(example["file_path_"]) + # convert it to RGB format if not image.mode == "RGB": image = image.convert("RGB") # default to score-sde preprocessing - img = np.array(image).astype(np.uint8) - crop = min(img.shape[0], img.shape[1]) - h, w, = img.shape[0], img.shape[1] + + img = np.array(image).astype(np.uint8) # convert image to numpy array + crop = min(img.shape[0], img.shape[1]) # crop the image to a square shape + h, w, = img.shape[0], img.shape[1] # get the height and width of image img = img[(h - crop) // 2:(h + crop) // 2, - (w - crop) // 2:(w + crop) // 2] + (w - crop) // 2:(w + crop) // 2] # crop the image to a square shape - image = Image.fromarray(img) - if self.size is not None: + image = Image.fromarray(img) # create an image from numpy array + if self.size is not None: # if image size is provided, resize the image image = image.resize((self.size, self.size), resample=self.interpolation) - image = self.flip(image) - image = np.array(image).astype(np.uint8) - example["image"] = (image / 127.5 - 1.0).astype(np.float32) - return example - + image = self.flip(image) # flip the image horizontally with the given probability + image = np.array(image).astype(np.uint8) + example["image"] = (image / 127.5 - 1.0).astype(np.float32) # normalize the image values and convert to float32 + return example # return the example dictionary containing the image and its file paths +#A dataset class for LSUN Churches training set. +# It initializes by calling the constructor of LSUNBase class and passing the appropriate arguments. +# The text file containing the paths to the images and the root directory where the images are stored are passed as arguments. Any additional keyword arguments passed to this class will be forwarded to the constructor of the parent class. class LSUNChurchesTrain(LSUNBase): def __init__(self, **kwargs): super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs) - +#A dataset class for LSUN Churches validation set. +# It is similar to LSUNChurchesTrain except that it uses a different text file and sets the flip probability to zero by default. class LSUNChurchesValidation(LSUNBase): def __init__(self, flip_p=0., **kwargs): super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches", flip_p=flip_p, **kwargs) - +# A dataset class for LSUN Bedrooms training set. +# It initializes by calling the constructor of LSUNBase class and passing the appropriate arguments. class LSUNBedroomsTrain(LSUNBase): def __init__(self, **kwargs): super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs) - +# A dataset class for LSUN Bedrooms validation set. +# It is similar to LSUNBedroomsTrain except that it uses a different text file and sets the flip probability to zero by default. class LSUNBedroomsValidation(LSUNBase): def __init__(self, flip_p=0.0, **kwargs): super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms", flip_p=flip_p, **kwargs) - +# A dataset class for LSUN Cats training set. +# It initializes by calling the constructor of LSUNBase class and passing the appropriate arguments. +# The text file containing the paths to the images and the root directory where the images are stored are passed as arguments. class LSUNCatsTrain(LSUNBase): def __init__(self, **kwargs): super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs) - +# A dataset class for LSUN Cats validation set. +# It is similar to LSUNCatsTrain except that it uses a different text file and sets the flip probability to zero by default. class LSUNCatsValidation(LSUNBase): def __init__(self, flip_p=0., **kwargs): super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats", diff --git a/examples/images/diffusion/ldm/data/teyvat.py b/examples/images/diffusion/ldm/data/teyvat.py index 61dc29d56e7c..eb5d3ea469d4 100644 --- a/examples/images/diffusion/ldm/data/teyvat.py +++ b/examples/images/diffusion/ldm/data/teyvat.py @@ -13,7 +13,7 @@ def make_multi_folder_data(paths, caption_files=None, **kwargs): """Make a concat dataset from multiple folders - Don't suport captions yet + Don't support captions yet If paths is a list, that's ok, if it's a Dict interpret it as: k=folder v=n_times to repeat that """ diff --git a/examples/images/diffusion/ldm/models/autoencoder.py b/examples/images/diffusion/ldm/models/autoencoder.py index b1bd8377835b..f0a69fe63a8c 100644 --- a/examples/images/diffusion/ldm/models/autoencoder.py +++ b/examples/images/diffusion/ldm/models/autoencoder.py @@ -1,16 +1,13 @@ import torch -try: - import lightning.pytorch as pl -except: - import pytorch_lightning as pl +import lightning.pytorch as pl -import torch.nn.functional as F +from torch import nn +from torch.nn import functional as F +from torch.nn import Identity from contextlib import contextmanager from ldm.modules.diffusionmodules.model import Encoder, Decoder from ldm.modules.distributions.distributions import DiagonalGaussianDistribution - -from ldm.util import instantiate_from_config from ldm.modules.ema import LitEma @@ -32,7 +29,7 @@ def __init__(self, self.image_key = image_key self.encoder = Encoder(**ddconfig) self.decoder = Decoder(**ddconfig) - self.loss = instantiate_from_config(lossconfig) + self.loss = Identity() assert ddconfig["double_z"] self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1) self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) diff --git a/examples/images/diffusion/ldm/models/diffusion/classifier.py b/examples/images/diffusion/ldm/models/diffusion/classifier.py index 612a8371bf20..3cf12f093bea 100644 --- a/examples/images/diffusion/ldm/models/diffusion/classifier.py +++ b/examples/images/diffusion/ldm/models/diffusion/classifier.py @@ -9,9 +9,10 @@ from einops import rearrange from glob import glob from natsort import natsorted - +from ldm.models.diffusion.ddpm import LatentDiffusion +from ldm.lr_scheduler import LambdaLinearScheduler from ldm.modules.diffusionmodules.openaimodel import EncoderUNetModel, UNetModel -from ldm.util import log_txt_as_img, default, ismap, instantiate_from_config +from ldm.util import log_txt_as_img, default, ismap __models__ = { 'class_label': EncoderUNetModel, @@ -86,7 +87,7 @@ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False): print(f"Unexpected Keys: {unexpected}") def load_diffusion(self): - model = instantiate_from_config(self.diffusion_config) + model = LatentDiffusion(**self.diffusion_config.get('params',dict())) self.diffusion_model = model.eval() self.diffusion_model.train = disabled_train for param in self.diffusion_model.parameters(): @@ -221,7 +222,7 @@ def configure_optimizers(self): optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay) if self.use_scheduler: - scheduler = instantiate_from_config(self.scheduler_config) + scheduler = LambdaLinearScheduler(**self.scheduler_config.get('params',dict())) print("Setting up LambdaLR scheduler...") scheduler = [ diff --git a/examples/images/diffusion/ldm/models/diffusion/ddpm.py b/examples/images/diffusion/ldm/models/diffusion/ddpm.py index b7315b048c66..842ec1371ea0 100644 --- a/examples/images/diffusion/ldm/models/diffusion/ddpm.py +++ b/examples/images/diffusion/ldm/models/diffusion/ddpm.py @@ -22,19 +22,22 @@ from functools import partial from einops import rearrange, repeat +from ldm.lr_scheduler import LambdaLinearScheduler from ldm.models.autoencoder import * from ldm.models.autoencoder import AutoencoderKL, IdentityFirstStage from ldm.models.diffusion.ddim import * from ldm.models.diffusion.ddim import DDIMSampler +from ldm.modules.midas.api import MiDaSInference from ldm.modules.diffusionmodules.model import * from ldm.modules.diffusionmodules.model import Decoder, Encoder, Model from ldm.modules.diffusionmodules.openaimodel import * -from ldm.modules.diffusionmodules.openaimodel import AttentionPool2d +from ldm.modules.diffusionmodules.openaimodel import AttentionPool2d, UNetModel from ldm.modules.diffusionmodules.util import extract_into_tensor, make_beta_schedule, noise_like from ldm.modules.distributions.distributions import DiagonalGaussianDistribution, normal_kl +from ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation from ldm.modules.ema import LitEma from ldm.modules.encoders.modules import * -from ldm.util import count_params, default, exists, instantiate_from_config, isimage, ismap, log_txt_as_img, mean_flat +from ldm.util import count_params, default, exists, isimage, ismap, log_txt_as_img, mean_flat from omegaconf import ListConfig from torch.optim.lr_scheduler import LambdaLR from torchvision.utils import make_grid @@ -690,7 +693,7 @@ def register_schedule(self, self.make_cond_schedule() def instantiate_first_stage(self, config): - model = instantiate_from_config(config) + model = AutoencoderKL(**config) self.first_stage_model = model.eval() self.first_stage_model.train = disabled_train for param in self.first_stage_model.parameters(): @@ -706,15 +709,13 @@ def instantiate_cond_stage(self, config): self.cond_stage_model = None # self.be_unconditional = True else: - model = instantiate_from_config(config) + model = FrozenOpenCLIPEmbedder(**config) self.cond_stage_model = model.eval() self.cond_stage_model.train = disabled_train for param in self.cond_stage_model.parameters(): param.requires_grad = False else: - assert config != '__is_first_stage__' - assert config != '__is_unconditional__' - model = instantiate_from_config(config) + model = FrozenOpenCLIPEmbedder(**config) self.cond_stage_model = model def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False): @@ -1479,8 +1480,7 @@ def configure_optimizers(self): # opt = torch.optim.AdamW(params, lr=lr) if self.use_scheduler: - assert 'target' in self.scheduler_config - scheduler = instantiate_from_config(self.scheduler_config) + scheduler = LambdaLinearScheduler(**self.scheduler_config) rank_zero_info("Setting up LambdaLR scheduler...") scheduler = [{'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), 'interval': 'step', 'frequency': 1}] @@ -1502,7 +1502,7 @@ class DiffusionWrapper(pl.LightningModule): def __init__(self, diff_model_config, conditioning_key): super().__init__() self.sequential_cross_attn = diff_model_config.pop("sequential_crossattn", False) - self.diffusion_model = instantiate_from_config(diff_model_config) + self.diffusion_model = UNetModel(**diff_model_config) self.conditioning_key = conditioning_key assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm'] @@ -1551,7 +1551,7 @@ def __init__(self, *args, low_scale_config, low_scale_key="LR", noise_level_key= self.noise_level_key = noise_level_key def instantiate_low_stage(self, config): - model = instantiate_from_config(config) + model = ImageConcatWithNoiseAugmentation(**config) self.low_scale_model = model.eval() self.low_scale_model.train = disabled_train for param in self.low_scale_model.parameters(): @@ -1933,7 +1933,7 @@ class LatentDepth2ImageDiffusion(LatentFinetuneDiffusion): def __init__(self, depth_stage_config, concat_keys=("midas_in",), *args, **kwargs): super().__init__(concat_keys=concat_keys, *args, **kwargs) - self.depth_model = instantiate_from_config(depth_stage_config) + self.depth_model = MiDaSInference(**depth_stage_config) self.depth_stage_key = concat_keys[0] @torch.no_grad() @@ -2006,7 +2006,7 @@ def __init__(self, self.low_scale_key = low_scale_key def instantiate_low_stage(self, config): - model = instantiate_from_config(config) + model = ImageConcatWithNoiseAugmentation(**config) self.low_scale_model = model.eval() self.low_scale_model.train = disabled_train for param in self.low_scale_model.parameters(): diff --git a/examples/images/diffusion/main.py b/examples/images/diffusion/main.py index 4dd88a5eca44..713029fc677d 100644 --- a/examples/images/diffusion/main.py +++ b/examples/images/diffusion/main.py @@ -10,11 +10,8 @@ import numpy as np import torch import torchvision +import lightning.pytorch as pl -try: - import lightning.pytorch as pl -except: - import pytorch_lightning as pl from functools import partial @@ -23,19 +20,15 @@ from PIL import Image from prefetch_generator import BackgroundGenerator from torch.utils.data import DataLoader, Dataset, Subset, random_split +from ldm.models.diffusion.ddpm import LatentDiffusion -try: - from lightning.pytorch import seed_everything - from lightning.pytorch.callbacks import Callback, LearningRateMonitor, ModelCheckpoint - from lightning.pytorch.trainer import Trainer - from lightning.pytorch.utilities import rank_zero_info, rank_zero_only - LIGHTNING_PACK_NAME = "lightning.pytorch." -except: - from pytorch_lightning import seed_everything - from pytorch_lightning.callbacks import Callback, LearningRateMonitor, ModelCheckpoint - from pytorch_lightning.trainer import Trainer - from pytorch_lightning.utilities import rank_zero_info, rank_zero_only - LIGHTNING_PACK_NAME = "pytorch_lightning." +from lightning.pytorch import seed_everything +from lightning.pytorch.callbacks import Callback, LearningRateMonitor, ModelCheckpoint +from lightning.pytorch.trainer import Trainer +from lightning.pytorch.utilities import rank_zero_info, rank_zero_only +from lightning.pytorch.loggers import WandbLogger, TensorBoardLogger +from lightning.pytorch.strategies import ColossalAIStrategy,DDPStrategy +LIGHTNING_PACK_NAME = "lightning.pytorch." from ldm.data.base import Txt2ImgIterableBaseDataset from ldm.util import instantiate_from_config @@ -44,14 +37,18 @@ class DataLoaderX(DataLoader): - +# A custom data loader class that inherits from DataLoader def __iter__(self): + # Overriding the __iter__ method of DataLoader to return a BackgroundGenerator + #This is to enable data loading in the background to improve training performance return BackgroundGenerator(super().__iter__()) def get_parser(**parser_kwargs): + #A function to create an ArgumentParser object and add arguments to it def str2bool(v): + # A helper function to parse boolean values from command line arguments if isinstance(v, bool): return v if v.lower() in ("yes", "true", "t", "y", "1"): @@ -60,8 +57,10 @@ def str2bool(v): return False else: raise argparse.ArgumentTypeError("Boolean value expected.") - + # Create an ArgumentParser object with specifies kwargs parser = argparse.ArgumentParser(**parser_kwargs) + + # Add various command line arguments with their default values and descriptions parser.add_argument( "-n", "--name", @@ -161,14 +160,18 @@ def str2bool(v): return parser - +# A function that returns the non-default arguments between two objects def nondefault_trainer_args(opt): + # create an argument parser parser = argparse.ArgumentParser() + # add pytorch lightning trainer default arguments parser = Trainer.add_argparse_args(parser) + # parse the empty arguments to obtain the default values args = parser.parse_args([]) + # return all non-default arguments return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k)) - +# A dataset wrapper class to create a pytorch dataset from an arbitrary object class WrappedDataset(Dataset): """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset""" @@ -181,7 +184,7 @@ def __len__(self): def __getitem__(self, idx): return self.data[idx] - +# A function to initialize worker processes def worker_init_fn(_): worker_info = torch.utils.data.get_worker_info() @@ -189,15 +192,18 @@ def worker_init_fn(_): worker_id = worker_info.id if isinstance(dataset, Txt2ImgIterableBaseDataset): + #divide the dataset into equal parts for each worker split_size = dataset.num_records // worker_info.num_workers + #set the sample IDs for the current worker # reset num_records to the true number to retain reliable length information dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size] + # set the seed for the current worker current_id = np.random.choice(len(np.random.get_state()[1]), 1) return np.random.seed(np.random.get_state()[1][current_id] + worker_id) else: return np.random.seed(np.random.get_state()[1][0] + worker_id) - +#Provide functionality for creating data loaders based on provided dataset configurations class DataModuleFromConfig(pl.LightningDataModule): def __init__(self, @@ -212,10 +218,12 @@ def __init__(self, use_worker_init_fn=False, shuffle_val_dataloader=False): super().__init__() + # Set data module attributes self.batch_size = batch_size self.dataset_configs = dict() self.num_workers = num_workers if num_workers is not None else batch_size * 2 self.use_worker_init_fn = use_worker_init_fn + # If a dataset is passed, add it to the dataset configs and create a corresponding dataloader method if train is not None: self.dataset_configs["train"] = train self.train_dataloader = self._train_dataloader @@ -231,21 +239,28 @@ def __init__(self, self.wrap = wrap def prepare_data(self): + # Instantiate datasets for data_cfg in self.dataset_configs.values(): instantiate_from_config(data_cfg) def setup(self, stage=None): + # Instantiate datasets from the dataset configs self.datasets = dict((k, instantiate_from_config(self.dataset_configs[k])) for k in self.dataset_configs) + + # If wrap is true, create a WrappedDataset for each dataset if self.wrap: for k in self.datasets: self.datasets[k] = WrappedDataset(self.datasets[k]) def _train_dataloader(self): + #Check if the train dataset is iterable is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset) + #Set the worker initialization function of the dataset is iterable or use_worker_init_fn is True if is_iterable_dataset or self.use_worker_init_fn: init_fn = worker_init_fn else: init_fn = None + # Return a DataLoaderX object for the train dataset return DataLoaderX(self.datasets["train"], batch_size=self.batch_size, num_workers=self.num_workers, @@ -253,10 +268,12 @@ def _train_dataloader(self): worker_init_fn=init_fn) def _val_dataloader(self, shuffle=False): + #Check if the validation dataset is iterable if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn: init_fn = worker_init_fn else: init_fn = None + # Return a DataLoaderX object for the validation dataset return DataLoaderX(self.datasets["validation"], batch_size=self.batch_size, num_workers=self.num_workers, @@ -264,7 +281,9 @@ def _val_dataloader(self, shuffle=False): shuffle=shuffle) def _test_dataloader(self, shuffle=False): + # Check if the test dataset is iterable is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset) + # Set the worker initialization function if the dataset is iterable or use_worker_init_fn is True if is_iterable_dataset or self.use_worker_init_fn: init_fn = worker_init_fn else: @@ -291,6 +310,7 @@ def _predict_dataloader(self, shuffle=False): class SetupCallback(Callback): + # Initialize the callback with the necessary parameters def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config): super().__init__() @@ -302,12 +322,14 @@ def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_confi self.config = config self.lightning_config = lightning_config + # Save a checkpoint if training is interrupted with keyboard interrupt def on_keyboard_interrupt(self, trainer, pl_module): if trainer.global_rank == 0: print("Summoning checkpoint.") ckpt_path = os.path.join(self.ckptdir, "last.ckpt") trainer.save_checkpoint(ckpt_path) + # Create necessary directories and save configuration files before training starts # def on_pretrain_routine_start(self, trainer, pl_module): def on_fit_start(self, trainer, pl_module): if trainer.global_rank == 0: @@ -316,6 +338,7 @@ def on_fit_start(self, trainer, pl_module): os.makedirs(self.ckptdir, exist_ok=True) os.makedirs(self.cfgdir, exist_ok=True) + #Create trainstep checkpoint directory if necessary if "callbacks" in self.lightning_config: if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']: os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True) @@ -323,11 +346,13 @@ def on_fit_start(self, trainer, pl_module): print(OmegaConf.to_yaml(self.config)) OmegaConf.save(self.config, os.path.join(self.cfgdir, "{}-project.yaml".format(self.now))) + # Save project config and lightning config as YAML files print("Lightning config") print(OmegaConf.to_yaml(self.lightning_config)) OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}), os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now))) + # Remove log directory if resuming training and directory already exists else: # ModelCheckpoint callback created log directory --- remove it if not self.resume and os.path.exists(self.logdir): @@ -346,25 +371,28 @@ def on_fit_start(self, trainer, pl_module): # trainer.save_checkpoint(ckpt_path) +# PyTorch Lightning callback for logging images during training and validation of a deep learning model class ImageLogger(Callback): def __init__(self, - batch_frequency, - max_images, - clamp=True, - increase_log_steps=True, - rescale=True, - disabled=False, - log_on_batch_idx=False, - log_first_step=False, - log_images_kwargs=None): + batch_frequency, # Frequency of batches on which to log images + max_images, # Maximum number of images to log + clamp=True, # Whether to clamp pixel values to [-1,1] + increase_log_steps=True, # Whether to increase frequency of log steps exponentially + rescale=True, # Whether to rescale pixel values to [0,1] + disabled=False, # Whether to disable logging + log_on_batch_idx=False, # Whether to log on batch index instead of global step + log_first_step=False, # Whether to log on the first step + log_images_kwargs=None): # Additional keyword arguments to pass to log_images method super().__init__() self.rescale = rescale self.batch_freq = batch_frequency self.max_images = max_images self.logger_log_images = { - pl.loggers.CSVLogger: self._testtube, + # Dictionary of logger classes and their corresponding logging methods + pl.loggers.CSVLogger: self._testtube, } + # Create a list of exponentially increasing log steps, starting from 1 and ending at batch_frequency self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)] if not increase_log_steps: self.log_steps = [self.batch_freq] @@ -374,17 +402,32 @@ def __init__(self, self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} self.log_first_step = log_first_step - @rank_zero_only - def _testtube(self, pl_module, images, batch_idx, split): + @rank_zero_only # Ensure that only the first process in distributed training executes this method + def _testtube(self, # The PyTorch Lightning module + pl_module, # A dictionary of images to log. + images, # + batch_idx, # The batch index. + split # The split (train/val) on which to log the images + ): + # Method for logging images using test-tube logger for k in images: grid = torchvision.utils.make_grid(images[k]) grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w tag = f"{split}/{k}" + # Add image grid to logger's experiment pl_module.logger.experiment.add_image(tag, grid, global_step=pl_module.global_step) @rank_zero_only - def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): + def log_local(self, + save_dir, + split, # The split (train/val) on which to log the images + images, # A dictionary of images to log + global_step, # The global step + current_epoch, # The current epoch. + batch_idx + ): + # Method for saving image grids to local file system root = os.path.join(save_dir, "images", split) for k in images: grid = torchvision.utils.make_grid(images[k], nrow=4) @@ -396,12 +439,16 @@ def log_local(self, save_dir, split, images, global_step, current_epoch, batch_i filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx) path = os.path.join(root, filename) os.makedirs(os.path.split(path)[0], exist_ok=True) + # Save image grid as PNG file Image.fromarray(grid).save(path) def log_img(self, pl_module, batch, batch_idx, split="train"): + #Function for logging images to both the logger and local file system. check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step + # check if it's time to log an image batch if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 hasattr(pl_module, "log_images") and callable(pl_module.log_images) and self.max_images > 0): + # Get logger type and check if training mode is on logger = type(pl_module.logger) is_train = pl_module.training @@ -409,8 +456,10 @@ def log_img(self, pl_module, batch, batch_idx, split="train"): pl_module.eval() with torch.no_grad(): + # Get images from log_images method of the pl_module images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) + # Clip images if specified and convert to CPU tensor for k in images: N = min(images[k].shape[0], self.max_images) images[k] = images[k][:N] @@ -419,15 +468,19 @@ def log_img(self, pl_module, batch, batch_idx, split="train"): if self.clamp: images[k] = torch.clamp(images[k], -1., 1.) + # Log images locally to file system self.log_local(pl_module.logger.save_dir, split, images, pl_module.global_step, pl_module.current_epoch, batch_idx) + # log the images using the logger logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None) logger_log_images(pl_module, images, pl_module.global_step, split) + # switch back to training mode if necessary if is_train: pl_module.train() + # The function checks if it's time to log an image batch def check_frequency(self, check_idx): if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (check_idx > 0 or self.log_first_step): @@ -439,14 +492,17 @@ def check_frequency(self, check_idx): return True return False + # Log images on train batch end if logging is not disabled def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): # if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): # self.log_img(pl_module, batch, batch_idx, split="train") pass + # Log images on validation batch end if logging is not disabled and in validation mode def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): if not self.disabled and pl_module.global_step > 0: self.log_img(pl_module, batch, batch_idx, split="val") + # log gradients during calibration if necessary if hasattr(pl_module, 'calibrate_grad_norm'): if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0: self.log_gradients(trainer, pl_module, batch_idx=batch_idx) @@ -458,6 +514,7 @@ class CUDACallback(Callback): def on_train_start(self, trainer, pl_module): rank_zero_info("Training is starting") + #the method is called at the end of each training epoch def on_train_end(self, trainer, pl_module): rank_zero_info("Training is ending") @@ -524,6 +581,7 @@ def on_train_epoch_end(self, trainer, pl_module): # params: # key: value + # get the current time to create a new logging directory now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") # add cwd for convenience and to make classes in this file available when @@ -535,11 +593,13 @@ def on_train_epoch_end(self, trainer, pl_module): parser = Trainer.add_argparse_args(parser) opt, unknown = parser.parse_known_args() + # Verify the arguments are both specified if opt.name and opt.resume: raise ValueError("-n/--name and -r/--resume cannot be specified both." "If you want to resume training in a new log folder, " "use -n/--name in combination with --resume_from_checkpoint") + # Check if the "resume" option is specified, resume training from the checkpoint if it is true ckpt = None if opt.resume: rank_zero_info("Resuming from {}".format(opt.resume)) @@ -557,8 +617,10 @@ def on_train_epoch_end(self, trainer, pl_module): logdir = opt.resume.rstrip("/") ckpt = os.path.join(logdir, "checkpoints", "last.ckpt") + # Finds all ".yaml" configuration files in the log directory and adds them to the list of base configurations base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml"))) opt.base = base_configs + opt.base + # Gets the name of the current log directory by splitting the path and taking the last element. _tmp = logdir.split("/") nowname = _tmp[-1] else: @@ -574,13 +636,17 @@ def on_train_epoch_end(self, trainer, pl_module): nowname = now + name + opt.postfix logdir = os.path.join(opt.logdir, nowname) + # Sets the checkpoint path of the 'ckpt' option is specified if opt.ckpt: ckpt = opt.ckpt + # Create the checkpoint and configuration directories within the log directory. ckptdir = os.path.join(logdir, "checkpoints") cfgdir = os.path.join(logdir, "configs") + # Sets the seed for the random number generator to ensure reproducibility seed_everything(opt.seed) + # Initialize and save configuration using teh OmegaConf library. try: # init and save configs configs = [OmegaConf.load(cfg) for cfg in opt.base] @@ -593,6 +659,7 @@ def on_train_epoch_end(self, trainer, pl_module): for k in nondefault_trainer_args(opt): trainer_config[k] = getattr(opt, k) + # Check whether the accelerator is gpu if not trainer_config["accelerator"] == "gpu": del trainer_config["accelerator"] cpu = True @@ -609,157 +676,131 @@ def on_train_epoch_end(self, trainer, pl_module): config.model["params"].update({"use_fp16": False}) if ckpt is not None: + #If a checkpoint path is specified in the ckpt variable, the code updates the "ckpt" key in the "params" dictionary of the config.model configuration with the value of ckpt config.model["params"].update({"ckpt": ckpt}) rank_zero_info("Using ckpt_path = {}".format(config.model["params"]["ckpt"])) - model = instantiate_from_config(config.model) + model = LatentDiffusion(**config.model.get("params", dict())) # trainer and callbacks trainer_kwargs = dict() # config the logger - # default logger configs + # Default logger configs to log training metrics during the training process. default_logger_cfgs = { "wandb": { - "target": LIGHTNING_PACK_NAME + "loggers.WandbLogger", - "params": { "name": nowname, "save_dir": logdir, "offline": opt.debug, "id": nowname, } - }, + , "tensorboard": { - "target": LIGHTNING_PACK_NAME + "loggers.TensorBoardLogger", - "params": { "save_dir": logdir, "name": "diff_tb", "log_graph": True } - } } + # Set up the logger for TensorBoard default_logger_cfg = default_logger_cfgs["tensorboard"] if "logger" in lightning_config: logger_cfg = lightning_config.logger + trainer_kwargs["logger"] = WandbLogger(**logger_cfg) else: logger_cfg = default_logger_cfg - logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg) - trainer_kwargs["logger"] = instantiate_from_config(logger_cfg) + trainer_kwargs["logger"] = TensorBoardLogger(**logger_cfg) # config the strategy, defualt is ddp if "strategy" in trainer_config: strategy_cfg = trainer_config["strategy"] - strategy_cfg["target"] = LIGHTNING_PACK_NAME + strategy_cfg["target"] + trainer_kwargs["strategy"] = ColossalAIStrategy(**strategy_cfg) else: - strategy_cfg = { - "target": LIGHTNING_PACK_NAME + "strategies.DDPStrategy", - "params": { - "find_unused_parameters": False - } - } - - trainer_kwargs["strategy"] = instantiate_from_config(strategy_cfg) + strategy_cfg = {"find_unused_parameters": False} + trainer_kwargs["strategy"] = DDPStrategy(**strategy_cfg) + # Set up ModelCheckpoint callback to save best models # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to # specify which metric is used to determine best models default_modelckpt_cfg = { - "target": LIGHTNING_PACK_NAME + "callbacks.ModelCheckpoint", - "params": { "dirpath": ckptdir, "filename": "{epoch:06}", "verbose": True, "save_last": True, } - } if hasattr(model, "monitor"): - default_modelckpt_cfg["params"]["monitor"] = model.monitor - default_modelckpt_cfg["params"]["save_top_k"] = 3 + default_modelckpt_cfg["monitor"] = model.monitor + default_modelckpt_cfg["save_top_k"] = 3 if "modelcheckpoint" in lightning_config: - modelckpt_cfg = lightning_config.modelcheckpoint + modelckpt_cfg = lightning_config.modelcheckpoint["params"] else: modelckpt_cfg = OmegaConf.create() modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg) if version.parse(pl.__version__) < version.parse('1.4.0'): - trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg) - - # add callback which sets up log directory - default_callbacks_cfg = { - "setup_callback": { - "target": "main.SetupCallback", - "params": { - "resume": opt.resume, - "now": now, - "logdir": logdir, - "ckptdir": ckptdir, - "cfgdir": cfgdir, - "config": config, - "lightning_config": lightning_config, - } - }, - "image_logger": { - "target": "main.ImageLogger", - "params": { - "batch_frequency": 750, - "max_images": 4, - "clamp": True - } - }, - "learning_rate_logger": { - "target": "main.LearningRateMonitor", - "params": { - "logging_interval": "step", - # "log_momentum": True - } - }, - "cuda_callback": { - "target": "main.CUDACallback" - }, - } - - if "callbacks" in lightning_config: - callbacks_cfg = lightning_config.callbacks - else: - callbacks_cfg = OmegaConf.create() - - if 'metrics_over_trainsteps_checkpoint' in callbacks_cfg: - print( - 'Caution: Saving checkpoints every n train steps without deleting. This might require some free space.') - default_metrics_over_trainsteps_ckpt_dict = { - 'metrics_over_trainsteps_checkpoint': { - "target": LIGHTNING_PACK_NAME + 'callbacks.ModelCheckpoint', - 'params': { - "dirpath": os.path.join(ckptdir, 'trainstep_checkpoints'), - "filename": "{epoch:06}-{step:09}", - "verbose": True, - 'save_top_k': -1, - 'every_n_train_steps': 10000, - 'save_weights_only': True - } - } + trainer_kwargs["checkpoint_callback"] = ModelCheckpoint(**modelckpt_cfg) + + #Create an empty OmegaConf configuration object + + callbacks_cfg = OmegaConf.create() + + #Instantiate items according to the configs + trainer_kwargs.setdefault("callbacks", []) + setup_callback_config = { + "resume": opt.resume, # resume training if applicable + "now": now, + "logdir": logdir, # directory to save the log file + "ckptdir": ckptdir, # directory to save the checkpoint file + "cfgdir": cfgdir, # directory to save the configuration file + "config": config, # configuration dictionary + "lightning_config": lightning_config, # LightningModule configuration } - default_callbacks_cfg.update(default_metrics_over_trainsteps_ckpt_dict) - - callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg) - - trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg] + trainer_kwargs["callbacks"].append(SetupCallback(**setup_callback_config)) + + image_logger_config = { + + "batch_frequency": 750, # how frequently to log images + "max_images": 4, # maximum number of images to log + "clamp": True # whether to clamp pixel values to [0,1] + } + trainer_kwargs["callbacks"].append(ImageLogger(**image_logger_config)) + + learning_rate_logger_config = { + "logging_interval": "step", # logging frequency (either 'step' or 'epoch') + # "log_momentum": True # whether to log momentum (currently commented out) + } + trainer_kwargs["callbacks"].append(LearningRateMonitor(**learning_rate_logger_config)) + + metrics_over_trainsteps_checkpoint_config= { + "dirpath": os.path.join(ckptdir, 'trainstep_checkpoints'), + "filename": "{epoch:06}-{step:09}", + "verbose": True, + 'save_top_k': -1, + 'every_n_train_steps': 10000, + 'save_weights_only': True + } + trainer_kwargs["callbacks"].append(ModelCheckpoint(**metrics_over_trainsteps_checkpoint_config)) + trainer_kwargs["callbacks"].append(CUDACallback()) + # Create a Trainer object with the specified command-line arguments and keyword arguments, and set the log directory trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs) trainer.logdir = logdir - # data - data = instantiate_from_config(config.data) + # Create a data module based on the configuration file + data = DataModuleFromConfig(**config.data) + # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html # calling these ourselves should not be necessary but it is. # lightning still takes care of proper multiprocessing though data.prepare_data() data.setup() + # Print some information about the datasets in the data module for k in data.datasets: rank_zero_info(f"{k}, {data.datasets[k].__class__.__name__}, {len(data.datasets[k])}") - # configure learning rate - bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate + # Configure learning rate based on the batch size, base learning rate and number of GPUs + # If scale_lr is true, calculate the learning rate based on additional factors + bs, base_lr = config.data.batch_size, config.model.base_learning_rate if not cpu: ngpu = trainer_config["devices"] else: @@ -780,7 +821,7 @@ def on_train_epoch_end(self, trainer, pl_module): rank_zero_info("++++ NOT USING LR SCALING ++++") rank_zero_info(f"Setting learning rate to {model.learning_rate:.2e}") - # allow checkpointing via USR1 + # Allow checkpointing via USR1 def melk(*args, **kwargs): # run all checkpoint hooks if trainer.global_rank == 0: @@ -794,20 +835,23 @@ def divein(*args, **kwargs): pudb.set_trace() import signal - + # Assign melk to SIGUSR1 signal and divein to SIGUSR2 signal signal.signal(signal.SIGUSR1, melk) signal.signal(signal.SIGUSR2, divein) - # run + # Run the training and validation if opt.train: try: trainer.fit(model, data) except Exception: melk() raise + # Print the maximum GPU memory allocated during training + print(f"GPU memory usage: {torch.cuda.max_memory_allocated() / 1024**2:.0f} MB") # if not opt.no_test and not trainer.interrupted: # trainer.test(model, data) except Exception: + # If there's an exception, debug it if opt.debug is true and the trainer's global rank is 0 if opt.debug and trainer.global_rank == 0: try: import pudb as debugger @@ -816,7 +860,7 @@ def divein(*args, **kwargs): debugger.post_mortem() raise finally: - # move newly created debug project to debug_runs + # Move the log directory to debug_runs if opt.debug is true and the trainer's global if opt.debug and not opt.resume and trainer.global_rank == 0: dst, name = os.path.split(logdir) dst = os.path.join(dst, "debug_runs", name) diff --git a/examples/images/diffusion/requirements.txt b/examples/images/diffusion/requirements.txt index d0af35353b66..59d027fcf60f 100644 --- a/examples/images/diffusion/requirements.txt +++ b/examples/images/diffusion/requirements.txt @@ -1,10 +1,10 @@ albumentations==1.3.0 -opencv-python==4.6.0 +opencv-python==4.6.0.66 pudb==2019.2 prefetch_generator imageio==2.9.0 imageio-ffmpeg==0.4.2 -torchmetrics==0.6 +torchmetrics==0.7 omegaconf==2.1.1 test-tube>=0.7.5 streamlit>=0.73.1 diff --git a/examples/images/diffusion/scripts/tests/test_checkpoint.py b/examples/images/diffusion/scripts/tests/test_checkpoint.py index a32e66d44cf2..13622c4989fd 100644 --- a/examples/images/diffusion/scripts/tests/test_checkpoint.py +++ b/examples/images/diffusion/scripts/tests/test_checkpoint.py @@ -7,8 +7,9 @@ from diffusers import StableDiffusionPipeline import torch -from ldm.util import instantiate_from_config + from main import get_parser +from ldm.modules.diffusionmodules.openaimodel import UNetModel if __name__ == "__main__": with torch.no_grad(): @@ -17,7 +18,7 @@ config = f.read() base_config = yaml.load(config, Loader=yaml.FullLoader) unet_config = base_config['model']['params']['unet_config'] - diffusion_model = instantiate_from_config(unet_config).to("cuda:0") + diffusion_model = UNetModel(**unet_config).to("cuda:0") pipe = StableDiffusionPipeline.from_pretrained( "/data/scratch/diffuser/stable-diffusion-v1-4" diff --git a/examples/images/diffusion/train_colossalai.sh b/examples/images/diffusion/train_colossalai.sh index c56ed7876e5a..7f1a1bd14615 100755 --- a/examples/images/diffusion/train_colossalai.sh +++ b/examples/images/diffusion/train_colossalai.sh @@ -3,3 +3,4 @@ TRANSFORMERS_OFFLINE=1 DIFFUSERS_OFFLINE=1 python main.py --logdir /tmp --train --base configs/Teyvat/train_colossalai_teyvat.yaml --ckpt diffuser_root_dir/512-base-ema.ckpt + diff --git a/examples/images/dreambooth/README.md b/examples/images/dreambooth/README.md index 14ed66c8d45b..ba4c1a71034a 100644 --- a/examples/images/dreambooth/README.md +++ b/examples/images/dreambooth/README.md @@ -5,12 +5,12 @@ The `train_dreambooth_colossalai.py` script shows how to implement the training By accommodating model data in CPU and GPU and moving the data to the computing device when necessary, [Gemini](https://www.colossalai.org/docs/advanced_tutorials/meet_gemini), the Heterogeneous Memory Manager of [Colossal-AI](https://github.com/hpcaitech/ColossalAI) can breakthrough the GPU memory wall by using GPU and CPU memory (composed of CPU DRAM or nvme SSD memory) together at the same time. Moreover, the model scale can be further improved by combining heterogeneous training with the other parallel approaches, such as data parallel, tensor parallel and pipeline parallel. -## Installing the dependencies +## Installation -Before running the scripts, make sure to install the library's training dependencies: +To begin with, make sure your operating system has the cuda version suitable for this exciting training session, which is cuda11.6-11.8. Notice that you may want to make sure the module versions suitable for the whole environment. Before running the scripts, make sure to install the library's training dependencies: ```bash -pip install -r requirements_colossalai.txt +pip install -r requirements.txt ``` ### Install [colossalai](https://github.com/hpcaitech/ColossalAI.git) @@ -37,9 +37,7 @@ The `text` include the tag `Teyvat`, `Name`,`Element`, `Weapon`, `Region`, `Mode ## Training -The arguement `placement` can be `cpu`, `auto`, `cuda`, with `cpu` the GPU RAM required can be minimized to 4GB but will deceleration, with `cuda` you can also reduce GPU memory by half but accelerated training， with `auto` a more balanced solution for speed and memory can be obtained。 - -**___Note: Change the `resolution` to 768 if you are using the [stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) 768x768 model.___** +We provide the script `colossalai.sh` to run the training task with colossalai. Meanwhile, we also provided traditional training process of dreambooth, `dreambooth.sh`, for possible comparison. For instance, the script of training process for [stable-diffusion-v1-4] model can be modified into: ```bash export MODEL_NAME="CompVis/stable-diffusion-v1-4" @@ -59,12 +57,17 @@ torchrun --nproc_per_node 2 train_dreambooth_colossalai.py \ --max_train_steps=400 \ --placement="cuda" ``` - +- `MODEL_NAME` refers to the model you are training. +- `INSTANCE_DIR` refers to personalized path to instance images, you might need to insert information here. +- `OUTPUT_DIR` refers to local path to save the trained model, you might need to find a path with enough space. +- `resolution` refers to the corresponding resolution number of your target model. Note: Change the `resolution` to 768 if you are using the [stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) 768x768 model. +- `placement` refers to the training strategy supported by Colossal AI, default = 'cuda', which refers to loading all the parameters into cuda memory. On the other hand, 'cpu' refers to 'cpu offload' strategy while 'auto' enables 'Gemini', both featured by Colossal AI. ### Training with prior-preservation loss Prior-preservation is used to avoid overfitting and language-drift. Refer to the paper to learn more about it. For prior-preservation we first generate images using the model with a class prompt and then use those during training along with our data. -According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases. The `num_class_images` flag sets the number of images to generate with the class prompt. You can place existing images in `class_data_dir`, and the training script will generate any additional images so that `num_class_images` are present in `class_data_dir` during training time. + +According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases. The `num_class_images` flag sets the number of images to generate with the class prompt. You can place existing images in `class_data_dir`, and the training script will generate any additional images so that `num_class_images` are present in `class_data_dir` during training time. The general script can be then modified as the following. ```bash export MODEL_NAME="CompVis/stable-diffusion-v1-4" @@ -89,9 +92,32 @@ torchrun --nproc_per_node 2 train_dreambooth_colossalai.py \ --placement="cuda" ``` +## New API +We have modified our previous implementation of Dreambooth with our new Booster API, which offers a more flexible and efficient way to train your model. The new API is more user-friendly and easy to use. You can find the new API in `train_dreambooth_colossalai.py`. +We have also offer a shell script `test_ci.sh` for you to go through all our plugins for the booster. +For more information about the booster API you can refer to https://colossalai.org/docs/basics/booster_api/. + +## Performance + +| Strategy | #GPU | Batch Size | GPU RAM(GB) | speedup | +|:--------------:|:----:|:----------:|:-----------:|:-------:| +| Traditional | 1 | 16 | oom | \ | +| Traditional | 1 | 8 | 61.81 | 1 | +| torch_ddp | 4 | 16 | oom | \ | +| torch_ddp | 4 | 8 | 41.97 | 0.97 | +| gemini | 4 | 16 | 53.29 | \ | +| gemini | 4 | 8 | 29.36 | 2.00 | +| low_level_zero | 4 | 16 | 52.80 | \ | +| low_level_zero | 4 | 8 | 28.87 | 2.02 | + +The evaluation is performed on 4 Nvidia A100 GPUs with 80GB memory each, with GPU 0 & 1, 2 & 3 connected with NVLink. +We finetuned the [stable-diffusion-v1-4](https://huggingface.co/stabilityai/stable-diffusion-v1-4) model with 512x512 resolution on the [Teyvat](https://huggingface.co/datasets/Fazzie/Teyvat) dataset and compared +the memory cost and the throughput for the plugins. + + ## Inference -Once you have trained a model using above command, the inference can be done simply using the `StableDiffusionPipeline`. Make sure to include the `identifier`(e.g. sks in above example) in your prompt. +Once you have trained a model using above command, the inference can be done simply using the `StableDiffusionPipeline`. Make sure to include the `identifier`(e.g. `--instance_prompt="a photo of sks dog" ` in the above example) in your prompt. ```python from diffusers import StableDiffusionPipeline diff --git a/examples/images/dreambooth/colossalai.sh b/examples/images/dreambooth/colossalai.sh index 227d8b8bdb04..db4562dbc921 100755 --- a/examples/images/dreambooth/colossalai.sh +++ b/examples/images/dreambooth/colossalai.sh @@ -1,22 +1,18 @@ -export MODEL_NAME= -export INSTANCE_DIR= -export CLASS_DIR="path-to-class-images" -export OUTPUT_DIR="path-to-save-model" - -HF_DATASETS_OFFLINE=1 -TRANSFORMERS_OFFLINE=1 +HF_DATASETS_OFFLINE=1 +TRANSFORMERS_OFFLINE=1 DIFFUSERS_OFFLINE=1 -torchrun --nproc_per_node 2 --master_port=25641 train_dreambooth_colossalai.py \ - --pretrained_model_name_or_path=$MODEL_NAME \ - --instance_data_dir=$INSTANCE_DIR \ - --output_dir=$OUTPUT_DIR \ - --instance_prompt="a photo of a dog" \ +torchrun --nproc_per_node 4 --standalone train_dreambooth_colossalai.py \ + --pretrained_model_name_or_path="/data/dreambooth/diffuser/stable-diffusion-v1-4" \ + --instance_data_dir="/data/dreambooth/Teyvat/data" \ + --output_dir="./weight_output" \ + --instance_prompt="a picture of a dog" \ --resolution=512 \ + --plugin="gemini" \ --train_batch_size=1 \ - --gradient_accumulation_steps=1 \ --learning_rate=5e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --num_class_images=200 \ - --placement="cuda" \ + --test_run=True \ + --placement="auto" \ diff --git a/examples/images/dreambooth/debug.py b/examples/images/dreambooth/debug.py index c4adb48230be..33219b2caa29 100644 --- a/examples/images/dreambooth/debug.py +++ b/examples/images/dreambooth/debug.py @@ -5,7 +5,7 @@ from diffusers import AutoencoderKL import colossalai -from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx +from colossalai.zero import ColoInitContext, post_process_colo_init_ctx path = "/data/scratch/diffuser/stable-diffusion-v1-4" diff --git a/examples/images/dreambooth/dreambooth.sh b/examples/images/dreambooth/dreambooth.sh index e063bc8279c5..f6b8f5e1b87e 100644 --- a/examples/images/dreambooth/dreambooth.sh +++ b/examples/images/dreambooth/dreambooth.sh @@ -1,7 +1,7 @@ python train_dreambooth.py \ - --pretrained_model_name_or_path= ## Your Model Path \ - --instance_data_dir= ## Your Training Input Pics Path \ - --output_dir="path-to-save-model" \ + --pretrained_model_name_or_path="/data/dreambooth/diffuser/stable-diffusion-v1-4" \ + --instance_data_dir="/data/dreambooth/Teyvat/data" \ + --output_dir="./weight_output" \ --instance_prompt="a photo of a dog" \ --resolution=512 \ --train_batch_size=1 \ diff --git a/examples/images/dreambooth/requirements.txt b/examples/images/dreambooth/requirements.txt index 6c4f40fb5dd0..1ec828c630ef 100644 --- a/examples/images/dreambooth/requirements.txt +++ b/examples/images/dreambooth/requirements.txt @@ -5,4 +5,3 @@ transformers>=4.21.0 ftfy tensorboard modelcards -colossalai diff --git a/examples/images/dreambooth/requirements_colossalai.txt b/examples/images/dreambooth/requirements_colossalai.txt deleted file mode 100644 index c4a0e91703bb..000000000000 --- a/examples/images/dreambooth/requirements_colossalai.txt +++ /dev/null @@ -1,8 +0,0 @@ -diffusers -torch -torchvision -ftfy -tensorboard -modelcards -transformers -colossalai==0.2.0+torch1.12cu11.3 -f https://release.colossalai.org diff --git a/examples/images/dreambooth/test_ci.sh b/examples/images/dreambooth/test_ci.sh index e69de29bb2d1..21f45adae2a0 100644 --- a/examples/images/dreambooth/test_ci.sh +++ b/examples/images/dreambooth/test_ci.sh @@ -0,0 +1,25 @@ +#!/bin/bash +set -xe +pip install -r requirements.txt + +HF_DATASETS_OFFLINE=1 +TRANSFORMERS_OFFLINE=1 +DIFFUSERS_OFFLINE=1 + +# "torch_ddp" "torch_ddp_fp16" "low_level_zero" +for plugin in "gemini"; do + torchrun --nproc_per_node 4 --standalone train_dreambooth_colossalai.py \ + --pretrained_model_name_or_path="/data/dreambooth/diffuser/stable-diffusion-v1-4" \ + --instance_data_dir="/data/dreambooth/Teyvat/data" \ + --output_dir="./weight_output" \ + --instance_prompt="a picture of a dog" \ + --resolution=512 \ + --plugin=$plugin \ + --train_batch_size=1 \ + --learning_rate=5e-6 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --test_run=True \ + --num_class_images=200 \ + --placement="auto" # "cuda" +done diff --git a/examples/images/dreambooth/train_dreambooth_colossalai.py b/examples/images/dreambooth/train_dreambooth_colossalai.py index 5c4c86bc7073..888b28de8306 100644 --- a/examples/images/dreambooth/train_dreambooth_colossalai.py +++ b/examples/images/dreambooth/train_dreambooth_colossalai.py @@ -4,6 +4,7 @@ import os from pathlib import Path from typing import Optional +import shutil import torch import torch.nn.functional as F @@ -21,10 +22,12 @@ from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel.utils import get_static_torch_model +from colossalai.nn.optimizer import HybridAdam from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext +from colossalai.zero.gemini import get_static_torch_model +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin disable_existing_loggers() logger = get_dist_logger() @@ -59,6 +62,13 @@ def parse_args(input_args=None): required=True, help="Path to pretrained model or model identifier from huggingface.co/models.", ) + parser.add_argument( + "--externel_unet_path", + type=str, + default=None, + required=False, + help="Path to the externel unet model.", + ) parser.add_argument( "--revision", type=str, @@ -188,12 +198,19 @@ def parse_args(input_args=None): parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument("--test_run", default=False, help="Whether to use a smaller dataset for test run.") parser.add_argument( "--hub_model_id", type=str, default=None, help="The name of the repository to keep in sync with the local `output_dir`.", ) + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") parser.add_argument( "--logging_dir", type=str, @@ -251,6 +268,7 @@ def __init__( class_prompt=None, size=512, center_crop=False, + test=False, ): self.size = size self.center_crop = center_crop @@ -261,6 +279,8 @@ def __init__( raise ValueError("Instance images root doesn't exists.") self.instance_images_path = list(Path(instance_data_root).iterdir()) + if test: + self.instance_images_path = self.instance_images_path[:10] self.num_instance_images = len(self.instance_images_path) self.instance_prompt = instance_prompt self._length = self.num_instance_images @@ -340,18 +360,6 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: return f"{organization}/{model_id}" -# Gemini + ZeRO DDP -def gemini_zero_dpp(model: torch.nn.Module, placememt_policy: str = "auto"): - from colossalai.nn.parallel import GeminiDDP - - model = GeminiDDP(model, - device=get_current_device(), - placement_policy=placememt_policy, - pin_memory=True, - search_range_mb=64) - return model - - def main(args): if args.seed is None: colossalai.launch_from_torch(config={}) @@ -393,7 +401,7 @@ def main(args): images = pipeline(example["prompt"]).images for i, image in enumerate(images): - hash_image = hashlib.sha1(image.tobytes()).hexdigest() + hash_image = hashlib.sha256(image.tobytes()).hexdigest() image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" image.save(image_filename) @@ -453,12 +461,18 @@ def main(args): revision=args.revision, ) - logger.info(f"Loading UNet2DConditionModel from {args.pretrained_model_name_or_path}", ranks=[0]) - with ColoInitContext(device=get_current_device()): + + if args.externel_unet_path is None: + logger.info(f"Loading UNet2DConditionModel from {args.pretrained_model_name_or_path}", ranks=[0]) unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, - subfolder="unet", - revision=args.revision, - low_cpu_mem_usage=False) + subfolder="unet", + revision=args.revision, + low_cpu_mem_usage=False) + else: + logger.info(f"Loading UNet2DConditionModel from {args.externel_unet_path}", ranks=[0]) + unet = UNet2DConditionModel.from_pretrained(args.externel_unet_path, + revision=args.revision, + low_cpu_mem_usage=False) vae.requires_grad_(False) text_encoder.requires_grad_(False) @@ -469,10 +483,22 @@ def main(args): if args.scale_lr: args.learning_rate = args.learning_rate * args.train_batch_size * world_size - unet = gemini_zero_dpp(unet, args.placement) + # Use Booster API to use Gemini/Zero with ColossalAI + + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy=args.placement, strict_ddp_mode=True, initial_scale=2 ** 5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2 ** 5) + + booster = Booster(plugin=plugin, **booster_kwargs) # config optimizer for colossalai zero - optimizer = GeminiAdamOptimizer(unet, lr=args.learning_rate, initial_scale=2**5, clipping_norm=args.max_grad_norm) + optimizer = HybridAdam(unet.parameters(), lr=args.learning_rate, initial_scale=2**5, clipping_norm=args.max_grad_norm) # load noise_scheduler noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") @@ -487,6 +513,7 @@ def main(args): tokenizer=tokenizer, size=args.resolution, center_crop=args.center_crop, + test=args.test_run ) def collate_fn(examples): @@ -555,6 +582,8 @@ def collate_fn(examples): # Afterwards we recalculate our number of training epochs args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + unet, optimizer, _, _, lr_scheduler = booster.boost(unet, optimizer, lr_scheduler=lr_scheduler) + # Train! total_batch_size = args.train_batch_size * world_size @@ -643,36 +672,24 @@ def collate_fn(examples): if global_step % args.save_steps == 0: torch.cuda.synchronize() - torch_unet = get_static_torch_model(unet) + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + booster.save_model(unet, os.path.join(save_path, "diffusion_pytorch_model.bin")) if local_rank == 0: - pipeline = DiffusionPipeline.from_pretrained( - args.pretrained_model_name_or_path, - unet=torch_unet, - revision=args.revision, - ) - save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") - pipeline.save_pretrained(save_path) + if not os.path.exists(os.path.join(save_path, "config.json")): + shutil.copy(os.path.join(args.pretrained_model_name_or_path, "unet/config.json"), save_path) logger.info(f"Saving model checkpoint to {save_path}", ranks=[0]) if global_step >= args.max_train_steps: break - torch.cuda.synchronize() - unet = get_static_torch_model(unet) + booster.save_model(unet, os.path.join(args.output_dir, "diffusion_pytorch_model.bin")) + logger.info(f"Saving model checkpoint to {args.output_dir} on rank {local_rank}") if local_rank == 0: - pipeline = DiffusionPipeline.from_pretrained( - args.pretrained_model_name_or_path, - unet=unet, - revision=args.revision, - ) - - pipeline.save_pretrained(args.output_dir) - logger.info(f"Saving model checkpoint to {args.output_dir}", ranks=[0]) - + if not os.path.exists(os.path.join(args.output_dir, "config.json")): + shutil.copy(os.path.join(args.pretrained_model_name_or_path, "unet/config.json"), args.output_dir) if args.push_to_hub: repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True) - if __name__ == "__main__": args = parse_args() main(args) diff --git a/examples/images/dreambooth/train_dreambooth_colossalai_lora.py b/examples/images/dreambooth/train_dreambooth_colossalai_lora.py index 3d789ae2ce0f..dce65ff514b7 100644 --- a/examples/images/dreambooth/train_dreambooth_colossalai_lora.py +++ b/examples/images/dreambooth/train_dreambooth_colossalai_lora.py @@ -4,6 +4,7 @@ import os from pathlib import Path from typing import Optional +import shutil import torch import torch.nn.functional as F @@ -23,10 +24,12 @@ from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel.utils import get_static_torch_model +from colossalai.nn.optimizer import HybridAdam from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, GeminiAdamOptimizer +from colossalai.zero.gemini import get_static_torch_model +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin disable_existing_loggers() logger = get_dist_logger() @@ -61,6 +64,13 @@ def parse_args(input_args=None): required=True, help="Path to pretrained model or model identifier from huggingface.co/models.", ) + parser.add_argument( + "--externel_unet_path", + type=str, + default=None, + required=False, + help="Path to the externel unet model.", + ) parser.add_argument( "--revision", type=str, @@ -196,6 +206,12 @@ def parse_args(input_args=None): default=None, help="The name of the repository to keep in sync with the local `output_dir`.", ) + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") parser.add_argument( "--logging_dir", type=str, @@ -342,18 +358,6 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: return f"{organization}/{model_id}" -# Gemini + ZeRO DDP -def gemini_zero_dpp(model: torch.nn.Module, placememt_policy: str = "auto"): - from colossalai.nn.parallel import GeminiDDP - - model = GeminiDDP(model, - device=get_current_device(), - placement_policy=placememt_policy, - pin_memory=True, - search_range_mb=64) - return model - - def main(args): if args.seed is None: colossalai.launch_from_torch(config={}) @@ -395,7 +399,7 @@ def main(args): images = pipeline(example["prompt"]).images for i, image in enumerate(images): - hash_image = hashlib.sha1(image.tobytes()).hexdigest() + hash_image = hashlib.sha256(image.tobytes()).hexdigest() image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" image.save(image_filename) @@ -455,32 +459,42 @@ def main(args): revision=args.revision, ) - logger.info(f"Loading UNet2DConditionModel from {args.pretrained_model_name_or_path}", ranks=[0]) - with ColoInitContext(device=get_current_device()): + + if args.externel_unet_path is None: + logger.info(f"Loading UNet2DConditionModel from {args.pretrained_model_name_or_path}", ranks=[0]) unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, - subfolder="unet", - revision=args.revision, - low_cpu_mem_usage=False) - unet.requires_grad_(False) - - # Set correct lora layers - lora_attn_procs = {} - for name in unet.attn_processors.keys(): - cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim - if name.startswith("mid_block"): - hidden_size = unet.config.block_out_channels[-1] - elif name.startswith("up_blocks"): - block_id = int(name[len("up_blocks.")]) - hidden_size = list(reversed(unet.config.block_out_channels))[block_id] - elif name.startswith("down_blocks"): - block_id = int(name[len("down_blocks.")]) - hidden_size = unet.config.block_out_channels[block_id] - - lora_attn_procs[name] = LoRACrossAttnProcessor(hidden_size=hidden_size, - cross_attention_dim=cross_attention_dim) - - unet.set_attn_processor(lora_attn_procs) - lora_layers = AttnProcsLayers(unet.attn_processors) + subfolder="unet", + revision=args.revision, + low_cpu_mem_usage=False) + else: + logger.info(f"Loading UNet2DConditionModel from {args.externel_unet_path}", ranks=[0]) + unet = UNet2DConditionModel.from_pretrained(args.externel_unet_path, + revision=args.revision, + low_cpu_mem_usage=False) + unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, + subfolder="unet", + revision=args.revision, + low_cpu_mem_usage=False) + unet.requires_grad_(False) + + # Set correct lora layers + lora_attn_procs = {} + for name in unet.attn_processors.keys(): + cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim + if name.startswith("mid_block"): + hidden_size = unet.config.block_out_channels[-1] + elif name.startswith("up_blocks"): + block_id = int(name[len("up_blocks.")]) + hidden_size = list(reversed(unet.config.block_out_channels))[block_id] + elif name.startswith("down_blocks"): + block_id = int(name[len("down_blocks.")]) + hidden_size = unet.config.block_out_channels[block_id] + + lora_attn_procs[name] = LoRACrossAttnProcessor(hidden_size=hidden_size, + cross_attention_dim=cross_attention_dim) + + unet.set_attn_processor(lora_attn_procs) + lora_layers = AttnProcsLayers(unet.attn_processors) vae.requires_grad_(False) text_encoder.requires_grad_(False) @@ -491,10 +505,22 @@ def main(args): if args.scale_lr: args.learning_rate = args.learning_rate * args.train_batch_size * world_size - unet = gemini_zero_dpp(unet, args.placement) + # Use Booster API to use Gemini/Zero with ColossalAI + + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2 ** 5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2 ** 5) + + booster = Booster(plugin=plugin, **booster_kwargs) # config optimizer for colossalai zero - optimizer = GeminiAdamOptimizer(unet, lr=args.learning_rate, initial_scale=2**5, clipping_norm=args.max_grad_norm) + optimizer = HybridAdam(unet.parameters(), lr=args.learning_rate, initial_scale=2**5, clipping_norm=args.max_grad_norm) # load noise_scheduler noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") @@ -577,6 +603,8 @@ def collate_fn(examples): # Afterwards we recalculate our number of training epochs args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + unet, optimizer, _, _, lr_scheduler = booster.boost(unet, optimizer, lr_scheduler=lr_scheduler) + # Train! total_batch_size = args.train_batch_size * world_size @@ -665,27 +693,24 @@ def collate_fn(examples): if global_step % args.save_steps == 0: torch.cuda.synchronize() - torch_unet = get_static_torch_model(unet) + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + booster.save_model(unet, os.path.join(save_path, "diffusion_pytorch_model.bin")) if local_rank == 0: - save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") - torch_unet = torch_unet.to(torch.float32) - torch_unet.save_attn_procs(save_path) + if not os.path.exists(os.path.join(save_path, "config.json")): + shutil.copy(os.path.join(args.pretrained_model_name_or_path, "unet/config.json"), save_path) logger.info(f"Saving model checkpoint to {save_path}", ranks=[0]) if global_step >= args.max_train_steps: break - torch.cuda.synchronize() - torch_unet = get_static_torch_model(unet) + booster.save_model(unet, os.path.join(args.output_dir, "diffusion_pytorch_model.bin")) + logger.info(f"Saving model checkpoint to {args.output_dir} on rank {local_rank}") if local_rank == 0: - torch_unet = torch_unet.to(torch.float32) - torch_unet.save_attn_procs(save_path) - logger.info(f"Saving model checkpoint to {args.output_dir}", ranks=[0]) - + if not os.path.exists(os.path.join(args.output_dir, "config.json")): + shutil.copy(os.path.join(args.pretrained_model_name_or_path, "unet/config.json"), args.output_dir) if args.push_to_hub: repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True) - if __name__ == "__main__": args = parse_args() main(args) diff --git a/examples/images/resnet/.gitignore b/examples/images/resnet/.gitignore new file mode 100644 index 000000000000..a79cf5236c08 --- /dev/null +++ b/examples/images/resnet/.gitignore @@ -0,0 +1,4 @@ +data +checkpoint +ckpt-fp16 +ckpt-fp32 diff --git a/examples/images/resnet/README.md b/examples/images/resnet/README.md new file mode 100644 index 000000000000..c69828637269 --- /dev/null +++ b/examples/images/resnet/README.md @@ -0,0 +1,56 @@ +# Train ResNet on CIFAR-10 from scratch + +## 🚀 Quick Start + +This example provides a training script and an evaluation script. The training script provides an example of training ResNet on CIFAR10 dataset from scratch. + +- Training Arguments + - `-p`, `--plugin`: Plugin to use. Choices: `torch_ddp`, `torch_ddp_fp16`, `low_level_zero`. Defaults to `torch_ddp`. + - `-r`, `--resume`: Resume from checkpoint file path. Defaults to `-1`, which means not resuming. + - `-c`, `--checkpoint`: The folder to save checkpoints. Defaults to `./checkpoint`. + - `-i`, `--interval`: Epoch interval to save checkpoints. Defaults to `5`. If set to `0`, no checkpoint will be saved. + - `--target_acc`: Target accuracy. Raise exception if not reached. Defaults to `None`. + +- Eval Arguments + - `-e`, `--epoch`: select the epoch to evaluate + - `-c`, `--checkpoint`: the folder where checkpoints are found + +### Install requirements + +```bash +pip install -r requirements.txt +``` + +### Train +The folders will be created automatically. +```bash +# train with torch DDP with fp32 +colossalai run --nproc_per_node 2 train.py -c ./ckpt-fp32 + +# train with torch DDP with mixed precision training +colossalai run --nproc_per_node 2 train.py -c ./ckpt-fp16 -p torch_ddp_fp16 + +# train with low level zero +colossalai run --nproc_per_node 2 train.py -c ./ckpt-low_level_zero -p low_level_zero +``` + +### Eval + +```bash +# evaluate fp32 training +python eval.py -c ./ckpt-fp32 -e 80 + +# evaluate fp16 mixed precision training +python eval.py -c ./ckpt-fp16 -e 80 + +# evaluate low level zero training +python eval.py -c ./ckpt-low_level_zero -e 80 +``` + +Expected accuracy performance will be: + +| Model | Single-GPU Baseline FP32 | Booster DDP with FP32 | Booster DDP with FP16 | Booster Low Level Zero | +| --------- | ------------------------ | --------------------- | --------------------- | ---------------------- | +| ResNet-18 | 85.85% | 84.91% | 85.46% | 84.50% | + +**Note: the baseline is adapted from the [script](https://pytorch-tutorial.readthedocs.io/en/latest/tutorial/chapter03_intermediate/3_2_2_cnn_resnet_cifar10/) to use `torchvision.models.resnet18`** diff --git a/examples/images/resnet/eval.py b/examples/images/resnet/eval.py new file mode 100644 index 000000000000..657708ec3ff2 --- /dev/null +++ b/examples/images/resnet/eval.py @@ -0,0 +1,48 @@ +import argparse + +import torch +import torch.nn as nn +import torchvision +import torchvision.transforms as transforms + +# ============================== +# Parse Arguments +# ============================== +parser = argparse.ArgumentParser() +parser.add_argument('-e', '--epoch', type=int, default=80, help="resume from the epoch's checkpoint") +parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory") +args = parser.parse_args() + +# ============================== +# Prepare Test Dataset +# ============================== +# CIFAR-10 dataset +test_dataset = torchvision.datasets.CIFAR10(root='./data/', train=False, transform=transforms.ToTensor()) + +# Data loader +test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False) + +# ============================== +# Load Model +# ============================== +model = torchvision.models.resnet18(num_classes=10).cuda() +state_dict = torch.load(f'{args.checkpoint}/model_{args.epoch}.pth') +model.load_state_dict(state_dict) + +# ============================== +# Run Evaluation +# ============================== +model.eval() + +with torch.no_grad(): + correct = 0 + total = 0 + for images, labels in test_loader: + images = images.cuda() + labels = labels.cuda() + outputs = model(images) + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + print('Accuracy of the model on the test images: {} %'.format(100 * correct / total)) diff --git a/examples/images/resnet/requirements.txt b/examples/images/resnet/requirements.txt new file mode 100644 index 000000000000..3c7da7743702 --- /dev/null +++ b/examples/images/resnet/requirements.txt @@ -0,0 +1,5 @@ +colossalai +torch +torchvision +tqdm +pytest \ No newline at end of file diff --git a/examples/images/resnet/test_ci.sh b/examples/images/resnet/test_ci.sh new file mode 100755 index 000000000000..b3fb67830dda --- /dev/null +++ b/examples/images/resnet/test_ci.sh @@ -0,0 +1,12 @@ +#!/bin/bash +set -xe + +export DATA=/data/scratch/cifar-10 + +pip install -r requirements.txt + +# TODO: skip ci test due to time limits, train.py needs to be rewritten. + +# for plugin in "torch_ddp" "torch_ddp_fp16" "low_level_zero"; do +# colossalai run --nproc_per_node 4 train.py --interval 0 --target_acc 0.84 --plugin $plugin +# done diff --git a/examples/images/resnet/train.py b/examples/images/resnet/train.py new file mode 100644 index 000000000000..fe0dabf08377 --- /dev/null +++ b/examples/images/resnet/train.py @@ -0,0 +1,204 @@ +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.nn as nn +import torchvision +import torchvision.transforms as transforms +from torch.optim import Optimizer +from torch.optim.lr_scheduler import MultiStepLR +from torch.utils.data import DataLoader +from tqdm import tqdm + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 80 +LEARNING_RATE = 1e-3 + + +def build_dataloader(batch_size: int, coordinator: DistCoordinator, plugin: DPPluginBase): + # transform + transform_train = transforms.Compose( + [transforms.Pad(4), + transforms.RandomHorizontalFlip(), + transforms.RandomCrop(32), + transforms.ToTensor()]) + transform_test = transforms.ToTensor() + + # CIFAR-10 dataset + data_path = os.environ.get('DATA', './data') + with coordinator.priority_execution(): + train_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=True, + transform=transform_train, + download=True) + test_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=False, + transform=transform_test, + download=True) + + # Data loader + train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True) + test_dataloader = plugin.prepare_dataloader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False) + return train_dataloader, test_dataloader + + +@torch.no_grad() +def evaluate(model: nn.Module, test_dataloader: DataLoader, coordinator: DistCoordinator) -> float: + model.eval() + correct = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + total = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + for images, labels in test_dataloader: + images = images.cuda() + labels = labels.cuda() + outputs = model(images) + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + dist.all_reduce(correct) + dist.all_reduce(total) + accuracy = correct.item() / total.item() + if coordinator.is_master(): + print(f'Accuracy of the model on the test images: {accuracy * 100:.2f} %') + return accuracy + + +def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: nn.Module, train_dataloader: DataLoader, + booster: Booster, coordinator: DistCoordinator): + model.train() + with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar: + for images, labels in pbar: + images = images.cuda() + labels = labels.cuda() + # Forward pass + outputs = model(images) + loss = criterion(outputs, labels) + + # Backward and optimize + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + # Print log info + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + # FIXME(ver217): gemini is not supported resnet now + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'low_level_zero'], + help="plugin to use") + parser.add_argument('-r', '--resume', type=int, default=-1, help="resume from the epoch's checkpoint") + parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory") + parser.add_argument('-i', '--interval', type=int, default=5, help="interval of saving checkpoint") + parser.add_argument('--target_acc', + type=float, + default=None, + help="target accuracy. Raise exception if not reached") + args = parser.parse_args() + + # ============================== + # Prepare Checkpoint Directory + # ============================== + if args.interval > 0: + Path(args.checkpoint).mkdir(parents=True, exist_ok=True) + + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}) + coordinator = DistCoordinator() + + # update the learning rate with linear scaling + # old_gpu_num / old_lr = new_gpu_num / new_lr + global LEARNING_RATE + LEARNING_RATE *= coordinator.world_size + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + train_dataloader, test_dataloader = build_dataloader(100, coordinator, plugin) + + # ==================================== + # Prepare model, optimizer, criterion + # ==================================== + # resent50 + model = torchvision.models.resnet18(num_classes=10) + + # Loss and optimizer + criterion = nn.CrossEntropyLoss() + optimizer = HybridAdam(model.parameters(), lr=LEARNING_RATE) + + # lr scheduler + lr_scheduler = MultiStepLR(optimizer, milestones=[20, 40, 60, 80], gamma=1 / 3) + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, criterion, _, lr_scheduler = booster.boost(model, + optimizer, + criterion=criterion, + lr_scheduler=lr_scheduler) + + # ============================== + # Resume from checkpoint + # ============================== + if args.resume >= 0: + booster.load_model(model, f'{args.checkpoint}/model_{args.resume}.pth') + booster.load_optimizer(optimizer, f'{args.checkpoint}/optimizer_{args.resume}.pth') + booster.load_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{args.resume}.pth') + + # ============================== + # Train model + # ============================== + start_epoch = args.resume if args.resume >= 0 else 0 + for epoch in range(start_epoch, NUM_EPOCHS): + train_epoch(epoch, model, optimizer, criterion, train_dataloader, booster, coordinator) + lr_scheduler.step() + + # save checkpoint + if args.interval > 0 and (epoch + 1) % args.interval == 0: + booster.save_model(model, f'{args.checkpoint}/model_{epoch + 1}.pth') + booster.save_optimizer(optimizer, f'{args.checkpoint}/optimizer_{epoch + 1}.pth') + booster.save_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{epoch + 1}.pth') + + accuracy = evaluate(model, test_dataloader, coordinator) + if args.target_acc is not None: + assert accuracy >= args.target_acc, f'Accuracy {accuracy} is lower than target accuracy {args.target_acc}' + + +if __name__ == '__main__': + main() diff --git a/examples/images/vit/README.md b/examples/images/vit/README.md index 4423d85d19e0..7c4147b76457 100644 --- a/examples/images/vit/README.md +++ b/examples/images/vit/README.md @@ -1,61 +1,28 @@ -# Vision Transformer with ColoTensor +## Overview -# Overview +Vision Transformer is a class of Transformer model tailored for computer vision tasks. It was first proposed in paper [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929) and achieved SOTA results on various tasks at that time. -In this example, we will run Vision Transformer with ColoTensor. +In our example, we are using pretrained weights of ViT loaded from HuggingFace. +We adapt the ViT training code to ColossalAI by leveraging [Boosting API](https://colossalai.org/docs/basics/booster_api) loaded with a chosen plugin, where each plugin corresponds to a specific kind of training strategy. This example supports plugins including TorchDDPPlugin, LowLevelZeroPlugin, and GeminiPlugin. -We use model **ViTForImageClassification** from Hugging Face [Link](https://huggingface.co/docs/transformers/model_doc/vit) for unit test. -You can change world size or decide whether use DDP in our code. +## Run Demo -We use model **vision_transformer** from timm [Link](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py) for training example. - -(2022/6/28) The default configuration now supports 2DP+2TP with gradient accumulation and checkpoint support. Zero is not supported at present. - -# Requirement - -Install colossalai version >= 0.1.11 - -## Unit test -To run unit test, you should install pytest, transformers with: -```shell -pip install pytest transformers +By running the following script: +```bash +bash run_demo.sh ``` +You will finetune a a [ViT-base](https://huggingface.co/google/vit-base-patch16-224) model on this [dataset](https://huggingface.co/datasets/beans), with more than 8000 images of bean leaves. This dataset is for image classification task and there are 3 labels: ['angular_leaf_spot', 'bean_rust', 'healthy']. -## Training example -To run training example with ViT-S, you should install **NVIDIA DALI** from [Link](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/installation.html) for dataloader support. -You also need to install timm and titans for model/dataloader support with: -```shell -pip install timm titans -``` +The script can be modified if you want to try another set of hyperparameters or change to another ViT model with different size. -### Data preparation -You can download the ImageNet dataset from the [ImageNet official website](https://www.image-net.org/download.php). You should get the raw images after downloading the dataset. As we use **NVIDIA DALI** to read data, we use the TFRecords dataset instead of raw Imagenet dataset. This offers better speedup to IO. If you don't have TFRecords dataset, follow [imagenet-tools](https://github.com/ver217/imagenet-tools) to build one. +The demo code refers to this [blog](https://huggingface.co/blog/fine-tune-vit). -Before you start training, you need to set the environment variable `DATA` so that the script knows where to fetch the data for DALI dataloader. -```shell -export DATA=/path/to/ILSVRC2012 -``` -# How to run +## Run Benchmark -## Unit test -In your terminal -```shell -pytest test_vit.py +You can run benchmark for ViT model by running the following script: +```bash +bash run_benchmark.sh ``` - -This will evaluate models with different **world_size** and **use_ddp**. - -## Training example -Modify the settings in run.sh according to your environment. -For example, if you set `--nproc_per_node=8` in `run.sh` and `TP_WORLD_SIZE=2` in your config file, -data parallel size will be automatically calculated as 4. -Thus, the parallel strategy is set to 4DP+2TP. - -Then in your terminal -```shell -sh run.sh -``` - -This will start ViT-S training with ImageNet. +The script will test performance (throughput & peak memory usage) for each combination of hyperparameters. You can also play with this script to configure your own set of hyperparameters for testing. \ No newline at end of file diff --git a/examples/images/vit/args.py b/examples/images/vit/args.py new file mode 100644 index 000000000000..e4a873a9eb52 --- /dev/null +++ b/examples/images/vit/args.py @@ -0,0 +1,124 @@ +from colossalai import get_default_parser + +def parse_demo_args(): + + parser = get_default_parser() + parser.add_argument( + "--model_name_or_path", + type=str, + default="google/vit-base-patch16-224", + help="Path to pretrained model or model identifier from huggingface.co/models." + ) + parser.add_argument( + "--output_path", + type=str, + default="./output_model.bin", + help="The path of your saved model after finetuning." + ) + parser.add_argument( + "--plugin", + type=str, + default="gemini", + help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'." + ) + parser.add_argument( + "--num_epoch", + type=int, + default=3, + help="Number of epochs." + ) + parser.add_argument( + "--batch_size", + type=int, + default=32, + help="Batch size (per dp group) for the training dataloader." + ) + parser.add_argument( + "--learning_rate", + type=float, + default=3e-4, + help="Initial learning rate (after the potential warmup period) to use." + ) + parser.add_argument( + "--warmup_ratio", + type=float, + default=0.3, + help="Ratio of warmup steps against total training steps." + ) + parser.add_argument( + "--weight_decay", + type=float, + default=0.1, + help="Weight decay to use." + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="A seed for reproducible training." + ) + + args = parser.parse_args() + return args + +def parse_benchmark_args(): + + parser = get_default_parser() + + parser.add_argument( + "--model_name_or_path", + type=str, + default="google/vit-base-patch16-224", + help="Path to a pretrained model or model identifier from huggingface.co/models." + ) + parser.add_argument( + "--plugin", + type=str, + default="gemini", + help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'." + ) + parser.add_argument( + "--batch_size", + type=int, + default=8, + help="Batch size (per dp group) for the training dataloader." + ) + parser.add_argument( + "--num_labels", + type=int, + default=10, + help="Number of labels for classification." + ) + parser.add_argument( + "--learning_rate", + type=float, + default=5e-5, + help="Initial learning rate (after the potential warmup period) to use." + ) + parser.add_argument( + "--weight_decay", + type=float, + default=0.0, + help="Weight decay to use." + ) + parser.add_argument( + "--max_train_steps", + type=int, + default=20, + help="Total number of training steps to perform." + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="A seed for reproducible training." + ) + parser.add_argument( + "--mem_cap", + type=int, + default=0, + help="Limit on the usage of space for each GPU (in GB)." + ) + args = parser.parse_args() + + return args \ No newline at end of file diff --git a/examples/images/vit/configs/vit_1d_tp2.py b/examples/images/vit/configs/vit_1d_tp2.py deleted file mode 100644 index fbf399f2e50d..000000000000 --- a/examples/images/vit/configs/vit_1d_tp2.py +++ /dev/null @@ -1,32 +0,0 @@ -from colossalai.amp import AMP_TYPE - -# hyperparameters -# BATCH_SIZE is as per GPU -# global batch size = BATCH_SIZE x data parallel size -BATCH_SIZE = 256 -LEARNING_RATE = 3e-3 -WEIGHT_DECAY = 0.3 -NUM_EPOCHS = 300 -WARMUP_EPOCHS = 32 - -# model config -IMG_SIZE = 224 -PATCH_SIZE = 16 -HIDDEN_SIZE = 384 -DEPTH = 12 -NUM_HEADS = 6 -MLP_RATIO = 4 -NUM_CLASSES = 1000 -CHECKPOINT = False -SEQ_LENGTH = (IMG_SIZE // PATCH_SIZE)**2 + 1 # add 1 for cls token - -USE_DDP = True -TP_WORLD_SIZE = 2 -TP_TYPE = 'row' -parallel = dict(tensor=dict(mode="1d", size=TP_WORLD_SIZE),) - -fp16 = dict(mode=AMP_TYPE.NAIVE) -clip_grad_norm = 1.0 -gradient_accumulation = 8 - -LOG_PATH = "./log" diff --git a/examples/images/vit/configs/vit_1d_tp2_ci.py b/examples/images/vit/configs/vit_1d_tp2_ci.py deleted file mode 100644 index e491e4ada45e..000000000000 --- a/examples/images/vit/configs/vit_1d_tp2_ci.py +++ /dev/null @@ -1,32 +0,0 @@ -from colossalai.amp import AMP_TYPE - -# hyperparameters -# BATCH_SIZE is as per GPU -# global batch size = BATCH_SIZE x data parallel size -BATCH_SIZE = 8 -LEARNING_RATE = 3e-3 -WEIGHT_DECAY = 0.3 -NUM_EPOCHS = 3 -WARMUP_EPOCHS = 1 - -# model config -IMG_SIZE = 224 -PATCH_SIZE = 16 -HIDDEN_SIZE = 32 -DEPTH = 2 -NUM_HEADS = 4 -MLP_RATIO = 4 -NUM_CLASSES = 10 -CHECKPOINT = False -SEQ_LENGTH = (IMG_SIZE // PATCH_SIZE)**2 + 1 # add 1 for cls token - -USE_DDP = True -TP_WORLD_SIZE = 2 -TP_TYPE = 'row' -parallel = dict(tensor=dict(mode="1d", size=TP_WORLD_SIZE),) - -fp16 = dict(mode=AMP_TYPE.NAIVE) -clip_grad_norm = 1.0 -gradient_accumulation = 2 - -LOG_PATH = "./log_ci" diff --git a/examples/images/vit/data.py b/examples/images/vit/data.py new file mode 100644 index 000000000000..00fde707b173 --- /dev/null +++ b/examples/images/vit/data.py @@ -0,0 +1,32 @@ +import torch +from torch.utils.data import Dataset +from datasets import load_dataset + +class BeansDataset(Dataset): + + def __init__(self, image_processor, split='train'): + + super().__init__() + self.image_processor = image_processor + self.ds = load_dataset('beans')[split] + self.label_names = self.ds.features['labels'].names + self.num_labels = len(self.label_names) + self.inputs = [] + for example in self.ds: + self.inputs.append(self.process_example(example)) + + def __len__(self): + return len(self.inputs) + + def __getitem__(self, idx): + return self.inputs[idx] + + def process_example(self, example): + input = self.image_processor(example['image'], return_tensors='pt') + input['labels'] = example['labels'] + return input + + +def beans_collator(batch): + return {'pixel_values': torch.cat([data['pixel_values'] for data in batch], dim=0), + 'labels': torch.tensor([data['labels'] for data in batch], dtype=torch.int64)} diff --git a/examples/images/vit/requirements.txt b/examples/images/vit/requirements.txt index 1f69794ebe70..edad87ca380f 100644 --- a/examples/images/vit/requirements.txt +++ b/examples/images/vit/requirements.txt @@ -1,8 +1,6 @@ colossalai >= 0.1.12 torch >= 1.8.1 numpy>=1.24.1 -timm>=0.6.12 -titans>=0.0.7 tqdm>=4.61.2 -transformers>=4.25.1 -nvidia-dali-cuda110>=1.8.0 --extra-index-url https://developer.download.nvidia.com/compute/redist +transformers>=4.20.0 +datasets \ No newline at end of file diff --git a/examples/images/vit/run.sh b/examples/images/vit/run.sh deleted file mode 100644 index 84fe58f11a6a..000000000000 --- a/examples/images/vit/run.sh +++ /dev/null @@ -1,15 +0,0 @@ -export DATA=/data/scratch/imagenet/tf_records -export OMP_NUM_THREADS=4 - -# resume -# CUDA_VISIBLE_DEVICES=4,5,6,7 colossalai run \ -# --nproc_per_node 4 train.py \ -# --config configs/vit_1d_tp2.py \ -# --resume_from checkpoint/epoch_10 \ -# --master_port 29598 | tee ./out 2>&1 - -# train -CUDA_VISIBLE_DEVICES=4,5,6,7 colossalai run \ ---nproc_per_node 4 train.py \ ---config configs/vit_1d_tp2.py \ ---master_port 29598 | tee ./out 2>&1 diff --git a/examples/images/vit/run_benchmark.sh b/examples/images/vit/run_benchmark.sh new file mode 100644 index 000000000000..2487bf81ee2b --- /dev/null +++ b/examples/images/vit/run_benchmark.sh @@ -0,0 +1,27 @@ +set -xe +pip install -r requirements.txt + +export BS=8 +export MEMCAP=0 +export GPUNUM=1 + +for BS in 8 32 128 +do +for PLUGIN in "torch_ddp" "torch_ddp_fp16" "low_level_zero" "gemini" +do +for GPUNUM in 1 4 +do + +MODEL_PATH="google/vit-base-patch16-224" +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + vit_benchmark.py \ + --model_name_or_path ${MODEL_PATH} \ + --mem_cap ${MEMCAP} \ + --plugin ${PLUGIN} \ + --batch_size ${BS} + +done +done +done diff --git a/examples/images/vit/run_demo.sh b/examples/images/vit/run_demo.sh new file mode 100644 index 000000000000..2d140dd6e423 --- /dev/null +++ b/examples/images/vit/run_demo.sh @@ -0,0 +1,44 @@ +set -xe +pip install -r requirements.txt + +# model name or path +MODEL="google/vit-base-patch16-224" + +# path for saving model +OUTPUT_PATH="./output_model.bin" + +# plugin(training strategy) +# can only be one of "torch_ddp"/"torch_ddp_fp16"/"low_level_zero"/"gemini" +PLUGIN="gemini" + +# number of gpus to use +GPUNUM=4 + +# batch size per gpu +BS=16 + +# learning rate +LR="2e-4" + +# number of epoch +EPOCH=3 + +# weight decay +WEIGHT_DECAY=0.05 + +# ratio of warmup steps +WARMUP_RATIO=0.3 + +# run the script for demo +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + vit_train_demo.py \ + --model_name_or_path ${MODEL} \ + --output_path ${OUTPUT_PATH} \ + --plugin ${PLUGIN} \ + --batch_size ${BS} \ + --num_epoch ${EPOCH} \ + --learning_rate ${LR} \ + --weight_decay ${WEIGHT_DECAY} \ + --warmup_ratio ${WARMUP_RATIO} diff --git a/examples/images/vit/test_ci.sh b/examples/images/vit/test_ci.sh index 41d25ee23521..8606015c0397 100644 --- a/examples/images/vit/test_ci.sh +++ b/examples/images/vit/test_ci.sh @@ -1,9 +1,19 @@ -export OMP_NUM_THREADS=4 - +set -xe pip install -r requirements.txt -# train -colossalai run \ ---nproc_per_node 4 train.py \ ---config configs/vit_1d_tp2_ci.py \ ---dummy_data +BS=8 +for PLUGIN in "torch_ddp" "torch_ddp_fp16" "low_level_zero" "gemini" +do +for GPUNUM in 1 4 +do + +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + vit_benchmark.py \ + --model_name_or_path "google/vit-base-patch16-224" \ + --plugin ${PLUGIN} \ + --batch_size ${BS} + +done +done diff --git a/examples/images/vit/test_vit.py b/examples/images/vit/test_vit.py deleted file mode 100644 index 90f2475b885e..000000000000 --- a/examples/images/vit/test_vit.py +++ /dev/null @@ -1,164 +0,0 @@ -import os -import random -from functools import partial - -import numpy as np -import pytest -import torch -import torch.multiprocessing as mp -from torch.nn.parallel import DistributedDataParallel as DDP -from vit import get_training_components - -import colossalai -from colossalai.context import ParallelMode -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.nn.parallel.data_parallel import ColoDDP -from colossalai.tensor import ComputePattern, ComputeSpec, DistSpecManager, ProcessGroup, ShardSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext - - -def set_seed(seed): - random.seed(seed) - os.environ['PYTHONHASHSEED'] = str(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.backends.cudnn.deterministic = True - - -def tensor_equal(A, B): - return torch.allclose(A, B, rtol=1e-3, atol=1e-1) - - -def tensor_shard_equal(tensor: torch.Tensor, shard: torch.Tensor): - assert tensor.ndim == shard.ndim - if tensor.shape == shard.shape: - return tensor_equal(tensor, shard) - else: - dims_not_eq = torch.nonzero(torch.tensor(tensor.shape) != torch.tensor(shard.shape)) - if dims_not_eq.numel() == 1: - # 1D shard - dim = dims_not_eq.item() - world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D) - rank = gpc.get_local_rank(ParallelMode.PARALLEL_1D) - return tensor_equal(tensor.chunk(world_size, dim)[rank], shard) - else: - raise - - -# Only for all Linear, it's 1d_row split because Linear will be transposed when calculating. -# But for other layers, it's 1d_col split. -# Layernorm is not supported for now. -# patch_embeddings.projection has nn.Conv2d -# https://github.com/huggingface/transformers/blob/dcb08b99f44919425f8ba9be9ddcc041af8ec25e/src/transformers/models/vit/modeling_vit.py#L182 -def init_1d_row_for_linear_weight_spec(model, world_size: int): - pg = ProcessGroup(tp_degree=world_size) - spec = (ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) - with DistSpecManager.no_grad(): - for n, p in model.named_parameters(): - if 'weight' in n and 'layernorm' not in n and 'embeddings.patch_embeddings.projection.weight' not in n: - p.set_process_group(pg) - p.set_tensor_spec(*spec) - - -# Similarly, it's col split for Linear but row split for others. -def init_1d_col_for_linear_weight_bias_spec(model, world_size: int): - pg = ProcessGroup(tp_degree=world_size) - spec = (ShardSpec([0], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) - with DistSpecManager.no_grad(): - for n, p in model.named_parameters(): - if ('weight' in n - or 'bias' in n) and 'layernorm' not in n and 'embeddings.patch_embeddings.projection' not in n: - p.set_process_group(pg) - p.set_tensor_spec(*spec) - - -def check_param_equal(model, torch_model): - for p, torch_p in zip(model.parameters(), torch_model.parameters()): - assert tensor_shard_equal(torch_p, p) - - -def check_grad_equal(model, torch_model): - for p, torch_p in zip(model.parameters(), torch_model.parameters()): - if (torch_p.grad.shape == p.grad.shape): - assert torch.allclose(torch_p.grad, p.grad, rtol=1e-3, atol=2.0) == True - else: - dims_not_eq = torch.nonzero(torch.tensor(torch_p.grad.shape) != torch.tensor(p.grad.shape)) - dim = dims_not_eq.item() - world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D) - rank = gpc.get_local_rank(ParallelMode.PARALLEL_1D) - assert torch.allclose(torch_p.grad.chunk(world_size, dim)[rank], p.grad, rtol=1e-3, atol=2.0) == True - - -def run_vit(init_spec_func, use_ddp): - model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_training_components() - with ColoInitContext(device=get_current_device()): - model = model_builder() - model = model.cuda() - torch_model = model_builder().cuda() - if use_ddp: - model = ColoDDP(model) - torch_model = DDP(torch_model, - device_ids=[gpc.get_global_rank()], - process_group=gpc.get_group(ParallelMode.DATA)) - for torch_p, p in zip(torch_model.parameters(), model.parameters()): - torch_p.data.copy_(p) - - world_size = torch.distributed.get_world_size() - init_spec_func(model, world_size) - - check_param_equal(model, torch_model) - model.train() - torch_model.train() - set_seed(gpc.get_local_rank(ParallelMode.DATA)) - - optimizer = optimizer_class(model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) - torch_optimizer = optimizer_class(torch_model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) - - for i, image_dict in enumerate(train_dataloader): - if use_ddp: - model.zero_grad() - else: - optimizer.zero_grad() - logits = model(image_dict['pixel_values']) - torch_logits = torch_model(image_dict['pixel_values']) - assert tensor_equal(torch_logits.logits, logits.logits) - loss = criterion(logits.logits, image_dict['label']) - torch_loss = criterion(torch_logits.logits, image_dict['label']) - if use_ddp: - model.backward(loss) - else: - loss.backward() - torch_loss.backward() - check_grad_equal(model, torch_model) - optimizer.step() - torch_optimizer.step() - check_param_equal(model, torch_model) - break - - -def run_dist(rank, world_size, port, use_ddp): - if use_ddp and world_size == 1: - return - tp_world_size = world_size // 2 if use_ddp else world_size - config = dict(parallel=dict(tensor=dict(mode="1d", size=tp_world_size),)) - colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - run_vit(init_1d_row_for_linear_weight_spec, use_ddp) - run_vit(init_1d_col_for_linear_weight_bias_spec, use_ddp) - - -@pytest.mark.dist -@pytest.mark.parametrize('world_size', [1, 4]) -@pytest.mark.parametrize('use_ddp', [False, True]) -@rerun_if_address_is_in_use() -def test_vit(world_size, use_ddp): - run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_vit(1, False) diff --git a/examples/images/vit/train.py b/examples/images/vit/train.py deleted file mode 100644 index 0b4489244368..000000000000 --- a/examples/images/vit/train.py +++ /dev/null @@ -1,174 +0,0 @@ -import os - -import torch -import torch.distributed as dist -import torch.nn as nn -import torch.nn.functional as F -from timm.models.vision_transformer import _create_vision_transformer -from titans.dataloader.imagenet import build_dali_imagenet -from tqdm import tqdm -from vit import DummyDataLoader - -import colossalai -from colossalai.core import global_context as gpc -from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn import CrossEntropyLoss -from colossalai.nn._ops import * -from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR -from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel.data_parallel import ColoDDP -from colossalai.tensor import ComputePattern, ComputeSpec, DistSpecManager, ProcessGroup, ShardSpec -from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext - - -def init_1d_row_for_linear_weight_spec(model, world_size: int): - pg = ProcessGroup(tp_degree=world_size) - spec = (ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) - with DistSpecManager.no_grad(): - for n, p in model.named_parameters(): - if 'weight' in n and 'norm' not in n and 'patch_embed.proj.weight' not in n: - p.set_process_group(pg) - p.set_tensor_spec(*spec) - - -# Similarly, it's col split for Linear but row split for others. -def init_1d_col_for_linear_weight_bias_spec(model, world_size: int): - pg = ProcessGroup(tp_degree=world_size) - spec = (ShardSpec([0], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) - with DistSpecManager.no_grad(): - for n, p in model.named_parameters(): - if ('weight' in n or 'bias' in n) and 'norm' not in n and ('patch_embed.proj.weight' not in n - and 'patch_embed.proj.bias' not in n): - p.set_process_group(pg) - p.set_tensor_spec(*spec) - - -def init_spec_func(model, tp_type): - world_size = torch.distributed.get_world_size() - if tp_type == 'row': - init_1d_row_for_linear_weight_spec(model, world_size) - elif tp_type == 'col': - init_1d_col_for_linear_weight_bias_spec(model, world_size) - else: - raise NotImplemented - - -def train_imagenet(): - - parser = colossalai.get_default_parser() - parser.add_argument('--resume_from', default=False, action='store_true') - parser.add_argument('--dummy_data', default=False, action='store_true') - - args = parser.parse_args() - colossalai.launch_from_torch(config=args.config) - use_ddp = gpc.config.USE_DDP - - disable_existing_loggers() - - logger = get_dist_logger() - if hasattr(gpc.config, 'LOG_PATH'): - if gpc.get_global_rank() == 0: - log_path = gpc.config.LOG_PATH - if not os.path.exists(log_path): - os.mkdir(log_path) - logger.log_to_file(log_path) - - logger.info('Build data loader', ranks=[0]) - if not args.dummy_data: - root = os.environ['DATA'] - train_dataloader, test_dataloader = build_dali_imagenet(root, - train_batch_size=gpc.config.BATCH_SIZE, - test_batch_size=gpc.config.BATCH_SIZE) - else: - train_dataloader = DummyDataLoader(length=10, - batch_size=gpc.config.BATCH_SIZE, - category=gpc.config.NUM_CLASSES, - image_size=gpc.config.IMG_SIZE, - return_dict=False) - test_dataloader = DummyDataLoader(length=5, - batch_size=gpc.config.BATCH_SIZE, - category=gpc.config.NUM_CLASSES, - image_size=gpc.config.IMG_SIZE, - return_dict=False) - - logger.info('Build model', ranks=[0]) - - model_kwargs = dict(img_size=gpc.config.IMG_SIZE, - patch_size=gpc.config.PATCH_SIZE, - embed_dim=gpc.config.HIDDEN_SIZE, - depth=gpc.config.DEPTH, - num_heads=gpc.config.NUM_HEADS, - mlp_ratio=gpc.config.MLP_RATIO, - num_classes=gpc.config.NUM_CLASSES, - drop_rate=0.1, - attn_drop_rate=0.1, - weight_init='jax') - - with ColoInitContext(device=get_current_device()): - model = _create_vision_transformer('vit_small_patch16_224', pretrained=False, **model_kwargs) - init_spec_func(model, gpc.config.TP_TYPE) - - world_size = torch.distributed.get_world_size() - model = ColoDDP(module=model, process_group=ProcessGroup(tp_degree=world_size)) - logger.info('Build criterion, optimizer, lr_scheduler', ranks=[0]) - optimizer = HybridAdam(model.parameters(), lr=gpc.config.LEARNING_RATE, weight_decay=gpc.config.WEIGHT_DECAY) - - criterion = CrossEntropyLoss() - lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, - total_steps=gpc.config.NUM_EPOCHS, - warmup_steps=gpc.config.WARMUP_EPOCHS) - - start_epoch = 0 - if args.resume_from: - load_model = torch.load(args.resume_from + '_model.pth') - start_epoch = load_model['epoch'] - model.load_state_dict(load_model['model']) - load_optim = torch.load(args.resume_from + '_optim_rank_{}.pth'.format(dist.get_rank())) - optimizer.load_state_dict(load_optim['optim']) - - for epoch in range(start_epoch, gpc.config.NUM_EPOCHS): - model.train() - for index, (x, y) in tqdm(enumerate(train_dataloader), total=len(train_dataloader), leave=False): - x, y = x.cuda(), y.cuda() - output = model(x) - loss = criterion(output, y) - loss = loss / gpc.config.gradient_accumulation - if use_ddp: - model.backward(loss) - else: - loss.backward() - if (index + 1) % gpc.config.gradient_accumulation == 0: - optimizer.step() - if use_ddp: - model.zero_grad() - else: - optimizer.zero_grad() - - logger.info( - f"Finish Train Epoch [{epoch+1}/{gpc.config.NUM_EPOCHS}] loss: {loss.item():.3f} lr: {optimizer.state_dict()['param_groups'][0]['lr']}", - ranks=[0]) - - model.eval() - test_loss = 0 - correct = 0 - test_sum = 0 - with torch.no_grad(): - for index, (x, y) in tqdm(enumerate(test_dataloader), total=len(test_dataloader), leave=False): - x, y = x.cuda(), y.cuda() - output = model(x) - test_loss += F.cross_entropy(output, y, reduction='sum').item() - pred = output.argmax(dim=1, keepdim=True) - correct += pred.eq(y.view_as(pred)).sum().item() - test_sum += y.size(0) - - test_loss /= test_sum - logger.info( - f"Finish Test Epoch [{epoch+1}/{gpc.config.NUM_EPOCHS}] loss: {test_loss:.3f} Accuracy: [{correct}/{test_sum}]({correct/test_sum:.3f})", - ranks=[0]) - - lr_scheduler.step() - - -if __name__ == '__main__': - train_imagenet() diff --git a/examples/images/vit/vit.py b/examples/images/vit/vit.py deleted file mode 100644 index f22e8ea90cec..000000000000 --- a/examples/images/vit/vit.py +++ /dev/null @@ -1,95 +0,0 @@ -from abc import ABC, abstractmethod - -import torch -import torch.nn as nn -from transformers import ViTConfig, ViTForImageClassification - -from colossalai.utils.cuda import get_current_device - - -class DummyDataGenerator(ABC): - - def __init__(self, length=10): - self.length = length - - @abstractmethod - def generate(self): - pass - - def __iter__(self): - self.step = 0 - return self - - def __next__(self): - if self.step < self.length: - self.step += 1 - return self.generate() - else: - raise StopIteration - - def __len__(self): - return self.length - - -class DummyDataLoader(DummyDataGenerator): - - def __init__(self, length=10, batch_size=4, channel=3, category=8, image_size=224, return_dict=True): - super().__init__(length) - self.batch_size = batch_size - self.channel = channel - self.category = category - self.image_size = image_size - self.return_dict = return_dict - - def generate(self): - image_dict = {} - image_dict['pixel_values'] = torch.rand( - self.batch_size, self.channel, self.image_size, self.image_size, device=get_current_device()) * 2 - 1 - image_dict['label'] = torch.randint(self.category, (self.batch_size,), - dtype=torch.int64, - device=get_current_device()) - if not self.return_dict: - return image_dict['pixel_values'], image_dict['label'] - return image_dict - - -class ViTCVModel(nn.Module): - - def __init__(self, - hidden_size=768, - num_hidden_layers=12, - num_attention_heads=12, - image_size=224, - patch_size=16, - num_channels=3, - num_labels=8, - checkpoint=False): - super().__init__() - self.checkpoint = checkpoint - self.model = ViTForImageClassification( - ViTConfig(hidden_size=hidden_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - image_size=image_size, - patch_size=patch_size, - num_channels=num_channels, - num_labels=num_labels)) - if checkpoint: - self.model.gradient_checkpointing_enable() - - def forward(self, pixel_values): - return self.model(pixel_values=pixel_values) - - -def vit_base_s(checkpoint=True): - return ViTCVModel(checkpoint=checkpoint) - - -def vit_base_micro(checkpoint=True): - return ViTCVModel(hidden_size=32, num_hidden_layers=2, num_attention_heads=4, checkpoint=checkpoint) - - -def get_training_components(): - trainloader = DummyDataLoader() - testloader = DummyDataLoader() - return vit_base_micro, trainloader, testloader, torch.optim.Adam, torch.nn.functional.cross_entropy diff --git a/examples/images/vit/vit_benchmark.py b/examples/images/vit/vit_benchmark.py new file mode 100644 index 000000000000..11d480bba65f --- /dev/null +++ b/examples/images/vit/vit_benchmark.py @@ -0,0 +1,129 @@ +import time + +import torch +import transformers +from transformers import ViTConfig, ViTForImageClassification +import tqdm + +import colossalai +from colossalai.nn.optimizer import HybridAdam +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.utils import get_current_device +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator + +from args import parse_benchmark_args + +def format_num(num: int, bytes=False): + """Scale bytes to its proper format, e.g. 1253656 => '1.20MB'""" + factor = 1024 if bytes else 1000 + suffix = "B" if bytes else "" + for unit in ["", " K", " M", " G", " T", " P"]: + if num < factor: + return f"{num:.2f}{unit}{suffix}" + num /= factor + + +def get_data(batch_size, num_labels, num_channels=3, height=224, width=224): + pixel_values = torch.randn(batch_size, num_channels, height, width, device=torch.cuda.current_device(), dtype=torch.float) + labels = torch.randint(0, num_labels, (batch_size, ), device=torch.cuda.current_device(), dtype=torch.int64) + return pixel_values, labels + + +def colo_memory_cap(size_in_GB): + from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device + cuda_capacity = colo_device_memory_capacity(get_current_device()) + if size_in_GB * (1024**3) < cuda_capacity: + colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity) + print(f"Limiting GPU memory usage to {size_in_GB} GB") + + +def main(): + + args = parse_benchmark_args() + + # Launch ColossalAI + colossalai.launch_from_torch(config={}, seed=args.seed) + coordinator = DistCoordinator() + world_size = coordinator.world_size + + # Manage loggers + disable_existing_loggers() + logger = get_dist_logger() + if coordinator.is_master(): + transformers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + + # Whether to set limit on memory capacity + if args.mem_cap > 0: + colo_memory_cap(args.mem_cap) + + # Build ViT model + config = ViTConfig.from_pretrained(args.model_name_or_path) + model = ViTForImageClassification(config) + logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0]) + + # Enable gradient checkpointing + model.gradient_checkpointing_enable() + + # Set plugin + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(device=get_current_device(), + placement_policy='cpu', + pin_memory=True, + strict_ddp_mode=True, + initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + logger.info(f"Set plugin as {args.plugin}", ranks=[0]) + + # Set optimizer + optimizer = HybridAdam(model.parameters(), lr=(args.learning_rate * world_size)) + + # Set booster + booster = Booster(plugin=plugin, **booster_kwargs) + model, optimizer, _, _, _ = booster.boost(model, optimizer) + + + # Start training. + logger.info(f"Start testing", ranks=[0]) + progress_bar = tqdm.tqdm(total=args.max_train_steps, desc="Training Step", disable=not coordinator.is_master()) + + torch.cuda.synchronize() + model.train() + start_time = time.time() + + for _ in range(args.max_train_steps): + + pixel_values, labels = get_data(args.batch_size, args.num_labels, 3, 224, 224) + optimizer.zero_grad() + outputs = model(pixel_values=pixel_values, labels=labels) + loss = outputs['loss'] + booster.backward(loss, optimizer) + optimizer.step() + + torch.cuda.synchronize() + progress_bar.update(1) + + # Compute Statistics + end_time = time.time() + throughput = "{:.4f}".format((world_size * args.max_train_steps * args.batch_size) / (end_time - start_time)) + max_mem = format_num(torch.cuda.max_memory_allocated(device=torch.cuda.current_device()), bytes=True) + + logger.info(f"Testing finished, " + f"batch size per gpu: {args.batch_size}, " + f"plugin: {args.plugin}, " + f"throughput: {throughput}, " + f"maximum memory usage per gpu: {max_mem}.", + ranks=[0]) + + +if __name__ == "__main__": + main() diff --git a/examples/images/vit/vit_train_demo.py b/examples/images/vit/vit_train_demo.py new file mode 100644 index 000000000000..3a739f10b5d0 --- /dev/null +++ b/examples/images/vit/vit_train_demo.py @@ -0,0 +1,177 @@ +import torch +import torch.distributed as dist +import transformers +from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor +from tqdm import tqdm + +import colossalai +from colossalai.nn.optimizer import HybridAdam +from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.utils import get_current_device +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator + +from args import parse_demo_args +from data import BeansDataset, beans_collator + + +def move_to_cuda(batch, device): + return {k: v.to(device) for k, v in batch.items()} + + +def train_epoch(epoch, model, optimizer, lr_scheduler, dataloader, booster, coordinator): + + torch.cuda.synchronize() + model.train() + + with tqdm(dataloader, desc=f'Epoch [{epoch + 1}]', disable=not coordinator.is_master()) as pbar: + + for batch in pbar: + + # Foward + optimizer.zero_grad() + batch = move_to_cuda(batch, torch.cuda.current_device()) + outputs = model(**batch) + loss = outputs['loss'] + + # Backward + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() + + # Print batch loss + pbar.set_postfix({'loss': loss.item()}) + + +@torch.no_grad() +def evaluate_model(epoch, model, eval_dataloader, num_labels, coordinator): + + model.eval() + accum_loss = torch.zeros(1, device=get_current_device()) + total_num = torch.zeros(1, device=get_current_device()) + accum_correct = torch.zeros(1, device=get_current_device()) + + for batch in eval_dataloader: + batch = move_to_cuda(batch, torch.cuda.current_device()) + outputs = model(**batch) + val_loss, logits = outputs[:2] + accum_loss += (val_loss / len(eval_dataloader)) + if num_labels > 1: + preds = torch.argmax(logits, dim=1) + elif num_labels == 1: + preds = logits.squeeze() + + labels = batch["labels"] + total_num += batch["labels"].shape[0] + accum_correct += (torch.sum(preds == labels)) + + dist.all_reduce(accum_loss) + dist.all_reduce(total_num) + dist.all_reduce(accum_correct) + avg_loss = "{:.4f}".format(accum_loss.item()) + accuracy = "{:.4f}".format(accum_correct.item() / total_num.item()) + if coordinator.is_master(): + print(f"Evaluation result for epoch {epoch + 1}: \ + average_loss={avg_loss}, \ + accuracy={accuracy}.") + + + + +def main(): + + args = parse_demo_args() + + # Launch ColossalAI + colossalai.launch_from_torch(config={}, seed=args.seed) + coordinator = DistCoordinator() + world_size = coordinator.world_size + + # Manage loggers + disable_existing_loggers() + logger = get_dist_logger() + if coordinator.is_master(): + transformers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + + # Prepare Dataset + image_processor = ViTImageProcessor.from_pretrained(args.model_name_or_path) + train_dataset = BeansDataset(image_processor, split='train') + eval_dataset = BeansDataset(image_processor, split='validation') + + + # Load pretrained ViT model + config = ViTConfig.from_pretrained(args.model_name_or_path) + config.num_labels = train_dataset.num_labels + config.id2label = {str(i): c for i, c in enumerate(train_dataset.label_names)} + config.label2id = {c: str(i) for i, c in enumerate(train_dataset.label_names)} + model = ViTForImageClassification.from_pretrained(args.model_name_or_path, + config=config, + ignore_mismatched_sizes=True) + logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0]) + + # Enable gradient checkpointing + model.gradient_checkpointing_enable() + + # Set plugin + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(device=get_current_device(), + placement_policy='cpu', + pin_memory=True, + strict_ddp_mode=True, + initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + logger.info(f"Set plugin as {args.plugin}", ranks=[0]) + + # Prepare dataloader + train_dataloader = plugin.prepare_dataloader(train_dataset, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + collate_fn=beans_collator) + eval_dataloader = plugin.prepare_dataloader(eval_dataset, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + collate_fn=beans_collator) + + # Set optimizer + optimizer = HybridAdam(model.parameters(), lr=(args.learning_rate * world_size), weight_decay=args.weight_decay) + + # Set lr scheduler + total_steps = len(train_dataloader) * args.num_epoch + num_warmup_steps = int(args.warmup_ratio * total_steps) + lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, + total_steps=(len(train_dataloader) * args.num_epoch), + warmup_steps=num_warmup_steps) + + # Set booster + booster = Booster(plugin=plugin, **booster_kwargs) + model, optimizer, _, train_dataloader, lr_scheduler = booster.boost(model=model, + optimizer=optimizer, + dataloader=train_dataloader, + lr_scheduler=lr_scheduler) + + # Finetuning + logger.info(f"Start finetuning", ranks=[0]) + for epoch in range(args.num_epoch): + train_epoch(epoch, model, optimizer, lr_scheduler, train_dataloader, booster, coordinator) + evaluate_model(epoch, model, eval_dataloader, eval_dataset.num_labels, coordinator) + logger.info(f"Finish finetuning", ranks=[0]) + + # Save the finetuned model + booster.save_model(model, args.output_path) + logger.info(f"Saving model checkpoint to {args.output_path}", ranks=[0]) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/language/bert/README.md b/examples/language/bert/README.md new file mode 100644 index 000000000000..81c3f03fffca --- /dev/null +++ b/examples/language/bert/README.md @@ -0,0 +1,34 @@ +## Overview + +This directory includes two parts: Using the Booster API finetune Huggingface Bert and AlBert models and benchmarking Bert and AlBert models with different Booster Plugin. + +## Finetune +``` +bash test_ci.sh +``` + +## Benchmark +``` +bash benchmark.sh +``` + +Now include these metrics in benchmark: CUDA mem occupy, throughput and the number of model parameters. If you have custom metrics, you can add them to benchmark_util. + +## Results + +### Bert + +| | max cuda mem | throughput(sample/s) | params | +| :-----| -----------: | :--------: | :----: | +| ddp | 21.44 GB | 3.0 | 82M | +| ddp_fp16 | 16.26 GB | 11.3 | 82M | +| gemini | 11.0 GB | 12.9 | 82M | +| low_level_zero | 11.29 G | 14.7 | 82M | + +### AlBert +| | max cuda mem | throughput(sample/s) | params | +| :-----| -----------: | :--------: | :----: | +| ddp | OOM | | | +| ddp_fp16 | OOM | | | +| gemini | 69.39 G | 1.3 | 208M | +| low_level_zero | 56.89 G | 1.4 | 208M | \ No newline at end of file diff --git a/examples/language/bert/benchmark.py b/examples/language/bert/benchmark.py new file mode 100644 index 000000000000..ae8b2269a534 --- /dev/null +++ b/examples/language/bert/benchmark.py @@ -0,0 +1,174 @@ +import argparse + +import torch +from benchmark_utils import benchmark +from torch.utils.data import DataLoader, Dataset +from transformers import ( + AlbertConfig, + AlbertForSequenceClassification, + BertConfig, + BertForSequenceClassification, + get_linear_schedule_with_warmup, +) + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 3 +BATCH_SIZE = 32 +LEARNING_RATE = 2.4e-5 +WEIGHT_DECAY = 0.01 +WARMUP_FRACTION = 0.1 +SEQ_LEN = 512 +VOCAB_SIZE = 1000 +NUM_LABELS = 10 +DATASET_LEN = 1000 + + +class RandintDataset(Dataset): + + def __init__(self, dataset_length: int, sequence_length: int, vocab_size: int, n_class: int): + + self._sequence_length = sequence_length + self._vocab_size = vocab_size + self._n_class = n_class + self._dataset_length = dataset_length + self._datas = torch.randint( + low=0, + high=self._vocab_size, + size=(self._dataset_length, self._sequence_length,), + dtype=torch.long, + ) + self._labels = torch.randint(low=0, high=self._n_class, size=(self._dataset_length, 1), dtype=torch.long) + + def __len__(self): + return self._dataset_length + + def __getitem__(self, idx): + return self._datas[idx], self._labels[idx] + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + parser.add_argument('-t', '--task', default='mrpc', help="GLUE task to run") + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") + parser.add_argument( + "--model_type", + type=str, + default="bert", + help="bert or albert", + ) + + args = parser.parse_args() + + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}, seed=42) + coordinator = DistCoordinator() + + # local_batch_size = BATCH_SIZE // coordinator.world_size + lr = LEARNING_RATE * coordinator.world_size + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + + train_dataset = RandintDataset(dataset_length=DATASET_LEN, + sequence_length=SEQ_LEN, + vocab_size=VOCAB_SIZE, + n_class=NUM_LABELS) + train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE) + + # ==================================== + # Prepare model, optimizer + # ==================================== + # bert pretrained model + + if args.model_type == "bert": + cfg = BertConfig(vocab_size=VOCAB_SIZE, num_labels=NUM_LABELS) + model = BertForSequenceClassification(cfg) + elif args.model_type == "albert": + cfg = AlbertConfig(vocab_size=VOCAB_SIZE, num_labels=NUM_LABELS) + model = AlbertForSequenceClassification(cfg) + else: + raise RuntimeError + + # optimizer + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], + "weight_decay": WEIGHT_DECAY, + }, + { + "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], + "weight_decay": 0.0, + }, + ] + + optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, eps=1e-8) + + # lr scheduler + total_steps = len(train_dataloader) * NUM_EPOCHS + num_warmup_steps = int(WARMUP_FRACTION * total_steps) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_steps, + ) + + # criterion + criterion = lambda inputs: inputs[0] + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, _, _, lr_scheduler = booster.boost(model, optimizer, lr_scheduler=lr_scheduler) + + # ============================== + # Benchmark model + # ============================== + + results = benchmark(model, + booster, + optimizer, + lr_scheduler, + train_dataloader, + criterion=criterion, + epoch_num=NUM_EPOCHS) + + coordinator.print_on_master(results) + + +if __name__ == '__main__': + main() diff --git a/examples/language/bert/benchmark.sh b/examples/language/bert/benchmark.sh new file mode 100755 index 000000000000..9453d1373f2f --- /dev/null +++ b/examples/language/bert/benchmark.sh @@ -0,0 +1,9 @@ +#!/bin/bash +set -xe + +pip install -r requirements.txt + +for plugin in "torch_ddp" "torch_ddp_fp16" "gemini" "low_level_zero"; do + torchrun --standalone --nproc_per_node 2 benchmark.py --plugin $plugin --model_type "bert" + torchrun --standalone --nproc_per_node 2 benchmark.py --plugin $plugin --model_type "albert" +done diff --git a/examples/language/bert/benchmark_utils.py b/examples/language/bert/benchmark_utils.py new file mode 100644 index 000000000000..886017a41826 --- /dev/null +++ b/examples/language/bert/benchmark_utils.py @@ -0,0 +1,146 @@ +import inspect +from logging import getLogger +from time import time +from typing import Callable + +import torch +import yaml +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader +from tqdm import tqdm + +from colossalai.booster import Booster +from colossalai.cluster import DistCoordinator + +logger = getLogger("colossalai-booster-benchmark") +_INVALID = float("nan") + + +def format_num(num: int, bytes=False): + """Scale bytes to its proper format, e.g. 1253656 => '1.20MB'""" + factor = 1024 if bytes else 1000 + suffix = "B" if bytes else "" + for unit in ["", " K", " M", " G", " T", " P"]: + if num < factor: + return f"{num:.2f}{unit}{suffix}" + num /= factor + + +def _is_valid(val): + return val == val + + +def get_call_arg_names(module_or_fn): + if isinstance(module_or_fn, torch.nn.Module): + return inspect.getfullargspec(module_or_fn.forward)[0][1:] + return inspect.getfullargspec(module_or_fn)[0] + + +def measure_params(model): + num_params = _INVALID + + try: + num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + except AttributeError as e: + logger.error(f"Unable to measure model params due to error: {e}") + + return num_params + + +def warm_up( + model, + booster, + dataloader, + criterion, + optimizer, + lr_scheduler, + num_runs=10, +): + for i, data in enumerate(dataloader): + if i > num_runs: + break + inputs, labels = data[0].cuda(), data[1].cuda() + outputs = model(inputs, labels=labels) + loss = criterion(outputs) + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + +def fmt(d: dict): + return yaml.dump(d) + + +def benchmark( + model: torch.nn.Module, + booster: Booster, + optimizer: torch.optim.Optimizer, + lr_scheduler: LRScheduler, + dataloader: DataLoader, + criterion: Callable = None, + warm_up_fn=warm_up, + epoch_num: int = 3, + batch_size: int = 32, + warm_up_steps: int = 3, +): + results = {} + model_device = torch.cuda.current_device() + + # Warm up + warm_up_fn( + model, + booster, + dataloader, + criterion, + optimizer, + lr_scheduler, + num_runs=warm_up_steps, + ) + # Measure params + params = measure_params(model) + if _is_valid(params): + results["params"] = format_num(params) + logger.info(f"Model parameters: {params} ({format_num(params)})") + + # Measure Allocated Memory and Throughput + memory = {} + throughput = {} + torch.cuda.reset_peak_memory_stats(device=model_device) + pre_mem = torch.cuda.memory_allocated(device=model_device) + + start_time = time() + + for epoch in range(epoch_num): + with tqdm(dataloader, desc=f'Epoch [{epoch + 1}/{epoch_num}]', + disable=not DistCoordinator().is_master()) as pbar: + for data in pbar: + inputs, labels = data[0].cuda(), data[1].cuda() + outputs = model(inputs, labels=labels) + loss = criterion(outputs) + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + end_time = time() + + all_sample = epoch_num * len(dataloader) + + post_mem = torch.cuda.memory_allocated(device=model_device) + max_mem = torch.cuda.max_memory_allocated(device=model_device) + + memory[f"batch_size_{batch_size}"] = { + "cuda_pre_training_bytes": format_num(pre_mem, bytes=True), + "cuda_max_training_bytes": format_num(max_mem, bytes=True), + "cuda_post_training_bytes": format_num(post_mem, bytes=True), + } + logger.info(fmt({f"Memory results (batch_size={batch_size})": memory[f"batch_size_{batch_size}"]})) + + throughput[f"batch_size_{batch_size}"] = {"throughput:": "{:.1f}".format(all_sample * DistCoordinator().world_size / (end_time - start_time))} + logger.info(fmt({f"Throughput results (batch_size={batch_size})": throughput[f"batch_size_{batch_size}"]})) + + results["throughput"] = throughput + results["memory"] = memory + + return results diff --git a/examples/language/bert/data.py b/examples/language/bert/data.py new file mode 100644 index 000000000000..981cedcca8c2 --- /dev/null +++ b/examples/language/bert/data.py @@ -0,0 +1,127 @@ +import datasets +from transformers import AutoTokenizer, PreTrainedTokenizer + +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase + + +class GLUEDataBuilder: + + task_text_field_map = { + "cola": ["sentence"], + "sst2": ["sentence"], + "mrpc": ["sentence1", "sentence2"], + "qqp": ["question1", "question2"], + "stsb": ["sentence1", "sentence2"], + "mnli": ["premise", "hypothesis"], + "qnli": ["question", "sentence"], + "rte": ["sentence1", "sentence2"], + "wnli": ["sentence1", "sentence2"], + "ax": ["premise", "hypothesis"], + } + + glue_task_num_labels = { + "cola": 2, + "sst2": 2, + "mrpc": 2, + "qqp": 2, + "stsb": 1, + "mnli": 3, + "qnli": 2, + "rte": 2, + "wnli": 2, + "ax": 3, + } + + loader_columns = [ + "datasets_idx", + "input_ids", + "token_type_ids", + "attention_mask", + "start_positions", + "end_positions", + "labels", + ] + + def __init__( + self, + model_name_or_path: str, + plugin: DPPluginBase, + task_name: str = "mrpc", + max_seq_length: int = 128, + train_batch_size: int = 32, + eval_batch_size: int = 32, + **kwargs, + ): + super().__init__() + self.model_name_or_path = model_name_or_path + self.task_name = task_name + self.max_seq_length = max_seq_length + self.train_batch_size = train_batch_size + self.eval_batch_size = eval_batch_size + self.plugin = plugin + + self.text_fields = self.task_text_field_map[task_name] + self.num_labels = self.glue_task_num_labels[task_name] + self.tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + self.setup() + + def setup(self): + self.dataset = datasets.load_dataset("glue", self.task_name) + + for split in self.dataset.keys(): + self.dataset[split] = self.dataset[split].map( + self.convert_to_features, + batched=True, + remove_columns=["label"], + ) + self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns] + self.dataset[split].set_format(type="torch", columns=self.columns) + + self.eval_splits = [x for x in self.dataset.keys() if "validation" in x] + + def prepare_data(self): + datasets.load_dataset("glue", self.task_name) + AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + + def train_dataloader(self): + return self.plugin.prepare_dataloader(self.dataset["train"], + batch_size=self.train_batch_size, + shuffle=True, + drop_last=True) + + def val_dataloader(self): + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["validation"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def test_dataloader(self): + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["test"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def convert_to_features(self, example_batch): + + # Either encode single sentence or sentence pairs + if len(self.text_fields) > 1: + texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]])) + else: + texts_or_text_pairs = example_batch[self.text_fields[0]] + + # Tokenize the text/text pairs + features = self.tokenizer.batch_encode_plus(texts_or_text_pairs, + max_length=self.max_seq_length, + padding='max_length', + truncation=True) + + # Rename label to labels to make it easier to pass to model forward + features["labels"] = example_batch["label"] + + return features diff --git a/examples/language/bert/finetune.py b/examples/language/bert/finetune.py new file mode 100644 index 000000000000..b209ffde85a4 --- /dev/null +++ b/examples/language/bert/finetune.py @@ -0,0 +1,220 @@ +import argparse +from typing import List, Union + +import evaluate +import torch +import torch.distributed as dist +import torch.nn as nn +from data import GLUEDataBuilder +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import ( + AlbertForSequenceClassification, + AutoConfig, + BertForSequenceClassification, + get_linear_schedule_with_warmup, +) + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 3 +BATCH_SIZE = 32 +LEARNING_RATE = 2.4e-5 +WEIGHT_DECAY = 0.01 +WARMUP_FRACTION = 0.1 + + +def move_to_cuda(batch): + return {k: v.cuda() for k, v in batch.items()} + + +@torch.no_grad() +def evaluate_model(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoader]], num_labels: int, task_name: str, + eval_splits: List[str], coordinator: DistCoordinator): + metric = evaluate.load("glue", task_name, process_id=coordinator.rank, num_process=coordinator.world_size) + model.eval() + + def evaluate_subset(dataloader: DataLoader): + accum_loss = torch.zeros(1, device=get_current_device()) + for batch in dataloader: + batch = move_to_cuda(batch) + outputs = model(**batch) + val_loss, logits = outputs[:2] + accum_loss.add_(val_loss) + + if num_labels > 1: + preds = torch.argmax(logits, axis=1) + elif num_labels == 1: + preds = logits.squeeze() + + labels = batch["labels"] + + metric.add_batch(predictions=preds, references=labels) + + results = metric.compute() + dist.all_reduce(accum_loss.div_(len(dataloader))) + if coordinator.is_master(): + results['loss'] = accum_loss.item() / coordinator.world_size + return results + + if isinstance(test_dataloader, DataLoader): + return evaluate_subset(test_dataloader) + else: + assert len(test_dataloader) == len(eval_splits) + final_results = {} + for split, sub_loader in zip(eval_splits, test_dataloader): + results = evaluate_subset(sub_loader) + final_results.update({f'{k}_{split}': v for k, v in results.items()}) + return final_results + + +def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, lr_scheduler, train_dataloader: DataLoader, + booster: Booster, coordinator: DistCoordinator): + model.train() + with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar: + for batch in pbar: + # Forward pass + batch = move_to_cuda(batch) + outputs = model(**batch) + loss = outputs[0] + + # Backward and optimize + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + lr_scheduler.step() + + # Print log info + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + parser.add_argument('-t', '--task', default='mrpc', help="GLUE task to run") + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") + parser.add_argument( + "--model_type", + type=str, + default="bert", + help="bert or albert", + ) + parser.add_argument('--target_f1', type=float, default=None, help="target f1 score. Raise exception if not reached") + args = parser.parse_args() + + if args.model_type == 'bert': + model_name = "bert-base-uncased" + elif args.model_type == 'albert': + model_name = "albert-xxlarge-v2" + else: + raise RuntimeError + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}, seed=42) + coordinator = DistCoordinator() + + # local_batch_size = BATCH_SIZE // coordinator.world_size + lr = LEARNING_RATE * coordinator.world_size + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + data_builder = GLUEDataBuilder(model_name, + plugin, + args.task, + train_batch_size=BATCH_SIZE, + eval_batch_size=BATCH_SIZE) + train_dataloader = data_builder.train_dataloader() + test_dataloader = data_builder.test_dataloader() + + # ==================================== + # Prepare model, optimizer + # ==================================== + # bert pretrained model + + cfg = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels) + if model_name == "bert-base-uncased": + model = BertForSequenceClassification.from_pretrained(model_name, config=cfg) + elif model_name == "albert-xxlarge-v2": + model = AlbertForSequenceClassification.from_pretrained(model_name, config=cfg) + else: + raise RuntimeError + + # optimizer + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], + "weight_decay": WEIGHT_DECAY, + }, + { + "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], + "weight_decay": 0.0, + }, + ] + + optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, eps=1e-8) + + # lr scheduler + total_steps = len(train_dataloader) * NUM_EPOCHS + num_warmup_steps = int(WARMUP_FRACTION * total_steps) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_steps, + ) + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, _, _, lr_scheduler = booster.boost(model, optimizer, lr_scheduler=lr_scheduler) + + # ============================== + # Train model + # ============================== + for epoch in range(NUM_EPOCHS): + train_epoch(epoch, model, optimizer, lr_scheduler, train_dataloader, booster, coordinator) + + results = evaluate_model(model, test_dataloader, data_builder.num_labels, args.task, data_builder.eval_splits, + coordinator) + + if coordinator.is_master(): + print(results) + if args.target_f1 is not None and 'f1' in results: + assert results['f1'] >= args.target_f1, f'f1 score {results["f1"]} is lower than target {args.target_f1}' + + +if __name__ == '__main__': + main() diff --git a/examples/language/bert/requirements.txt b/examples/language/bert/requirements.txt new file mode 100644 index 000000000000..377422c260ad --- /dev/null +++ b/examples/language/bert/requirements.txt @@ -0,0 +1,9 @@ +colossalai +evaluate +datasets +torch +tqdm +transformers +scipy +scikit-learn +ptflops diff --git a/examples/language/bert/run_gemini.sh b/examples/language/bert/run_gemini.sh deleted file mode 100644 index d791334e8c97..000000000000 --- a/examples/language/bert/run_gemini.sh +++ /dev/null @@ -1,22 +0,0 @@ -set -x -# distplan in ["CAI_ZeRO1", "CAI_ZeRO2", "CAI_Gemini", "Pytorch_DDP", "Pytorch_ZeRO"] -export DISTPLAN=${DISTPLAN:-"CAI_Gemini"} - -# The following options only valid when DISTPLAN="colossalai" -export GPUNUM=${GPUNUM:-1} -export PLACEMENT=${PLACEMENT:-"cpu"} -export BATCH_SIZE=${BATCH_SIZE:-16} - -# bert | albert -export MODEL_TYPE=${MODEL_TYPE:-"bert"} -export TRAIN_STEP=${TRAIN_STEP:-10} - -mkdir -p gemini_logs - -env CUDA_LAUNCH_BLOCKING=1 torchrun --standalone --nproc_per_node=${GPUNUM} ./train_bert_demo.py \ ---model_type=${MODEL_TYPE} \ ---batch_size=${BATCH_SIZE} \ ---placement=${PLACEMENT} \ ---distplan=${DISTPLAN} \ ---train_step=${TRAIN_STEP} \ -2>&1 | tee ./gemini_logs/${MODEL_TYPE}_${DISTPLAN}_gpu_${GPUNUM}_bs_${BATCH_SIZE}_${PLACEMENT}.log diff --git a/examples/language/bert/test_ci.sh b/examples/language/bert/test_ci.sh old mode 100644 new mode 100755 index 42c63fec50c0..7fc6daabb2f3 --- a/examples/language/bert/test_ci.sh +++ b/examples/language/bert/test_ci.sh @@ -1,2 +1,8 @@ -set -x -env GPUNUM=1 bash run_gemini.sh +#!/bin/bash +set -xe + +pip install -r requirements.txt + +for plugin in "torch_ddp" "torch_ddp_fp16" "gemini" "low_level_zero"; do + torchrun --standalone --nproc_per_node 4 finetune.py --target_f1 0.86 --plugin $plugin --model_type "bert" +done diff --git a/examples/language/bert/train_bert_demo.py b/examples/language/bert/train_bert_demo.py deleted file mode 100644 index b690ff787d01..000000000000 --- a/examples/language/bert/train_bert_demo.py +++ /dev/null @@ -1,332 +0,0 @@ -import os -from functools import partial -from time import time - -import psutil -import torch -from packaging import version -from torch import nn -from torch.nn.parallel import DistributedDataParallel as DDP -from transformers import AlbertConfig, AlbertForSequenceClassification, BertConfig, BertForSequenceClassification - -import colossalai -from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper -from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec -from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext - -CAI_VERSION = colossalai.__version__ - - -def get_tflops(model_numel, batch_size, seq_len, step_time): - return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12) - - -def get_profile_context(enable_flag, warmup_steps, active_steps, save_dir): - from contextlib import nullcontext - - from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler - if enable_flag: - return profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - schedule=schedule(wait=0, warmup=warmup_steps, active=active_steps), - on_trace_ready=tensorboard_trace_handler(save_dir), - record_shapes=True, - profile_memory=True) - else: - - class DummyProfiler: - - def __init__(self): - self.step_number = 0 - - def step(self): - self.step_number += 1 - - return nullcontext(DummyProfiler()) - - -def get_time_stamp(): - import time - cur_time = time.strftime("%d-%H:%M", time.localtime()) - return cur_time - - -def get_bert_data(batch_size: int, sequence_length: int, vacob_size: int, n_class: int, device: torch.device): - input = torch.randint( - low=0, - high=vacob_size, - size=(batch_size, sequence_length), - device=device, - dtype=torch.long, - ) - label = torch.randint(low=0, high=n_class, size=(batch_size,), device=device, dtype=torch.long) - return input, label - - -def parse_args(): - parser = colossalai.get_default_parser() - parser.add_argument( - "--distplan", - type=str, - default='CAI_Gemini', - help="The distributed plan [colossalai, zero1, zero2, torch_ddp, torch_zero].", - ) - parser.add_argument( - "--placement", - type=str, - default='cpu', - help="Placement Policy for Gemini. Valid when using colossalai as dist plan.", - ) - parser.add_argument( - "--batch_size", - type=int, - default=8, - help="batch size per DP group of training.", - ) - parser.add_argument( - "--model_type", - type=str, - default="bert", - help="bert or albert", - ) - parser.add_argument( - "--train_step", - type=int, - default=10, - help="training iterations for test", - ) - - args = parser.parse_args() - return args - - -SEQ_LEN = 512 -VOCAB_SIZE = 1000 -NUM_LABELS = 10 - - -# Parameter Sharding Strategies for Tensor Parallelism -def split_param_single_dim_tp1d(dim: int, param: ColoParameter, pg: ProcessGroup): - spec = (ShardSpec([dim], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) - param.set_tensor_spec(*spec) - - -def split_param_row_tp1d(param: ColoParameter, pg: ProcessGroup): - split_param_single_dim_tp1d(0, param, pg) - - -def split_param_col_tp1d(param: ColoParameter, pg: ProcessGroup): - split_param_single_dim_tp1d(-1, param, pg) - - -def get_cpu_mem(): - return psutil.Process().memory_info().rss / 1024**2 - - -def get_gpu_mem(): - return torch.cuda.memory_allocated() / 1024**2 - - -def get_mem_info(prefix=''): - return f'{prefix}GPU memory usage: {get_gpu_mem():.2f} MB, CPU memory usage: {get_cpu_mem():.2f} MB' - - -def get_model_size(model: nn.Module): - total_numel = 0 - for module in model.modules(): - for p in module.parameters(recurse=False): - total_numel += p.numel() - return total_numel - - -def model_builder(args): - if args.model_type == "bert": - cfg = BertConfig(vocab_size=VOCAB_SIZE, num_labels=NUM_LABELS) - return BertForSequenceClassification(cfg) - elif args.model_type == "albert": - cfg = AlbertConfig(vocab_size=VOCAB_SIZE, num_labels=NUM_LABELS) - return AlbertForSequenceClassification(cfg) - else: - raise RuntimeError - - -def model_size_formatter(numel: int) -> str: - GB_SIZE = 10**9 - MB_SIZE = 10**6 - KB_SIZE = 10**3 - if numel >= GB_SIZE: - return f'{numel / GB_SIZE:.1f}B' - elif numel >= MB_SIZE: - return f'{numel / MB_SIZE:.1f}M' - elif numel >= KB_SIZE: - return f'{numel / KB_SIZE:.1f}K' - else: - return str(numel) - - -def set_cpu_maximum_parallelism(): - conf_str = torch.__config__.parallel_info() - inter_str = conf_str.split("hardware_concurrency() : ")[1] - max_concurrency = inter_str.split('\n')[0] - os.environ["OMP_NUM_THREADS"] = max_concurrency - print(f"environmental variable OMP_NUM_THREADS is set to {max_concurrency}.") - - -def main(): - # version check - # this example is supposed to work for versions greater than 0.2.0 - assert version.parse(CAI_VERSION) >= version.parse("0.2.0") - - set_cpu_maximum_parallelism() - args = parse_args() - - # if args.distplan not in ["colossalai", "torch_ddp", "torch_zero", "zero1", "zero2"]: - if args.distplan not in ["CAI_ZeRO1", "CAI_ZeRO2", "CAI_Gemini", "Pytorch_DDP", "Pytorch_ZeRO"]: - raise TypeError(f"{args.distplan} is error") - - # batch size per DP degree - BATCH_SIZE = args.batch_size - - NUM_STEPS = args.train_step - - WARMUP_STEPS = 1 - assert WARMUP_STEPS < NUM_STEPS, "warmup steps should smaller than the total steps" - assert (NUM_STEPS - WARMUP_STEPS) % 2 == 1, "the number of valid steps should be odd to take the median" - PROF_FLAG = False # The flag of profiling, False by default - - disable_existing_loggers() - colossalai.launch_from_torch(config={}) - - logger = get_dist_logger() - logger.info(f" {args.distplan}, batch size {BATCH_SIZE}", ranks=[0]) - - torch.manual_seed(123) - if args.distplan.startswith("CAI"): - # all param must use the same process group. - world_size = torch.distributed.get_world_size() - - # build a base-bert model - with ColoInitContext(device=get_current_device(), dtype=torch.half): - model = model_builder(args) - # model = BertForSequenceClassification(BertConfig(vocal_size = VOCAB_SIZE)) - - # asign running configurations - gemini_config = None - if args.distplan.startswith("CAI_ZeRO"): - optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True) - elif args.distplan == "CAI_Gemini": - gemini_config = dict(strict_ddp_mode=True, - device=get_current_device(), - placement_policy=args.placement, - pin_memory=True, - hidden_dim=model.config.hidden_size, - search_range_mb=128) - optim_config = dict(gpu_margin_mem_ratio=0.) - else: - raise RuntimeError - - # build a highly optimized gpu/cpu optimizer - optimizer = HybridAdam(model.parameters(), lr=1e-3) - - if args.distplan == "CAI_ZeRO1": - zero_stage = 1 - elif args.distplan == "CAI_ZeRO2": - zero_stage = 2 - elif args.distplan == "CAI_Gemini": - zero_stage = 3 - else: - raise RuntimeError - - # wrap your model and optimizer - model = zero_model_wrapper(model, zero_stage, gemini_config) - optimizer = zero_optim_wrapper(model, optimizer, optim_config=optim_config) - - logger.info(get_mem_info(prefix='After init optim, '), ranks=[0]) - elif args.distplan.startswith("Pytorch"): - model = model_builder(args).cuda() - model = DDP(model) - if args.distplan.endswith("DDP"): - optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) - elif args.distplan.endswith("ZeRO"): - from torch.distributed.optim import ZeroRedundancyOptimizer - optimizer = ZeroRedundancyOptimizer(model.parameters(), optimizer_class=torch.optim.Adam, lr=1e-3) - else: - raise RuntimeError - - # model is shared after TP - numel = get_model_size(model) - logger.info(f"the size of testing model size is {model_size_formatter(numel)}.") - logger.info(get_mem_info(prefix='After init model, '), ranks=[0]) - - # Tflops_per_GPU = global_batch * global_numel * seq_len * 8 / #gpu - # = (batch_per_DP_group * dp_degree) * (numel * tp_degree) * seq_len * 8 / (tp_degree * dp_degree) - # = batch_per_DP_group * numel * seq_len * 8 - get_tflops_func = partial(get_tflops, numel, BATCH_SIZE, SEQ_LEN) - - torch.cuda.synchronize() - model.train() - tflops_list = [] - - def train_step(): - # we just use randomly generated data here - input_ids, labels = get_bert_data(BATCH_SIZE, - SEQ_LEN, - VOCAB_SIZE, - NUM_LABELS, - device=torch.cuda.current_device()) - optimizer.zero_grad() - - start = time() - outputs = model(input_ids, labels=labels) - loss, logits = outputs[:2] - torch.cuda.synchronize() - fwd_end = time() - fwd_time = fwd_end - start - logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Forward '), ranks=[0]) - - if args.distplan.startswith("CAI"): - optimizer.backward(loss) - elif args.distplan.startswith("Pytorch"): - loss.backward() - else: - raise RuntimeError - - torch.cuda.synchronize() - bwd_end = time() - bwd_time = bwd_end - fwd_end - logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Backward '), ranks=[0]) - - optimizer.step() - torch.cuda.synchronize() - optim_time = time() - bwd_end - step_time = time() - start - logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Optimizer step '), ranks=[0]) - - step_tflops = get_tflops_func(step_time) - logger.info( - f"[{n + 1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}, FWD time: {fwd_time:.3f}s, BWD time: {bwd_time:.3f}s, OPTIM time: {optim_time:.3f}s", - ranks=[0], - ) - if n >= WARMUP_STEPS: - tflops_list.append(step_tflops) - - demo_profiler = get_profile_context(PROF_FLAG, - WARMUP_STEPS, - NUM_STEPS - WARMUP_STEPS, - save_dir=f"profile/{get_time_stamp()}-demo") - - with demo_profiler as prof: - for n in range(NUM_STEPS): - train_step() - prof.step() - - tflops_list.sort() - median_index = ((NUM_STEPS - WARMUP_STEPS) >> 1) + WARMUP_STEPS - logger.info(f"Median TFLOPS is {tflops_list[median_index]:.3f}") - torch.cuda.synchronize() - - -if __name__ == '__main__': - main() diff --git a/examples/language/gpt/README.md b/examples/language/gpt/README.md index 10d6c2ddd5d7..47d24a4d69cb 100644 --- a/examples/language/gpt/README.md +++ b/examples/language/gpt/README.md @@ -40,7 +40,7 @@ We provide two stable solutions. One utilizes the Gemini to implement hybrid parallel strategies of Gemini, DDP/ZeRO, and Tensor Parallelism for a huggingface GPT model. The other one use [Titans](https://github.com/hpcaitech/Titans), a distributed executed model zoo maintained by ColossalAI,to implement the hybrid parallel strategies of TP + ZeRO + PP. -We recommend using Gemini to qucikly run your model in a distributed manner. +We recommend using Gemini to quickly run your model in a distributed manner. It doesn't require significant changes to the model structures, therefore you can apply it on a new model easily. And use Titans as an advanced weapon to pursue a more extreme performance. Titans has included the some typical models, such as Vit and GPT. diff --git a/examples/language/gpt/experiments/auto_offload/README.md b/examples/language/gpt/experiments/auto_offload/README.md new file mode 100644 index 000000000000..535aa76541cc --- /dev/null +++ b/examples/language/gpt/experiments/auto_offload/README.md @@ -0,0 +1,37 @@ +# Auto-Offload Demo with GPT2 + +## Requirements + +Before you can launch training, you need to install the following requirements. + +### Install PyTorch + +```bash +#conda +conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch +#pip +pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113 +``` + +### Install [Colossal-AI v0.2.0](https://colossalai.org/download/) From Official Website + +```bash +pip install colossalai==0.2.0+torch1.12cu11.3 -f https://release.colossalai.org +``` + +### Install transformers + +```bash +pip install transformers +``` + +## Dataset + +For simplicity, the input data is randomly generated here. + +## Training + +```bash +#Run the auto offload on GPT with default setting and a dummy dataset. +bash run.sh +``` diff --git a/examples/language/gpt/experiments/auto_offload/model_zoo.py b/examples/language/gpt/experiments/auto_offload/model_zoo.py new file mode 100644 index 000000000000..35e44608f810 --- /dev/null +++ b/examples/language/gpt/experiments/auto_offload/model_zoo.py @@ -0,0 +1,65 @@ +import torch +import torch.nn as nn +from transformers import GPT2Config, GPT2LMHeadModel + +class GPTLMModel(nn.Module): + + def __init__(self, + hidden_size=768, + num_layers=12, + num_attention_heads=12, + max_seq_len=1024, + vocab_size=50257): + super().__init__() + self.model = GPT2LMHeadModel( + GPT2Config(n_embd=hidden_size, + n_layer=num_layers, + n_head=num_attention_heads, + n_positions=max_seq_len, + n_ctx=max_seq_len, + vocab_size=vocab_size)) + + def forward(self, input_ids, attention_mask): + # Only return lm_logits + return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)[0] + + +class GPTLMLoss(nn.Module): + + def __init__(self): + super().__init__() + self.loss_fn = nn.CrossEntropyLoss() + + def forward(self, logits, labels): + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + +def get_gpt2_components(model_type: str, batch_size: int): + vocab_size = 1024 + seq_len = 8 + + def gpt2_model_builder(): + if model_type == "gpt2_medium": + return GPTLMModel(hidden_size=1024, num_layers=24, num_attention_heads=16) + elif model_type == "gpt2_xl": + return GPTLMModel(hidden_size=1600, num_layers=48, num_attention_heads=32) + elif model_type == "gpt2_10b": + return GPTLMModel(hidden_size=4096, num_layers=50, num_attention_heads=16) + elif model_type == "gpt2_14b": + return GPTLMModel(hidden_size=4096, num_layers=70, num_attention_heads=16) + elif model_type == "gpt2_20b": + return GPTLMModel(hidden_size=8192, num_layers=25, num_attention_heads=16) + elif model_type == "gpt2_24b": + return GPTLMModel(hidden_size=8192, num_layers=30, num_attention_heads=16) + else: + raise TypeError(f"model_builder {model_type}") + + def gpt2_data_gen(device="cuda"): + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device) + attention_mask = torch.ones_like(input_ids, device=device) + kwargs = dict(input_ids=input_ids, attention_mask=attention_mask) + return kwargs + + return gpt2_model_builder, gpt2_data_gen \ No newline at end of file diff --git a/examples/language/roberta/requirements.txt b/examples/language/gpt/experiments/auto_offload/requirements.txt similarity index 58% rename from examples/language/roberta/requirements.txt rename to examples/language/gpt/experiments/auto_offload/requirements.txt index 137a69e80498..3ebde8d460aa 100644 --- a/examples/language/roberta/requirements.txt +++ b/examples/language/gpt/experiments/auto_offload/requirements.txt @@ -1,2 +1,2 @@ colossalai >= 0.1.12 -torch >= 1.8.1 +torch >= 1.8.1 \ No newline at end of file diff --git a/examples/language/gpt/experiments/auto_offload/run.sh b/examples/language/gpt/experiments/auto_offload/run.sh new file mode 100644 index 000000000000..6a272ec442ab --- /dev/null +++ b/examples/language/gpt/experiments/auto_offload/run.sh @@ -0,0 +1,8 @@ +export BATCH_SIZE=${BATCH_SIZE:-64} +export MODEL_TYPE=${MODEL_TYPE:-"gpt2_medium"} +export MEMORY_BUDGET=${MEMORY_BUDGET:-16} +export SOLVER_TYPE=${SOLVER_TYPE:-"asyn"} + +mkdir -p offload_logs + +python train_gpt_offload.py --model_type=${MODEL_TYPE} --memory_budget=${MEMORY_BUDGET} --solver_type=${SOLVER_TYPE} --batch_size=${BATCH_SIZE} 2>&1 | tee ./offload_logs/${MODEL_TYPE}_bs_${BATCH_SIZE}_st_${SOLVER_TYPE}.log diff --git a/examples/language/gpt/experiments/auto_offload/train_gpt_offload.py b/examples/language/gpt/experiments/auto_offload/train_gpt_offload.py new file mode 100644 index 000000000000..89415c23f93c --- /dev/null +++ b/examples/language/gpt/experiments/auto_offload/train_gpt_offload.py @@ -0,0 +1,97 @@ +import argparse +import time + +import pytest +import torch +from model_zoo import GPTLMLoss, get_gpt2_components +from torch.utils._pytree import tree_map + +import colossalai +from colossalai.auto_parallel.offload.amp_optimizer import AMPOptimizer +from colossalai.auto_parallel.offload.mem_optimize import memory_optimize +from colossalai.auto_parallel.offload.solver import NOT_NVML +from colossalai.fx.profiler import parameter_size +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import spawn +from colossalai.utils import get_current_device + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--model_type', type=str, default="gpt2_medium") + parser.add_argument('--batch_size', type=int, default=64) + parser.add_argument('--solver_type', type=str, default='asyn') + parser.add_argument('--memory_budget', type=float, default=16) + return parser.parse_args() + + +@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed') +def train_gpt(args): + memory_budget = args.memory_budget * 1024 * 1024 * 1024 + solver_type = args.solver_type + model_type = args.model_type + batch_size = args.batch_size + + # build model + model_builder, data_gen = get_gpt2_components(model_type=model_type, batch_size=batch_size) + label = torch.randint(low=0, high=128, size=( + 64, + 8, + ), device=get_current_device()) + criterion = GPTLMLoss() + + start_time = time.time() + model = model_builder() + model.train() + param_size = parameter_size(model) / 1024**2 / 2 + init_time = time.time() - start_time + print(f"init_param_size={param_size:.3f} MB | init_model_time={init_time:.3f} s") + + data_args = data_gen(device="cpu") + wrap_fn = lambda x: x.to(dtype=torch.half) if isinstance(x, torch.Tensor) and torch.is_floating_point(x) else x + data_args = tree_map(wrap_fn, data_args) + start_time = time.time() + model = memory_optimize(model, data_args, memory_budget, solver_type) + solver_time = time.time() - start_time + print(f"solver_time={solver_time:.3f} s") + + hybrid_optimizer = HybridAdam(model.model.parameters(), lr=1e-3) + optim = AMPOptimizer(hybrid_optimizer, model) + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + time_list = [] + data_args = data_gen(device="cuda") + data_args = tree_map(wrap_fn, data_args) + for step in range(10): + optim.zero_grad() + torch.cuda.synchronize() + start_time = time.time() + loss = criterion(model(**data_args), label) + optim.backward(loss) + torch.cuda.synchronize() + time_list.append(time.time() - start_time) + optim.step() + + torch.cuda.synchronize() + + exec_time = sum(sorted(time_list)[:5]) / 5 + runtime_peak_mem_alc = torch.cuda.max_memory_allocated() / 1024**2 + runtime_peak_mem_res = torch.cuda.max_memory_reserved() / 1024**2 + print(f'solver_type: {solver_type} | model_type: {model_type}') + print(f'| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB ' + f'| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|') + print(time_list) + + +def run(rank, world_size, port, args): + config = {} + colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + train_gpt(args) + + +if __name__ == '__main__': + args = parse_args() + spawn(run, 1, args=args) diff --git a/examples/language/gpt/experiments/auto_parallel/README.md b/examples/language/gpt/experiments/auto_parallel/README.md index 404c8391109e..32688873f8f1 100644 --- a/examples/language/gpt/experiments/auto_parallel/README.md +++ b/examples/language/gpt/experiments/auto_parallel/README.md @@ -13,10 +13,10 @@ conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113 ``` -### Install [Colossal-AI v0.2.0](https://colossalai.org/download/) From Official Website +### Install Colossal-AI ```bash -pip install colossalai==0.2.0+torch1.12cu11.3 -f https://release.colossalai.org +pip install colossalai==0.2.0 ``` ### Install transformers @@ -34,7 +34,7 @@ conda install -c conda-forge coin-or-cbc ## Dataset -For simplicity, the input data is randonly generated here. +For simplicity, the input data is randomly generated here. ## Training diff --git a/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py b/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py index 6ceb7fd87c0a..e331fc8fcf10 100644 --- a/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py +++ b/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py @@ -1,18 +1,13 @@ from functools import partial from time import time -from typing import Dict, Optional, Tuple, Union import psutil import torch -import torch.multiprocessing as mp -import torch.nn as nn import transformers from gpt_modules import GPT2LMHeadModel, GPTLMLoss -from torch.fx import GraphModule -from colossalai.auto_parallel.tensor_shard.initialize import autoparallelize, initialize_model +from colossalai.auto_parallel.tensor_shard.initialize import autoparallelize from colossalai.core import global_context as gpc -from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch_from_torch from colossalai.logging import disable_existing_loggers, get_dist_logger diff --git a/examples/language/gpt/experiments/auto_parallel/requirements.txt b/examples/language/gpt/experiments/auto_parallel/requirements.txt index ff046ad1cae9..1b2561f098d5 100644 --- a/examples/language/gpt/experiments/auto_parallel/requirements.txt +++ b/examples/language/gpt/experiments/auto_parallel/requirements.txt @@ -1,4 +1,4 @@ colossalai >= 0.1.12 torch >= 1.8.1 -transformers >= 4.231 +transformers >= 4.23.1 PuLP >= 2.7.0 diff --git a/examples/language/gpt/experiments/pipeline_parallel/README.md b/examples/language/gpt/experiments/pipeline_parallel/README.md index 702e3c8d6540..5af994a00665 100644 --- a/examples/language/gpt/experiments/pipeline_parallel/README.md +++ b/examples/language/gpt/experiments/pipeline_parallel/README.md @@ -27,7 +27,7 @@ pip install transformers ## Dataset -For simplicity, the input data is randonly generated here. +For simplicity, the input data is randomly generated here. ## Training diff --git a/examples/language/gpt/gemini/test_ci.sh b/examples/language/gpt/gemini/test_ci.sh index 6079d5ed615b..0ddfd3a6211c 100644 --- a/examples/language/gpt/gemini/test_ci.sh +++ b/examples/language/gpt/gemini/test_ci.sh @@ -3,7 +3,7 @@ $(cd `dirname $0`;pwd) export TRAIN_STEP=4 for MODEL_TYPE in "gpt2_medium"; do - for DISTPLAN in "colossalai"; do + for DISTPLAN in "CAI_Gemini"; do for BATCH_SIZE in 2; do for GPUNUM in 1 4; do for TPDEGREE in 1 2; do diff --git a/examples/language/gpt/gemini/train_gpt_demo.py b/examples/language/gpt/gemini/train_gpt_demo.py index f46226bce2b5..9e61779a1dbf 100644 --- a/examples/language/gpt/gemini/train_gpt_demo.py +++ b/examples/language/gpt/gemini/train_gpt_demo.py @@ -11,12 +11,13 @@ from torch.nn.parallel import DistributedDataParallel as DDP import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin from colossalai.logging import disable_existing_loggers, get_dist_logger from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext CAI_VERSION = colossalai.__version__ @@ -161,7 +162,7 @@ def tensor_parallelize(model: torch.nn.Module, pg: ProcessGroup): # shard it w.r.t tp pattern if 'mlp.c_fc' in mn: if 'weight' in pn or 'bias' in pn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice # keep the shape of the output from c_fc param.compute_spec.set_output_replicate(False) else: @@ -172,9 +173,9 @@ def tensor_parallelize(model: torch.nn.Module, pg: ProcessGroup): else: param.set_dist_spec(ReplicaSpec()) elif 'wte' in mn or 'wpe' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif 'c_attn' in mn or 'c_proj' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice else: param.set_dist_spec(ReplicaSpec()) param.visited = True @@ -236,24 +237,7 @@ def main(): if args.tp_degree > 1: tensor_parallelize(model, tp_pg) - # asign running configurations - gemini_config = None - if args.distplan.startswith("CAI_ZeRO"): - optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True) - elif args.distplan == "CAI_Gemini": - gemini_config = dict(strict_ddp_mode=args.tp_degree == 1, - device=get_current_device(), - placement_policy=args.placement, - pin_memory=True, - hidden_dim=model.config.n_embd, - search_range_mb=128) - optim_config = dict(gpu_margin_mem_ratio=0.) - else: - raise RuntimeError - - # build a highly optimized gpu/cpu optimizer - optimizer = HybridAdam(model.parameters(), lr=1e-3) - + # assign running configurations if args.distplan == "CAI_ZeRO1": zero_stage = 1 elif args.distplan == "CAI_ZeRO2": @@ -263,22 +247,42 @@ def main(): else: raise RuntimeError - # wrap your model and optimizer - model = zero_model_wrapper(model, zero_stage, gemini_config) - optimizer = zero_optim_wrapper(model, optimizer, optim_config=optim_config) + plugin = None + if args.distplan.startswith("CAI_ZeRO"): + plugin = LowLevelZeroPlugin(stage=zero_stage, + reduce_bucket_size_in_m=12, + overlap_communication=True, + verbose=True) + elif args.distplan == "CAI_Gemini": + plugin = GeminiPlugin(device=get_current_device(), + placement_policy=args.placement, + pin_memory=True, + strict_ddp_mode=args.tp_degree == 1, + search_range_m=128, + hidden_dim=model.config.n_embd, + gpu_margin_mem_ratio=0.) + else: + raise RuntimeError + + # build a highly optimized gpu/cpu optimizer + optimizer = HybridAdam(model.parameters(), lr=1e-3) logger.info(get_mem_info(prefix='After init optim, '), ranks=[0]) elif args.distplan.startswith("Pytorch"): assert args.tp_degree == 1, "The degree of TP should be 1 for DDP examples." model = model_builder(args.model_type)(checkpoint=True).cuda() - model = DDP(model) + plugin = TorchDDPPlugin() if args.distplan.endswith("DDP"): optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) elif args.distplan.endswith("ZeRO"): from torch.distributed.optim import ZeroRedundancyOptimizer optimizer = ZeroRedundancyOptimizer(model.parameters(), optimizer_class=torch.optim.Adam, lr=1e-3) + else: raise RuntimeError + # wrap your model and optimizer + booster = Booster(plugin=plugin) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) # model is shared after TP numel = get_model_size(model) @@ -306,13 +310,7 @@ def train_step(): fwd_end = time() fwd_time = fwd_end - start logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Forward '), ranks=[0]) - - if args.distplan.startswith("CAI"): - optimizer.backward(loss) - elif args.distplan.startswith("Pytorch"): - loss.backward() - else: - raise RuntimeError + booster.backward(loss, optimizer) torch.cuda.synchronize() bwd_end = time() diff --git a/examples/language/gpt/titans/model/embed.py b/examples/language/gpt/titans/model/embed.py index 6369b9f8c5a1..d825ae92a285 100644 --- a/examples/language/gpt/titans/model/embed.py +++ b/examples/language/gpt/titans/model/embed.py @@ -305,7 +305,7 @@ def forward(ctx, vocab_parallel_logits, target): @staticmethod def backward(ctx, grad_output): - # Retreive tensors from the forward path. + # Retrieve tensors from the forward path. softmax, target_mask, masked_target_1d = ctx.saved_tensors # All the inputs have softmax as their gradient. diff --git a/examples/language/gpt/titans/train_gpt.py b/examples/language/gpt/titans/train_gpt.py index 66225d6c8044..6be0b9e8da30 100644 --- a/examples/language/gpt/titans/train_gpt.py +++ b/examples/language/gpt/titans/train_gpt.py @@ -15,7 +15,7 @@ from colossalai.trainer import Trainer, hooks from colossalai.utils import colo_set_process_memory_fraction, is_using_pp from colossalai.utils.timer import MultiTimer -from colossalai.zero.init_ctx import ZeroInitContext +from colossalai.zero.legacy.init_ctx import ZeroInitContext def calc_local_model_size(model: torch.nn.Module): diff --git a/examples/language/llama/README.md b/examples/language/llama/README.md new file mode 100644 index 000000000000..871804f2ca86 --- /dev/null +++ b/examples/language/llama/README.md @@ -0,0 +1,11 @@ +# Pretraining LLaMA: best practices for building LLaMA-like base models + +

      + +

      + +- 65-billion-parameter large model pretraining accelerated by 38% +[[code]](https://github.com/hpcaitech/ColossalAI/tree/example/llama/examples/language/llama) +[[blog]](https://www.hpc-ai.tech/blog/large-model-pretraining) + +> Since the main branch is being updated, in order to maintain the stability of the code, this example is temporarily kept as an [independent branch](https://github.com/hpcaitech/ColossalAI/tree/example/llama/examples/language/llama). diff --git a/examples/language/opt/README.md b/examples/language/opt/README.md index c2fd254571c7..37e1ff4d9008 100644 --- a/examples/language/opt/README.md +++ b/examples/language/opt/README.md @@ -19,15 +19,35 @@ Meta recently released [Open Pretrained Transformer (OPT)](https://github.com/fa The following example of [Colossal-AI](https://github.com/hpcaitech/ColossalAI) demonstrates fine-tuning Casual Language Modelling at low cost. -We are using the pre-training weights of the OPT model provided by Hugging Face Hub on the raw WikiText-2 (no tokens were replaced before -the tokenization). This training script is adapted from the [HuggingFace Language Modelling examples](https://github.com/huggingface/transformers/tree/main/examples/pytorch/language-modeling). ## Our Modifications -We adapt the OPT training code to ColossalAI by leveraging Gemini and ZeRO DDP. -## Quick Start -You can launch training by using the following bash script +We are using the pre-training weights of the OPT model provided by Hugging Face Hub on the raw WikiText-2 (no tokens were replaced before +the tokenization). + +We adapt the OPT training code to ColossalAI by leveraging [Boosting API](https://colossalai.org/docs/basics/booster_api) loaded with a chosen plugin, where each plugin corresponds to a specific kind of training strategy. This example supports plugins including TorchDDPPlugin, LowLevelZeroPlugin, and GeminiPlugin. + +## Run Demo +By running the following script: ```bash -bash ./run_gemini.sh +bash run_demo.sh ``` +You will finetune a [facebook/opt-350m](https://huggingface.co/facebook/opt-350m) model on this [dataset](https://huggingface.co/datasets/hugginglearners/netflix-shows), which contains more than 8000 comments on Netflix shows. + +The script can be modified if you want to try another set of hyperparameters or change to another OPT model with different size. + +The demo code is adapted from this [blog](https://medium.com/geekculture/fine-tune-eleutherai-gpt-neo-to-generate-netflix-movie-descriptions-in-only-47-lines-of-code-40c9b4c32475) and the [HuggingFace Language Modelling examples](https://github.com/huggingface/transformers/tree/main/examples/pytorch/language-modeling). + + + +## Run Benchmark + +You can run benchmark for OPT model by running the following script: +```bash +bash run_benchmark.sh +``` +The script will test performance (throughput & peak memory usage) for each combination of hyperparameters. You can also play with this script to configure your set of hyperparameters for testing. + + + diff --git a/examples/language/opt/args.py b/examples/language/opt/args.py new file mode 100644 index 000000000000..16730be7ebea --- /dev/null +++ b/examples/language/opt/args.py @@ -0,0 +1,120 @@ +from colossalai import get_default_parser + + +def parse_demo_args(): + + parser = get_default_parser() + parser.add_argument( + "--model_name_or_path", + type=str, + default="facebook/opt-350m", + help="Path to pretrained model or model identifier from huggingface.co/models." + ) + parser.add_argument( + "--output_path", + type=str, + default="./output_model.bin", + help="The path of your saved model after finetuning." + ) + parser.add_argument( + "--plugin", + type=str, + default="gemini", + help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'." + ) + parser.add_argument( + "--num_epoch", + type=int, + default=10, + help="Number of epochs." + ) + parser.add_argument( + "--batch_size", + type=int, + default=32, + help="Batch size (per dp group) for the training dataloader." + ) + parser.add_argument( + "--learning_rate", + type=float, + default=5e-5, + help="Initial learning rate (after the potential warmup period) to use." + ) + parser.add_argument( + "--warmup_ratio", + type=float, + default=0.1, + help="Ratio of warmup steps against total training steps." + ) + parser.add_argument( + "--weight_decay", + type=float, + default=0.01, + help="Weight decay to use." + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="A seed for reproducible training." + ) + + args = parser.parse_args() + return args + + + +def parse_benchmark_args(): + + parser = get_default_parser() + parser.add_argument( + "--model_name_or_path", + type=str, + default="facebook/opt-125m", + help="Path to pretrained model or model identifier from huggingface.co/models." + ) + parser.add_argument( + "--plugin", + type=str, + default="gemini", + help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'." + ) + parser.add_argument( + "--batch_size", + type=int, + default=32, + help="Batch size (per dp group) for the training dataloader." + ) + parser.add_argument( + "--learning_rate", + type=float, + default=5e-5, + help="Initial learning rate (after the potential warmup period) to use." + ) + parser.add_argument( + "--weight_decay", + type=float, + default=0.0, + help="Weight decay to use." + ) + parser.add_argument( + "--max_train_steps", + type=int, + default=20, + help="Total number of training steps to perform." + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="A seed for reproducible training." + ) + parser.add_argument( + "--mem_cap", + type=int, + default=0, + help="Limit on the usage of space for each GPU (in GB)." + ) + args = parser.parse_args() + + return args \ No newline at end of file diff --git a/examples/language/opt/benchmark.sh b/examples/language/opt/benchmark.sh deleted file mode 100644 index 0d04b5e9b33c..000000000000 --- a/examples/language/opt/benchmark.sh +++ /dev/null @@ -1,21 +0,0 @@ -export BS=16 -export MEMCAP=0 -export MODEL="6.7b" -export GPUNUM=1 - -for MODEL in "6.7b" "13b" "1.3b" -do -for GPUNUM in 8 1 -do -for BS in 16 24 32 8 -do -for MEMCAP in 0 40 -do -pkill -9 torchrun -pkill -9 python - -env BS=$BS MEM_CAP=$MEMCAP MODEL=$MODEL GPUNUM=$GPUNUM bash ./run_gemini.sh -done -done -done -done diff --git a/examples/language/opt/data.py b/examples/language/opt/data.py new file mode 100644 index 000000000000..6cfffb5fc95b --- /dev/null +++ b/examples/language/opt/data.py @@ -0,0 +1,37 @@ +import torch +from torch.utils.data import Dataset +from datasets import load_dataset + + +class NetflixDataset(Dataset): + + def __init__(self, tokenizer): + + super().__init__() + + self.tokenizer = tokenizer + self.input_ids = [] + self.attn_masks = [] + self.labels = [] + self.txt_list = netflix_descriptions = load_dataset("hugginglearners/netflix-shows", split="train")['description'] + self.max_length = max([len(self.tokenizer.encode(description)) for description in netflix_descriptions]) + + for txt in self.txt_list: + encodings_dict = self.tokenizer('' + txt + '', + truncation=True, + max_length=self.max_length, + padding="max_length") + self.input_ids.append(torch.tensor(encodings_dict['input_ids'])) + self.attn_masks.append(torch.tensor(encodings_dict['attention_mask'])) + + def __len__(self): + return len(self.input_ids) + + def __getitem__(self, idx): + return self.input_ids[idx], self.attn_masks[idx] + + +def netflix_collator(data): + return {'input_ids': torch.stack([x[0] for x in data]), + 'attention_mask': torch.stack([x[1] for x in data]), + 'labels': torch.stack([x[0] for x in data])} diff --git a/examples/language/opt/opt_benchmark.py b/examples/language/opt/opt_benchmark.py new file mode 100755 index 000000000000..2d69036b50c6 --- /dev/null +++ b/examples/language/opt/opt_benchmark.py @@ -0,0 +1,137 @@ +import time + +import torch +import transformers +from transformers import AutoConfig, OPTForCausalLM +from transformers.utils.versions import require_version +import tqdm + +import colossalai +from colossalai.nn.optimizer import HybridAdam +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.tensor import ProcessGroup, ShardSpec +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator + +from args import parse_benchmark_args + +require_version("transformers>=4.20.0", "To fix: pip install -r requirements.txt") + + +def format_num(num: int, bytes=False): + """Scale bytes to its proper format, e.g. 1253656 => '1.20MB'""" + factor = 1024 if bytes else 1000 + suffix = "B" if bytes else "" + for unit in ["", " K", " M", " G", " T", " P"]: + if num < factor: + return f"{num:.2f}{unit}{suffix}" + num /= factor + + +def get_data(batch_size, seq_len, vocab_size): + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device()) + attention_mask = torch.ones_like(input_ids) + return input_ids, attention_mask + + +def colo_memory_cap(size_in_GB): + from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device + cuda_capacity = colo_device_memory_capacity(get_current_device()) + if size_in_GB * (1024**3) < cuda_capacity: + colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity) + print(f"Limiting GPU memory usage to {size_in_GB} GB") + + +def main(): + + args = parse_benchmark_args() + + # Launch ColossalAI + colossalai.launch_from_torch(config={}, seed=args.seed) + coordinator = DistCoordinator() + world_size = coordinator.world_size + + # Manage loggers + disable_existing_loggers() + logger = get_dist_logger() + if coordinator.is_master(): + transformers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + + # Whether to set limit of memory capacity + if args.mem_cap > 0: + colo_memory_cap(args.mem_cap) + + # Build OPT model + config = AutoConfig.from_pretrained(args.model_name_or_path) + model = OPTForCausalLM(config=config) + logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0]) + + # Enable gradient checkpointing + model.gradient_checkpointing_enable() + + # Set plugin + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(device=get_current_device(), + placement_policy='cpu', + pin_memory=True, + strict_ddp_mode=True, + initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + logger.info(f"Set plugin as {args.plugin}", ranks=[0]) + + # Set optimizer + optimizer = HybridAdam(model.parameters(), lr=args.learning_rate) + + # Set booster + booster = Booster(plugin=plugin, **booster_kwargs) + model, optimizer, _, _, _ = booster.boost(model, optimizer) + + SEQ_LEN = 1024 + VOCAB_SIZE = 50257 + + # Start training. + logger.info(f"Start testing", ranks=[0]) + progress_bar = tqdm.tqdm(total=args.max_train_steps, desc="Training Step", disable=not coordinator.is_master()) + + torch.cuda.synchronize() + model.train() + start_time = time.time() + + for _ in range(args.max_train_steps): + + input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE) + optimizer.zero_grad() + outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False) + loss = outputs['loss'] + booster.backward(loss, optimizer) + optimizer.step() + + torch.cuda.synchronize() + progress_bar.update(1) + + # Compute Statistics + end_time = time.time() + throughput = "{:.4f}".format((world_size * args.max_train_steps * args.batch_size) / (end_time - start_time)) + max_mem = format_num(torch.cuda.max_memory_allocated(device=torch.cuda.current_device()), bytes=True) + + logger.info(f"Testing finished, " + f"batch size per gpu: {args.batch_size}, " + f"plugin: {args.plugin}, " + f"throughput: {throughput}, " + f"maximum memory usage per gpu: {max_mem}.", + ranks=[0]) + + +if __name__ == "__main__": + main() diff --git a/examples/language/opt/opt_train_demo.py b/examples/language/opt/opt_train_demo.py new file mode 100644 index 000000000000..fa7feca9c9a9 --- /dev/null +++ b/examples/language/opt/opt_train_demo.py @@ -0,0 +1,142 @@ +import time + +import torch +import datasets +import transformers +from transformers import AutoConfig, OPTForCausalLM, AutoTokenizer +from transformers import get_linear_schedule_with_warmup +from transformers.utils.versions import require_version +from tqdm import tqdm + +import colossalai +from colossalai.nn.optimizer import HybridAdam +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.tensor import ProcessGroup, ShardSpec +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator + +from args import parse_demo_args +from data import NetflixDataset, netflix_collator + +require_version("datasets>=1.8.0", "To fix: pip install -r requirements.txt") +require_version("transformers>=4.20.0", "To fix: pip install -r requirements.txt") + + +def move_to_cuda(batch, device): + return {k: v.to(device) for k, v in batch.items()} + + +def train_epoch(epoch, model, optimizer, lr_scheduler, dataloader, booster, coordinator): + + torch.cuda.synchronize() + model.train() + + with tqdm(dataloader, desc=f'Epoch [{epoch + 1}]', disable=not coordinator.is_master()) as pbar: + + for batch in pbar: + + # Forward + optimizer.zero_grad() + batch = move_to_cuda(batch, torch.cuda.current_device()) + + outputs = model(use_cache=False, **batch) + loss = outputs['loss'] + + # Backward + booster.backward(loss, optimizer) + optimizer.step() + lr_scheduler.step() + + # Print batch loss + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + + args = parse_demo_args() + + # Launch ColossalAI + colossalai.launch_from_torch(config={}, seed=args.seed) + coordinator = DistCoordinator() + world_size = coordinator.world_size + + # Manage loggers + disable_existing_loggers() + logger = get_dist_logger() + if coordinator.is_master(): + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + + # Build OPT model + config = AutoConfig.from_pretrained(args.model_name_or_path) + model = OPTForCausalLM.from_pretrained(args.model_name_or_path, config=config) + logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0]) + + # Enable gradient checkpointing + model.gradient_checkpointing_enable() + + # Set plugin + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(device=get_current_device(), + placement_policy='cpu', + pin_memory=True, + strict_ddp_mode=True, + initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + logger.info(f"Set plugin as {args.plugin}", ranks=[0]) + + # Prepare tokenizer and dataloader + tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path) + dataset = NetflixDataset(tokenizer) + dataloader = plugin.prepare_dataloader(dataset, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + collate_fn=netflix_collator) + + # Set optimizer + optimizer = HybridAdam(model.parameters(), + lr=(args.learning_rate * world_size), + weight_decay=args.weight_decay) + + # Set lr scheduler + total_steps = len(dataloader) * args.num_epoch + num_warmup_steps = int(args.warmup_ratio * total_steps) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=len(dataloader) * args.num_epoch + ) + + # Set booster + booster = Booster(plugin=plugin, **booster_kwargs) + model, optimizer, _, dataloader, lr_scheduler = booster.boost(model=model, + optimizer=optimizer, + dataloader=dataloader, + lr_scheduler=lr_scheduler) + + # Start finetuning + logger.info(f"Start finetuning", ranks=[0]) + for epoch in range(args.num_epoch): + train_epoch(epoch, model, optimizer, lr_scheduler, dataloader, booster, coordinator) + + # Finish training and evaluate + logger.info(f"Finish finetuning", ranks=[0]) + booster.save_model(model, args.output_path) + logger.info(f"Saving model checkpoint to {args.output_path}", ranks=[0]) + + +if __name__ == "__main__": + main() diff --git a/examples/language/opt/requirements.txt b/examples/language/opt/requirements.txt index 137a69e80498..4422216e6a1c 100644 --- a/examples/language/opt/requirements.txt +++ b/examples/language/opt/requirements.txt @@ -1,2 +1,4 @@ colossalai >= 0.1.12 torch >= 1.8.1 +datasets >= 1.8.0 +transformers >= 4.20.0 \ No newline at end of file diff --git a/examples/language/opt/run_benchmark.sh b/examples/language/opt/run_benchmark.sh new file mode 100644 index 000000000000..76c5e8601989 --- /dev/null +++ b/examples/language/opt/run_benchmark.sh @@ -0,0 +1,30 @@ +set -xe +pip install -r requirements.txt + +export BS=32 +export MEMCAP=0 +export GPUNUM=1 + +# acceptable values include `125m`, `350m`, `1.3b`, `2.7b`, `6.7b`, `13b`, `30b`, `66b` +export MODEL="125m" + +for BS in 8 32 128 +do +for PLUGIN in "torch_ddp" "torch_ddp_fp16" "low_level_zero" "gemini" +do +for GPUNUM in 1 4 +do + +MODLE_PATH="facebook/opt-${MODEL}" +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + opt_benchmark.py \ + --model_name_or_path ${MODLE_PATH} \ + --mem_cap ${MEMCAP} \ + --plugin ${PLUGIN} \ + --batch_size ${BS} + +done +done +done diff --git a/examples/language/opt/run_demo.sh b/examples/language/opt/run_demo.sh new file mode 100644 index 000000000000..0c9759c34039 --- /dev/null +++ b/examples/language/opt/run_demo.sh @@ -0,0 +1,44 @@ +set -xe +pip install -r requirements.txt + +# model name or path +MODEL="facebook/opt-350m" + +# path for saving model +OUTPUT_PATH="./output_model.bin" + +# plugin(training strategy) +# can only be one of "torch_ddp"/"torch_ddp_fp16"/"low_level_zero"/"gemini" +PLUGIN="gemini" + +# number of gpus to use +GPUNUM=4 + +# batch size per gpu +BS=16 + +# learning rate +LR="5e-5" + +# number of epoch +EPOCH=10 + +# weight decay +WEIGHT_DECAY=0.01 + +# ratio of warmup steps +WARMUP_RATIO=0.1 + +# run the script for demo +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + opt_train_demo.py \ + --model_name_or_path ${MODEL} \ + --output_path ${OUTPUT_PATH} \ + --plugin ${PLUGIN} \ + --batch_size ${BS} \ + --num_epoch ${EPOCH} \ + --learning_rate ${LR} \ + --weight_decay ${WEIGHT_DECAY} \ + --warmup_ratio ${WARMUP_RATIO} diff --git a/examples/language/opt/run_gemini.sh b/examples/language/opt/run_gemini.sh deleted file mode 100644 index 73f231292a13..000000000000 --- a/examples/language/opt/run_gemini.sh +++ /dev/null @@ -1,28 +0,0 @@ -set -x -export BS=${BS:-16} -export MEMCAP=${MEMCAP:-0} -# Acceptable values include `125m`, `350m`, `1.3b`, `2.7b`, `6.7b`, `13b`, `30b`, `66b`. For `175b` -export MODEL=${MODEL:-"125m"} -export GPUNUM=${GPUNUM:-1} -export USE_SHARD_INIT=${USE_SHARD_INIT:-"false"} - -# make directory for logs -mkdir -p ./logs - -if [ ${USE_SHARD_INIT} = "true" ]; then - USE_SHARD_INIT="--shardinit" -else - USE_SHARD_INIT="" -fi - -export MODLE_PATH="facebook/opt-${MODEL}" - -# HF_DATASETS_OFFLINE=1 TRANSFORMERS_OFFLINE=1 -torchrun \ - --nproc_per_node ${GPUNUM} \ - --master_port 19198 \ - train_gemini_opt.py \ - --mem_cap ${MEMCAP} \ - --model_name_or_path ${MODLE_PATH} \ - ${USE_SHARD_INIT} \ - --batch_size ${BS} 2>&1 | tee ./logs/colo_${MODEL}_bs_${BS}_cap_${MEMCAP}_gpu_${GPUNUM}.log diff --git a/examples/language/opt/test_ci.sh b/examples/language/opt/test_ci.sh index 317f602cda3c..fa14f52b70d2 100644 --- a/examples/language/opt/test_ci.sh +++ b/examples/language/opt/test_ci.sh @@ -1,4 +1,19 @@ -for GPUNUM in 2 1 +set -xe +pip install -r requirements.txt + +BS=4 +for PLUGIN in "torch_ddp" "torch_ddp_fp16" "low_level_zero" "gemini" do -env BS=2 MODEL="125m" GPUNUM=$GPUNUM bash ./run_gemini.sh +for GPUNUM in 1 4 +do + +torchrun \ + --standalone \ + --nproc_per_node ${GPUNUM} \ + opt_benchmark.py \ + --model_name_or_path "facebook/opt-125m" \ + --plugin ${PLUGIN} \ + --batch_size ${BS} + +done done diff --git a/examples/language/opt/train_gemini_opt.py b/examples/language/opt/train_gemini_opt.py deleted file mode 100755 index 4993ce25db17..000000000000 --- a/examples/language/opt/train_gemini_opt.py +++ /dev/null @@ -1,231 +0,0 @@ -#!/usr/bin/env python -# coding=utf-8 -# Copyright 2021 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) -on a text file or a dataset without using HuggingFace Trainer. - -Here is the full list of checkpoints on the hub that can be fine-tuned by this script: -https://huggingface.co/models?filter=text-generation -""" -# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments. - -import time -from functools import partial - -import datasets -import torch -import torch.distributed as dist -import transformers -from transformers import CONFIG_MAPPING, MODEL_MAPPING, AutoConfig, OPTForCausalLM -from transformers.utils.versions import require_version - -import colossalai -from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel import GeminiDDP -from colossalai.utils import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext - -from colossalai.tensor import ProcessGroup, ShardSpec - - -def get_data(batch_size, seq_len, vocab_size): - input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device()) - attention_mask = torch.ones_like(input_ids) - return input_ids, attention_mask - - -require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") - -MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys()) -MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES) - - -def get_time_stamp(): - torch.cuda.synchronize() - return time.time() - - -def get_tflops(model_numel, batch_size, seq_len, step_time): - return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12) - - -def parse_args(): - parser = colossalai.get_default_parser() - parser.add_argument( - "--model_name_or_path", - type=str, - help="Path to pretrained model or model identifier from huggingface.co/models.", - required=True, - ) - parser.add_argument( - "--config_name", - type=str, - default=None, - help="Pretrained config name or path if not the same as model_name", - ) - parser.add_argument( - "--batch_size", - type=int, - default=8, - help="Batch size (per dp group) for the training dataloader.", - ) - parser.add_argument( - "--learning_rate", - type=float, - default=5e-5, - help="Initial learning rate (after the potential warmup period) to use.", - ) - parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.") - parser.add_argument( - "--max_train_steps", - type=int, - default=20, - help="Total number of training steps to perform.", - ) - parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") - parser.add_argument( - "--model_type", - type=str, - default=None, - help="Model type to use if training from scratch.", - choices=MODEL_TYPES, - ) - parser.add_argument( - "--shardinit", - action="store_true", - help="Initialize the model with tensor parallel", - ) - parser.add_argument("--mem_cap", type=int, default=0, help="use mem cap") - parser.add_argument("--init_in_cpu", action='store_true', default=False, help="init training model in cpu") - args = parser.parse_args() - - return args - - -def colo_memory_cap(size_in_GB): - from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device - cuda_capacity = colo_device_memory_capacity(get_current_device()) - if size_in_GB * (1024**3) < cuda_capacity: - colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity) - print("Using {} GB of GPU memory".format(size_in_GB)) - - -def main(): - args = parse_args() - disable_existing_loggers() - colossalai.launch_from_torch({}) - logger = get_dist_logger() - is_main_process = dist.get_rank() == 0 - - if is_main_process: - datasets.utils.logging.set_verbosity_warning() - transformers.utils.logging.set_verbosity_info() - else: - datasets.utils.logging.set_verbosity_error() - transformers.utils.logging.set_verbosity_error() - - if args.mem_cap > 0: - colo_memory_cap(args.mem_cap) - - # If passed along, set the training seed now. - if args.seed is not None: - torch.mannul_seed(args.seed) - logger.info(f"Rank {dist.get_rank()}: random seed is set to {args.seed}") - - # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at - # https://huggingface.co/docs/datasets/loading_datasets.html. - - # Load pretrained model - # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently - # download model & vocab. - if args.config_name: - config = AutoConfig.from_pretrained(args.config_name) - elif args.model_name_or_path: - config = AutoConfig.from_pretrained(args.model_name_or_path) - else: - config = CONFIG_MAPPING[args.model_type]() - logger.warning("You are instantiating a new config instance from scratch.") - logger.info("Model config has been created", ranks=[0]) - - if args.init_in_cpu: - init_dev = torch.device('cpu') - else: - init_dev = get_current_device() - - # shard init prameters - if args.shardinit: - logger.info("Sharding initialization !", ranks=[0]) - else: - logger.info("Skipping sharding initialization", ranks=[0]) - - world_size = torch.distributed.get_world_size() - shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None - default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None - - # build model - if args.model_name_or_path is None: - logger.info("Train a new model from scratch", ranks=[0]) - with ColoInitContext(device=init_dev, dtype=torch.half, - default_dist_spec=default_dist_spec, - default_pg=shard_pg): - model = OPTForCausalLM(config) - else: - logger.info("Finetune a pre-trained model", ranks=[0]) - with ColoInitContext(device=init_dev, dtype=torch.half, - default_dist_spec=default_dist_spec, - default_pg=shard_pg): - model = OPTForCausalLM.from_pretrained(args.model_name_or_path, - from_tf=bool(".ckpt" in args.model_name_or_path), - config=config, - local_files_only=False) - - # enable graident checkpointing - model.gradient_checkpointing_enable() - - numel = sum([p.numel() for p in model.parameters()]) - PLACEMENT_POLICY = 'cpu' - model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, - pin_memory=True, strict_ddp_mode=args.shardinit) - optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0) - - SEQ_LEN = 1024 - VOCAB_SIZE = 50257 - - get_tflops_func = partial(get_tflops, numel, args.batch_size, SEQ_LEN) - - model.train() - for step in range(args.max_train_steps): - st_time = time.time() - input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE) - - outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False) - loss = outputs['loss'] - optimizer.backward(loss) - - optimizer.step() - optimizer.zero_grad() - torch.cuda.synchronize() - step_time = time.time() - st_time - step_tflops = get_tflops_func(step_time) - - logger.info("step {} finished, Tflops {}".format(step, step_tflops), ranks=[0]) - - logger.info("Training finished", ranks=[0]) - - -if __name__ == "__main__": - main() diff --git a/examples/language/palm/README.md b/examples/language/palm/README.md index 486bf240f89c..3ff3939d63d4 100644 --- a/examples/language/palm/README.md +++ b/examples/language/palm/README.md @@ -43,6 +43,9 @@ palm = PaLM( ) ``` +## New API +We have modified our previous implementation of PaLM with our new Booster API, which offers a more flexible and efficient way to train your model. The new API is more user-friendly and easy to use. You can find the new API in train.py. We have also offer a shell script test_ci.sh for you to go through all our plugins for the booster. For more information about the booster API you can refer to https://colossalai.org/docs/basics/booster_api/. + ## Test on Enwik8 ```bash diff --git a/examples/language/palm/run.sh b/examples/language/palm/run.sh index 7a533509e009..2a846e81a9a7 100644 --- a/examples/language/palm/run.sh +++ b/examples/language/palm/run.sh @@ -3,9 +3,11 @@ export DISTPAN="colossalai" # The following options only valid when DISTPAN="colossalai" export TPDEGREE=1 -export GPUNUM=1 +export GPUNUM=4 export PLACEMENT='cpu' export USE_SHARD_INIT=False -export BATCH_SIZE=4 +export BATCH_SIZE=1 -env OMP_NUM_THREADS=12 torchrun --standalone --nproc_per_node=${GPUNUM} --master_port 29501 train.py --tp_degree=${TPDEGREE} --batch_size=${BATCH_SIZE} --placement ${PLACEMENT} --shardinit ${USE_SHARD_INIT} --distplan ${DISTPAN} 2>&1 | tee run.log +env OMP_NUM_THREADS=12 torchrun --standalone --nproc_per_node=${GPUNUM} --master_port 29501 train.py \ +--dummy_data=True --tp_degree=${TPDEGREE} --batch_size=${BATCH_SIZE} --plugin='gemini' \ +--placement ${PLACEMENT} --shardinit ${USE_SHARD_INIT} --distplan ${DISTPAN} 2>&1 | tee run.log diff --git a/examples/language/palm/test_ci.sh b/examples/language/palm/test_ci.sh index f21095578077..4de6a44e5bf7 100644 --- a/examples/language/palm/test_ci.sh +++ b/examples/language/palm/test_ci.sh @@ -4,6 +4,6 @@ for BATCH_SIZE in 2 do for GPUNUM in 1 4 do -env OMP_NUM_THREADS=12 torchrun --standalone --nproc_per_node=${GPUNUM} --master_port 29501 train.py --dummy_data=True --batch_size=${BATCH_SIZE} 2>&1 | tee run.log +env OMP_NUM_THREADS=12 torchrun --standalone --nproc_per_node=${GPUNUM} --standalone train.py --dummy_data=True --batch_size=${BATCH_SIZE} --plugin='gemini' 2>&1 | tee run.log done done diff --git a/examples/language/palm/train.py b/examples/language/palm/train.py index 2f012780da77..a0600db1bc5b 100644 --- a/examples/language/palm/train.py +++ b/examples/language/palm/train.py @@ -9,17 +9,19 @@ import torch.optim as optim import tqdm from packaging import version + +from colossalai.nn import HybridAdam from palm_pytorch import PaLM from palm_pytorch.autoregressive_wrapper import AutoregressiveWrapper from torch.utils.data import DataLoader, Dataset import colossalai from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel import ZeroDDP from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec from colossalai.utils import MultiTimer, get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, GeminiAdamOptimizer, ZeroDDP +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin # constants @@ -60,6 +62,12 @@ def parse_args(): help= "Shard the tensors when init the model to shrink peak memory size on the assigned device. Valid when using colossalai as dist plan.", ) + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") parser.add_argument( "--batch_size", type=int, @@ -103,31 +111,9 @@ def get_model_size(model: nn.Module): return total_numel -# Gemini + ZeRO DDP -def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placememt_policy: str = "auto"): - cai_version = colossalai.__version__ - if version.parse(cai_version) > version.parse("0.1.10"): - from colossalai.nn.parallel import GeminiDDP - model = GeminiDDP(model, - device=get_current_device(), - placement_policy=placememt_policy, - pin_memory=True, - search_range_mb=32) - elif version.parse(cai_version) <= version.parse("0.1.10") and version.parse(cai_version) >= version.parse("0.1.9"): - from colossalai.gemini import ChunkManager, GeminiManager - chunk_size = ChunkManager.search_chunk_size(model, 64 * 1024**2, 32) - gemini_manager = GeminiManager(placememt_policy, chunk_manager) - chunk_manager = ChunkManager(chunk_size, - pg, - enable_distributed_storage=True, - init_device=GeminiManager.get_default_device(placememt_policy)) - model = ZeroDDP(model, gemini_manager) - else: - raise NotImplemented(f"CAI version {cai_version} is not supported") - return model -## Parameter Sharding Strategies for Tensor Parallelism +# Parameter Sharding Strategies for Tensor Parallelism def split_param_single_dim_tp1d(dim: int, param: ColoParameter, pg: ProcessGroup): spec = (ShardSpec([dim], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)) param.set_tensor_spec(*spec) @@ -154,15 +140,15 @@ def tensor_parallelize(model: torch.nn.Module, pg: ProcessGroup): continue param.set_dist_spec(ReplicaSpec()) if 'net.0' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif 'to_q' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif 'to_kv' in mn: split_param_row_tp1d(param, pg) # row slice elif 'to_out' in mn: split_param_row_tp1d(param, pg) # row slice elif '1.1' in mn: - split_param_col_tp1d(param, pg) # colmn slice + split_param_col_tp1d(param, pg) # column slice elif '1.2' in mn: split_param_row_tp1d(param, pg) # row slice else: @@ -220,6 +206,18 @@ def __len__(self): if args.distplan == "colossalai": # instantiate GPT-like decoder model + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy=args.placement, strict_ddp_mode=True, initial_scale=2 ** 5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2 ** 5) + logger.info(f"plugin: {plugin}") + booster = Booster(plugin=plugin, **booster_kwargs) + default_pg = ProcessGroup(tp_degree=args.tp_degree) default_dist_spec = ShardSpec([-1], [args.tp_degree]) if args.shardinit else None ctx = ColoInitContext(device='cpu', default_dist_spec=default_dist_spec, default_pg=default_pg) @@ -230,12 +228,12 @@ def __len__(self): pg = default_pg tensor_parallelize(model, pg) - model = gemini_zero_dpp(model, pg, args.placement) - #optimizer + # optimizer + + optimizer = HybridAdam(model.parameters(), lr=LEARNING_RATE, initial_scale=2**5) + model, optimizer, _, _, _ = booster.boost(model, optimizer) - #optimizer = GeminiAdamOptimizer(model, lr=1e-7, initial_scale=2**5) - optimizer = GeminiAdamOptimizer(model, lr=LEARNING_RATE, initial_scale=2**5) else: model = PaLM(num_tokens=256, dim=512, depth=8) model = AutoregressiveWrapper(model, max_seq_len=2048) diff --git a/examples/language/roberta/README.md b/examples/language/roberta/README.md deleted file mode 100644 index a42b1935dd85..000000000000 --- a/examples/language/roberta/README.md +++ /dev/null @@ -1,58 +0,0 @@ -# Introduction -This repo introduce how to pretrain a chinese roberta-large from scratch, including preprocessing, pretraining, finetune. The repo can help you quickly train a high-quality bert. - -## 0. Prerequisite -- Install Colossal-AI -- Editing the port from /etc/ssh/sshd_config and /etc/ssh/ssh_config, every host expose the same ssh port of server and client. If you are a root user, you also set the **PermitRootLogin** from /etc/ssh/sshd_config to "yes" -- Ensure that each host can log in to each other without password. If you have n hosts, need to execute n2 times - -``` -ssh-keygen -ssh-copy-id -i ~/.ssh/id_rsa.pub ip_destination -``` - -- In all hosts, edit /etc/hosts to record all hosts' name and ip.The example is shown below. - -```bash -192.168.2.1 GPU001 -192.168.2.2 GPU002 -192.168.2.3 GPU003 -192.168.2.4 GPU004 -192.168.2.5 GPU005 -192.168.2.6 GPU006 -192.168.2.7 GPU007 -... -``` - -- restart ssh -``` -service ssh restart -``` - -## 1. Corpus Preprocessing -```bash -cd preprocessing -``` -following the `README.md`, preprocess original corpus to h5py+numpy - -## 2. Pretrain - -```bash -cd pretraining -``` -following the `README.md`, load the h5py generated by preprocess of step 1 to pretrain the model - -## 3. Finetune - -The checkpoint produced by this repo can replace `pytorch_model.bin` from [hfl/chinese-roberta-wwm-ext-large](https://huggingface.co/hfl/chinese-roberta-wwm-ext-large/tree/main) directly. Then use transfomers from Hugging Face to finetune downstream application. - -## Contributors -The repo is contributed by AI team from [Moore Threads](https://www.mthreads.com/). If you find any problems for pretraining, please file an issue or send an email to yehua.zhang@mthreads.com. At last, welcome any form of contribution! - -``` -@misc{ - title={A simple Chinese RoBERTa Example for Whole Word Masked}, - author={Yehua Zhang, Chen Zhang}, - year={2022} -} -``` diff --git a/examples/language/roberta/configs/colossalai_ddp.py b/examples/language/roberta/configs/colossalai_ddp.py deleted file mode 100644 index c3c59aa4079c..000000000000 --- a/examples/language/roberta/configs/colossalai_ddp.py +++ /dev/null @@ -1,4 +0,0 @@ -from colossalai.zero.shard_utils import TensorShardStrategy -from colossalai.nn.optimizer import FusedAdam - -clip_grad_norm = 1.0 diff --git a/examples/language/roberta/configs/colossalai_zero.py b/examples/language/roberta/configs/colossalai_zero.py deleted file mode 100644 index c5debdce0988..000000000000 --- a/examples/language/roberta/configs/colossalai_zero.py +++ /dev/null @@ -1,32 +0,0 @@ -from colossalai.zero.shard_utils import TensorShardStrategy -from colossalai.nn.optimizer import FusedAdam - -# fp16 = dict( -# mode=AMP_TYPE.TORCH, -# ) - -# seed = 2 -zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(), - reduce_scatter_bucket_size_mb=25, - fp32_reduce_scatter=False, - tensor_placement_policy="cuda", - gradient_predivide_factor=1.0, - reuse_fp16_shard=False), - optimizer_config=dict(gpu_margin_mem_ratio=0.8, - initial_scale=2**5, - min_scale=1, - growth_factor=2, - backoff_factor=0.5, - growth_interval=1000, - hysteresis=2, - max_scale=2**32)) - -# gradient_accumulation = 4 -clip_grad_norm = 1.0 -optimizer = dict( - type=FusedAdam, - lr=0.00015, - weight_decay=1e-2, -) - -# 64433 \ No newline at end of file diff --git a/examples/language/roberta/preprocessing/mask.cpp b/examples/language/roberta/preprocessing/mask.cpp deleted file mode 100644 index 8355c45cff0a..000000000000 --- a/examples/language/roberta/preprocessing/mask.cpp +++ /dev/null @@ -1,184 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace py = pybind11; - -const int32_t LONG_SENTENCE_LEN = 512; - -struct MaskedLMInstance { - int index; - std::string label; - MaskedLMInstance(int index, std::string label) { - this->index = index; - this->label = label; - } -}; - -auto get_new_segment(std::vector segment, std::vector segment_jieba, const std::vector chinese_vocab) { // const std::unordered_set &chinese_vocab - std::unordered_set seq_cws_dict; - for (auto word : segment_jieba) { - seq_cws_dict.insert(word); - } - int i = 0; - std::vector new_segment; - int segment_size = segment.size(); - while (i < segment_size) { - if (!chinese_vocab[i]) { //chinese_vocab.find(segment[i]) == chinese_vocab.end() - new_segment.emplace_back(segment[i]); - i += 1; - continue; - } - bool has_add = false; - for (int length = 3; length >= 1; length--) { - if (i + length > segment_size) { - continue; - } - std::string chinese_word = ""; - for (int j = i; j < i + length; j++) { - chinese_word += segment[j]; - } - if (seq_cws_dict.find(chinese_word) != seq_cws_dict.end()) { - new_segment.emplace_back(segment[i]); - for (int j = i + 1; j < i + length; j++) { - new_segment.emplace_back("##" + segment[j]); - } - i += length; - has_add = true; - break; - } - } - if (!has_add) { - new_segment.emplace_back(segment[i]); - i += 1; - } - } - - return new_segment; -} - -bool startsWith(const std::string& s, const std::string& sub) { - return s.find(sub) == 0 ? true : false; -} - -auto create_whole_masked_lm_predictions(std::vector &tokens, - const std::vector &original_tokens, - const std::vector &vocab_words, - std::map &vocab, - const int max_predictions_per_seq, - const double masked_lm_prob) { - // for (auto item : vocab) { - // std::cout << "key=" << std::string(py::str(item.first)) << ", " - // << "value=" << std::string(py::str(item.second)) << std::endl; - // } - std::vector > cand_indexes; - std::vector cand_temp; - int tokens_size = tokens.size(); - std::string prefix = "##"; - bool do_whole_masked = true; - - for (int i = 0; i < tokens_size; i++) { - if (tokens[i] == "[CLS]" || tokens[i] == "[SEP]") { - continue; - } - if (do_whole_masked && (cand_indexes.size() > 0) && (tokens[i].rfind(prefix, 0) == 0)) { - cand_temp.emplace_back(i); - } - else { - if (cand_temp.size() > 0) { - cand_indexes.emplace_back(cand_temp); - } - cand_temp.clear(); - cand_temp.emplace_back(i); - } - } - auto seed = std::chrono::system_clock::now().time_since_epoch().count(); - std::shuffle(cand_indexes.begin(), cand_indexes.end(), std::default_random_engine(seed)); - // for (auto i : cand_indexes) { - // for (auto j : i) { - // std::cout << tokens[j] << " "; - // } - // std::cout << std::endl; - // } - // for (auto i : output_tokens) { - // std::cout << i; - // } - // std::cout << std::endl; - - int num_to_predict = std::min(max_predictions_per_seq, - std::max(1, int(tokens_size * masked_lm_prob))); - // std::cout << num_to_predict << std::endl; - - std::set covered_indexes; - std::vector masked_lm_output(tokens_size, -1); - int vocab_words_len = vocab_words.size(); - std::default_random_engine e(seed); - std::uniform_real_distribution u1(0.0, 1.0); - std::uniform_int_distribution u2(0, vocab_words_len - 1); - int mask_cnt = 0; - std::vector output_tokens; - output_tokens = original_tokens; - - for (auto index_set : cand_indexes) { - if (mask_cnt > num_to_predict) { - break; - } - int index_set_size = index_set.size(); - if (mask_cnt + index_set_size > num_to_predict) { - continue; - } - bool is_any_index_covered = false; - for (auto index : index_set) { - if (covered_indexes.find(index) != covered_indexes.end()) { - is_any_index_covered = true; - break; - } - } - if (is_any_index_covered) { - continue; - } - for (auto index : index_set) { - - covered_indexes.insert(index); - std::string masked_token; - if (u1(e) < 0.8) { - masked_token = "[MASK]"; - } - else { - if (u1(e) < 0.5) { - masked_token = output_tokens[index]; - } - else { - int random_index = u2(e); - masked_token = vocab_words[random_index]; - } - } - // masked_lms.emplace_back(MaskedLMInstance(index, output_tokens[index])); - masked_lm_output[index] = vocab[output_tokens[index]]; - output_tokens[index] = masked_token; - mask_cnt++; - } - } - - // for (auto p : masked_lms) { - // masked_lm_output[p.index] = vocab[p.label]; - // } - return std::make_tuple(output_tokens, masked_lm_output); -} - -PYBIND11_MODULE(mask, m) { - m.def("create_whole_masked_lm_predictions", &create_whole_masked_lm_predictions); - m.def("get_new_segment", &get_new_segment); -} diff --git a/examples/language/roberta/pretraining/arguments.py b/examples/language/roberta/pretraining/arguments.py deleted file mode 100644 index 3a9370e00b0c..000000000000 --- a/examples/language/roberta/pretraining/arguments.py +++ /dev/null @@ -1,152 +0,0 @@ -import colossalai -from numpy import require - -__all__ = ['parse_args'] - - -def parse_args(): - parser = colossalai.get_default_parser() - - parser.add_argument( - '--lr', - type=float, - required=True, - help='initial learning rate') - parser.add_argument( - '--epoch', - type=int, - required=True, - help='number of epoch') - parser.add_argument( - '--data_path_prefix', - type=str, - required=True, - help="location of the train data corpus") - parser.add_argument( - '--eval_data_path_prefix', - type=str, - required=True, - help='location of the evaluation data corpus') - parser.add_argument( - '--tokenizer_path', - type=str, - required=True, - help='location of the tokenizer') - parser.add_argument( - '--max_seq_length', - type=int, - default=512, - help='sequence length') - parser.add_argument( - '--refresh_bucket_size', - type=int, - default=1, - help= - "This param makes sure that a certain task is repeated for this time steps to \ - optimise on the back propogation speed with APEX's DistributedDataParallel") - parser.add_argument( - "--max_predictions_per_seq", - "--max_pred", - default=80, - type=int, - help= - "The maximum number of masked tokens in a sequence to be predicted.") - parser.add_argument( - "--gradient_accumulation_steps", - default=1, - type=int, - help="accumulation_steps") - parser.add_argument( - "--train_micro_batch_size_per_gpu", - default=2, - type=int, - required=True, - help="train batch size") - parser.add_argument( - "--eval_micro_batch_size_per_gpu", - default=2, - type=int, - required=True, - help="eval batch size") - parser.add_argument( - "--num_workers", - default=8, - type=int, - help="") - parser.add_argument( - "--async_worker", - action='store_true', - help="") - parser.add_argument( - "--bert_config", - required=True, - type=str, - help="location of config.json") - parser.add_argument( - "--wandb", - action='store_true', - help="use wandb to watch model") - parser.add_argument( - "--wandb_project_name", - default='roberta', - help="wandb project name") - parser.add_argument( - "--log_interval", - default=100, - type=int, - help="report interval") - parser.add_argument( - "--log_path", - type=str, - required=True, - help="log file which records train step") - parser.add_argument( - "--tensorboard_path", - type=str, - required=True, - help="location of tensorboard file") - parser.add_argument( - "--colossal_config", - type=str, - required=True, - help="colossal config, which contains zero config and so on") - parser.add_argument( - "--ckpt_path", - type=str, - required=True, - help="location of saving checkpoint, which contains model and optimizer") - parser.add_argument( - '--seed', - type=int, - default=42, - help="random seed for initialization") - parser.add_argument( - '--vscode_debug', - action='store_true', - help="use vscode to debug") - parser.add_argument( - '--load_pretrain_model', - default='', - type=str, - help="location of model's checkpoin") - parser.add_argument( - '--load_optimizer_lr', - default='', - type=str, - help="location of checkpoint, which contains optimerzier, learning rate, epoch, shard and global_step") - parser.add_argument( - '--resume_train', - action='store_true', - help="whether resume training from a early checkpoint") - parser.add_argument( - '--mlm', - default='bert', - type=str, - help="model type, bert or deberta") - parser.add_argument( - '--checkpoint_activations', - action='store_true', - help="whether to use gradient checkpointing") - - args = parser.parse_args() - return args diff --git a/examples/tutorial/README.md b/examples/tutorial/README.md index f4843331fd54..0664d41fd359 100644 --- a/examples/tutorial/README.md +++ b/examples/tutorial/README.md @@ -29,7 +29,11 @@ quickly deploy large AI model training and inference, reducing large AI model tr - Fine-tuning and Inference for OPT [[code]](https://github.com/hpcaitech/ColossalAI/tree/main/examples/tutorial/opt) [[video]](https://www.youtube.com/watch?v=jbEFNVzl67Y) - Optimized AlphaFold [[code]](https://github.com/hpcaitech/ColossalAI/tree/main/examples/tutorial/fastfold) [[video]](https://www.youtube.com/watch?v=-zP13LfJP7w) - Optimized Stable Diffusion [[code]](https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion) [[video]](https://www.youtube.com/watch?v=8KHeUjjc-XQ) - + - ColossalChat: Cloning ChatGPT with a Complete RLHF Pipeline +[[code]](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat) +[[blog]](https://medium.com/@yangyou_berkeley/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b) +[[demo]](https://www.youtube.com/watch?v=HcTiHzApHm0) +[[video]](https://www.youtube.com/watch?v=-qFBZFmOJfg) ## Discussion diff --git a/examples/tutorial/auto_parallel/README.md b/examples/tutorial/auto_parallel/README.md index bb014b9067b2..6a12e0dd5a48 100644 --- a/examples/tutorial/auto_parallel/README.md +++ b/examples/tutorial/auto_parallel/README.md @@ -45,6 +45,7 @@ colossalai run --nproc_per_node 4 auto_parallel_with_resnet.py You should expect to the log like this. This log shows the edge cost on the computation graph as well as the sharding strategy for an operation. For example, `layer1_0_conv1 S01R = S01R X RR` means that the first dimension (batch) of the input and output is sharded while the weight is not sharded (S means sharded, R means replicated), simply equivalent to data parallel training. ![](https://raw.githubusercontent.com/hpcaitech/public_assets/main/examples/tutorial/auto-parallel%20demo.png) +**Note: This experimental feature has been tested on torch 1.12.1 and transformer 4.22.2. If you are using other versions, you may need to modify the code to make it work.** ### Auto-Checkpoint Tutorial diff --git a/examples/tutorial/auto_parallel/auto_ckpt_batchsize_test.py b/examples/tutorial/auto_parallel/auto_ckpt_batchsize_test.py index 5decfc695f6f..5a68aae18041 100644 --- a/examples/tutorial/auto_parallel/auto_ckpt_batchsize_test.py +++ b/examples/tutorial/auto_parallel/auto_ckpt_batchsize_test.py @@ -1,19 +1,14 @@ -import time -from argparse import ArgumentParser from copy import deepcopy from functools import partial -import matplotlib.pyplot as plt -import numpy as np import torch -import torch.multiprocessing as mp import torchvision.models as tm from bench_utils import bench, data_gen_resnet import colossalai from colossalai.auto_parallel.checkpoint import CheckpointSolverRotor from colossalai.fx import metainfo_trace, symbolic_trace -from colossalai.utils import free_port +from colossalai.testing import spawn def _benchmark(rank, world_size, port): @@ -50,9 +45,7 @@ def _benchmark(rank, world_size, port): def auto_activation_checkpoint_batchsize_benchmark(): - world_size = 1 - run_func_module = partial(_benchmark, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_benchmark, 1) if __name__ == "__main__": diff --git a/examples/tutorial/auto_parallel/auto_ckpt_solver_test.py b/examples/tutorial/auto_parallel/auto_ckpt_solver_test.py index ab0f2ef661df..aa5c47294a82 100644 --- a/examples/tutorial/auto_parallel/auto_ckpt_solver_test.py +++ b/examples/tutorial/auto_parallel/auto_ckpt_solver_test.py @@ -4,14 +4,13 @@ import matplotlib.pyplot as plt import torch -import torch.multiprocessing as mp import torchvision.models as tm from bench_utils import GPTLMLoss, bench_rotor, data_gen_gpt2, data_gen_resnet, gpt2_medium import colossalai from colossalai.auto_parallel.checkpoint import CheckpointSolverRotor from colossalai.fx import metainfo_trace, symbolic_trace -from colossalai.utils import free_port +from colossalai.testing import spawn def _benchmark(rank, world_size, port, args): @@ -77,8 +76,7 @@ def _benchmark(rank, world_size, port, args): def auto_activation_checkpoint_benchmark(args): world_size = 1 - run_func_module = partial(_benchmark, world_size=world_size, port=free_port(), args=args) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_benchmark, world_size, args=args) if __name__ == "__main__": diff --git a/examples/tutorial/auto_parallel/requirements.txt b/examples/tutorial/auto_parallel/requirements.txt index ce89e7c80070..cc61362ba6f9 100644 --- a/examples/tutorial/auto_parallel/requirements.txt +++ b/examples/tutorial/auto_parallel/requirements.txt @@ -1,7 +1,7 @@ -torch +torch==1.12.1 colossalai titans pulp datasets matplotlib -transformers +transformers==4.22.1 diff --git a/examples/tutorial/fastfold/FastFold b/examples/tutorial/fastfold/FastFold index 867587b3aa4e..eba496808a91 160000 --- a/examples/tutorial/fastfold/FastFold +++ b/examples/tutorial/fastfold/FastFold @@ -1 +1 @@ -Subproject commit 867587b3aa4e43bdaf64f9910127842f1dfbfebd +Subproject commit eba496808a91bbcd9661cf832349a418b197015f diff --git a/examples/tutorial/new_api/README.md b/examples/tutorial/new_api/README.md new file mode 100644 index 000000000000..cec88f41caf1 --- /dev/null +++ b/examples/tutorial/new_api/README.md @@ -0,0 +1,5 @@ +# New API Features + +**The New API is not officially released yet.** + +This folder contains some of the demonstrations of the new API. The new API is still under intensive development and will be released soon. diff --git a/examples/tutorial/new_api/cifar_resnet/.gitignore b/examples/tutorial/new_api/cifar_resnet/.gitignore new file mode 100644 index 000000000000..a79cf5236c08 --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/.gitignore @@ -0,0 +1,4 @@ +data +checkpoint +ckpt-fp16 +ckpt-fp32 diff --git a/examples/tutorial/new_api/cifar_resnet/README.md b/examples/tutorial/new_api/cifar_resnet/README.md new file mode 100644 index 000000000000..4ed86aa7a0ad --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/README.md @@ -0,0 +1,56 @@ +# Train ResNet on CIFAR-10 from scratch + +## 🚀 Quick Start + +This example provides a training script and an evaluation script. The training script provides an example of training ResNet on CIFAR10 dataset from scratch. + +- Training Arguments + - `-p`, `--plugin`: Plugin to use. Choices: `torch_ddp`, `torch_ddp_fp16`, `low_level_zero`. Defaults to `torch_ddp`. + - `-r`, `--resume`: Resume from checkpoint file path. Defaults to `-1`, which means not resuming. + - `-c`, `--checkpoint`: The folder to save checkpoints. Defaults to `./checkpoint`. + - `-i`, `--interval`: Epoch interval to save checkpoints. Defaults to `5`. If set to `0`, no checkpoint will be saved. + - `--target_acc`: Target accuracy. Raise exception if not reached. Defaults to `None`. + +- Eval Arguments + - `-e`, `--epoch`: select the epoch to evaluate + - `-c`, `--checkpoint`: the folder where checkpoints are found + +### Install requirements + +```bash +pip install -r requirements.txt +``` + +### Train + +```bash +# train with torch DDP with fp32 +colossalai run --nproc_per_node 2 train.py -c ./ckpt-fp32 + +# train with torch DDP with mixed precision training +colossalai run --nproc_per_node 2 train.py -c ./ckpt-fp16 -p torch_ddp_fp16 + +# train with low level zero +colossalai run --nproc_per_node 2 train.py -c ./ckpt-low_level_zero -p low_level_zero +``` + +### Eval + +```bash +# evaluate fp32 training +python eval.py -c ./ckpt-fp32 -e 80 + +# evaluate fp16 mixed precision training +python eval.py -c ./ckpt-fp16 -e 80 + +# evaluate low level zero training +python eval.py -c ./ckpt-low_level_zero -e 80 +``` + +Expected accuracy performance will be: + +| Model | Single-GPU Baseline FP32 | Booster DDP with FP32 | Booster DDP with FP16 | Booster Low Level Zero | +| --------- | ------------------------ | --------------------- | --------------------- | ---------------------- | +| ResNet-18 | 85.85% | 84.91% | 85.46% | 84.50% | + +**Note: the baseline is adapted from the [script](https://pytorch-tutorial.readthedocs.io/en/latest/tutorial/chapter03_intermediate/3_2_2_cnn_resnet_cifar10/) to use `torchvision.models.resnet18`** diff --git a/examples/tutorial/new_api/cifar_resnet/eval.py b/examples/tutorial/new_api/cifar_resnet/eval.py new file mode 100644 index 000000000000..657708ec3ff2 --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/eval.py @@ -0,0 +1,48 @@ +import argparse + +import torch +import torch.nn as nn +import torchvision +import torchvision.transforms as transforms + +# ============================== +# Parse Arguments +# ============================== +parser = argparse.ArgumentParser() +parser.add_argument('-e', '--epoch', type=int, default=80, help="resume from the epoch's checkpoint") +parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory") +args = parser.parse_args() + +# ============================== +# Prepare Test Dataset +# ============================== +# CIFAR-10 dataset +test_dataset = torchvision.datasets.CIFAR10(root='./data/', train=False, transform=transforms.ToTensor()) + +# Data loader +test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False) + +# ============================== +# Load Model +# ============================== +model = torchvision.models.resnet18(num_classes=10).cuda() +state_dict = torch.load(f'{args.checkpoint}/model_{args.epoch}.pth') +model.load_state_dict(state_dict) + +# ============================== +# Run Evaluation +# ============================== +model.eval() + +with torch.no_grad(): + correct = 0 + total = 0 + for images, labels in test_loader: + images = images.cuda() + labels = labels.cuda() + outputs = model(images) + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + print('Accuracy of the model on the test images: {} %'.format(100 * correct / total)) diff --git a/examples/tutorial/new_api/cifar_resnet/requirements.txt b/examples/tutorial/new_api/cifar_resnet/requirements.txt new file mode 100644 index 000000000000..85522f4129c4 --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/requirements.txt @@ -0,0 +1,4 @@ +colossalai +torch +torchvision +tqdm diff --git a/examples/tutorial/new_api/cifar_resnet/test_ci.sh b/examples/tutorial/new_api/cifar_resnet/test_ci.sh new file mode 100755 index 000000000000..3954b84ff1ba --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/test_ci.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -xe + +export DATA=/data/scratch/cifar-10 + +pip install -r requirements.txt + +for plugin in "torch_ddp" "torch_ddp_fp16" "low_level_zero"; do + colossalai run --nproc_per_node 4 train.py --interval 0 --target_acc 0.84 --plugin $plugin +done diff --git a/examples/tutorial/new_api/cifar_resnet/train.py b/examples/tutorial/new_api/cifar_resnet/train.py new file mode 100644 index 000000000000..fe0dabf08377 --- /dev/null +++ b/examples/tutorial/new_api/cifar_resnet/train.py @@ -0,0 +1,204 @@ +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.nn as nn +import torchvision +import torchvision.transforms as transforms +from torch.optim import Optimizer +from torch.optim.lr_scheduler import MultiStepLR +from torch.utils.data import DataLoader +from tqdm import tqdm + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 80 +LEARNING_RATE = 1e-3 + + +def build_dataloader(batch_size: int, coordinator: DistCoordinator, plugin: DPPluginBase): + # transform + transform_train = transforms.Compose( + [transforms.Pad(4), + transforms.RandomHorizontalFlip(), + transforms.RandomCrop(32), + transforms.ToTensor()]) + transform_test = transforms.ToTensor() + + # CIFAR-10 dataset + data_path = os.environ.get('DATA', './data') + with coordinator.priority_execution(): + train_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=True, + transform=transform_train, + download=True) + test_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=False, + transform=transform_test, + download=True) + + # Data loader + train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True) + test_dataloader = plugin.prepare_dataloader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False) + return train_dataloader, test_dataloader + + +@torch.no_grad() +def evaluate(model: nn.Module, test_dataloader: DataLoader, coordinator: DistCoordinator) -> float: + model.eval() + correct = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + total = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + for images, labels in test_dataloader: + images = images.cuda() + labels = labels.cuda() + outputs = model(images) + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + dist.all_reduce(correct) + dist.all_reduce(total) + accuracy = correct.item() / total.item() + if coordinator.is_master(): + print(f'Accuracy of the model on the test images: {accuracy * 100:.2f} %') + return accuracy + + +def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: nn.Module, train_dataloader: DataLoader, + booster: Booster, coordinator: DistCoordinator): + model.train() + with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar: + for images, labels in pbar: + images = images.cuda() + labels = labels.cuda() + # Forward pass + outputs = model(images) + loss = criterion(outputs, labels) + + # Backward and optimize + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + # Print log info + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + # FIXME(ver217): gemini is not supported resnet now + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'low_level_zero'], + help="plugin to use") + parser.add_argument('-r', '--resume', type=int, default=-1, help="resume from the epoch's checkpoint") + parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory") + parser.add_argument('-i', '--interval', type=int, default=5, help="interval of saving checkpoint") + parser.add_argument('--target_acc', + type=float, + default=None, + help="target accuracy. Raise exception if not reached") + args = parser.parse_args() + + # ============================== + # Prepare Checkpoint Directory + # ============================== + if args.interval > 0: + Path(args.checkpoint).mkdir(parents=True, exist_ok=True) + + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}) + coordinator = DistCoordinator() + + # update the learning rate with linear scaling + # old_gpu_num / old_lr = new_gpu_num / new_lr + global LEARNING_RATE + LEARNING_RATE *= coordinator.world_size + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + train_dataloader, test_dataloader = build_dataloader(100, coordinator, plugin) + + # ==================================== + # Prepare model, optimizer, criterion + # ==================================== + # resent50 + model = torchvision.models.resnet18(num_classes=10) + + # Loss and optimizer + criterion = nn.CrossEntropyLoss() + optimizer = HybridAdam(model.parameters(), lr=LEARNING_RATE) + + # lr scheduler + lr_scheduler = MultiStepLR(optimizer, milestones=[20, 40, 60, 80], gamma=1 / 3) + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, criterion, _, lr_scheduler = booster.boost(model, + optimizer, + criterion=criterion, + lr_scheduler=lr_scheduler) + + # ============================== + # Resume from checkpoint + # ============================== + if args.resume >= 0: + booster.load_model(model, f'{args.checkpoint}/model_{args.resume}.pth') + booster.load_optimizer(optimizer, f'{args.checkpoint}/optimizer_{args.resume}.pth') + booster.load_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{args.resume}.pth') + + # ============================== + # Train model + # ============================== + start_epoch = args.resume if args.resume >= 0 else 0 + for epoch in range(start_epoch, NUM_EPOCHS): + train_epoch(epoch, model, optimizer, criterion, train_dataloader, booster, coordinator) + lr_scheduler.step() + + # save checkpoint + if args.interval > 0 and (epoch + 1) % args.interval == 0: + booster.save_model(model, f'{args.checkpoint}/model_{epoch + 1}.pth') + booster.save_optimizer(optimizer, f'{args.checkpoint}/optimizer_{epoch + 1}.pth') + booster.save_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{epoch + 1}.pth') + + accuracy = evaluate(model, test_dataloader, coordinator) + if args.target_acc is not None: + assert accuracy >= args.target_acc, f'Accuracy {accuracy} is lower than target accuracy {args.target_acc}' + + +if __name__ == '__main__': + main() diff --git a/examples/tutorial/new_api/cifar_vit/README.md b/examples/tutorial/new_api/cifar_vit/README.md new file mode 100644 index 000000000000..fa76447c508f --- /dev/null +++ b/examples/tutorial/new_api/cifar_vit/README.md @@ -0,0 +1,37 @@ +# Train ViT on CIFAR-10 from scratch + +## 🚀 Quick Start + +This example provides a training script, which provides an example of training ViT on CIFAR10 dataset from scratch. + +- Training Arguments + - `-p`, `--plugin`: Plugin to use. Choices: `torch_ddp`, `torch_ddp_fp16`, `low_level_zero`. Defaults to `torch_ddp`. + - `-r`, `--resume`: Resume from checkpoint file path. Defaults to `-1`, which means not resuming. + - `-c`, `--checkpoint`: The folder to save checkpoints. Defaults to `./checkpoint`. + - `-i`, `--interval`: Epoch interval to save checkpoints. Defaults to `5`. If set to `0`, no checkpoint will be saved. + - `--target_acc`: Target accuracy. Raise exception if not reached. Defaults to `None`. + +### Install requirements + +```bash +pip install -r requirements.txt +``` + +### Train + +```bash +# train with torch DDP with fp32 +colossalai run --nproc_per_node 4 train.py -c ./ckpt-fp32 + +# train with torch DDP with mixed precision training +colossalai run --nproc_per_node 4 train.py -c ./ckpt-fp16 -p torch_ddp_fp16 + +# train with low level zero +colossalai run --nproc_per_node 4 train.py -c ./ckpt-low_level_zero -p low_level_zero +``` + +Expected accuracy performance will be: + +| Model | Single-GPU Baseline FP32 | Booster DDP with FP32 | Booster DDP with FP16 | Booster Low Level Zero | +| --------- | ------------------------ | --------------------- | --------------------- | ---------------------- | +| ViT | 83.00% | 84.03% | 84.00% | 84.43% | diff --git a/examples/tutorial/new_api/cifar_vit/requirements.txt b/examples/tutorial/new_api/cifar_vit/requirements.txt new file mode 100644 index 000000000000..6d53ce7b5a7d --- /dev/null +++ b/examples/tutorial/new_api/cifar_vit/requirements.txt @@ -0,0 +1,5 @@ +colossalai +timm +torch +torchvision +tqdm diff --git a/examples/tutorial/new_api/cifar_vit/test_ci.sh b/examples/tutorial/new_api/cifar_vit/test_ci.sh new file mode 100755 index 000000000000..43239d400586 --- /dev/null +++ b/examples/tutorial/new_api/cifar_vit/test_ci.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -xe + +export DATA=/data/scratch/cifar-10 + +pip install -r requirements.txt + +for plugin in "torch_ddp" "torch_ddp_fp16" "low_level_zero"; do + colossalai run --nproc_per_node 4 train.py --interval 0 --target_acc 0.83 --plugin $plugin +done diff --git a/examples/tutorial/new_api/cifar_vit/train.py b/examples/tutorial/new_api/cifar_vit/train.py new file mode 100644 index 000000000000..82a8f2ed97e4 --- /dev/null +++ b/examples/tutorial/new_api/cifar_vit/train.py @@ -0,0 +1,219 @@ +import argparse +import os +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.nn as nn +import torchvision +import torchvision.transforms as transforms +from timm.models.vision_transformer import _cfg, _create_vision_transformer +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from tqdm import tqdm + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase +from colossalai.cluster import DistCoordinator +from colossalai.nn.lr_scheduler import LinearWarmupLR +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 60 +WARMUP_EPOCHS = 5 +LEARNING_RATE = 1e-3 + + +def vit_cifar(**kwargs): + pretrained_cfg = _cfg(num_classes=10, input_size=(3, 32, 32), crop_pct=1.0) + model_kwargs = dict(patch_size=4, embed_dim=512, depth=6, num_heads=8, drop_rate=0.1, mlp_ratio=1.0, **kwargs) + model = _create_vision_transformer('vit_cifar', pretrained_cfg=pretrained_cfg, **model_kwargs) + return model + + +def build_dataloader(batch_size: int, coordinator: DistCoordinator, plugin: DPPluginBase): + # transform + transform_train = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)), + ]) + transform_test = transforms.Compose([ + transforms.Resize(32), + transforms.ToTensor(), + transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)), + ]) + + # CIFAR-10 dataset + data_path = os.environ.get('DATA', './data') + with coordinator.priority_execution(): + train_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=True, + transform=transform_train, + download=True) + test_dataset = torchvision.datasets.CIFAR10(root=data_path, + train=False, + transform=transform_test, + download=True) + + # Data loader + train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True) + test_dataloader = plugin.prepare_dataloader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False) + return train_dataloader, test_dataloader + + +@torch.no_grad() +def evaluate(model: nn.Module, test_dataloader: DataLoader, coordinator: DistCoordinator) -> float: + model.eval() + correct = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + total = torch.zeros(1, dtype=torch.int64, device=get_current_device()) + for images, labels in test_dataloader: + images = images.cuda() + labels = labels.cuda() + outputs = model(images) + _, predicted = torch.max(outputs.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + dist.all_reduce(correct) + dist.all_reduce(total) + accuracy = correct.item() / total.item() + if coordinator.is_master(): + print(f'Accuracy of the model on the test images: {accuracy * 100:.2f} %') + return accuracy + + +def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: nn.Module, train_dataloader: DataLoader, + booster: Booster, coordinator: DistCoordinator): + model.train() + with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar: + for images, labels in pbar: + images = images.cuda() + labels = labels.cuda() + # Forward pass + outputs = model(images) + loss = criterion(outputs, labels) + + # Backward and optimize + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + # Print log info + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + # FIXME(ver217): gemini is not supported resnet now + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'low_level_zero'], + help="plugin to use") + parser.add_argument('-r', '--resume', type=int, default=-1, help="resume from the epoch's checkpoint") + parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory") + parser.add_argument('-i', '--interval', type=int, default=5, help="interval of saving checkpoint") + parser.add_argument('--target_acc', + type=float, + default=None, + help="target accuracy. Raise exception if not reached") + args = parser.parse_args() + + # ============================== + # Prepare Checkpoint Directory + # ============================== + if args.interval > 0: + Path(args.checkpoint).mkdir(parents=True, exist_ok=True) + + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}) + coordinator = DistCoordinator() + + # update the learning rate with linear scaling + # old_gpu_num / old_lr = new_gpu_num / new_lr + global LEARNING_RATE + LEARNING_RATE *= coordinator.world_size + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + train_dataloader, test_dataloader = build_dataloader(512, coordinator, plugin) + + # ==================================== + # Prepare model, optimizer, criterion + # ==================================== + # resent50 + model = torchvision.models.resnet18(num_classes=10) + + # Loss and optimizer + criterion = nn.CrossEntropyLoss() + optimizer = HybridAdam(model.parameters(), lr=LEARNING_RATE) + + # lr scheduler + lr_scheduler = LinearWarmupLR(optimizer, NUM_EPOCHS, WARMUP_EPOCHS) + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, criterion, train_dataloader, lr_scheduler = booster.boost(model, + optimizer, + criterion=criterion, + dataloader=train_dataloader, + lr_scheduler=lr_scheduler) + + # ============================== + # Resume from checkpoint + # ============================== + if args.resume >= 0: + booster.load_model(model, f'{args.checkpoint}/model_{args.resume}.pth') + booster.load_optimizer(optimizer, f'{args.checkpoint}/optimizer_{args.resume}.pth') + booster.load_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{args.resume}.pth') + + # ============================== + # Train model + # ============================== + start_epoch = args.resume if args.resume >= 0 else 0 + for epoch in range(start_epoch, NUM_EPOCHS): + train_epoch(epoch, model, optimizer, criterion, train_dataloader, booster, coordinator) + lr_scheduler.step() + + # save checkpoint + if args.interval > 0 and (epoch + 1) % args.interval == 0: + booster.save_model(model, f'{args.checkpoint}/model_{epoch + 1}.pth') + booster.save_optimizer(optimizer, f'{args.checkpoint}/optimizer_{epoch + 1}.pth') + booster.save_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{epoch + 1}.pth') + + accuracy = evaluate(model, test_dataloader, coordinator) + if args.target_acc is not None: + assert accuracy >= args.target_acc, f'Accuracy {accuracy} is lower than target accuracy {args.target_acc}' + + +if __name__ == '__main__': + main() diff --git a/examples/tutorial/new_api/glue_bert/README.md b/examples/tutorial/new_api/glue_bert/README.md new file mode 100644 index 000000000000..0030eead9f5b --- /dev/null +++ b/examples/tutorial/new_api/glue_bert/README.md @@ -0,0 +1,39 @@ +# Finetune BERT on GLUE + +## 🚀 Quick Start + +This example provides a training script, which provides an example of finetuning BERT on GLUE dataset. + +- Training Arguments + - `-t`, `--task`: GLUE task to run. Defaults to `mrpc`. + - `-p`, `--plugin`: Plugin to use. Choices: `torch_ddp`, `torch_ddp_fp16`, `gemini`, `low_level_zero`. Defaults to `torch_ddp`. + - `--target_f1`: Target f1 score. Raise exception if not reached. Defaults to `None`. + + +### Install requirements + +```bash +pip install -r requirements.txt +``` + +### Train + +```bash +# train with torch DDP with fp32 +colossalai run --nproc_per_node 4 finetune.py + +# train with torch DDP with mixed precision training +colossalai run --nproc_per_node 4 finetune.py -p torch_ddp_fp16 + +# train with gemini +colossalai run --nproc_per_node 4 finetune.py -p gemini + +# train with low level zero +colossalai run --nproc_per_node 4 finetune.py -p low_level_zero +``` + +Expected F1-score will be: + +| Model | Single-GPU Baseline FP32 | Booster DDP with FP32 | Booster DDP with FP16 | Booster Gemini | Booster Low Level Zero | +| ----------------- | ------------------------ | --------------------- | --------------------- |--------------- | ---------------------- | +| bert-base-uncased | 0.86 | 0.88 | 0.87 | 0.88 | 0.89 | diff --git a/examples/tutorial/new_api/glue_bert/data.py b/examples/tutorial/new_api/glue_bert/data.py new file mode 100644 index 000000000000..981cedcca8c2 --- /dev/null +++ b/examples/tutorial/new_api/glue_bert/data.py @@ -0,0 +1,127 @@ +import datasets +from transformers import AutoTokenizer, PreTrainedTokenizer + +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase + + +class GLUEDataBuilder: + + task_text_field_map = { + "cola": ["sentence"], + "sst2": ["sentence"], + "mrpc": ["sentence1", "sentence2"], + "qqp": ["question1", "question2"], + "stsb": ["sentence1", "sentence2"], + "mnli": ["premise", "hypothesis"], + "qnli": ["question", "sentence"], + "rte": ["sentence1", "sentence2"], + "wnli": ["sentence1", "sentence2"], + "ax": ["premise", "hypothesis"], + } + + glue_task_num_labels = { + "cola": 2, + "sst2": 2, + "mrpc": 2, + "qqp": 2, + "stsb": 1, + "mnli": 3, + "qnli": 2, + "rte": 2, + "wnli": 2, + "ax": 3, + } + + loader_columns = [ + "datasets_idx", + "input_ids", + "token_type_ids", + "attention_mask", + "start_positions", + "end_positions", + "labels", + ] + + def __init__( + self, + model_name_or_path: str, + plugin: DPPluginBase, + task_name: str = "mrpc", + max_seq_length: int = 128, + train_batch_size: int = 32, + eval_batch_size: int = 32, + **kwargs, + ): + super().__init__() + self.model_name_or_path = model_name_or_path + self.task_name = task_name + self.max_seq_length = max_seq_length + self.train_batch_size = train_batch_size + self.eval_batch_size = eval_batch_size + self.plugin = plugin + + self.text_fields = self.task_text_field_map[task_name] + self.num_labels = self.glue_task_num_labels[task_name] + self.tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + self.setup() + + def setup(self): + self.dataset = datasets.load_dataset("glue", self.task_name) + + for split in self.dataset.keys(): + self.dataset[split] = self.dataset[split].map( + self.convert_to_features, + batched=True, + remove_columns=["label"], + ) + self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns] + self.dataset[split].set_format(type="torch", columns=self.columns) + + self.eval_splits = [x for x in self.dataset.keys() if "validation" in x] + + def prepare_data(self): + datasets.load_dataset("glue", self.task_name) + AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True) + + def train_dataloader(self): + return self.plugin.prepare_dataloader(self.dataset["train"], + batch_size=self.train_batch_size, + shuffle=True, + drop_last=True) + + def val_dataloader(self): + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["validation"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def test_dataloader(self): + if len(self.eval_splits) == 1: + return self.plugin.prepare_dataloader(self.dataset["test"], batch_size=self.eval_batch_size) + elif len(self.eval_splits) > 1: + return [ + self.plugin.prepare_dataloader(self.dataset[x], batch_size=self.eval_batch_size) + for x in self.eval_splits + ] + + def convert_to_features(self, example_batch): + + # Either encode single sentence or sentence pairs + if len(self.text_fields) > 1: + texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]])) + else: + texts_or_text_pairs = example_batch[self.text_fields[0]] + + # Tokenize the text/text pairs + features = self.tokenizer.batch_encode_plus(texts_or_text_pairs, + max_length=self.max_seq_length, + padding='max_length', + truncation=True) + + # Rename label to labels to make it easier to pass to model forward + features["labels"] = example_batch["label"] + + return features diff --git a/examples/tutorial/new_api/glue_bert/finetune.py b/examples/tutorial/new_api/glue_bert/finetune.py new file mode 100644 index 000000000000..63bdfc5d02cf --- /dev/null +++ b/examples/tutorial/new_api/glue_bert/finetune.py @@ -0,0 +1,198 @@ +import argparse +from typing import List, Union + +import datasets +import torch +import torch.distributed as dist +import torch.nn as nn +from data import GLUEDataBuilder +from torch.optim import Optimizer +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoConfig, BertForSequenceClassification, get_linear_schedule_with_warmup + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin +from colossalai.cluster import DistCoordinator +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +# ============================== +# Prepare Hyperparameters +# ============================== +NUM_EPOCHS = 3 +BATCH_SIZE = 32 +LEARNING_RATE = 2.4e-5 +WEIGHT_DECAY = 0.01 +WARMUP_FRACTION = 0.1 + + +def move_to_cuda(batch): + return {k: v.cuda() for k, v in batch.items()} + + +@torch.no_grad() +def evaluate(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoader]], num_labels: int, task_name: str, + eval_splits: List[str], coordinator: DistCoordinator): + metric = datasets.load_metric("glue", task_name, process_id=coordinator.rank, num_process=coordinator.world_size) + model.eval() + + def evaluate_subset(dataloader: DataLoader): + accum_loss = torch.zeros(1, device=get_current_device()) + for batch in dataloader: + batch = move_to_cuda(batch) + outputs = model(**batch) + val_loss, logits = outputs[:2] + accum_loss.add_(val_loss) + + if num_labels > 1: + preds = torch.argmax(logits, axis=1) + elif num_labels == 1: + preds = logits.squeeze() + + labels = batch["labels"] + + metric.add_batch(predictions=preds, references=labels) + + results = metric.compute() + dist.all_reduce(accum_loss.div_(len(dataloader))) + if coordinator.is_master(): + results['loss'] = accum_loss.item() / coordinator.world_size + return results + + if isinstance(test_dataloader, DataLoader): + return evaluate_subset(test_dataloader) + else: + assert len(test_dataloader) == len(eval_splits) + final_results = {} + for split, sub_loader in zip(eval_splits, test_dataloader): + results = evaluate_subset(sub_loader) + final_results.update({f'{k}_{split}': v for k, v in results.items()}) + return final_results + + +def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, lr_scheduler, train_dataloader: DataLoader, + booster: Booster, coordinator: DistCoordinator): + model.train() + with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar: + for batch in pbar: + # Forward pass + batch = move_to_cuda(batch) + outputs = model(**batch) + loss = outputs[0] + + # Backward and optimize + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + lr_scheduler.step() + + # Print log info + pbar.set_postfix({'loss': loss.item()}) + + +def main(): + # ============================== + # Parse Arguments + # ============================== + parser = argparse.ArgumentParser() + parser.add_argument('-t', '--task', default='mrpc', help="GLUE task to run") + parser.add_argument('-p', + '--plugin', + type=str, + default='torch_ddp', + choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'], + help="plugin to use") + parser.add_argument('--target_f1', type=float, default=None, help="target f1 score. Raise exception if not reached") + args = parser.parse_args() + + # ============================== + # Launch Distributed Environment + # ============================== + colossalai.launch_from_torch(config={}, seed=42) + coordinator = DistCoordinator() + + # local_batch_size = BATCH_SIZE // coordinator.world_size + lr = LEARNING_RATE * coordinator.world_size + model_name = 'bert-base-uncased' + + # ============================== + # Instantiate Plugin and Booster + # ============================== + booster_kwargs = {} + if args.plugin == 'torch_ddp_fp16': + booster_kwargs['mixed_precision'] = 'fp16' + if args.plugin.startswith('torch_ddp'): + plugin = TorchDDPPlugin() + elif args.plugin == 'gemini': + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5) + elif args.plugin == 'low_level_zero': + plugin = LowLevelZeroPlugin(initial_scale=2**5) + + booster = Booster(plugin=plugin, **booster_kwargs) + + # ============================== + # Prepare Dataloader + # ============================== + data_builder = GLUEDataBuilder(model_name, + plugin, + args.task, + train_batch_size=BATCH_SIZE, + eval_batch_size=BATCH_SIZE) + train_dataloader = data_builder.train_dataloader() + test_dataloader = data_builder.test_dataloader() + + # ==================================== + # Prepare model, optimizer + # ==================================== + # bert pretrained model + config = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels) + model = BertForSequenceClassification.from_pretrained(model_name, config=config) + + # optimizer + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], + "weight_decay": WEIGHT_DECAY, + }, + { + "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], + "weight_decay": 0.0, + }, + ] + + optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, eps=1e-8) + + # lr scheduler + total_steps = len(train_dataloader) * NUM_EPOCHS + num_warmup_steps = int(WARMUP_FRACTION * total_steps) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_steps, + ) + + # ============================== + # Boost with ColossalAI + # ============================== + model, optimizer, _, _, lr_scheduler = booster.boost(model, optimizer, lr_scheduler=lr_scheduler) + + # ============================== + # Train model + # ============================== + for epoch in range(NUM_EPOCHS): + train_epoch(epoch, model, optimizer, lr_scheduler, train_dataloader, booster, coordinator) + + results = evaluate(model, test_dataloader, data_builder.num_labels, args.task, data_builder.eval_splits, + coordinator) + + if coordinator.is_master(): + print(results) + if args.target_f1 is not None and 'f1' in results: + assert results['f1'] >= args.target_f1, f'f1 score {results["f1"]} is lower than target {args.target_f1}' + + +if __name__ == '__main__': + main() diff --git a/examples/tutorial/new_api/glue_bert/requirements.txt b/examples/tutorial/new_api/glue_bert/requirements.txt new file mode 100644 index 000000000000..950c2d378f08 --- /dev/null +++ b/examples/tutorial/new_api/glue_bert/requirements.txt @@ -0,0 +1,7 @@ +colossalai +datasets +torch +tqdm +transformers +scipy +scikit-learn diff --git a/examples/tutorial/new_api/glue_bert/test_ci.sh b/examples/tutorial/new_api/glue_bert/test_ci.sh new file mode 100755 index 000000000000..c2c097f8d026 --- /dev/null +++ b/examples/tutorial/new_api/glue_bert/test_ci.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -xe + +pip install -r requirements.txt + +for plugin in "torch_ddp" "torch_ddp_fp16" "gemini" "low_level_zero"; do + torchrun --standalone --nproc_per_node 4 finetune.py --target_f1 0.86 --plugin $plugin +done diff --git a/examples/tutorial/new_api/test_ci.sh b/examples/tutorial/new_api/test_ci.sh new file mode 100644 index 000000000000..a08844dbe5fa --- /dev/null +++ b/examples/tutorial/new_api/test_ci.sh @@ -0,0 +1,6 @@ +#!/bin/bash +set -xe + +# FIXME(ver217): only run bert finetune to save time + +cd glue_bert && bash ./test_ci.sh && cd .. diff --git a/examples/tutorial/opt/opt/colossalai_zero.py b/examples/tutorial/opt/opt/colossalai_zero.py index 833745f3e8d8..7c2c152450c5 100644 --- a/examples/tutorial/opt/opt/colossalai_zero.py +++ b/examples/tutorial/opt/opt/colossalai_zero.py @@ -1,4 +1,8 @@ -from colossalai.zero.shard_utils import TensorShardStrategy +try: + from colossalai.zero.shard_utils import TensorShardStrategy +except ImportError: + # colossalai > 0.2.8 + from colossalai.zero.legacy import TensorShardStrategy zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(), tensor_placement_policy="auto", diff --git a/examples/tutorial/opt/opt/requirements.txt b/examples/tutorial/opt/opt/requirements.txt index c34df7992d3f..d0ed2c717aee 100644 --- a/examples/tutorial/opt/opt/requirements.txt +++ b/examples/tutorial/opt/opt/requirements.txt @@ -4,3 +4,4 @@ datasets >= 1.8.0 sentencepiece != 0.1.92 protobuf accelerate == 0.13.2 +transformers diff --git a/examples/tutorial/opt/opt/run_clm.py b/examples/tutorial/opt/opt/run_clm.py index c4f576cb18aa..fdc86adab665 100755 --- a/examples/tutorial/opt/opt/run_clm.py +++ b/examples/tutorial/opt/opt/run_clm.py @@ -30,24 +30,13 @@ import datasets import torch import torch.distributed as dist +import transformers from accelerate.utils import set_seed from context import barrier_context from datasets import load_dataset from packaging import version from torch.utils.data import DataLoader from tqdm.auto import tqdm - -import colossalai -import transformers -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.logging import disable_existing_loggers, get_dist_logger -from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP -from colossalai.tensor import ProcessGroup -from colossalai.utils import get_current_device, get_dataloader -from colossalai.utils.model.colo_init_context import ColoInitContext from transformers import ( CONFIG_MAPPING, MODEL_MAPPING, @@ -61,6 +50,15 @@ ) from transformers.utils.versions import require_version +import colossalai +from colossalai.context import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.nn.optimizer import HybridAdam +from colossalai.tensor import ProcessGroup +from colossalai.utils import get_current_device, get_dataloader +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer + require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys()) @@ -415,7 +413,11 @@ def main(): cai_version = colossalai.__version__ logger.info(f'using Colossal-AI version {cai_version}') if version.parse(cai_version) > version.parse("0.1.10"): - from colossalai.nn.parallel import GeminiDDP + try: + from colossalai.nn.parallel import GeminiDDP + except ImportError: + # this works for unreleased main branch, and this may be released on 0.2.9 + from colossalai.zero import GeminiDDP model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, pin_memory=True) elif version.parse(cai_version) <= version.parse("0.1.10") and version.parse(cai_version) >= version.parse("0.1.9"): from colossalai.gemini import ChunkManager, GeminiManager diff --git a/examples/tutorial/opt/opt/test_ci.sh b/examples/tutorial/opt/opt/test_ci.sh new file mode 100755 index 000000000000..e505da1364de --- /dev/null +++ b/examples/tutorial/opt/opt/test_ci.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +set -xue + +pip install -r requirements.txt + +BS=8 +MEMCAP=0 +GPUNUM=2 +MODLE="facebook/opt-125m" + +torchrun \ + --nproc_per_node ${GPUNUM} \ + --master_port 19198 \ + run_clm.py \ + -s \ + --output_dir $PWD \ + --mem_cap ${MEMCAP} \ + --model_name_or_path ${MODLE} \ + --per_device_train_batch_size ${BS} \ + --num_train_epochs 1 diff --git a/examples/tutorial/opt/test_ci.sh b/examples/tutorial/opt/test_ci.sh new file mode 100755 index 000000000000..8341bb10510f --- /dev/null +++ b/examples/tutorial/opt/test_ci.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +cd opt && bash test_ci.sh diff --git a/op_builder/README.md b/op_builder/README.md index b7ac6107300c..9c33a4a328d7 100644 --- a/op_builder/README.md +++ b/op_builder/README.md @@ -15,8 +15,8 @@ Method 2 is good because it allows the user to only build the kernel they actual ## PyTorch Extensions in Colossal-AI -The project DeepSpeed (https://github.com/microsoft/DeepSpeed) has proposed a [solution](https://github.com/microsoft/DeepSpeed/tree/master/op_builder)) to support kernel-build during either installation or runtime. -We have adapted from DeepSpeed's solution to build extensions. The extension build requries two main functions from PyTorch: +The project [DeepSpeed](https://github.com/microsoft/DeepSpeed) has proposed a [solution](https://github.com/microsoft/DeepSpeed/tree/master/op_builder) to support kernel-build during either installation or runtime. +We have adapted from DeepSpeed's solution to build extensions. The extension build requires two main functions from PyTorch: 1. `torch.utils.cpp_extension.CUDAExtension`: used to build extensions in `setup.py` during `pip install`. 2. `torch.utils.cpp_extension.load`: used to build and load extension during runtime diff --git a/op_builder/builder.py b/op_builder/builder.py index b9f44decc119..8396235e5cfe 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -7,7 +7,7 @@ import time from abc import ABC, abstractmethod from pathlib import Path -from typing import List +from typing import List, Optional from .utils import check_cuda_availability, check_system_pytorch_cuda_match, print_rank_0 @@ -78,7 +78,7 @@ def sources_files(self) -> List[str]: @abstractmethod def include_dirs(self) -> List[str]: """ - This function should return a list of inlcude files for extensions. + This function should return a list of include files for extensions. """ pass @@ -127,18 +127,18 @@ def check_runtime_build_environment(self): if CUDA_HOME is None: raise RuntimeError( - "CUDA_HOME is not found. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions" + "CUDA_HOME is not found. You need to export CUDA_HOME environment variable or install CUDA Toolkit first in order to build CUDA extensions" ) # make sure CUDA is available for compilation during cuda_available = check_cuda_availability() if not cuda_available: - raise RuntimeError("CUDA is not available on your system as torch.cuda.is_avaible() returns False.") + raise RuntimeError("CUDA is not available on your system as torch.cuda.is_available() returns False.") # make sure system CUDA and pytorch CUDA match, an error will raised inside the function if not check_system_pytorch_cuda_match(CUDA_HOME) - def load(self, verbose=True): + def load(self, verbose: Optional[bool] = None): """ load the kernel during runtime. If the kernel is not built during pip install, it will build the kernel. If the kernel is built during runtime, it will be stored in `~/.cache/colossalai/torch_extensions/`. If the @@ -149,6 +149,8 @@ def load(self, verbose=True): Args: verbose (bool, optional): show detailed info. Defaults to True. """ + if verbose is None: + verbose = os.environ.get('CAI_KERNEL_VERBOSE', '0') == '1' # if the kernel has be compiled and cached, we directly use it if self.cached_op_module is not None: return self.cached_op_module @@ -159,7 +161,7 @@ def load(self, verbose=True): op_module = self.import_op() if verbose: print_rank_0( - f"[extension] OP {self.prebuilt_import_path} has been compileed ahead of time, skip building.") + f"[extension] OP {self.prebuilt_import_path} has been compiled ahead of time, skip building.") except ImportError: # check environment self.check_runtime_build_environment() diff --git a/op_builder/utils.py b/op_builder/utils.py index 4029703e4829..cb528eea66a1 100644 --- a/op_builder/utils.py +++ b/op_builder/utils.py @@ -36,7 +36,7 @@ def get_cuda_version_in_pytorch() -> List[int]: torch_cuda_minor = torch.version.cuda.split(".")[1] except: raise ValueError( - "[extension] Cannot retrive the CUDA version in the PyTorch binary given by torch.version.cuda") + "[extension] Cannot retrieve the CUDA version in the PyTorch binary given by torch.version.cuda") return torch_cuda_major, torch_cuda_minor @@ -90,7 +90,6 @@ def check_system_pytorch_cuda_match(cuda_dir): 'Please make sure you have set the CUDA_HOME correctly and installed the correct PyTorch in https://pytorch.org/get-started/locally/ .' ) - print(bare_metal_minor != torch_cuda_minor) if bare_metal_minor != torch_cuda_minor: warnings.warn( f"[extension] The CUDA version on the system ({bare_metal_major}.{bare_metal_minor}) does not match with the version ({torch_cuda_major}.{torch_cuda_minor}) torch was compiled with. " @@ -111,7 +110,7 @@ def get_pytorch_version() -> List[int]: torch_version = torch.__version__.split('+')[0] TORCH_MAJOR = int(torch_version.split('.')[0]) TORCH_MINOR = int(torch_version.split('.')[1]) - TORCH_PATCH = int(torch_version.split('.')[2]) + TORCH_PATCH = int(torch_version.split('.')[2], 16) return TORCH_MAJOR, TORCH_MINOR, TORCH_PATCH @@ -156,16 +155,15 @@ def set_cuda_arch_list(cuda_dir): # we only need to set this when CUDA is not available for cross-compilation if not cuda_available: - warnings.warn( - '\n[extension] PyTorch did not find available GPUs on this system.\n' - 'If your intention is to cross-compile, this is not an error.\n' - 'By default, Colossal-AI will cross-compile for \n' - '1. Pascal (compute capabilities 6.0, 6.1, 6.2),\n' - '2. Volta (compute capability 7.0)\n' - '3. Turing (compute capability 7.5),\n' - '4. Ampere (compute capability 8.0, 8.6)if the CUDA version is >= 11.0\n' - '\nIf you wish to cross-compile for a single specific architecture,\n' - 'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n') + warnings.warn('\n[extension] PyTorch did not find available GPUs on this system.\n' + 'If your intention is to cross-compile, this is not an error.\n' + 'By default, Colossal-AI will cross-compile for \n' + '1. Pascal (compute capabilities 6.0, 6.1, 6.2),\n' + '2. Volta (compute capability 7.0)\n' + '3. Turing (compute capability 7.5),\n' + '4. Ampere (compute capability 8.0, 8.6)if the CUDA version is >= 11.0\n' + '\nIf you wish to cross-compile for a single specific architecture,\n' + 'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n') if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None: bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir) diff --git a/pytest.ini b/pytest.ini index ac31ace4bfae..01e5cd217c5d 100644 --- a/pytest.ini +++ b/pytest.ini @@ -3,4 +3,4 @@ markers = cpu: tests which can run on CPU gpu: tests which requires a single GPU dist: tests which are run in a multi-GPU or multi-machine environment - experiment: tests for experimental features \ No newline at end of file + experiment: tests for experimental features diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt index 05c0e6ac5e5c..9f6580c72d1b 100644 --- a/requirements/requirements-test.txt +++ b/requirements/requirements-test.txt @@ -1,14 +1,18 @@ diffusers fbgemm-gpu==0.2.0 pytest -pytest-cov +coverage==7.2.3 +git+https://github.com/hpcaitech/pytest-testmon torchvision -transformers +transformers==4.30.2 timm titans torchaudio +torchx-nightly==2022.6.29 # torchrec 0.2.0 requires torchx-nightly. This package is updated every day. We fix the version to a specific date to avoid breaking changes. torchrec==0.2.0 contexttimer einops triton==2.0.0.dev20221202 git+https://github.com/HazyResearch/flash-attention.git@c422fee3776eb3ea24e011ef641fd5fbeb212623#egg=flash_attn +requests==2.27.1 # downgrade to avoid huggingface error https://github.com/huggingface/transformers/issues/17611 +SentencePiece diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 8e619ac24477..b34dc2e223ae 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -8,4 +8,5 @@ click fabric contexttimer ninja -torch +torch>=1.11 +safetensors diff --git a/setup.py b/setup.py index 89a7b0de461b..5d8f831218d9 100644 --- a/setup.py +++ b/setup.py @@ -46,7 +46,7 @@ def environment_check_for_cuda_extension_build(): if not CUDA_HOME: raise RuntimeError( - "[extension] CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions" + "[extension] CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment variable or install CUDA Toolkit first in order to build CUDA extensions" ) check_system_pytorch_cuda_match(CUDA_HOME) diff --git a/tests/components_to_test/__init__.py b/tests/components_to_test/__init__.py index 106f4e61c7e1..f29efefce4a4 100644 --- a/tests/components_to_test/__init__.py +++ b/tests/components_to_test/__init__.py @@ -9,11 +9,11 @@ resnet, simple_net, ) -from .utils import run_fwd_bwd +from .utils import run_fwd, run_fwd_bwd from . import albert # isort:skip __all__ = [ 'bert', 'gpt2', 'hanging_param_model', 'inline_op_model', 'nested_model', 'repeated_computed_layers', 'resnet', - 'simple_net', 'run_fwd_bwd', 'albert', 'beit' + 'simple_net', 'run_fwd_bwd', 'albert', 'beit', 'run_fwd' ] diff --git a/tests/components_to_test/albert.py b/tests/components_to_test/albert.py index d5b6bc89a83e..8924eb2fbc92 100644 --- a/tests/components_to_test/albert.py +++ b/tests/components_to_test/albert.py @@ -27,8 +27,8 @@ def bert_model_builder(checkpoint: bool = False): attention_probs_dropout_prob=0.) print('building AlbertForSequenceClassification model') - # adapting huggingface BertForSequenceClassification for single unitest calling interface - class ModelAaptor(AlbertForSequenceClassification): + # adapting huggingface BertForSequenceClassification for single unittest calling interface + class ModelAdaptor(AlbertForSequenceClassification): def forward(self, input_ids, labels): """ @@ -37,23 +37,23 @@ def forward(self, input_ids, labels): """ return super().forward(input_ids=input_ids, labels=labels)[0] - model = ModelAaptor(config) + model = ModelAdaptor(config) # if checkpoint and version.parse(transformers.__version__) >= version.parse("4.11.0"): # model.gradient_checkpointing_enable() return model - is_distrbuted = torch.distributed.is_initialized() + is_distributed = torch.distributed.is_initialized() trainloader = get_bert_data_loader(n_class=vocab_size, batch_size=2, total_samples=10000, sequence_length=sequence_length, - is_distrbuted=is_distrbuted) + is_distributed=is_distributed) testloader = get_bert_data_loader(n_class=vocab_size, batch_size=2, total_samples=10000, sequence_length=sequence_length, - is_distrbuted=is_distrbuted) + is_distributed=is_distributed) criterion = None return bert_model_builder, trainloader, testloader, torch.optim.Adam, criterion diff --git a/tests/components_to_test/beit.py b/tests/components_to_test/beit.py index 1252071f4075..2021ae6f6e35 100644 --- a/tests/components_to_test/beit.py +++ b/tests/components_to_test/beit.py @@ -27,7 +27,7 @@ def generate(self): @non_distributed_component_funcs.register(name='beit') def get_training_components(): - def model_buider(checkpoint=False): + def model_builder(checkpoint=False): model = Beit(img_size=DummyDataLoader.img_size, num_classes=DummyDataLoader.num_class, embed_dim=32, @@ -39,4 +39,4 @@ def model_buider(checkpoint=False): testloader = DummyDataLoader() criterion = torch.nn.CrossEntropyLoss() - return model_buider, trainloader, testloader, torch.optim.Adam, criterion + return model_builder, trainloader, testloader, torch.optim.Adam, criterion diff --git a/tests/components_to_test/bert.py b/tests/components_to_test/bert.py index c1faa6f9d892..e7d1d50806b8 100644 --- a/tests/components_to_test/bert.py +++ b/tests/components_to_test/bert.py @@ -13,7 +13,7 @@ def get_bert_data_loader( total_samples, sequence_length, device=torch.device('cpu:0'), - is_distrbuted=False, + is_distributed=False, ): train_data = torch.randint( low=0, @@ -24,7 +24,7 @@ def get_bert_data_loader( ) train_label = torch.randint(low=0, high=2, size=(total_samples,), device=device, dtype=torch.long) train_dataset = torch.utils.data.TensorDataset(train_data, train_label) - if is_distrbuted: + if is_distributed: sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) else: sampler = SequentialSampler(train_dataset) @@ -52,8 +52,8 @@ def bert_model_builder(checkpoint: bool = False): attention_probs_dropout_prob=0.) print('building BertForSequenceClassification model') - # adapting huggingface BertForSequenceClassification for single unitest calling interface - class ModelAaptor(BertForSequenceClassification): + # adapting huggingface BertForSequenceClassification for single unittest calling interface + class ModelAdaptor(BertForSequenceClassification): def forward(self, input_ids, labels): """ @@ -62,23 +62,23 @@ def forward(self, input_ids, labels): """ return super().forward(input_ids=input_ids, labels=labels)[0] - model = ModelAaptor(config) + model = ModelAdaptor(config) if checkpoint and version.parse(transformers.__version__) >= version.parse("4.11.0"): model.gradient_checkpointing_enable() return model - is_distrbuted = torch.distributed.is_initialized() + is_distributed = torch.distributed.is_initialized() trainloader = get_bert_data_loader(n_class=vocab_size, batch_size=2, total_samples=10000, sequence_length=sequence_length, - is_distrbuted=is_distrbuted) + is_distributed=is_distributed) testloader = get_bert_data_loader(n_class=vocab_size, batch_size=2, total_samples=10000, sequence_length=sequence_length, - is_distrbuted=is_distrbuted) + is_distributed=is_distributed) criterion = None return bert_model_builder, trainloader, testloader, torch.optim.Adam, criterion diff --git a/tests/components_to_test/registry.py b/tests/components_to_test/registry.py index 728ed9eba6ea..edfcaaa7275b 100644 --- a/tests/components_to_test/registry.py +++ b/tests/components_to_test/registry.py @@ -9,10 +9,10 @@ def __init__(self): def register(self, name): assert name not in self._registry - def _regsiter(callable_): + def _register(callable_): self._registry[name] = callable_ - return _regsiter + return _register def get_callable(self, name: str): return self._registry[name] @@ -34,6 +34,6 @@ def __next__(self): non_distributed_component_funcs = Registry() -model_paralle_component_funcs = Registry() +model_parallel_component_funcs = Registry() -__all__ = ['non_distributed_component_funcs', 'model_paralle_component_funcs'] +__all__ = ['non_distributed_component_funcs', 'model_parallel_component_funcs'] diff --git a/tests/components_to_test/utils/__init__.py b/tests/components_to_test/utils/__init__.py index f223f7d322cb..150124b58800 100644 --- a/tests/components_to_test/utils/__init__.py +++ b/tests/components_to_test/utils/__init__.py @@ -1,2 +1,2 @@ from .dummy_data_generator import DummyDataGenerator -from .executor import run_fwd_bwd +from .executor import run_fwd, run_fwd_bwd diff --git a/tests/components_to_test/utils/executor.py b/tests/components_to_test/utils/executor.py index e77152561e6c..631401e022e6 100644 --- a/tests/components_to_test/utils/executor.py +++ b/tests/components_to_test/utils/executor.py @@ -1,9 +1,9 @@ import torch -def run_fwd_bwd(model, data, label, criterion, optimizer=None) -> torch.Tensor: - """run_fwd_bwd - run fwd and bwd for the model +def run_fwd(model, data, label, criterion) -> torch.Tensor: + """run_fwd + run fwd for the model Args: model (torch.nn.Module): a PyTorch model @@ -22,6 +22,23 @@ def run_fwd_bwd(model, data, label, criterion, optimizer=None) -> torch.Tensor: loss = model(data, label) loss = loss.float() + return loss + + +def run_fwd_bwd(model, data, label, criterion, optimizer=None) -> torch.Tensor: + """run_fwd_bwd + run fwd and bwd for the model + + Args: + model (torch.nn.Module): a PyTorch model + data (torch.Tensor): input data + label (torch.Tensor): label + criterion (Optional[Callable]): a function of criterion + + Returns: + torch.Tensor: loss of fwd + """ + loss = run_fwd(model, data, label, criterion) if optimizer: optimizer.backward(loss) else: diff --git a/tests/kit/model_zoo/__init__.py b/tests/kit/model_zoo/__init__.py index 710038ffa387..466a2a558829 100644 --- a/tests/kit/model_zoo/__init__.py +++ b/tests/kit/model_zoo/__init__.py @@ -1,5 +1,4 @@ from . import diffusers, timm, torchaudio, torchrec, torchvision, transformers - from .registry import model_zoo __all__ = ['model_zoo'] diff --git a/tests/kit/model_zoo/diffusers/diffusers.py b/tests/kit/model_zoo/diffusers/diffusers.py index 8aa3f4c6741f..204c1d7773ca 100644 --- a/tests/kit/model_zoo/diffusers/diffusers.py +++ b/tests/kit/model_zoo/diffusers/diffusers.py @@ -18,6 +18,7 @@ data_unet_fn = lambda: dict(sample=torch.randn(2, 3, 32, 32), timestep=3) identity_output = lambda x: x +clip_vision_model_output = lambda x: dict(pooler_output=x[1]) def data_clip_model(): @@ -65,7 +66,7 @@ def data_clip_vision(): model_zoo.register(name='diffusers_clip_vision_model', model_fn=partial(transformers.CLIPVisionModel, config=transformers.CLIPVisionConfig()), data_gen_fn=data_clip_vision, - output_transform_fn=identity_output) + output_transform_fn=clip_vision_model_output) model_zoo.register(name='diffusers_unet2d_model', model_fn=diffusers.UNet2DModel, diff --git a/tests/kit/model_zoo/registry.py b/tests/kit/model_zoo/registry.py index 7470327a65b6..1e7ef3b62736 100644 --- a/tests/kit/model_zoo/registry.py +++ b/tests/kit/model_zoo/registry.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from typing import Callable -__all__ = ['ModelZooRegistry', 'ModelAttributem', 'model_zoo'] +__all__ = ['ModelZooRegistry', 'ModelAttribute', 'model_zoo'] @dataclass @@ -28,27 +28,35 @@ def register(self, model_fn: Callable, data_gen_fn: Callable, output_transform_fn: Callable, + loss_fn: Callable = None, model_attribute: ModelAttribute = None): """ Register a model and data generation function. Examples: - >>> # Register - >>> model_zoo = ModelZooRegistry() - >>> model_zoo.register('resnet18', resnet18, resnet18_data_gen) - >>> # Run the model - >>> data = resnresnet18_data_gen() # do not input any argument - >>> model = resnet18() # do not input any argument - >>> out = model(**data) + + ```python + # normal forward workflow + model = resnet18() + data = resnet18_data_gen() + output = model(**data) + transformed_output = output_transform_fn(output) + loss = loss_fn(transformed_output) + + # Register + model_zoo = ModelZooRegistry() + model_zoo.register('resnet18', resnet18, resnet18_data_gen, output_transform_fn, loss_fn) + ``` Args: name (str): Name of the model. - model_fn (callable): A function that returns a model. **It must not contain any arguments.** - output_transform_fn (callable): A function that transforms the output of the model into Dict. - data_gen_fn (callable): A function that returns a data sample in the form of Dict. **It must not contain any arguments.** + model_fn (Callable): A function that returns a model. **It must not contain any arguments.** + data_gen_fn (Callable): A function that returns a data sample in the form of Dict. **It must not contain any arguments.** + output_transform_fn (Callable): A function that transforms the output of the model into Dict. + loss_fn (Callable): a function to compute the loss from the given output. Defaults to None model_attribute (ModelAttribute): Attributes of the model. Defaults to None. """ - self[name] = (model_fn, data_gen_fn, output_transform_fn, model_attribute) + self[name] = (model_fn, data_gen_fn, output_transform_fn, loss_fn, model_attribute) def get_sub_registry(self, keyword: str): """ @@ -62,6 +70,8 @@ def get_sub_registry(self, keyword: str): for k, v in self.items(): if keyword in k: new_dict[k] = v + + assert len(new_dict) > 0, f'No model found with keyword {keyword}' return new_dict diff --git a/tests/kit/model_zoo/torchaudio/torchaudio.py b/tests/kit/model_zoo/torchaudio/torchaudio.py index 74611720292f..9a244ac312c0 100644 --- a/tests/kit/model_zoo/torchaudio/torchaudio.py +++ b/tests/kit/model_zoo/torchaudio/torchaudio.py @@ -1,3 +1,5 @@ +from functools import partial + import torch import torchaudio.models as tm @@ -101,13 +103,11 @@ def tacotron_data_gen_fn(): mel_specgram_lengths=mel_specgram_lengths) -model_zoo.register( - name='torchaudio_tacotron', - model_fn=lambda: tm.Tacotron2(n_mels=N_MELS), - data_gen_fn=tacotron_data_gen_fn, - output_transform_fn=lambda outputs: dict( - spectrogram_before=outputs[0], spectrogram_after=outputs[1], stop_tokens=outputs[2], attn_weights=outputs[3]), - model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='torchaudio_tacotron', + model_fn=lambda: tm.Tacotron2(n_mels=N_MELS), + data_gen_fn=tacotron_data_gen_fn, + output_transform_fn=lambda outputs: dict(summed_output=sum(x.sum() for x in outputs)), + model_attribute=ModelAttribute(has_control_flow=True)) def wav2vec_data_gen_fn(): @@ -118,7 +118,7 @@ def wav2vec_data_gen_fn(): model_zoo.register(name='torchaudio_wav2vec2_base', - model_fn=tm.wav2vec2_base, + model_fn=partial(tm.wav2vec2_base, encoder_layer_drop=0.0), data_gen_fn=wav2vec_data_gen_fn, output_transform_fn=transformer_output_transform_fn, model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/torchrec/torchrec.py b/tests/kit/model_zoo/torchrec/torchrec.py index 014e9218b226..dda563155fca 100644 --- a/tests/kit/model_zoo/torchrec/torchrec.py +++ b/tests/kit/model_zoo/torchrec/torchrec.py @@ -2,96 +2,141 @@ from functools import partial import torch - -try: - from torchrec.models import deepfm, dlrm - from torchrec.modules.embedding_configs import EmbeddingBagConfig - from torchrec.modules.embedding_modules import EmbeddingBagCollection - from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor - NO_TORCHREC = False -except ImportError: - NO_TORCHREC = True +from torchrec.models import deepfm, dlrm +from torchrec.modules.embedding_configs import EmbeddingBagConfig +from torchrec.modules.embedding_modules import EmbeddingBagCollection +from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor from ..registry import ModelAttribute, model_zoo +BATCH = 2 +SHAPE = 10 + -def register_torchrec_models(): - BATCH = 2 - SHAPE = 10 - # KeyedTensor +def gen_kt(): KT = KeyedTensor(keys=["f1", "f2"], length_per_key=[SHAPE, SHAPE], values=torch.rand((BATCH, 2 * SHAPE))) + return KT + - # KeyedJaggedTensor +# KeyedJaggedTensor +def gen_kjt(): KJT = KeyedJaggedTensor.from_offsets_sync(keys=["f1", "f2"], values=torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]), offsets=torch.tensor([0, 2, 4, 6, 8])) + return KJT + + +data_gen_fn = lambda: dict(features=torch.rand((BATCH, SHAPE))) + + +def interaction_arch_data_gen_fn(): + KT = gen_kt() + return dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KT) + + +def simple_dfm_data_gen_fn(): + KJT = gen_kjt() + return dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KJT) + + +def sparse_arch_data_gen_fn(): + KJT = gen_kjt() + return dict(features=KJT) + + +def output_transform_fn(x): + if isinstance(x, KeyedTensor): + output = dict() + for key in x.keys(): + output[key] = x[key] + return output + else: + return dict(output=x) + + +def output_transform_fn(x): + if isinstance(x, KeyedTensor): + output = dict() + for key in x.keys(): + output[key] = x[key] + return output + else: + return dict(output=x) + + +def get_ebc(): + # EmbeddingBagCollection + eb1_config = EmbeddingBagConfig(name="t1", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f1"]) + eb2_config = EmbeddingBagConfig(name="t2", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f2"]) + return EmbeddingBagCollection(tables=[eb1_config, eb2_config], device=torch.device('cpu')) + - data_gen_fn = lambda: dict(features=torch.rand((BATCH, SHAPE))) +def sparse_arch_model_fn(): + ebc = get_ebc() + return deepfm.SparseArch(ebc) - interaction_arch_data_gen_fn = lambda: dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KT) - simple_dfm_data_gen_fn = lambda: dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KJT) +def simple_deep_fmnn_model_fn(): + ebc = get_ebc() + return deepfm.SimpleDeepFMNN(SHAPE, ebc, SHAPE, SHAPE) - sparse_arch_data_gen_fn = lambda: dict(features=KJT) - output_transform_fn = lambda x: dict(output=x) +def dlrm_model_fn(): + ebc = get_ebc() + return dlrm.DLRM(ebc, SHAPE, [SHAPE, SHAPE], [5, 1]) - def get_ebc(): - # EmbeddingBagCollection - eb1_config = EmbeddingBagConfig(name="t1", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f1"]) - eb2_config = EmbeddingBagConfig(name="t2", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f2"]) - return EmbeddingBagCollection(tables=[eb1_config, eb2_config]) - model_zoo.register(name='deepfm_densearch', - model_fn=partial(deepfm.DenseArch, SHAPE, SHAPE, SHAPE), - data_gen_fn=data_gen_fn, - output_transform_fn=output_transform_fn) +def dlrm_sparsearch_model_fn(): + ebc = get_ebc() + return dlrm.SparseArch(ebc) - model_zoo.register(name='deepfm_interactionarch', - model_fn=partial(deepfm.FMInteractionArch, SHAPE * 3, ["f1", "f2"], SHAPE), - data_gen_fn=interaction_arch_data_gen_fn, - output_transform_fn=output_transform_fn) - model_zoo.register(name='deepfm_overarch', - model_fn=partial(deepfm.OverArch, SHAPE), - data_gen_fn=data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='deepfm_densearch', + model_fn=partial(deepfm.DenseArch, SHAPE, SHAPE, SHAPE), + data_gen_fn=data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='deepfm_simpledeepfmnn', - model_fn=partial(deepfm.SimpleDeepFMNN, SHAPE, get_ebc(), SHAPE, SHAPE), - data_gen_fn=simple_dfm_data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='deepfm_interactionarch', + model_fn=partial(deepfm.FMInteractionArch, SHAPE * 3, ["f1", "f2"], SHAPE), + data_gen_fn=interaction_arch_data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='deepfm_sparsearch', - model_fn=partial(deepfm.SparseArch, get_ebc()), - data_gen_fn=sparse_arch_data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='deepfm_overarch', + model_fn=partial(deepfm.OverArch, SHAPE), + data_gen_fn=data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='dlrm', - model_fn=partial(dlrm.DLRM, get_ebc(), SHAPE, [SHAPE, SHAPE], [5, 1]), - data_gen_fn=simple_dfm_data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='deepfm_simpledeepfmnn', + model_fn=simple_deep_fmnn_model_fn, + data_gen_fn=simple_dfm_data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='dlrm_densearch', - model_fn=partial(dlrm.DenseArch, SHAPE, [SHAPE, SHAPE]), - data_gen_fn=data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='deepfm_sparsearch', + model_fn=sparse_arch_model_fn, + data_gen_fn=sparse_arch_data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='dlrm_interactionarch', - model_fn=partial(dlrm.InteractionArch, 2), - data_gen_fn=interaction_arch_data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='dlrm', + model_fn=dlrm_model_fn, + data_gen_fn=simple_dfm_data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='dlrm_overarch', - model_fn=partial(dlrm.OverArch, SHAPE, [5, 1]), - data_gen_fn=data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='dlrm_densearch', + model_fn=partial(dlrm.DenseArch, SHAPE, [SHAPE, SHAPE]), + data_gen_fn=data_gen_fn, + output_transform_fn=output_transform_fn) - model_zoo.register(name='dlrm_sparsearch', - model_fn=partial(dlrm.SparseArch, get_ebc()), - data_gen_fn=sparse_arch_data_gen_fn, - output_transform_fn=output_transform_fn) +model_zoo.register(name='dlrm_interactionarch', + model_fn=partial(dlrm.InteractionArch, 2), + data_gen_fn=interaction_arch_data_gen_fn, + output_transform_fn=output_transform_fn) +model_zoo.register(name='dlrm_overarch', + model_fn=partial(dlrm.OverArch, SHAPE, [5, 1]), + data_gen_fn=data_gen_fn, + output_transform_fn=output_transform_fn) -if not NO_TORCHREC: - register_torchrec_models() +model_zoo.register(name='dlrm_sparsearch', + model_fn=dlrm_sparsearch_model_fn, + data_gen_fn=sparse_arch_data_gen_fn, + output_transform_fn=output_transform_fn) diff --git a/tests/kit/model_zoo/torchvision/torchvision.py b/tests/kit/model_zoo/torchvision/torchvision.py index 62bda93d5a75..ddc3ec24b2ff 100644 --- a/tests/kit/model_zoo/torchvision/torchvision.py +++ b/tests/kit/model_zoo/torchvision/torchvision.py @@ -36,12 +36,12 @@ def swin_s(): # special output transform fn -google_net_output_transform_fn = lambda x: dict(output=x.logits) if isinstance(x, torchvision.models.GoogLeNetOutputs - ) else dict(output=x) +google_net_output_transform_fn = lambda x: dict(output=sum(x)) if isinstance(x, torchvision.models.GoogLeNetOutputs + ) else dict(output=x) swin_s_output_output_transform_fn = lambda x: {f'output{idx}': val for idx, val in enumerate(x)} if isinstance(x, tuple) else dict(output=x) -inception_v3_output_transform_fn = lambda x: dict(output=x.logits) if isinstance(x, torchvision.models.InceptionOutputs - ) else dict(output=x) +inception_v3_output_transform_fn = lambda x: dict(output=sum(x)) if isinstance(x, torchvision.models.InceptionOutputs + ) else dict(output=x) model_zoo.register(name='torchvision_alexnet', model_fn=tm.alexnet, diff --git a/tests/kit/model_zoo/transformers/__init__.py b/tests/kit/model_zoo/transformers/__init__.py index f56ff7ad84eb..4aa01abe13ee 100644 --- a/tests/kit/model_zoo/transformers/__init__.py +++ b/tests/kit/model_zoo/transformers/__init__.py @@ -1,5 +1,7 @@ from .albert import * from .bert import * +from .bloom import * from .gpt import * +from .llama import * from .opt import * from .t5 import * diff --git a/tests/kit/model_zoo/transformers/bert.py b/tests/kit/model_zoo/transformers/bert.py index 99135704da70..d2d3de7b7bee 100644 --- a/tests/kit/model_zoo/transformers/bert.py +++ b/tests/kit/model_zoo/transformers/bert.py @@ -6,83 +6,147 @@ # =============================== # Register single-sentence BERT # =============================== -BATCH_SIZE = 2 -SEQ_LENGTH = 16 -def data_gen_fn(): - input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - token_type_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) +# define data gen function +def data_gen(): + # Generated from following code snippet + # + # from transformers import BertTokenizer + # input = 'Hello, my dog is cute' + # tokenized_input = tokenizer(input, return_tensors='pt') + # input_ids = tokenized_input['input_ids'] + # attention_mask = tokenized_input['attention_mask'] + # token_type_ids = tokenized_input['token_type_ids'] + input_ids = torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]], dtype=torch.int64) + token_type_ids = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 0]], dtype=torch.int64) + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64) return dict(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) +def data_gen_for_lm(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + data['labels'] = data['input_ids'].clone() + return data + + +def data_gen_for_pretraining(): + # pretraining data gen + # `next_sentence_label` is the label for next sentence prediction, 0 or 1 + data = data_gen_for_lm() + data['next_sentence_label'] = torch.tensor([1], dtype=torch.int64) + return data + + +def data_gen_for_sequence_classification(): + # sequence classification data gen + # `labels` is the label for sequence classification, 0 or 1 + data = data_gen() + data['labels'] = torch.tensor([1], dtype=torch.int64) + return data + + +def data_gen_for_token_classification(): + # token classification data gen + # `labels` is the type not the token id for token classification, 0 or 1 + data = data_gen() + data['labels'] = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 0]], dtype=torch.int64) + return data + + +def data_gen_for_mcq(): + # multiple choice question data gen + # Generated from following code snippet + # + # tokenizer = transformers.BertTokenizer.from_pretrained("bert-base-uncased") + # prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + # choice0 = "It is eaten with a fork and a knife." + # choice1 = "It is eaten while held in the hand." + # data = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True) + # data = {k: v.unsqueeze(0) for k, v in encoding.items()} + # data['labels'] = torch.tensor([0], dtype=torch.int64) + input_ids = torch.tensor([[[ + 101, 1999, 3304, 1010, 10733, 2366, 1999, 5337, 10906, 1010, 2107, 2004, 2012, 1037, 4825, 1010, 2003, 3591, + 4895, 14540, 6610, 2094, 1012, 102, 2009, 2003, 8828, 2007, 1037, 9292, 1998, 1037, 5442, 1012, 102 + ], + [ + 101, 1999, 3304, 1010, 10733, 2366, 1999, 5337, 10906, 1010, 2107, 2004, 2012, 1037, + 4825, 1010, 2003, 3591, 4895, 14540, 6610, 2094, 1012, 102, 2009, 2003, 8828, 2096, + 2218, 1999, 1996, 2192, 1012, 102, 0 + ]]]) + token_type_ids = torch.tensor( + [[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]]]) + attention_mask = torch.tensor( + [[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]]]) + labels = torch.tensor([0], dtype=torch.int64) + + return dict(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, labels=labels) + + +# define output transform function output_transform_fn = lambda x: x -config = transformers.BertConfig(hidden_size=128, num_hidden_layers=2, num_attention_heads=4, intermediate_size=256) +# define loss funciton +loss_fn_for_bert_model = lambda x: x.pooler_output.mean() +loss_fn = lambda x: x.loss + +config = transformers.BertConfig(hidden_size=128, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=256, + hidden_dropout_prob=0, + attention_probs_dropout_prob=0) # register the BERT variants model_zoo.register(name='transformers_bert', model_fn=lambda: transformers.BertModel(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_bert_model, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_for_pretraining', model_fn=lambda: transformers.BertForPreTraining(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen_for_pretraining, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_lm_head_model', model_fn=lambda: transformers.BertLMHeadModel(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen_for_lm, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_for_masked_lm', model_fn=lambda: transformers.BertForMaskedLM(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen_for_lm, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_for_sequence_classification', model_fn=lambda: transformers.BertForSequenceClassification(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen_for_sequence_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_for_token_classification', model_fn=lambda: transformers.BertForTokenClassification(config), - data_gen_fn=data_gen_fn, + data_gen_fn=data_gen_for_token_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) - - -# =============================== -# Register multi-sentence BERT -# =============================== -def data_gen_for_next_sentence(): - tokenizer = transformers.BertTokenizer.from_pretrained("bert-base-uncased") - prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." - next_sentence = "The sky is blue due to the shorter wavelength of blue light." - encoding = tokenizer(prompt, next_sentence, return_tensors="pt") - return encoding - - -def data_gen_for_mcq(): - tokenizer = transformers.BertTokenizer.from_pretrained("bert-base-uncased") - prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." - choice0 = "It is eaten with a fork and a knife." - choice1 = "It is eaten while held in the hand." - encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True) - encoding = {k: v.unsqueeze(0) for k, v in encoding.items()} - return encoding - - -# register the following models model_zoo.register(name='transformers_bert_for_next_sentence', model_fn=lambda: transformers.BertForNextSentencePrediction(config), - data_gen_fn=data_gen_for_next_sentence, + data_gen_fn=data_gen_for_sequence_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_bert_for_mcq', model_fn=lambda: transformers.BertForMultipleChoice(config), data_gen_fn=data_gen_for_mcq, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/transformers/bloom.py b/tests/kit/model_zoo/transformers/bloom.py new file mode 100644 index 000000000000..71146c0b9819 --- /dev/null +++ b/tests/kit/model_zoo/transformers/bloom.py @@ -0,0 +1,107 @@ +import torch +import transformers + +from ..registry import ModelAttribute, model_zoo + +# =============================== +# Register Bloom +# =============================== + + +def data_gen(): + # Generated from following code snippet + # + # from transformers import BloomTokenizer + # input = 'Hello, my dog is cute' + # tokenized_input = tokenizer(input, return_tensors='pt') + # input_ids = tokenized_input['input_ids'] + # attention_mask = tokenized_input['attention_mask'] + input_ids = torch.tensor([[59414, 15, 2670, 35433, 632, 207595]], dtype=torch.int64) + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1]], dtype=torch.int64) + return dict(input_ids=input_ids, attention_mask=attention_mask) + + +def data_gen_for_lm(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + data['labels'] = data['input_ids'].clone() + return data + + +def data_gen_for_token_classification(): + # token classification data gen + # `labels` is the type not the token id for token classification, 0 or 1 + data = data_gen() + data['labels'] = torch.tensor([[0, 0, 0, 0, 0, 0]], dtype=torch.int64) + return data + + +def data_gen_for_sequence_classification(): + # sequence classification data gen + data = data_gen() + data['labels'] = torch.tensor([0], dtype=torch.int64) + return data + + +def data_gen_for_question_answering(): + # obtained with the following code + # + # from transformers import AutoTokenizer + # tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m") + # question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" + # inputs = tokenizer(question, text, return_tensors="pt") + + input_ids = torch.tensor( + [[57647, 1620, 23967, 620, 107373, 34, 91514, 620, 107373, 1620, 267, 35378, 48946, 18161]], dtype=torch.int64) + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64) + return dict(input_ids=input_ids, attention_mask=attention_mask) + + +# define output transform function +output_transform_fn = lambda x: x + +# define loss function +loss_fn_for_bloom_model = lambda x: x.last_hidden_state.mean() +loss_fn_for_causal_lm = lambda x: x.loss +loss_fn_for_classification = lambda x: x.logits.mean() +loss_fn_for_question_answering = lambda x: x.end_logits.mean() + +config = transformers.BloomConfig(n_layer=1, + n_head=4, + vocab_size=250880, + hidden_dropout=0, + attention_dropout=0, + hidden_size=64) + +# register the following models +model_zoo.register(name='transformers_bloom', + model_fn=lambda: transformers.BloomModel(config), + data_gen_fn=data_gen, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_bloom_model, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_bloom_for_causal_lm', + model_fn=lambda: transformers.BloomForCausalLM(config), + data_gen_fn=data_gen_for_lm, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_causal_lm, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_bloom_for_sequence_classification', + model_fn=lambda: transformers.BloomForSequenceClassification(config), + data_gen_fn=data_gen_for_sequence_classification, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_classification, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_bloom_for_token_classification', + model_fn=lambda: transformers.BloomForTokenClassification(config), + data_gen_fn=data_gen_for_token_classification, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_classification, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_bloom_for_question_answering', + model_fn=lambda: transformers.BloomForQuestionAnswering(config), + data_gen_fn=data_gen_for_question_answering, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_question_answering, + model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/transformers/gpt.py b/tests/kit/model_zoo/transformers/gpt.py index 2a100c981dea..b9e0310780af 100644 --- a/tests/kit/model_zoo/transformers/gpt.py +++ b/tests/kit/model_zoo/transformers/gpt.py @@ -6,44 +6,89 @@ # =============================== # Register single-sentence GPT # =============================== -BATCH_SIZE = 1 # it can only be 1 as GPT cannot handle batch sizes > 1 if no padding token is defined. -SEQ_LENGTH = 16 def data_gen(): - input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - token_type_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - return dict(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) + # Generated from following code snippet + # + # from transformers import GPT2Tokenizer + # input = 'Hello, my dog is cute' + # tokenized_input = tokenizer(input, return_tensors='pt') + # input_ids = tokenized_input['input_ids'] + # attention_mask = tokenized_input['attention_mask'] + input_ids = torch.tensor([[15496, 11, 616, 3290, 318, 13779]], dtype=torch.int64) + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1]], dtype=torch.int64) + return dict(input_ids=input_ids, attention_mask=attention_mask) +def data_gen_for_lm(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + data['labels'] = data['input_ids'].clone() + return data + + +def data_gen_for_token_classification(): + # token classification data gen + # `labels` is the type not the token id for token classification, 0 or 1 + data = data_gen() + data['labels'] = torch.tensor([[0, 0, 0, 0, 0, 0]], dtype=torch.int64) + return data + + +def data_gen_for_sequence_classification(): + # sequence classification data gen + data = data_gen() + data['labels'] = torch.tensor([0], dtype=torch.int64) + return data + + +# define output transform function output_transform_fn = lambda x: x -config = transformers.GPT2Config(n_position=64, n_layer=2, n_head=4) +# define loss function +loss_fn_for_gpt2_model = lambda x: x.last_hidden_state.mean() +loss_fn = lambda x: x.loss + +config = transformers.GPT2Config(n_layer=2, + n_head=4, + vocab_size=50258, + attn_pdrop=0, + embd_pdrop=0, + resid_pdrop=0, + summary_first_dropout=0, + hidden_dropout=0, + problem_type="single_label_classification") # register the following models model_zoo.register(name='transformers_gpt', model_fn=lambda: transformers.GPT2Model(config), data_gen_fn=data_gen, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_gpt2_model, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_gpt_lm', model_fn=lambda: transformers.GPT2LMHeadModel(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_lm, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_gpt_double_heads', model_fn=lambda: transformers.GPT2DoubleHeadsModel(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_lm, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_gpt_for_token_classification', model_fn=lambda: transformers.GPT2ForTokenClassification(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_token_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_gpt_for_sequence_classification', model_fn=lambda: transformers.GPT2ForSequenceClassification(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_sequence_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn, model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/transformers/llama.py b/tests/kit/model_zoo/transformers/llama.py new file mode 100644 index 000000000000..705bbc7364ba --- /dev/null +++ b/tests/kit/model_zoo/transformers/llama.py @@ -0,0 +1,76 @@ +import torch +import transformers + +from ..registry import ModelAttribute, model_zoo + +try: + from transformers import LlamaConfig, LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel + HAS_LLAMA = True +except ImportError: + HAS_LLAMA = False + +if HAS_LLAMA: + # =============================== + # Register LLaMA + # =============================== + + def data_gen(): + # the input ids are corresponding to the sentence + # 'Hello, my dog is cute' + # + # the code is give below: + # ----------------------------------- + # from transformers import LlamaTokenizerFast + # tokenizer = LlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer") + # input = 'Hello, my dog is cute' + # tokenized_input = tokenizer(input, return_tensors='pt').to('cuda') + # ----------------------------------- + + input_ids = torch.Tensor([[1, 15043, 29892, 590, 11203, 338, 274, 1082]]).long() + attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1, 1, 1]]).long() + return dict(input_ids=input_ids, attention_mask=attention_mask) + + # label is needed for casual lm + def data_gen_for_casual_lm(): + data = data_gen() + labels = data['input_ids'].clone() + data['labels'] = labels + return data + + # transform the output to a dict + output_transform_fn = lambda x: x + + # function to get the loss + loss_fn = lambda output: output.last_hidden_state.mean() + loss_fn_for_casual_lm = lambda output: output.loss + loss_fn_for_seq_classification = lambda output: output.logits.mean() + + config = LlamaConfig(num_hidden_layers=4, + hidden_size=128, + intermediate_size=256, + num_attention_heads=4, + max_position_embeddings=128, + num_labels=16) + + # register the following models + # transformers.LlamaModel, + # transformers.LlamaForCausalLM, + # transformers.LlamaForSequenceClassification, + model_zoo.register(name='transformers_llama', + model_fn=lambda: transformers.LlamaModel(config), + data_gen_fn=data_gen, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn, + model_attribute=ModelAttribute(has_control_flow=True)) + model_zoo.register(name='transformers_llama_for_casual_lm', + model_fn=lambda: transformers.LlamaForCausalLM(config), + data_gen_fn=data_gen_for_casual_lm, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_casual_lm, + model_attribute=ModelAttribute(has_control_flow=True)) + model_zoo.register(name='transformers_llama_for_sequence_classification', + model_fn=lambda: transformers.LlamaForSequenceClassification(config), + data_gen_fn=data_gen, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_seq_classification, + model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/transformers/opt.py b/tests/kit/model_zoo/transformers/opt.py index d9c4a0b3c23c..4463ae12b901 100644 --- a/tests/kit/model_zoo/transformers/opt.py +++ b/tests/kit/model_zoo/transformers/opt.py @@ -11,14 +11,47 @@ def data_gen(): - input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) + input_ids = torch.Tensor([[1, 15043, 29892, 590, 11203, 338, 274, 1082]]).long() + attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1, 1, 1]]).long() return dict(input_ids=input_ids, attention_mask=attention_mask) -output_transform_fn = lambda x: x +def data_gen_for_causal_lm(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + labels = data['input_ids'].clone() + data['labels'] = labels + return data + + +def data_gen_for_sequence_classification(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + labels = data['input_ids'].clone() + data['labels'] = torch.tensor([1]) + return data + + +def data_gen_for_question_answering(): + # LM data gen + # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels` + data = data_gen() + data['start_positions'] = torch.tensor([0]) + data['end_positions'] = torch.tensor([1]) + return data + -config = transformers.OPTConfig(hidden_size=128, num_hidden_layers=2, num_attention_heads=4) +output_transform_fn = lambda x: x +loss_fn_for_opt_model = lambda x: x.last_hidden_state.mean() +loss_fn_for_lm = lambda x: x.loss +config = transformers.OPTConfig( + hidden_size=128, + num_hidden_layers=2, + num_attention_heads=4, + dropout=0, +) # register the following models # transformers.OPTModel, @@ -27,9 +60,23 @@ def data_gen(): model_fn=lambda: transformers.OPTModel(config), data_gen_fn=data_gen, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_opt_model, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_opt_for_causal_lm', model_fn=lambda: transformers.OPTForCausalLM(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_causal_lm, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_lm, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_opt_for_question_answering', + model_fn=lambda: transformers.OPTForQuestionAnswering(config), + data_gen_fn=data_gen_for_question_answering, + output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_lm, + model_attribute=ModelAttribute(has_control_flow=True)) +model_zoo.register(name='transformers_opt_for_sequence_classification', + model_fn=lambda: transformers.OPTForSequenceClassification(config), + data_gen_fn=data_gen_for_sequence_classification, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_lm, model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/kit/model_zoo/transformers/t5.py b/tests/kit/model_zoo/transformers/t5.py index b81bcad90db8..689db2c40abb 100644 --- a/tests/kit/model_zoo/transformers/t5.py +++ b/tests/kit/model_zoo/transformers/t5.py @@ -6,24 +6,50 @@ # =============================== # Register single-sentence T5 # =============================== -BATCH_SIZE = 2 -SEQ_LENGTH = 16 - - -def data_gen(): - input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - decoder_input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) - return dict(input_ids=input_ids, decoder_input_ids=decoder_input_ids) +# define data gen function def data_gen_for_encoder_only(): - input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) + # Generated from following code snippet + # + # from transformers import T5Config, T5Tokenizer + # config = T5Config(decoder_start_token_id=0) + # tokenizer = T5Tokenizer.from_pretrained("t5-small") + # input_ids = tokenizer("translate English to German: The house is wonderful.", return_tensors="pt").input_ids + input_ids = torch.Tensor([[13959, 1566, 12, 2968, 10, 37, 629, 19, 1627, 5, 1]]).long() return dict(input_ids=input_ids) +def data_gen_for_conditional_generation(): + # labels is generated with the following code + # + # labels = tokenizer("Das Haus ist wunderbar.", return_tensors="pt").input_ids + data = data_gen_for_encoder_only() + labels = torch.Tensor([[644, 4598, 229, 19250, 5, 1]]).long() + data['labels'] = labels + return data + + +def data_gen_for_t5_model(): + # decoder_inputs_ids is obtained with the following code + # + # decoder_input_ids = model._shift_right(input_ids) + data = data_gen_for_encoder_only() + decoder_input_ids = torch.Tensor([[0, 13959, 1566, 12, 2968, 10, 37, 629, 19, 1627, 5]]).long() + data['decoder_input_ids'] = decoder_input_ids + return data + + +# output transform function output_transform_fn = lambda x: x -config = transformers.T5Config(d_model=128, num_layers=2) +# define loss funciton +loss_fn_for_t5_model = lambda x: x.last_hidden_state.mean() +loss_fn_for_encoder_only = lambda x: x.last_hidden_state.mean() +loss_fn_for_conditional_generation = lambda x: x.loss + +# define model config +config = transformers.T5Config(d_model=128, num_layers=2, dropout_rate=0, decoder_start_token_id=0) # register the following models # transformers.T5Model, @@ -31,16 +57,19 @@ def data_gen_for_encoder_only(): # transformers.T5EncoderModel, model_zoo.register(name='transformers_t5', model_fn=lambda: transformers.T5Model(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_t5_model, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_t5_model, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_t5_for_conditional_generation', model_fn=lambda: transformers.T5ForConditionalGeneration(config), - data_gen_fn=data_gen, + data_gen_fn=data_gen_for_conditional_generation, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_conditional_generation, model_attribute=ModelAttribute(has_control_flow=True)) model_zoo.register(name='transformers_t5_encoder_model', model_fn=lambda: transformers.T5EncoderModel(config), data_gen_fn=data_gen_for_encoder_only, output_transform_fn=output_transform_fn, + loss_fn=loss_fn_for_encoder_only, model_attribute=ModelAttribute(has_control_flow=True)) diff --git a/tests/test_amp/test_naive_fp16.py b/tests/test_amp/test_naive_fp16.py index c01de469b8f1..6ce4c7f49725 100644 --- a/tests/test_amp/test_naive_fp16.py +++ b/tests/test_amp/test_naive_fp16.py @@ -1,14 +1,11 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.amp import convert_to_apex_amp, convert_to_naive_amp -from colossalai.testing import assert_close_loose, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import assert_close_loose, clear_cache_before_run, rerun_if_address_is_in_use, spawn from tests.components_to_test.registry import non_distributed_component_funcs @@ -87,10 +84,9 @@ def run_dist(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() +@clear_cache_before_run() def test_naive_amp(): - world_size = 1 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 1) if __name__ == '__main__': diff --git a/tests/test_amp/test_torch_fp16.py b/tests/test_amp/test_torch_fp16.py index e65dd8cded26..6451aa6264a3 100644 --- a/tests/test_amp/test_torch_fp16.py +++ b/tests/test_amp/test_torch_fp16.py @@ -1,14 +1,11 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.amp import convert_to_apex_amp, convert_to_torch_amp -from colossalai.testing import assert_close_loose, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import assert_close_loose, clear_cache_before_run, rerun_if_address_is_in_use, spawn from tests.components_to_test.registry import non_distributed_component_funcs @@ -87,10 +84,9 @@ def run_dist(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() +@clear_cache_before_run() def test_torch_amp(): - world_size = 1 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 1) if __name__ == '__main__': diff --git a/tests/test_analyzer/__init__.py b/tests/test_analyzer/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_analyzer/test_fx/test_bias_addition.py b/tests/test_analyzer/test_fx/test_bias_addition.py index 5c9ec7cc3477..f7b5eb140f24 100644 --- a/tests/test_analyzer/test_fx/test_bias_addition.py +++ b/tests/test_analyzer/test_fx/test_bias_addition.py @@ -1,7 +1,10 @@ import pytest import torch +from packaging import version from torch.utils.checkpoint import checkpoint +from colossalai.testing.utils import clear_cache_before_run, parameterize + try: from colossalai._analyzer.fx import symbolic_trace except: @@ -55,9 +58,13 @@ def __init__(self, bias) -> None: self.linear = LinearModel(3, 3, bias) self.conv = ConvModel(3, 6, 3, bias) - def forward(self, x, select=0): + def forward(self, x, select=torch.Tensor([0])): x = self.linear(x) - x = checkpoint(self.conv, x, select) + if select: + x = checkpoint(self.conv, x, 0) + else: + x = checkpoint(self.conv, x, 1) + return x @@ -73,11 +80,12 @@ def forward(self, x): return x -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize("bias", [True, False]) -@pytest.mark.parametrize("bias_addition_split", [True, False]) -@pytest.mark.parametrize("shape", [(3, 3, 3), (3, 3, 3, 3)]) -@pytest.mark.parametrize("select", [0, 1]) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize("bias", [True, False]) +@parameterize("bias_addition_split", [True, False]) +@parameterize("shape", [(3, 3, 3), (3, 3, 3, 3)]) +@parameterize("select", [torch.Tensor([0]), torch.Tensor([1])]) def test_siu_model(bias, bias_addition_split, shape, select): model = SiuModel(bias=bias) x = torch.rand(shape) @@ -86,18 +94,18 @@ def test_siu_model(bias, bias_addition_split, shape, select): concrete_args={'select': select}, trace_act_ckpt=True, bias_addition_split=bias_addition_split) - assert torch.allclose(model(x, select), gm(x, select)), 'original model and traced model should be the same!' + assert torch.allclose(model(x, select), gm(x)), 'original model and traced model should be the same!' if bias and bias_addition_split: assert '+' in gm.code, 'bias addition should be split!' else: assert '+' not in gm.code, 'bias addition should not be split!' -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize("alpha", [1, 2]) -@pytest.mark.parametrize("beta", [1, 2]) -@pytest.mark.parametrize("bias_addition_split", [True, False]) -@pytest.mark.parametrize("shape", [(3, 3), (5, 5)]) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@parameterize("alpha", [1, 2]) +@parameterize("beta", [1, 2]) +@parameterize("bias_addition_split", [True, False]) +@parameterize("shape", [(3, 3), (5, 5)]) def test_addmm_model(alpha, beta, bias_addition_split, shape): model = AddmmModel(alpha=alpha, beta=beta) x = torch.rand(shape) @@ -110,4 +118,5 @@ def test_addmm_model(alpha, beta, bias_addition_split, shape): if __name__ == '__main__': - test_siu_model(True, True, (3, 3, 3)) + test_siu_model() + test_addmm_model() diff --git a/tests/test_analyzer/test_fx/test_mod_dir.py b/tests/test_analyzer/test_fx/test_mod_dir.py index 15e0c2ec21c7..f62147b297a2 100644 --- a/tests/test_analyzer/test_fx/test_mod_dir.py +++ b/tests/test_analyzer/test_fx/test_mod_dir.py @@ -1,6 +1,8 @@ import pytest import torch +from colossalai.testing import clear_cache_before_run, parameterize + try: from colossalai._analyzer.fx import symbolic_trace except: @@ -62,9 +64,10 @@ def forward(self, x): @pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize("bias", [True, False]) -@pytest.mark.parametrize("bias_addition_split", [True, False]) -@pytest.mark.parametrize("shape", [(3, 3, 3), (3, 3, 3, 3)]) +@clear_cache_before_run() +@parameterize("bias", [True, False]) +@parameterize("bias_addition_split", [True, False]) +@parameterize("shape", [(3, 3, 3), (3, 3, 3, 3)]) def test_mod_dir(bias, bias_addition_split, shape): model = AModel(bias=bias) x = torch.rand(shape) @@ -75,4 +78,4 @@ def test_mod_dir(bias, bias_addition_split, shape): if __name__ == '__main__': - test_mod_dir(True, True, (3, 3, 3)) + test_mod_dir(bias=True, bias_addition_split=True, shape=(3, 3, 3)) diff --git a/tests/test_analyzer/test_fx/test_nested_ckpt.py b/tests/test_analyzer/test_fx/test_nested_ckpt.py index c31aab6752f8..bd16f5a4f95d 100644 --- a/tests/test_analyzer/test_fx/test_nested_ckpt.py +++ b/tests/test_analyzer/test_fx/test_nested_ckpt.py @@ -1,7 +1,9 @@ +import pytest import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint -import pytest + +from colossalai.testing import clear_cache_before_run try: from colossalai._analyzer.fx import symbolic_trace @@ -42,6 +44,7 @@ def forward(self, x): @pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') +@clear_cache_before_run() def test_nested_ckpt(): model = MyModule() x = torch.rand(10, 10) diff --git a/tests/test_analyzer/test_fx/test_shape_prop.py b/tests/test_analyzer/test_fx/test_shape_prop.py index b19884a70fb2..a849feb795e5 100644 --- a/tests/test_analyzer/test_fx/test_shape_prop.py +++ b/tests/test_analyzer/test_fx/test_shape_prop.py @@ -1,16 +1,17 @@ import pytest -import timm.models as tmm import torch import torchvision.models as tm -from .zoo import tm_models, tmm_models +from packaging import version + +from colossalai.testing.utils import clear_cache_before_run, parameterize +from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models try: from colossalai._analyzer._subclasses import MetaTensorMode from colossalai._analyzer.fx import symbolic_trace from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass from colossalai._analyzer.fx.symbolic_profile import register_shape_impl - - + @register_shape_impl(torch.nn.functional.linear) def linear_impl(*args, **kwargs): assert True @@ -23,15 +24,16 @@ def _check_gm_validity(gm: torch.fx.GraphModule): for node in gm.graph.nodes: assert node.meta['info'].outputs, f'In {gm.__class__.__name__}, {node} has no output shape.' if node.op in [ - # 'call_module', # can apply to params - # 'call_function', # can apply to params - # 'call_method', # can apply to params + 'call_module', # can apply to params + 'call_function', # can apply to params + 'call_method', # can apply to params ]: - assert node.meta['info'].inputs, f'In {gm.__class__.__name__}, {node} has no input shape.' + assert hasattr(node.meta['info'], 'inputs'), f'In {gm.__class__.__name__}, {node} has no input shape.' -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize('m', tm_models) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize('m', tm_models) def test_torchvision_shape_prop(m): with MetaTensorMode(): model = m() @@ -44,8 +46,9 @@ def test_torchvision_shape_prop(m): _check_gm_validity(gm) -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize('m', tmm_models) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize('m', tmm_models) def test_timm_shape_prop(m): with MetaTensorMode(): model = m() @@ -53,11 +56,12 @@ def test_timm_shape_prop(m): meta_args = { "x": data, } + gm = symbolic_trace(model, meta_args=meta_args) shape_prop_pass(gm, data) _check_gm_validity(gm) if __name__ == "__main__": - test_torchvision_shape_prop(tm.resnet18) - test_timm_shape_prop(tmm.vgg11) + test_torchvision_shape_prop() + test_timm_shape_prop() diff --git a/tests/test_analyzer/test_fx/test_symbolic_profile.py b/tests/test_analyzer/test_fx/test_symbolic_profile.py index 5f749e6f3c50..17deee7a7118 100644 --- a/tests/test_analyzer/test_fx/test_symbolic_profile.py +++ b/tests/test_analyzer/test_fx/test_symbolic_profile.py @@ -1,8 +1,10 @@ import pytest -import timm.models as tmm import torch import torchvision.models as tm -from .zoo import tm_models, tmm_models +from packaging import version + +from colossalai.testing.utils import clear_cache_before_run, parameterize +from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models try: from colossalai._analyzer._subclasses import MetaTensorMode @@ -16,8 +18,9 @@ def _check_gm_validity(gm: torch.fx.GraphModule): assert len(node.meta['info'].global_ctx), f'In {gm.__class__.__name__}, {node} has empty global context.' -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize('m', tm_models) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize('m', tm_models) def test_torchvision_profile(m, verbose=False, bias_addition_split=False): with MetaTensorMode(): model = m() @@ -30,8 +33,9 @@ def test_torchvision_profile(m, verbose=False, bias_addition_split=False): _check_gm_validity(gm) -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize('m', tmm_models) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize('m', tmm_models) def test_timm_profile(m, verbose=False, bias_addition_split=False): with MetaTensorMode(): model = m() @@ -45,5 +49,5 @@ def test_timm_profile(m, verbose=False, bias_addition_split=False): if __name__ == "__main__": - test_torchvision_profile(tm.vit_b_16, verbose=True, bias_addition_split=False) - test_timm_profile(tmm.gmlp_b16_224, verbose=True, bias_addition_split=False) + test_torchvision_profile() + test_timm_profile() diff --git a/tests/test_analyzer/test_fx/zoo.py b/tests/test_analyzer/test_fx/zoo.py index 925078d0dcbe..a96aa3949134 100644 --- a/tests/test_analyzer/test_fx/zoo.py +++ b/tests/test_analyzer/test_fx/zoo.py @@ -33,18 +33,18 @@ tmm.dm_nfnet_f0, tmm.eca_nfnet_l0, tmm.efficientformer_l1, - tmm.ese_vovnet19b_dw, + # tmm.ese_vovnet19b_dw, tmm.gmixer_12_224, tmm.gmlp_b16_224, - tmm.hardcorenas_a, + # tmm.hardcorenas_a, tmm.hrnet_w18_small, tmm.inception_v3, tmm.mixer_b16_224, tmm.nf_ecaresnet101, tmm.nf_regnet_b0, # tmm.pit_b_224, # pretrained only - tmm.regnetv_040, - tmm.skresnet18, + # tmm.regnetv_040, + # tmm.skresnet18, # tmm.swin_base_patch4_window7_224, # fx bad case # tmm.tnt_b_patch16_224, # bad case tmm.vgg11, diff --git a/tests/test_analyzer/test_subclasses/test_aten.py b/tests/test_analyzer/test_subclasses/test_aten.py index 591a8d617580..b7858110ac09 100644 --- a/tests/test_analyzer/test_subclasses/test_aten.py +++ b/tests/test_analyzer/test_subclasses/test_aten.py @@ -1,9 +1,11 @@ from typing import Any, Callable, Union -import pytest +import pytest import torch import torch.nn as nn +from colossalai.testing import clear_cache_before_run + try: from colossalai._analyzer._subclasses import MetaTensor except: @@ -72,6 +74,7 @@ def run_and_compare(f: Union[nn.Module, Callable], x: torch.Tensor, requires_bac @pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') +@clear_cache_before_run() def test_meta_aten(): for (aten_op, requires_backward), v in registered_meta.items(): for f, x in v: diff --git a/tests/test_analyzer/test_subclasses/test_flop_tensor.py b/tests/test_analyzer/test_subclasses/test_flop_tensor.py index 551628103325..4e9c9852649b 100644 --- a/tests/test_analyzer/test_subclasses/test_flop_tensor.py +++ b/tests/test_analyzer/test_subclasses/test_flop_tensor.py @@ -1,9 +1,11 @@ import pytest import torch -import torch.nn as nn import torch.nn.functional as F import torchvision.models as tm -from .zoo import tm_models, tmm_models +from packaging import version + +from colossalai.testing import clear_cache_before_run, parameterize +from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models try: from colossalai._analyzer._subclasses import MetaTensorMode, flop_count @@ -11,7 +13,7 @@ pass -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') @pytest.mark.parametrize('m', tm_models + tmm_models) def test_flop_count_module(m): x = torch.rand(2, 3, 224, 224) @@ -37,7 +39,7 @@ def test_flop_count_module(m): ] -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') @pytest.mark.parametrize('func, args, kwargs', odd_cases) def test_flop_count_function(func, args, kwargs): rs_fwd, rs_bwd = flop_count(func, *args, **kwargs, verbose=True) @@ -46,5 +48,5 @@ def test_flop_count_function(func, args, kwargs): if __name__ == '__main__': - test_flop_count_module(tm.resnet18, torch.rand(2, 3, 224, 224)) + test_flop_count_module(tm.resnet18) test_flop_count_function(F.relu, (torch.rand(2, 3, 224, 224, requires_grad=True),), {'inplace': True}) diff --git a/tests/test_analyzer/test_subclasses/test_meta_mode.py b/tests/test_analyzer/test_subclasses/test_meta_mode.py index d8122b019619..d2a0a1b9cfb5 100644 --- a/tests/test_analyzer/test_subclasses/test_meta_mode.py +++ b/tests/test_analyzer/test_subclasses/test_meta_mode.py @@ -1,12 +1,15 @@ import pytest import torch -import torch.distributed as dist import torchvision.models as tm +from packaging import version + +from colossalai.testing import clear_cache_before_run, parameterize + try: from colossalai._analyzer._subclasses import MetaTensor, MetaTensorMode except: pass -from .zoo import tm_models, tmm_models +from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models def compare_all(tensor: torch.Tensor, meta_tensor: torch.Tensor): @@ -28,8 +31,9 @@ def run_and_compare(model): compare_all(x.grad, meta_x.grad) -@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12') -@pytest.mark.parametrize('m', tm_models + tmm_models) +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() +@parameterize('m', tm_models + tmm_models) def test_meta_mode_shape(m): run_and_compare(m()) diff --git a/tests/test_analyzer/test_subclasses/zoo.py b/tests/test_analyzer/test_subclasses/zoo.py deleted file mode 100644 index 925078d0dcbe..000000000000 --- a/tests/test_analyzer/test_subclasses/zoo.py +++ /dev/null @@ -1,53 +0,0 @@ -import timm.models as tmm -import torchvision.models as tm - -# input shape: (batch_size, 3, 224, 224) -tm_models = [ - tm.alexnet, - tm.convnext_base, - tm.densenet121, - # tm.efficientnet_v2_s, - # tm.googlenet, # output bad case - # tm.inception_v3, # bad case - tm.mobilenet_v2, - tm.mobilenet_v3_small, - tm.mnasnet0_5, - tm.resnet18, - tm.regnet_x_16gf, - tm.resnext50_32x4d, - tm.shufflenet_v2_x0_5, - tm.squeezenet1_0, - # tm.swin_s, # fx bad case - tm.vgg11, - tm.vit_b_16, - tm.wide_resnet50_2, -] - -tmm_models = [ - tmm.beit_base_patch16_224, - tmm.beitv2_base_patch16_224, - tmm.cait_s24_224, - tmm.coat_lite_mini, - tmm.convit_base, - tmm.deit3_base_patch16_224, - tmm.dm_nfnet_f0, - tmm.eca_nfnet_l0, - tmm.efficientformer_l1, - tmm.ese_vovnet19b_dw, - tmm.gmixer_12_224, - tmm.gmlp_b16_224, - tmm.hardcorenas_a, - tmm.hrnet_w18_small, - tmm.inception_v3, - tmm.mixer_b16_224, - tmm.nf_ecaresnet101, - tmm.nf_regnet_b0, - # tmm.pit_b_224, # pretrained only - tmm.regnetv_040, - tmm.skresnet18, - # tmm.swin_base_patch4_window7_224, # fx bad case - # tmm.tnt_b_patch16_224, # bad case - tmm.vgg11, - tmm.vit_base_patch16_18x2_224, - tmm.wide_resnet50_2, -] diff --git a/tests/test_auto_parallel/test_ckpt_solvers/test_C_solver_consistency.py b/tests/test_auto_parallel/test_ckpt_solvers/test_C_solver_consistency.py index f8dd0b16b7f6..f184f64b35d0 100644 --- a/tests/test_auto_parallel/test_ckpt_solvers/test_C_solver_consistency.py +++ b/tests/test_auto_parallel/test_ckpt_solvers/test_C_solver_consistency.py @@ -3,7 +3,6 @@ import pytest import torch import torch.fx -import torch.multiprocessing as mp import torchvision.models as tm import colossalai @@ -13,7 +12,7 @@ # from colossalai.fx.passes.algorithms import solver_rotor # from colossalai.fx.passes.algorithms.operation import Sequence from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn if is_compatible_with_meta(): from colossalai.fx.profiler.tensor import MetaTensor @@ -26,8 +25,8 @@ withcodegen = False -def _run_C_solver_consistency_test(rank=0): - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_C_solver_consistency_test(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') for M, mem_budget in [(tm.resnet50, 4000), (tm.densenet121, 8080)]: model = M() @@ -70,8 +69,9 @@ def _run_C_solver_consistency_test(rank=0): @pytest.mark.skip("TODO(lyl): refactor all tests.") @pytest.mark.skipif(not withcodegen, reason="torch version is less than 1.12.0") +@rerun_if_address_is_in_use() def test_C_solver_consistency(): - mp.spawn(_run_C_solver_consistency_test, nprocs=1) + spawn(_run_C_solver_consistency_test, 1) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_ckpt_solvers/test_ckpt_torchvision.py b/tests/test_auto_parallel/test_ckpt_solvers/test_ckpt_torchvision.py index 89600ea098a9..db268b91d0a0 100644 --- a/tests/test_auto_parallel/test_ckpt_solvers/test_ckpt_torchvision.py +++ b/tests/test_auto_parallel/test_ckpt_solvers/test_ckpt_torchvision.py @@ -4,7 +4,6 @@ import pytest import torch -import torch.multiprocessing as mp import torchvision.models as tm from torch.fx import GraphModule @@ -15,7 +14,7 @@ from colossalai.fx.graph_module import ColoGraphModule # from colossalai.fx.passes.algorithms import chen_greedy, solver_rotor from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn if is_compatible_with_meta(): from colossalai.fx.profiler.tensor import MetaTensor @@ -68,8 +67,8 @@ def check_backward_consistency(m: torch.nn.Module, gm: GraphModule, solver: Call assert _is_all_gradient_close(m, gm), f'Solver {solver} did not work correctly in backward pass on {model_cls}' -def _run_ckpt_solver(rank): - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_ckpt_solver(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') MODEL_LIST = [tm.densenet121] torch.backends.cudnn.deterministic = True @@ -98,12 +97,13 @@ def _run_ckpt_solver(rank): @pytest.mark.skip("TODO(super-dainiu): refactor all tests.") @pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0') +@rerun_if_address_is_in_use() def test_ckpt_solver(): - mp.spawn(_run_ckpt_solver, nprocs=1) + spawn(_run_ckpt_solver, 1) -def _run_ckpt_solver_torch11(rank): - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_ckpt_solver_torch11(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') MODEL_LIST = [tm.densenet121] torch.backends.cudnn.deterministic = True @@ -131,8 +131,9 @@ def _run_ckpt_solver_torch11(rank): @pytest.mark.skipif(with_codegen, reason='torch version is equal to or higher than 1.12.0') @pytest.mark.skip(reason="currently torch11 ColoGraphModule is not done") +@rerun_if_address_is_in_use() def test_ckpt_solver_torch11(): - mp.spawn(_run_ckpt_solver_torch11, nprocs=1) + spawn(_run_ckpt_solver_torch11, 1) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_ckpt_solvers/test_linearize.py b/tests/test_auto_parallel/test_ckpt_solvers/test_linearize.py index 0f90ba0b0989..59880815dc5e 100644 --- a/tests/test_auto_parallel/test_ckpt_solvers/test_linearize.py +++ b/tests/test_auto_parallel/test_ckpt_solvers/test_linearize.py @@ -8,6 +8,7 @@ # from colossalai.fx.passes.algorithms import linearize, solver_rotor # from colossalai.fx.passes.algorithms.operation import (ForwardCheck, ForwardEnable, ForwardNograd, Loss) from colossalai.fx.passes.meta_info_prop import MetaInfoProp +from colossalai.testing import clear_cache_before_run if is_compatible_with_meta(): from colossalai.fx.profiler.tensor import MetaTensor @@ -24,6 +25,7 @@ @pytest.mark.skip(reason='TODO: modify the logger') @pytest.mark.skip("TODO(lyl): refactor all tests.") @pytest.mark.skipif(not with_codegen, reason="torch version is lower than 1.12.0") +@clear_cache_before_run() def test_linearize(): MODEL_DICT = {tm.resnet18: [2100, 3000], tm.densenet121: [8100, 17000]} tracer = ColoTracer() @@ -84,6 +86,7 @@ def test_linearize(): @pytest.mark.skip("TODO(lyl): refactor all tests.") @pytest.mark.skip(reason="torch11 meta tensor not implemented") @pytest.mark.skipif(with_codegen, reason="torch version is equal to or higher than 1.12.0") +@clear_cache_before_run() def test_linearize_torch11(): MODEL_DICT = {tm.resnet18: [2100, 3000], tm.densenet121: [8100, 17000]} tracer = ColoTracer() diff --git a/tests/test_auto_parallel/test_offload/model_utils.py b/tests/test_auto_parallel/test_offload/model_utils.py new file mode 100644 index 000000000000..c22b17ae42ba --- /dev/null +++ b/tests/test_auto_parallel/test_offload/model_utils.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn +from transformers import GPT2Config, GPT2LMHeadModel +from transformers import BertConfig, BertLMHeadModel +from tests.components_to_test.registry import non_distributed_component_funcs + +class GPTLMModel(nn.Module): + + def __init__(self, + hidden_size=768, + num_layers=12, + num_attention_heads=12, + max_seq_len=1024, + vocab_size=50257): + super().__init__() + self.model = GPT2LMHeadModel( + GPT2Config(n_embd=hidden_size, + n_layer=num_layers, + n_head=num_attention_heads, + n_positions=max_seq_len, + n_ctx=max_seq_len, + vocab_size=vocab_size)) + + def forward(self, input_ids, attention_mask): + # Only return lm_logits + return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)[0] + + +class LMLoss(nn.Module): + + def __init__(self): + super().__init__() + self.loss_fn = nn.CrossEntropyLoss() + + def forward(self, logits, labels): + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + +class BertLMModel(nn.Module): + def __init__(self, hidden_size=768, num_layers=12, num_attention_heads=32, vocab_size=30522): + super().__init__() + self.model = BertLMHeadModel(BertConfig(n_embd=hidden_size, num_hidden_layers=num_layers, hidden_size=hidden_size, + num_attention_heads=num_attention_heads, max_position_embeddings=hidden_size, + vocab_size=vocab_size)) + + def forward(self, input_ids, attention_mask): + # Only return lm_logits + return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)[0] + +@non_distributed_component_funcs.register(name='bert_') +def get_bert_components(): + vocab_size = 1024 + seq_len = 64 + batchSize = 64 + + def bert_model_builder(): + model = BertLMModel(hidden_size=8192, num_layers=4, num_attention_heads=32, vocab_size=vocab_size) + return model + + def bert_data_gen(device="meta"): + input_ids = torch.randint(0, vocab_size, (batchSize, seq_len), device=device) + attention_mask = torch.ones_like(input_ids, device=device) + kwargs = dict(input_ids=input_ids, attention_mask=attention_mask) + return kwargs + + return bert_model_builder, bert_data_gen + +@non_distributed_component_funcs.register(name='gpt2_') +def get_gpt2_components(): + vocab_size = 1024 + seq_len = 8 + batchSize = 64 + + def gpt2_model_builder(): + model = GPTLMModel(hidden_size=8192, num_layers=2, num_attention_heads=32, vocab_size=vocab_size) + return model + + def gpt2_data_gen(device="meta"): + input_ids = torch.randint(0, vocab_size, (batchSize, seq_len), device=device) + attention_mask = torch.ones_like(input_ids, device=device) + kwargs = dict(input_ids=input_ids, attention_mask=attention_mask) + return kwargs + + return gpt2_model_builder, gpt2_data_gen \ No newline at end of file diff --git a/tests/test_auto_parallel/test_offload/test_perf.py b/tests/test_auto_parallel/test_offload/test_perf.py new file mode 100644 index 000000000000..45c22efc4127 --- /dev/null +++ b/tests/test_auto_parallel/test_offload/test_perf.py @@ -0,0 +1,147 @@ +import time + +import pytest +import torch +from torch.utils._pytree import tree_map + +import colossalai +from colossalai.auto_parallel.offload.amp_optimizer import AMPOptimizer +from colossalai.auto_parallel.offload.mem_optimize import memory_optimize +from colossalai.auto_parallel.offload.solver import NOT_NVML +from colossalai.fx.profiler import parameter_size +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext, zero_model_wrapper, zero_optim_wrapper +from tests.test_auto_parallel.test_offload.model_utils import * +from tests.test_tensor.common_utils import set_seed + + +@parameterize('model_name', ['gpt2_']) +@parameterize('memory_budget', [5000]) +@parameterize('solver_name', ['asyn']) +def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str): + + # build model + get_components_func = non_distributed_component_funcs.get_callable(model_name) + model_builder, data_gen = get_components_func() + label = torch.randint(low=0, high=128, size=( + 64, + 8, + ), device=get_current_device()) + criterion = LMLoss() + + set_seed(42) + start_time = time.time() + model = model_builder() + model.train() + param_size = parameter_size(model) / 1024**2 / 2 + init_time = time.time() - start_time + print(f"init_param_size={param_size:.3f} MB | init_model_time={init_time:.3f} s") + + data_args = data_gen(device="cpu") + wrap_fn = lambda x: x.to(dtype=torch.half) if isinstance(x, torch.Tensor) and torch.is_floating_point(x) else x + data_args = tree_map(wrap_fn, data_args) + start_time = time.time() + model = memory_optimize(model, data_args, memory_budget * 1024 * 1024, solver_name) + solver_time = time.time() - start_time + print(f"solver_time={solver_time:.3f} s") + + hybrid_optimizer = HybridAdam(model.model.parameters(), lr=1e-3) + optim = AMPOptimizer(hybrid_optimizer, model) + + with ColoInitContext(device=torch.device('cpu')): + gemini_model = model_builder() + gemini_model.train() + + hybrid_optimizer = HybridAdam(gemini_model.parameters(), lr=1e-3) + gemini_config = dict(strict_ddp_mode=False, + device=torch.device('cpu'), + placement_policy='cpu', + pin_memory=True, + hidden_dim=8192, + search_range_m=128) + gemini_model = zero_model_wrapper(gemini_model, 3, gemini_config) + optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True) + gemini_optim = zero_optim_wrapper(gemini_model, hybrid_optimizer, optim_config=optim_config) + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + # test gemini + time_list = [] + set_seed(42) + data_args = data_gen(device="cuda") + for step in range(10): + gemini_optim.zero_grad() + torch.cuda.synchronize() + start_time = time.time() + gemini_out = gemini_model(**data_args) + gemini_loss = criterion(gemini_out, label) + gemini_optim.backward(gemini_loss) + torch.cuda.synchronize() + time_list.append(time.time() - start_time) + gemini_optim.step() + + torch.cuda.synchronize() + + exec_time = sum(sorted(time_list)[:5]) / 5 + runtime_peak_mem_alc = torch.cuda.max_memory_allocated() / 1024**2 + runtime_peak_mem_res = torch.cuda.max_memory_reserved() / 1024**2 + print(f'gemini | model_name: {model_name}') + print(f'| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB ' + f'| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|') + print(time_list) + + del data_args + del gemini_model + del gemini_optim + del gemini_out + del gemini_loss + + # test asyn offload + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + time_list = [] + set_seed(42) + data_args = data_gen(device="cuda") + data_args = tree_map(wrap_fn, data_args) + for step in range(10): + optim.zero_grad() + torch.cuda.synchronize() + start_time = time.time() + loss = criterion(model(**data_args), label) + optim.backward(loss) + torch.cuda.synchronize() + time_list.append(time.time() - start_time) + optim.step() + + torch.cuda.synchronize() + + exec_time = sum(sorted(time_list)[:5]) / 5 + runtime_peak_mem_alc = torch.cuda.max_memory_allocated() / 1024**2 + runtime_peak_mem_res = torch.cuda.max_memory_reserved() / 1024**2 + print(f'solver_name: {solver_name} | model_name: {model_name}') + print(f'| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB ' + f'| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|') + print(time_list) + + +def run_dist(rank, world_size, port): + config = {} + colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + exam_fwd_bwd() + + +@pytest.mark.skip("this test failed") +@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed') +@rerun_if_address_is_in_use() +def test_perf(): + spawn(run_dist, 1) + + +if __name__ == '__main__': + test_perf() diff --git a/tests/test_auto_parallel/test_offload/test_solver.py b/tests/test_auto_parallel/test_offload/test_solver.py new file mode 100644 index 000000000000..aa2c9a36849f --- /dev/null +++ b/tests/test_auto_parallel/test_offload/test_solver.py @@ -0,0 +1,67 @@ +import pytest +import torch.fx +from torch.fx import GraphModule +from torch.utils._pytree import tree_map + +from colossalai.auto_parallel.offload.region_manager import RegionManager +from colossalai.auto_parallel.offload.solver import NOT_NVML, SolverFactory +from colossalai.fx import ColoTracer, is_compatible_with_meta +from colossalai.fx.passes.meta_info_prop import MetaInfoProp +from colossalai.testing import clear_cache_before_run, parameterize +from tests.test_auto_parallel.test_offload.model_utils import * + + +@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed') +@clear_cache_before_run() +@parameterize('model_name', ['gpt2_', 'bert_']) +@parameterize('memory_budget', [4000]) +@parameterize('solver_name', ['syn', 'asyn']) +def solver_test(model_name: str, memory_budget: float, solver_name: str): + + get_components_func = non_distributed_component_funcs.get_callable(model_name) + model_builder, data_gen = get_components_func() + data_args = data_gen(device="cpu") + wrap_fn = lambda x: x.to(dtype=torch.half) if isinstance(x, torch.Tensor) and torch.is_floating_point(x) else x + data_args = tree_map(wrap_fn, data_args) + model = model_builder() + model.train() + model = model.cpu().half() + + tracer = ColoTracer() + assert is_compatible_with_meta() + wrap_fn = lambda x: x.to("meta") if isinstance(x, torch.Tensor) else x + meta_args = tree_map(wrap_fn, data_args) + graph = tracer.trace(model, meta_args=meta_args) + gm = GraphModule(model, graph, model.__class__.__name__) + + interp = MetaInfoProp(gm) + interp.propagate(*meta_args.values()) + + region_manager = RegionManager(graph, solver_name=solver_name) + region_manager._pre_process() + region_list = region_manager.region_list + + solver_cls = SolverFactory.create(solver_name) + memory_budget = memory_budget * 1024 * 1024 + solver = solver_cls(region_list, memory_budget) + solver._call_solver() + + assert solver.best_ts.peak_mem < memory_budget + + print("****************** execution plan *******************") + for region in region_list: + need_offload = region.need_offload + to_prefetch = region.fwd_prefetch_region.r_id if region.fwd_prefetch_region is not None else None + print( + f'| {model_name} forward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}' + ) + for region in region_list.__reversed__(): + need_offload = region.need_offload + to_prefetch = region.bwd_prefetch_region.r_id if region.bwd_prefetch_region is not None else None + print( + f'| {model_name} backward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}' + ) + + +if __name__ == '__main__': + solver_test() diff --git a/tests/test_auto_parallel/test_pass/test_node_converting_pass.py b/tests/test_auto_parallel/test_pass/test_node_converting_pass.py index d0d107610f7a..429e89aae5d3 100644 --- a/tests/test_auto_parallel/test_pass/test_node_converting_pass.py +++ b/tests/test_auto_parallel/test_pass/test_node_converting_pass.py @@ -6,6 +6,7 @@ from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.tracer import ColoTracer from colossalai.tensor.sharding_spec import ShardingSpec +from colossalai.testing import clear_cache_before_run class TestModule(torch.nn.Module): @@ -26,6 +27,7 @@ def insert_narrow(gm, x_node): return gm +@clear_cache_before_run() def test_node_args_converting_pass(): model = TestModule() physical_mesh_id = torch.arange(0, 4) diff --git a/tests/test_auto_parallel/test_pass/test_size_value_converting_pass.py b/tests/test_auto_parallel/test_pass/test_size_value_converting_pass.py index 3494830080ff..bca81201c6ef 100644 --- a/tests/test_auto_parallel/test_pass/test_size_value_converting_pass.py +++ b/tests/test_auto_parallel/test_pass/test_size_value_converting_pass.py @@ -1,11 +1,14 @@ +import pytest import torch import torch.nn.functional as F +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.passes.runtime_preparation_pass import size_value_converting_pass from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.graph_module import ColoGraphModule -from colossalai.fx.tracer import ColoTracer from colossalai.tensor.sharding_spec import ShardingSpec +from colossalai.testing import clear_cache_before_run class TestModule(torch.nn.Module): @@ -33,6 +36,8 @@ def recover_narrow(gm, narrow_node): return gm +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') +@clear_cache_before_run() def test_size_value_converting_pass(): model = TestModule() physical_mesh_id = torch.arange(0, 4) @@ -40,14 +45,14 @@ def test_size_value_converting_pass(): device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) meta_args = {'x': torch.rand(4, 8).to('meta')} input = torch.rand(4, 8) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) graph = tracer.trace(root=model, meta_args=meta_args) - x_node = list(graph.nodes)[0] x_sharding_spec = ShardingSpec(device_mesh, entire_shape=(4, 8), dim_partition_dict={0: [0]}) setattr(x_node, 'sharding_spec', x_sharding_spec) gm = ColoGraphModule(model, graph) gm = insert_narrow(gm, x_node) + shape_prop_pass(gm, *meta_args.values()) gm.recompile() size = gm(input) assert size == torch.Size([2, 8]) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py b/tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py index f43885a6ac44..9fbe674ef4f4 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py @@ -2,15 +2,17 @@ import pytest import torch -import torch.multiprocessing as mp -from colossalai.auto_parallel.tensor_shard.initialize import initialize_model +try: + from colossalai.auto_parallel.tensor_shard.initialize import initialize_model + NO_CODEGEN = False +except: + NO_CODEGEN = True + from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import assert_close, rerun_if_address_is_in_use, run_on_environment_flag, spawn class LinearModel(torch.nn.Module): @@ -77,14 +79,12 @@ def check_conv_module(rank, world_size, port): @run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.skipif(NO_CODEGEN, reason='No codegen found') @pytest.mark.dist @rerun_if_address_is_in_use() def test_bias_addition_module(): - world_size = 4 - run_func_linear = partial(check_linear_module, world_size=world_size, port=free_port()) - mp.spawn(run_func_linear, nprocs=world_size) - run_func_conv = partial(check_conv_module, world_size=world_size, port=free_port()) - mp.spawn(run_func_conv, nprocs=world_size) + spawn(check_linear_module, 4) + spawn(check_conv_module, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_checkpoint.py b/tests/test_auto_parallel/test_tensor_shard/test_checkpoint.py index 0b42722fec5f..398458306e3d 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_checkpoint.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_checkpoint.py @@ -1,23 +1,21 @@ -from functools import partial -from typing import Optional, Tuple, Union +from typing import Optional, Tuple import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from torch.utils.checkpoint import checkpoint from transformers.pytorch_utils import Conv1D -from colossalai.auto_parallel.tensor_shard.initialize import initialize_model +try: + from colossalai.auto_parallel.tensor_shard.initialize import initialize_model + NO_CODEGEN = False +except: + NO_CODEGEN = True + from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.graph_module import ColoGraphModule -from colossalai.fx.tracer import ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.tensor.shape_consistency import ShapeConsistencyManager -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn HIDDEN_SIZE = 16 @@ -43,6 +41,7 @@ def check_act_ckpt(rank, world_size, port): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = GPT2MLPWithCkpt(intermediate_size=4 * HIDDEN_SIZE, hidden_size=HIDDEN_SIZE) + input = torch.rand(1, 64, HIDDEN_SIZE) input_sample = { 'hidden_states': torch.rand(1, 64, HIDDEN_SIZE).to('meta'), } @@ -54,16 +53,15 @@ def check_act_ckpt(rank, world_size, port): gm = initialize_model(model, input_sample, device_mesh) code = gm.module.graph.python_code('self').src assert "runtime_comm_spec_apply_1 = colossalai_auto_parallel_passes_runtime_apply_pass_runtime_comm_spec_apply(linear_1, comm_actions_dict, 12, 'linear_1')" in code - assert "view_3 = colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0, False, view_1, comm_actions_dict, use_reentrant=True)" in code + assert "view_3 = torch.utils.checkpoint.checkpoint(self.checkpoint_0, view_1, comm_actions_dict, use_reentrant=False)" in code @run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.skipif(NO_CODEGEN, reason='No codegen found') @pytest.mark.dist @rerun_if_address_is_in_use() def test_mlp_layer(): - world_size = 4 - run_func = partial(check_act_ckpt, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_act_ckpt, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_ddp.py b/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_ddp.py index e4982a5d7f5a..6908a1781869 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_ddp.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_ddp.py @@ -1,18 +1,19 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP -from colossalai.auto_parallel.tensor_shard.initialize import initialize_model +try: + from colossalai.auto_parallel.tensor_shard.initialize import initialize_model + NO_CODEGEN = False +except: + NO_CODEGEN = True + from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import assert_close, rerun_if_address_is_in_use, run_on_environment_flag, spawn class MLP(torch.nn.Module): @@ -93,12 +94,11 @@ def check_compatibility_with_ddp(rank, world_size, port): @run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.skipif(NO_CODEGEN, reason='No codegen found') @pytest.mark.dist @rerun_if_address_is_in_use() def test_compatibility_with_ddp(): - world_size = 4 - run_func = partial(check_compatibility_with_ddp, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_compatibility_with_ddp, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py b/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py index 760401c3f2c2..4e3c26c1ba9c 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_compatibility_with_gemini.py @@ -1,22 +1,22 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp -from torch.nn.parallel import DistributedDataParallel as DDP -from colossalai.auto_parallel.tensor_shard.initialize import initialize_model +try: + from colossalai.auto_parallel.tensor_shard.initialize import initialize_model + NO_CODEGEN = False +except: + NO_CODEGEN = True + from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper from colossalai.tensor.process_group import ProcessGroup -from colossalai.testing import assert_close, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port, get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx +from colossalai.testing import assert_close, rerun_if_address_is_in_use, run_on_environment_flag, spawn +from colossalai.utils import get_current_device +from colossalai.zero import post_process_colo_init_ctx, zero_model_wrapper, zero_optim_wrapper class MLP(torch.nn.Module): @@ -75,7 +75,7 @@ def check_auto_parallel_with_gemini(rank, world_size, port): device=get_current_device(), placement_policy='cpu', pin_memory=True, - search_range_mb=128) + search_range_m=128) post_process_colo_init_ctx(gm, device=get_current_device(), default_pg=dp_process_group) gm = zero_model_wrapper(gm, zero_stage=3, gemini_config=gemini_config) @@ -102,12 +102,11 @@ def check_auto_parallel_with_gemini(rank, world_size, port): @run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.skipif(NO_CODEGEN, reason='No codegen found') @pytest.mark.dist @rerun_if_address_is_in_use() def test_auto_parallel_with_gemini(): - world_size = 4 - run_func = partial(check_auto_parallel_with_gemini, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_auto_parallel_with_gemini, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_find_repeat_block.py b/tests/test_auto_parallel/test_tensor_shard/test_find_repeat_block.py index 90301521f207..a0b407b240e1 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_find_repeat_block.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_find_repeat_block.py @@ -5,10 +5,12 @@ from torch.fx import GraphModule from transformers.pytorch_utils import Conv1D +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes import shape_prop_pass +# from colossalai.fx.tracer.tracer import ColoTracer +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.utils.factory import find_repeat_blocks -from colossalai.fx.tracer.tracer import ColoTracer -from colossalai.testing import parameterize -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, parameterize, run_on_environment_flag NUM_REPEAT_BLOCKS = 4 BATCH_SIZE = 1 @@ -78,16 +80,18 @@ def forward(self, x): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() @parameterize('model_cls', [RepeatModel, NonRepeatModel]) def test_repeat_blocks(model_cls): model = model_cls(4 * HIDDEN_DIM, HIDDEN_DIM, NUM_REPEAT_BLOCKS) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) input_sample = {'x': torch.rand(BATCH_SIZE, SEQ_LENGTH, HIDDEN_DIM).to('meta')} graph = tracer.trace(root=model, meta_args=input_sample) gm = GraphModule(model, graph, model.__class__.__name__) + shape_prop_pass(gm, *input_sample.values()) gm.recompile() node_list = list(graph.nodes) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_runtime_with_gpt_modules.py b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_runtime_with_gpt_modules.py index ebeef9870fe9..48d2672c6571 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_runtime_with_gpt_modules.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_runtime_with_gpt_modules.py @@ -1,30 +1,35 @@ import copy import random -from functools import partial from typing import Dict import numpy as np import pytest import torch -import torch.multiprocessing as mp import transformers from torch.fx import GraphModule -from colossalai.auto_parallel.tensor_shard.initialize import ( - ModuleWrapper, - build_strategy_constructor, - solve_solution, - transform_to_sharded_model, -) +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +# from colossalai.fx.tracer.tracer import ColoTracer +from colossalai._analyzer.fx.tracer.tracer import ColoTracer + +try: + from colossalai.auto_parallel.tensor_shard.initialize import ( + ModuleWrapper, + build_strategy_constructor, + solve_solution, + transform_to_sharded_model, + ) + NO_CODEGEN = False +except: + NO_CODEGEN = True + from colossalai.auto_parallel.tensor_shard.sharding_strategy import ShardingSpec from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.tracer.tracer import ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.tensor.shape_consistency import to_global -from colossalai.testing import assert_close, assert_close_loose, parameterize, rerun_if_address_is_in_use +from colossalai.testing import assert_close, assert_close_loose, parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_gpt.gpt_modules import GPT2MLP, GPT2Attention, GPT2Block, GPT2Model BATCH_SIZE = 1 @@ -52,9 +57,8 @@ def _check_module_grad(module: torch.nn.Module, origin_param_dict: Dict[str, tor param_sharding_spec = best_sharding_spec_dict[new_name] grad_to_compare = copy.deepcopy(param_grad) param_grad_global = to_global(grad_to_compare, param_sharding_spec) - try: - assert_close_loose(param_grad_global, origin_param_grad, rtol=1e-03, atol=1e-03) + assert_close_loose(param_grad_global, origin_param_grad, rtol=1e-03, atol=1e-05) except: difference = param_grad_global - origin_param_grad avg_diff = difference.abs().sum() / difference.numel() @@ -66,7 +70,7 @@ def check_attention_layer(rank, model_cls, world_size, port): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - config = transformers.GPT2Config(n_position=64, n_layer=1, n_head=16, n_embd=HIDDEN_DIM) + config = transformers.GPT2Config(n_position=64, n_layer=2, n_head=16, n_embd=HIDDEN_DIM) if model_cls == GPT2MLP: model = model_cls(intermediate_size=4 * config.hidden_size, config=config).to('cuda') @@ -111,15 +115,17 @@ def check_attention_layer(rank, model_cls, world_size, port): # [[0, 1] # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) graph = tracer.trace(root=model, meta_args=meta_input_sample) gm = GraphModule(model, graph, model.__class__.__name__) + shape_prop_pass(gm, *meta_input_sample.values()) gm.recompile() strategies_constructor = build_strategy_constructor(graph, device_mesh, 'standard', 'replicated', 'standard') solution = solve_solution(gm, strategies_constructor, memory_budget=-1) - gm, sharding_spec_dicts = transform_to_sharded_model(gm, solution, device_mesh, strategies_constructor) + gm, sharding_spec_dicts = transform_to_sharded_model(gm, meta_input_sample, solution, device_mesh, + strategies_constructor) gm = ModuleWrapper(gm, *sharding_spec_dicts) nodes = [strategies_vector.node for strategies_vector in strategies_constructor.leaf_strategies] @@ -176,13 +182,12 @@ def check_attention_layer(rank, model_cls, world_size, port): @run_on_environment_flag(name='AUTO_PARALLEL') +@pytest.mark.skipif(NO_CODEGEN, reason="no codegen module") @pytest.mark.dist @parameterize('model_cls', [GPT2MLP, GPT2Block, GPT2Attention, GPT2Model]) @rerun_if_address_is_in_use() def test_mlp_layer(model_cls): - world_size = 4 - run_func = partial(check_attention_layer, model_cls=model_cls, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_attention_layer, 4, model_cls=model_cls) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_solver_with_gpt_module.py b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_solver_with_gpt_module.py index 4adb4fbaf047..5a8c3c4bf5a0 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_solver_with_gpt_module.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_solver_with_gpt_module.py @@ -1,15 +1,15 @@ import torch -import torch.nn as nn import transformers from torch.fx import GraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.constants import BATCHNORM_MODULE_OP from colossalai.auto_parallel.tensor_shard.options import SolverOptions -from colossalai.auto_parallel.tensor_shard.solver import CostGraph, GraphAnalyser, Solver, StrategiesConstructor +from colossalai.auto_parallel.tensor_shard.solver import CostGraph, Solver, StrategiesConstructor from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.tracer.tracer import ColoTracer from colossalai.tensor.shape_consistency import ShapeConsistencyManager -from colossalai.testing import parameterize +from colossalai.testing import clear_cache_before_run, parameterize from colossalai.testing.pytest_wrapper import run_on_environment_flag from tests.test_auto_parallel.test_tensor_shard.test_gpt.gpt_modules import GPT2MLP, GPT2Attention, GPT2Block, GPT2Model @@ -19,9 +19,10 @@ @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() @parameterize('model_cls', [GPT2Block, GPT2Attention, GPT2MLP, GPT2Model]) def test_self_attention_block(model_cls): - config = transformers.GPT2Config(n_position=64, n_layer=12, n_head=16, n_embd=HIDDEN_DIM) + config = transformers.GPT2Config(n_position=64, n_layer=2, n_head=16, n_embd=HIDDEN_DIM) if model_cls == GPT2MLP: model = model_cls(intermediate_size=4 * config.hidden_size, config=config) else: @@ -33,7 +34,7 @@ def test_self_attention_block(model_cls): device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) shape_consistency_manager = ShapeConsistencyManager() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) if model_cls == GPT2MLP: input_sample = { 'hidden_states': torch.rand(BATCH_SIZE, SEQ_LENGTH, HIDDEN_DIM).to('meta'), @@ -52,6 +53,7 @@ def test_self_attention_block(model_cls): graph = tracer.trace(root=model, meta_args=input_sample) gm = GraphModule(model, graph, model.__class__.__name__) + shape_prop_pass(gm, *input_sample.values()) print(gm.graph) gm.recompile() solver_options = SolverOptions() diff --git a/tests/test_auto_parallel/test_tensor_shard/test_liveness_analysis.py b/tests/test_auto_parallel/test_tensor_shard/test_liveness_analysis.py index f5de7bf702ff..d10b222c060d 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_liveness_analysis.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_liveness_analysis.py @@ -1,8 +1,13 @@ +import pytest import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.solver import GraphAnalyser from colossalai.fx import ColoGraphModule, ColoTracer +from colossalai.testing import clear_cache_before_run class LinearModel(nn.Module): @@ -22,15 +27,15 @@ def forward(self, x1, x2): return out +@pytest.mark.skip('meta tensor has some bugs in 1.11') +@clear_cache_before_run() def test_liveness_analysis(): model = LinearModel() - tracer = ColoTracer() - graph = tracer.trace(model, - meta_args={ - 'x1': torch.rand(4, 4, device='meta'), - 'x2': torch.rand(4, 4, device='meta') - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'x1': torch.rand(4, 4, device='meta'), 'x2': torch.rand(4, 4, device='meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(root=model, graph=graph, class_name=model.__class__.__name__) + shape_prop_pass(gm, *meta_args.values()) graph_analyser = GraphAnalyser(gm) liveness_list = graph_analyser.liveness_analysis() diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_activation_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_activation_metainfo.py index e41ac4fa690b..e0a2133e654e 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_activation_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_activation_metainfo.py @@ -1,23 +1,14 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn from colossalai.auto_parallel.meta_profiler import meta_register from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy, print_results +from colossalai.testing.utils import clear_cache_before_run, parameterize +from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import print_results @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() @parameterize('func', [ torch.nn.functional.softmax, torch.nn.functional.relu, diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_binary_elementwise_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_binary_elementwise_metainfo.py index 1b745d8906b0..68ccc7835bc3 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_binary_elementwise_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_binary_elementwise_metainfo.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from colossalai.device.device_mesh import DeviceMesh @@ -10,8 +7,7 @@ from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing.utils import rerun_if_address_is_in_use, spawn from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy @@ -62,9 +58,7 @@ def _binary_elementwise_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_binary_elementwise_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_binary_elementwise_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_binary_elementwise_mem_test, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_conv_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_conv_metainfo.py index a973a8182cf3..c6f7b88f44a5 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_conv_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_conv_metainfo.py @@ -1,17 +1,12 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing.utils import rerun_if_address_is_in_use, spawn from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy @@ -25,7 +20,7 @@ def forward(self, input): return nn.functional.conv2d(input, self.conv_weight) -def _conv_module_mem_test(rank, bias, world_size, port): +def _conv_module_mem_test(rank, world_size, port, bias): """This function is for conv memory test Test and print real memory cost and estimated, this test will not be executed except with the tag AUTO_PARALLEL @@ -62,9 +57,7 @@ def _conv_module_mem_test(rank, bias, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_conv_meta_concrete_info_match(bias=False): - world_size = 4 - run_func_module = partial(_conv_module_mem_test, bias=bias, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_conv_module_mem_test, 4, bias=bias) def _conv_function_mem_test(rank, world_size, port): @@ -103,9 +96,7 @@ def _conv_function_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_conv_function_concrete_info_match(): - world_size = 4 - run_func_module = partial(_conv_function_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_conv_function_mem_test, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_embedding_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_embedding_metainfo.py index 2fb1306546ca..e3f76a95c4a5 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_embedding_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_embedding_metainfo.py @@ -1,33 +1,16 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn - -from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - MemoryCost, - OperationData, - OperationDataType, - ShardingStrategy, - StrategiesVector, - TrainCycleItem, -) -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port + +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType +from colossalai.testing.utils import clear_cache_before_run from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import print_results if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + from colossalai.auto_parallel.meta_profiler import meta_register @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() def test_embedding_meta_info(): meta_func = meta_register.get(torch.nn.Embedding) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_linear_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_linear_metainfo.py index e9c0601eb1e4..fb3ded339ddf 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_linear_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_linear_metainfo.py @@ -1,24 +1,14 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ShardingStrategy, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing.utils import rerun_if_address_is_in_use, spawn from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy -if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register - class MyModule(nn.Module): @@ -63,9 +53,7 @@ def _linear_module_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_linear_module_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_linear_module_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_linear_module_mem_test, 4) def _linear_function_mem_test(rank, world_size, port): @@ -101,9 +89,7 @@ def _linear_function_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_linear_function_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_linear_function_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_linear_function_mem_test, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_matmul_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_matmul_metainfo.py index fd29c63fb522..2d2d77f0c637 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_matmul_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_matmul_metainfo.py @@ -1,33 +1,16 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn - -from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - MemoryCost, - OperationData, - OperationDataType, - ShardingStrategy, - StrategiesVector, - TrainCycleItem, -) -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port + +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, TrainCycleItem +from colossalai.testing.utils import clear_cache_before_run, parameterize from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import print_results if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + from colossalai.auto_parallel.meta_profiler import ShardMetaInfo, meta_register @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() @parameterize( 'tensor_shapes', [ diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_norm_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_norm_metainfo.py index 9d3ab9c82670..808172977b60 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_norm_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_norm_metainfo.py @@ -1,29 +1,17 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - MemoryCost, - OperationData, - OperationDataType, - ShardingStrategy, - StrategiesVector, - TrainCycleItem, -) +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, TrainCycleItem from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use, spawn from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy, print_results if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + from colossalai.auto_parallel.meta_profiler import meta_register def _batchnorm_module_mem_test(rank, world_size, port): @@ -62,9 +50,7 @@ def _batchnorm_module_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_batchnorm_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_batchnorm_module_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_batchnorm_module_mem_test, 4) @pytest.mark.skipif(torch.__version__ < '1.12.0', reason='need pytorch 1.12.0 or higher for aten level operations') diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_pooling_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_pooling_metainfo.py index 529686d27d19..4cddf4e19fca 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_pooling_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_pooling_metainfo.py @@ -1,17 +1,12 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing.utils import rerun_if_address_is_in_use, spawn from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy @@ -51,9 +46,7 @@ def _adaptiveavgpool_module_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_adaptiveavgpool_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_adaptiveavgpool_module_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_adaptiveavgpool_module_mem_test, 4) def _maxpool_module_mem_test(rank, world_size, port): @@ -92,9 +85,7 @@ def _maxpool_module_mem_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_maxpool_meta_concrete_info_match(): - world_size = 4 - run_func_module = partial(_maxpool_module_mem_test, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(_maxpool_module_mem_test, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py index a0ab66fdc060..6e8145885d67 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_tensor_metainfo.py @@ -1,30 +1,13 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - MemoryCost, - OperationData, - OperationDataType, - ShardingStrategy, - StrategiesVector, - TrainCycleItem, -) -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType +from colossalai.testing.utils import clear_cache_before_run from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import print_results if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + from colossalai.auto_parallel.meta_profiler import ShardMetaInfo, meta_register class SplitModule(nn.Module): @@ -37,6 +20,7 @@ def forward(self, x): @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() def test_tensor_meta_info(): """test tensor related meta information We will just use torch.Tensor.split for the test diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_where_metainfo.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_where_metainfo.py index 20156f9ab4d5..b4564312eeb4 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_where_metainfo.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_where_metainfo.py @@ -1,31 +1,16 @@ import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn - -from colossalai.auto_parallel.tensor_shard.node_handler import LinearModuleHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - MemoryCost, - OperationData, - OperationDataType, - ShardingStrategy, - StrategiesVector, - TrainCycleItem, -) -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port + +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, TrainCycleItem +from colossalai.testing.utils import clear_cache_before_run from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import print_results if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo, meta_register + from colossalai.auto_parallel.meta_profiler import ShardMetaInfo, meta_register @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() def test_where_meta_info(): meta_func = meta_register.get(torch.where) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/utils.py b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/utils.py index 60ecd1dd9801..4ca85d34da30 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_metainfo/utils.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_metainfo/utils.py @@ -5,16 +5,19 @@ import torch from torch.fx import GraphModule +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes import shape_prop_pass +# from colossalai.fx.tracer.tracer import ColoTracer +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass from colossalai.auto_parallel.tensor_shard.options import SolverOptions from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationDataType, TrainCycleItem from colossalai.auto_parallel.tensor_shard.solver import StrategiesConstructor from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.tracer.tracer import ColoTracer if torch.__version__ >= '1.12.0': - from colossalai.auto_parallel.meta_profiler import MetaInfo + from colossalai.auto_parallel.meta_profiler import ShardMetaInfo def mem_test_for_node_strategy(rank: int, @@ -30,14 +33,16 @@ def mem_test_for_node_strategy(rank: int, model_to_shard, args_to_shard, kwargs_to_shard = copy.deepcopy(model), copy.deepcopy(input_args), copy.deepcopy( input_kwargs) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) input_sample = {} for input_arg, meta_arg_name in zip(input_args, meta_arg_names): input_sample[meta_arg_name] = torch.rand(input_arg.shape).to('meta') for meta_kwarg_name, input_kwarg in input_kwargs.items(): input_sample[meta_kwarg_name] = torch.rand(input_kwarg.shape).to('meta') graph = tracer.trace(root=model_to_shard, meta_args=input_sample) - gm = GraphModule(model_to_shard, graph, model_to_shard.__class__.__name__) + gm = ColoGraphModule(model_to_shard, graph, model_to_shard.__class__.__name__) + shape_prop_pass(gm, *input_sample.values()) + gm.recompile() solver_options = SolverOptions() strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options) strategies_constructor.build_strategies_and_cost() @@ -108,10 +113,10 @@ def mem_test_for_node_strategy(rank: int, # estimated memory if target_node.op == "call_module": - metainfo = MetaInfo(target_node.strategies_vector[strategy_index], - target_node.graph.owning_module.get_submodule(target_node.target)) + metainfo = ShardMetaInfo(target_node.strategies_vector[strategy_index], + target_node.graph.owning_module.get_submodule(target_node.target)) else: - metainfo = MetaInfo(target_node.strategies_vector[strategy_index], target_node.target) + metainfo = ShardMetaInfo(target_node.strategies_vector[strategy_index], target_node.target) print("estimated memory:") print( diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addbmm_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addbmm_handler.py index ffc15e403f35..80e6a6c1460c 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addbmm_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addbmm_handler.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from colossalai.auto_parallel.tensor_shard.node_handler import BMMFunctionHandler @@ -11,9 +8,7 @@ from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -45,7 +40,7 @@ def forward(self, bias, x1, x2): return output -def check_2d_device_mesh(rank, module, bias_shape, using_kwargs, world_size, port): +def check_2d_device_mesh(rank, world_size, port, module, bias_shape, using_kwargs): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = module(using_kwargs).cuda() @@ -249,14 +244,13 @@ def check_1d_device_mesh(rank, module, bias_shape, using_kwargs, world_size, por @parameterize('using_kwargs', [True, False]) @rerun_if_address_is_in_use() def test_2d_device_mesh(module, bias_shape, using_kwargs): - world_size = 4 - run_func = partial(check_2d_device_mesh, - module=module, - bias_shape=bias_shape, - world_size=world_size, - using_kwargs=using_kwargs, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn( + check_2d_device_mesh, + 4, + module=module, + bias_shape=bias_shape, + using_kwargs=using_kwargs, + ) @pytest.mark.skip("skip due to bias cases not ready") @@ -267,14 +261,13 @@ def test_2d_device_mesh(module, bias_shape, using_kwargs): @parameterize('using_kwargs', [True, False]) @rerun_if_address_is_in_use() def test_1d_device_mesh(module, bias_shape, using_kwargs): - world_size = 4 - run_func = partial(check_1d_device_mesh, - module=module, - bias_shape=bias_shape, - using_kwargs=using_kwargs, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn( + check_1d_device_mesh, + 4, + module=module, + bias_shape=bias_shape, + using_kwargs=using_kwargs, + ) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addmm_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addmm_handler.py index aa5a57474335..fe6554cd81ee 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addmm_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_addmm_handler.py @@ -1,27 +1,20 @@ -from faulthandler import disable -from functools import partial -from xml.dom import WrongDocumentErr - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from typing_extensions import Self +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( - OperationData, OperationDataType, ShardingStrategy, StrategiesVector, ) from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -47,7 +40,7 @@ def forward(self, m1): return x -def check_addmm_function_handler(rank, input_shape, model_cls, world_size, port): +def check_addmm_function_handler(rank, world_size, port, input_shape, model_cls): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') if model_cls == AddmmModel: @@ -96,7 +89,7 @@ def check_addmm_function_handler(rank, input_shape, model_cls, world_size, port) meta_arg_names=meta_arg_names, node_type='bias_module') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %m1 : torch.Tensor [#users=1] = placeholder[target=m1] @@ -109,6 +102,7 @@ def check_addmm_function_handler(rank, input_shape, model_cls, world_size, port) # return add graph = tracer.trace(model, meta_args=meta_args_for_tracer) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args_for_tracer.values()) # [input_1, m1, m2, addmm, output] node_list = list(graph.nodes) linear_node = node_list[4] @@ -190,13 +184,7 @@ def check_addmm_function_handler(rank, input_shape, model_cls, world_size, port) @parameterize('model_cls', [AddmmModel, AddmmModel_with_param]) @rerun_if_address_is_in_use() def test_addmm_handler(input_shape, model_cls): - world_size = 4 - run_func_function = partial(check_addmm_function_handler, - input_shape=input_shape, - model_cls=model_cls, - world_size=world_size, - port=free_port()) - mp.spawn(run_func_function, nprocs=world_size) + spawn(check_addmm_function_handler, 4, input_shape=input_shape, model_cls=model_cls) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_batch_norm_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_batch_norm_handler.py index 0ab70abffb4c..c3ceef4c7adf 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_batch_norm_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_batch_norm_handler.py @@ -1,19 +1,16 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.batch_norm_handler import BatchNormModuleHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -30,7 +27,7 @@ def check_bn_module_handler(rank, world_size, port): # the index of bn node in computation graph node_index = 1 # the total number of bn strategies without sync bn mode - # TODO: add sync bn stategies after related passes ready + # TODO: add sync bn strategies after related passes ready strategy_number = 4 numerical_test_for_node_strategy(model=model, device_mesh=device_mesh, @@ -38,13 +35,15 @@ def check_bn_module_handler(rank, world_size, port): strategy_number=strategy_number, input_args=[input], meta_arg_names=['input']) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %_0 : [#users=1] = call_module[target=0](args = (%input_1,), kwargs = {}) # return _0 - graph = tracer.trace(model, meta_args={"input": torch.rand(4, 16, 64, 64).to('meta')}) + meta_args = {"input": torch.rand(4, 16, 64, 64).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) bn_mod_node = list(graph.nodes)[1] strategies_vector = StrategiesVector(bn_mod_node) @@ -110,9 +109,7 @@ def check_bn_module_handler(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_bn_module_handler(): - world_size = 4 - run_func = partial(check_bn_module_handler, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_bn_module_handler, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_function_node.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_function_node.py index 162d1fbba295..800bc11a50e4 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_function_node.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_function_node.py @@ -1,14 +1,10 @@ -from faulthandler import disable -from functools import partial -from xml.dom import WrongDocumentErr - import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn import torch.nn.functional as F -from typing_extensions import Self +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler, LinearModuleHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( OperationData, @@ -17,13 +13,9 @@ StrategiesVector, ) from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy WEIGHT_SHAPE = (32, 16) @@ -66,7 +58,7 @@ def check_linear_module_handler(rank, world_size, port): meta_arg_names=meta_arg_names, node_type='bias_module') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %x : torch.Tensor [#users=1] = placeholder[target=x] # %weight : [#users=1] = get_attr[target=weight] @@ -74,8 +66,10 @@ def check_linear_module_handler(rank, world_size, port): # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%x, %weight), kwargs = {}) # %add : [#users=1] = call_function[target=operator.add](args = (%linear, %bias), kwargs = {}) # return add - graph = tracer.trace(model, meta_args={"x": torch.rand(4, 4, 4, 16).to('meta')}) + meta_args = {"x": torch.rand(4, 4, 4, 16).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_mod_node = list(graph.nodes)[3] strategies_vector = StrategiesVector(linear_mod_node) @@ -168,9 +162,7 @@ def check_linear_module_handler(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_linear_handler(): - world_size = 4 - run_func_module = partial(check_linear_module_handler, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(check_linear_module_handler) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_module_node.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_module_node.py index c5c3f378197e..c29a065d10ba 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_module_node.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bias_linear_module_node.py @@ -1,14 +1,10 @@ -from faulthandler import disable -from functools import partial -from xml.dom import WrongDocumentErr - import pytest import torch -import torch.multiprocessing as mp -import torch.nn as nn -from typing_extensions import Self -from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler, LinearModuleHandler +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer +from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( OperationData, OperationDataType, @@ -16,13 +12,9 @@ StrategiesVector, ) from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -37,7 +29,7 @@ def forward(self, x): return x -def check_linear_module_handler(rank, bias, world_size, port): +def check_linear_module_handler(rank, world_size, port, bias): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = LinearModule(16, 32, bias=bias).cuda() @@ -62,9 +54,11 @@ def check_linear_module_handler(rank, bias, world_size, port): meta_arg_names=meta_arg_names, node_type='bias_module') - tracer = ColoTracer() - graph = tracer.trace(model, meta_args={"x": torch.rand(4, 4, 4, 16).to('meta')}) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {"x": torch.rand(4, 4, 4, 16).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_mod_node = list(graph.nodes)[3] strategies_vector = StrategiesVector(linear_mod_node) @@ -157,9 +151,7 @@ def check_linear_module_handler(rank, bias, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_linear_handler(bias=True): - world_size = 4 - run_func_module = partial(check_linear_module_handler, bias=bias, world_size=world_size, port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) + spawn(check_linear_module_handler, bias=bias) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_binary_elementwise_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_binary_elementwise_handler.py index 50385c0450a8..83f3aafe220e 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_binary_elementwise_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_binary_elementwise_handler.py @@ -1,23 +1,20 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import BinaryElementwiseHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy -def check_binary_elementwise_handler_with_tensor(rank, op, other_dim, world_size, port): +def check_binary_elementwise_handler_with_tensor(rank, world_size, port, op, other_dim): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') @@ -52,10 +49,11 @@ def forward(self, x1, x2): input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) meta_args = {'x1': torch.rand(4, 4).to('meta'), 'x2': torch.rand([4] * other_dim).to('meta')} graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) op_node = list(graph.nodes)[2] strategies_vector = StrategiesVector(op_node) @@ -146,7 +144,7 @@ def forward(self, x1): return out -def check_binary_elementwise_handler_with_int(rank, op, other_dim, model_cls, world_size, port): +def check_binary_elementwise_handler_with_int(rank, world_size, port, op, other_dim, model_cls): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') @@ -172,12 +170,11 @@ def check_binary_elementwise_handler_with_int(rank, op, other_dim, model_cls, wo strategy_number=strategy_number, input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) meta_args = {'x1': torch.rand(4, 4).to('meta')} graph = tracer.trace(model, meta_args=meta_args) - print(graph) - # assert False gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) if model_cls == BEOpModelWithNodeConst: op_node = list(graph.nodes)[2] @@ -234,13 +231,12 @@ def check_binary_elementwise_handler_with_int(rank, op, other_dim, model_cls, wo @pytest.mark.dist @rerun_if_address_is_in_use() def test_binary_elementwise_handler_with_tensor(op, other_dim): - world_size = 4 - run_func_tensor = partial(check_binary_elementwise_handler_with_tensor, - op=op, - other_dim=other_dim, - world_size=world_size, - port=free_port()) - mp.spawn(run_func_tensor, nprocs=world_size) + spawn( + check_binary_elementwise_handler_with_tensor, + 4, + op=op, + other_dim=other_dim, + ) @run_on_environment_flag(name='AUTO_PARALLEL') @@ -250,14 +246,13 @@ def test_binary_elementwise_handler_with_tensor(op, other_dim): @pytest.mark.dist @rerun_if_address_is_in_use() def test_binary_elementwise_handler_with_int(op, model_cls, other_dim): - world_size = 4 - run_func_int = partial(check_binary_elementwise_handler_with_int, - op=op, - model_cls=model_cls, - other_dim=other_dim, - world_size=world_size, - port=free_port()) - mp.spawn(run_func_int, nprocs=world_size) + spawn( + check_binary_elementwise_handler_with_int, + 4, + op=op, + model_cls=model_cls, + other_dim=other_dim, + ) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bmm_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bmm_handler.py index 02c7e0671149..f4fdc458f80e 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bmm_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_bmm_handler.py @@ -1,19 +1,16 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import BMMFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -52,13 +49,11 @@ def check_2d_device_mesh(rank, module, world_size, port): strategy_number=strategy_number, input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() - graph = tracer.trace(model, - meta_args={ - "x1": torch.rand(4, 8, 16).to('meta'), - 'x2': torch.rand(4, 16, 8).to('meta') - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'x1': torch.rand(4, 8, 16).to('meta'), 'x2': torch.rand(4, 16, 8).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_mod_node = list(graph.nodes)[2] strategies_vector = StrategiesVector(linear_mod_node) @@ -147,13 +142,11 @@ def check_1d_device_mesh(rank, module, world_size, port): strategy_number=strategy_number, input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() - graph = tracer.trace(model, - meta_args={ - "x1": torch.rand(4, 8, 16).to('meta'), - 'x2': torch.rand(4, 16, 8).to('meta') - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'x1': torch.rand(4, 8, 16).to('meta'), 'x2': torch.rand(4, 16, 8).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_mod_node = list(graph.nodes)[2] strategies_vector = StrategiesVector(linear_mod_node) @@ -205,14 +198,12 @@ def check_1d_device_mesh(rank, module, world_size, port): @run_on_environment_flag(name='AUTO_PARALLEL') @parameterize('module', [BMMTensorMethodModule, BMMTorchFunctionModule]) +@parameterize('module', [BMMTensorMethodModule, BMMTorchFunctionModule]) @pytest.mark.dist @rerun_if_address_is_in_use() def test_bmm_handler(module): - world_size = 4 - run_func_2d = partial(check_2d_device_mesh, module=module, world_size=world_size, port=free_port()) - mp.spawn(run_func_2d, nprocs=world_size) - run_func_1d = partial(check_1d_device_mesh, module=module, world_size=world_size, port=free_port()) - mp.spawn(run_func_1d, nprocs=world_size) + spawn(check_2d_device_mesh, 4, module=module) + spawn(check_1d_device_mesh, 4, module=module) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_conv_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_conv_handler.py index 2acd015c8f59..f9632b1cd8f9 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_conv_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_conv_handler.py @@ -1,23 +1,20 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler, ConvModuleHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy -def check_conv_module_handler(rank, bias, world_size, port): +def check_conv_module_handler(rank, world_size, port, bias): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = nn.Sequential(nn.Conv2d(4, 16, 3, padding=1, bias=bias)).cuda() @@ -41,9 +38,11 @@ def check_conv_module_handler(rank, bias, world_size, port): strategy_number=strategy_number, input_args=[input], meta_arg_names=['input']) - tracer = ColoTracer() - graph = tracer.trace(model, meta_args={"input": torch.rand(4, 4, 64, 64).to('meta')}) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'input': torch.rand(4, 4, 64, 64).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) conv_mod_node = list(graph.nodes)[1] strategies_vector = StrategiesVector(conv_mod_node) @@ -151,7 +150,7 @@ def forward(self, input, others, bias=None): return x -def check_conv_function_handler(rank, bias, world_size, port): +def check_conv_function_handler(rank, world_size, port, bias): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = ConvModel().cuda() @@ -178,7 +177,7 @@ def check_conv_function_handler(rank, bias, world_size, port): meta_arg_names=meta_arg_names, input_kwargs=input_kwargs) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %others : torch.Tensor [#users=1] = placeholder[target=others] @@ -189,6 +188,7 @@ def check_conv_function_handler(rank, bias, world_size, port): meta_args['bias'] = torch.rand(16).to('meta') graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) if bias: conv_mod_node = list(graph.nodes)[3] @@ -297,9 +297,7 @@ def check_conv_function_handler(rank, bias, world_size, port): # @parameterize('bias', [True, False]) @rerun_if_address_is_in_use() def test_conv_module_handler(bias=False): - world_size = 4 - run_func = partial(check_conv_module_handler, bias=bias, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_conv_module_handler, 4, bias=bias) @run_on_environment_flag(name='AUTO_PARALLEL') @@ -309,9 +307,7 @@ def test_conv_module_handler(bias=False): # @parameterize('bias', [True, False]) @rerun_if_address_is_in_use() def test_conv_function_handler(bias=False): - world_size = 4 - run_func = partial(check_conv_function_handler, bias=bias, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_conv_function_handler, 4, bias=bias) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_default_reshape_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_default_reshape_handler.py index ea7c2b729635..64f56ba98e2b 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_default_reshape_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_default_reshape_handler.py @@ -1,12 +1,14 @@ import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import DefaultReshapeHandler from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag class ReshapeModel(nn.Module): @@ -21,21 +23,23 @@ def forward(self, input, other): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_reshape_handler(): model = ReshapeModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %other : torch.Tensor [#users=1] = placeholder[target=other] # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %other), kwargs = {}) # %view : [#users=1] = call_method[target=view](args = (%conv2d, 2, -1), kwargs = {}) # return view - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(4, 4, 64, 64).to('meta'), - "other": torch.rand(4, 16, 3, 3).to('meta'), - }) + meta_args = { + "input": torch.rand(4, 4, 64, 64).to('meta'), + "other": torch.rand(16, 4, 3, 3).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) @@ -67,13 +71,13 @@ def test_reshape_handler(): assert mapping['input'].name == "conv2d" assert mapping['input'].data.is_meta - assert mapping['input'].data.shape == torch.Size([4, 4, 62, 62]) + assert mapping['input'].data.shape == torch.Size([4, 16, 62, 62]) assert mapping['input'].type == OperationDataType.ARG - assert mapping['input'].logical_shape == torch.Size([4, 4, 62, 62]) + assert mapping['input'].logical_shape == torch.Size([4, 16, 62, 62]) assert mapping['output'].name == "view" assert mapping['output'].data.is_meta - assert mapping['output'].data.shape == torch.Size([2, 30752]) + assert mapping['output'].data.shape == torch.Size([2, 123008]) assert mapping['output'].type == OperationDataType.OUTPUT # reshape handler is a following strategy handler, so the number of strategies is equal to the predecessor node. diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_embedding_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_embedding_handler.py index 5bce383dd0ab..4fa0313b1cb5 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_embedding_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_embedding_handler.py @@ -1,22 +1,20 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.embedding_handler import ( EmbeddingFunctionHandler, EmbeddingModuleHandler, ) from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy NUM_EMBEDDINGS = 16 @@ -60,9 +58,11 @@ def check_embedding_module_handler(rank, world_size, port): input_args=[input], meta_arg_names=['input']) - tracer = ColoTracer() - graph = tracer.trace(model, meta_args={"input": torch.rand(4, 16, 16).to('meta')}) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {"input": torch.randint(NUM_EMBEDDINGS, (4, 16, 16)).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) embedding_node = list(graph.nodes)[1] strategies_vector = StrategiesVector(embedding_node) @@ -171,18 +171,19 @@ def check_embedding_function_handler(rank, world_size, port): input_args=input_args, meta_arg_names=meta_arg_names, input_kwargs=input_kwargs) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %others : torch.Tensor [#users=1] = placeholder[target=others] # %embedding : [#users=1] = call_function[target=torch.nn.functional.embedding](args = (%input_1, %others), kwargs = {padding_idx: None, max_norm: None, norm_type: 2.0, scale_grad_by_freq: False, sparse: False}) # return embedding meta_args = { - "input": torch.rand(4, 16, 16).to('meta'), + "input": torch.randint(NUM_EMBEDDINGS, (4, 16, 16)).to('meta'), "others": torch.rand(NUM_EMBEDDINGS, EMBEDDING_DIMS).to('meta') } graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) embedding_node = list(graph.nodes)[2] strategies_vector = StrategiesVector(embedding_node) @@ -267,18 +268,14 @@ def check_embedding_function_handler(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_embedding_module_handler(): - world_size = 4 - run_func = partial(check_embedding_module_handler, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_embedding_module_handler, 4) @run_on_environment_flag(name='AUTO_PARALLEL') @pytest.mark.dist @rerun_if_address_is_in_use() def test_embedding_function_handler(): - world_size = 4 - run_func = partial(check_embedding_function_handler, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_embedding_function_handler, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getattr_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getattr_handler.py index 681e93a5fe16..a089df743ec0 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getattr_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getattr_handler.py @@ -1,10 +1,14 @@ +import pytest import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.getattr_handler import GetattrHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer +from colossalai.testing import clear_cache_before_run class GetattrModel(nn.Module): @@ -18,15 +22,19 @@ def forward(self, input): return weight +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') +@clear_cache_before_run() def test_getattr_handler(): model = GetattrModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=0] = placeholder[target=input] # %conv_weight : [#users=1] = get_attr[target=conv.weight] # return conv_weight - graph = tracer.trace(model, meta_args={'input': torch.rand(4, 4, 64, 64).to('meta')}) + meta_args = {'input': torch.rand(4, 4, 64, 64).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getitem_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getitem_handler.py index c72d2a6a80e8..a2e0968b18bb 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getitem_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_getitem_handler.py @@ -2,22 +2,21 @@ import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.default_reshape_handler import DefaultReshapeHandler from colossalai.auto_parallel.tensor_shard.node_handler.getitem_handler import GetItemHandler from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.placeholder_handler import PlaceholderHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.fx.tracer.meta_patch.patched_module import linear from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -58,15 +57,15 @@ def check_getitem_from_tensor_handler(rank, getitem_index, world_size, port): meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() - - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = { + "input": torch.rand(8, 16, 64, 32).to('meta'), + "other": torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *list(meta_args.values())) linear_mod_node = list(graph.nodes)[2] getitem_mod_node = list(graph.nodes)[3] getitem_strategies_vector = StrategiesVector(getitem_mod_node) @@ -101,12 +100,7 @@ def check_getitem_from_tensor_handler(rank, getitem_index, world_size, port): # @parameterize('getitem_index', [slice(0, 2), (slice(None), slice(None))]) @parameterize('getitem_index', [1, (1, 4), slice(0, 2), (slice(None), slice(None))]) def test_getitem_from_tensor_handler(getitem_index): - world_size = 4 - run_func = partial(check_getitem_from_tensor_handler, - getitem_index=getitem_index, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_getitem_from_tensor_handler, 4) class GetItemFromTupleModel(nn.Module): @@ -121,6 +115,7 @@ def forward(self, input): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_getitem_from_tuple_handler(): model = GetItemFromTupleModel() tracer = ColoTracer() @@ -129,10 +124,12 @@ def test_getitem_from_tuple_handler(): # %split : [#users=1] = call_function[target=torch.functional.split](args = (%conv2d, 2), kwargs = {dim: 0}) # %getitem : [#users=1] = call_function[target=operator.getitem](args = (%split, 1), kwargs = {}) # return getitem - graph = tracer.trace(model, meta_args={ + meta_args = { "input": torch.rand(4, 4, 64, 64).to('meta'), - }) + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_layer_norm_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_layer_norm_handler.py index f4d0063fd6b6..ad72c2026b9a 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_layer_norm_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_layer_norm_handler.py @@ -1,20 +1,17 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.layer_norm_handler import LayerNormModuleHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.fx.tracer.meta_patch.patched_module import linear from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -40,13 +37,15 @@ def check_ln_module_handler(rank, world_size, port): strategy_number=strategy_number, input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %_0 : [#users=1] = call_module[target=0](args = (%input_1,), kwargs = {}) # return _0 - graph = tracer.trace(model, meta_args={"input": torch.rand(4, 16).to('meta')}) + meta_args = {"input": torch.rand(4, 16).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) ln_mod_node = list(graph.nodes)[1] strategies_vector = StrategiesVector(ln_mod_node) @@ -100,9 +99,7 @@ def check_ln_module_handler(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_ln_module_handler(): - world_size = 4 - run_func = partial(check_ln_module_handler, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_ln_module_handler, 4) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_linear_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_linear_handler.py index 18afacf56b8e..ec695cd8f7b9 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_linear_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_linear_handler.py @@ -1,10 +1,10 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler, LinearModuleHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( OperationData, @@ -13,17 +13,15 @@ StrategiesVector, ) from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag from colossalai.testing.utils import parameterize -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy -def check_linear_module_handler(rank, bias, input_shape, world_size, port): +def check_linear_module_handler(rank, world_size, port, bias, input_shape): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = nn.Sequential(nn.Linear(16, 32, bias=bias)).cuda() @@ -49,9 +47,11 @@ def check_linear_module_handler(rank, bias, input_shape, world_size, port): input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() - graph = tracer.trace(model, meta_args={"input": torch.rand(input_shape).to('meta')}) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {"input": torch.rand(input_shape).cuda()} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_mod_node = list(graph.nodes)[1] strategies_vector = StrategiesVector(linear_mod_node) @@ -168,7 +168,7 @@ def forward(self, input, others, bias=None): return x -def check_linear_function_handler(rank, bias, input_shape, world_size, port): +def check_linear_function_handler(rank, world_size, port, bias, input_shape): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = LinearModel().cuda() @@ -196,13 +196,12 @@ def check_linear_function_handler(rank, bias, input_shape, world_size, port): input_args=input_args, meta_arg_names=meta_arg_names) - tracer = ColoTracer() - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(input_shape).to('meta'), - 'others': torch.rand(32, 16).to('meta') - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'input': torch.rand(input_shape).to('meta'), 'others': torch.rand(32, 16).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) + if bias: linear_func_node = list(graph.nodes)[3] else: @@ -310,19 +309,18 @@ def check_linear_function_handler(rank, bias, input_shape, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_linear_handler(input_shape, bias=False): - world_size = 4 - run_func_module = partial(check_linear_module_handler, - bias=bias, - input_shape=input_shape, - world_size=world_size, - port=free_port()) - mp.spawn(run_func_module, nprocs=world_size) - run_func_function = partial(check_linear_function_handler, - bias=bias, - input_shape=input_shape, - world_size=world_size, - port=free_port()) - mp.spawn(run_func_function, nprocs=world_size) + spawn( + check_linear_module_handler, + 4, + bias=bias, + input_shape=input_shape, + ) + spawn( + check_linear_function_handler, + 4, + bias=bias, + input_shape=input_shape, + ) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_matmul_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_matmul_handler.py index 91b3ae27d599..938acd3d1eea 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_matmul_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_matmul_handler.py @@ -2,6 +2,9 @@ import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.matmul_handler import ( MatMulHandler, MatMulType, @@ -15,8 +18,7 @@ StrategiesVector, ) from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing.utils import parameterize +from colossalai.testing.utils import clear_cache_before_run, parameterize class MatMulModule(nn.Module): @@ -26,6 +28,7 @@ def forward(self, x1, x2): @pytest.mark.skipif(torch.__version__ < '1.12.0', reason="need pytorch 1.12.0 or higher for aten level operations") +@clear_cache_before_run() @parameterize( 'tensor_shapes', [ @@ -57,9 +60,11 @@ def test_matmul_node_handler(tensor_shapes): model = MatMulModule() - tracer = ColoTracer() - graph = tracer.trace(model, meta_args={"x1": x1.to('meta'), 'x2': x2.to('meta')}) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {"x1": x1.to('meta'), 'x2': x2.to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) print(graph) @@ -124,7 +129,6 @@ def test_matmul_node_handler(tensor_shapes): input_sharding_spec = strategy.get_sharding_spec_by_name('x1') other_sharding_spec = strategy.get_sharding_spec_by_name('x2') output_sharding_spec = strategy.get_sharding_spec_by_name('matmul') - if matmul_type == MatMulType.DOT: # dot product will produce a scaler # results should fulfill: @@ -159,7 +163,10 @@ def test_matmul_node_handler(tensor_shapes): if len(other_shape) > 1: assert other_sharding_spec.sharding_sequence[-1] == output_sharding_spec.sharding_sequence[-1] if len(input_shape) > 1: - assert input_sharding_spec.sharding_sequence[-2] == output_sharding_spec.sharding_sequence[-2] + if len(other_shape) == 1: + assert input_sharding_spec.sharding_sequence[-2] == output_sharding_spec.sharding_sequence[-1] + else: + assert input_sharding_spec.sharding_sequence[-2] == output_sharding_spec.sharding_sequence[-2] if len(other_shape) > 2: assert other_sharding_spec.sharding_sequence[-2] == input_sharding_spec.sharding_sequence[-1] diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_norm_pooling_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_norm_pooling_handler.py index f219bc2f3976..6bff9f9648e2 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_norm_pooling_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_norm_pooling_handler.py @@ -1,26 +1,29 @@ -import pytest import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.normal_pooling_handler import NormPoolingHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.fx.tracer.meta_patch.patched_module import linear -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_norm_pool_handler(): model = nn.Sequential(nn.MaxPool2d(4, padding=1).to('meta')) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %_0 : [#users=1] = call_module[target=0](args = (%input_1,), kwargs = {}) # return _0 - graph = tracer.trace(model, meta_args={"input": torch.rand(4, 4, 64, 64).to('meta')}) + meta_args = {"input": torch.rand(4, 4, 64, 64).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_output_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_output_handler.py index 26376c429ebc..1703d5ded2f2 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_output_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_output_handler.py @@ -1,11 +1,14 @@ +import pytest import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.output_handler import OutputHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import clear_cache_before_run, parameterize class OutputModel(nn.Module): @@ -18,19 +21,20 @@ def forward(self, x): return x, y +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') @parameterize('output_option', ['distributed', 'replicated']) -@rerun_if_address_is_in_use() +@clear_cache_before_run() def test_output_handler(output_option): model = OutputModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %x : torch.Tensor [#users=2] = placeholder[target=x] # %mul : [#users=1] = call_function[target=operator.mul](args = (%x, 2), kwargs = {}) # return (x, mul) - graph = tracer.trace(model, meta_args={ - "x": torch.rand(4, 4, 64, 64).to('meta'), - }) + meta_args = {'x': torch.rand(4, 4, 64, 64).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) @@ -39,14 +43,14 @@ def test_output_handler(output_option): output_strategies_vector = StrategiesVector(output_node) # build handler - otuput_handler = OutputHandler(node=output_node, + output_handler = OutputHandler(node=output_node, device_mesh=device_mesh, strategies_vector=output_strategies_vector, output_option=output_option) - otuput_handler.register_strategy(compute_resharding_cost=False) + output_handler.register_strategy(compute_resharding_cost=False) # check operation data mapping - mapping = otuput_handler.get_operation_data_mapping() + mapping = output_handler.get_operation_data_mapping() for name, op_data in mapping.items(): op_data: OperationData @@ -55,7 +59,7 @@ def test_output_handler(output_option): assert mapping['output'].name == "output" assert mapping['output'].type == OperationDataType.OUTPUT - strategy_name_list = [val.name for val in otuput_handler.strategies_vector] + strategy_name_list = [val.name for val in output_handler.strategies_vector] if output_option == 'distributed': assert "Distributed Output" in strategy_name_list else: diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_permute_and_transpose_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_permute_and_transpose_handler.py index af03481d830e..f071cd120fb7 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_permute_and_transpose_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_permute_and_transpose_handler.py @@ -2,20 +2,20 @@ import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import PermuteHandler, TransposeHandler from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -53,7 +53,7 @@ def forward(self, input, other): return permute_node -def check_view_handler(rank, call_function, reshape_dims, model_cls, world_size, port): +def check_view_handler(rank, world_size, port, call_function, reshape_dims, model_cls): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') if call_function == torch.permute: @@ -88,7 +88,7 @@ def check_view_handler(rank, call_function, reshape_dims, model_cls, world_size, input_args=[input, other], meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) if model_cls.__name__ == 'ConvReshapeModel': # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] @@ -96,11 +96,11 @@ def check_view_handler(rank, call_function, reshape_dims, model_cls, world_size, # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %other), kwargs = {bias: None}) # %permute : [#users=1] = call_function[target=torch.permute](args = (%conv2d, (0, 2, 1, 3)), kwargs = {}) # return permute - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 8, 66, 66).to('meta'), - "other": torch.rand(16, 8, 3, 3).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 8, 66, 66).to('meta'), + 'other': torch.rand(16, 8, 3, 3).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) if model_cls.__name__ == 'LinearReshapeModel': # graph(): @@ -109,13 +109,14 @@ def check_view_handler(rank, call_function, reshape_dims, model_cls, world_size, # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%input_1, %other), kwargs = {bias: None}) # %permute : [#users=1] = call_method[target=view](args = (%linear, 32, 4, 32, 32, 4), kwargs = {}) # return permute - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 16, 64, 32).to('meta'), + 'other': torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) previous_mod_node = list(graph.nodes)[2] reshape_node = list(graph.nodes)[3] @@ -325,14 +326,13 @@ def check_view_handler(rank, call_function, reshape_dims, model_cls, world_size, @parameterize('reshape_dims', [((0, 2, 1, 3), (1, 2)), ((2, 0, 1, 3), (1, 3))]) @parameterize('model_cls', [ConvReshapeModel, LinearReshapeModel]) def test_view_handler(call_function, reshape_dims, model_cls): - world_size = 4 - run_func = partial(check_view_handler, - call_function=call_function, - reshape_dims=reshape_dims, - model_cls=model_cls, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn( + check_view_handler, + 4, + call_function=call_function, + reshape_dims=reshape_dims, + model_cls=model_cls, + ) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_placeholder_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_placeholder_handler.py index 9bc453a27cdc..6d02b0e0ba74 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_placeholder_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_placeholder_handler.py @@ -1,11 +1,14 @@ +import pytest import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.placeholder_handler import PlaceholderHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import clear_cache_before_run, parameterize class PlaceholderModel(nn.Module): @@ -17,18 +20,21 @@ def forward(self, input): return input +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') @parameterize('placeholder_option', ['distributed', 'replicated']) -@rerun_if_address_is_in_use() +@clear_cache_before_run() def test_placeholder_handler(placeholder_option): model = PlaceholderModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # return input_1 - graph = tracer.trace(model, meta_args={ + meta_args = { "input": torch.rand(4, 4, 64, 64).to('meta'), - }) + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_shard_option.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_shard_option.py index f6895d92ab03..14c364c45fc4 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_shard_option.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_shard_option.py @@ -1,17 +1,14 @@ -from functools import partial - import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.options import ShardOption from colossalai.auto_parallel.tensor_shard.sharding_strategy import StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing import parameterize -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.testing.utils import parameterize +from colossalai.testing import clear_cache_before_run, run_on_environment_flag class LinearModel(nn.Module): @@ -30,13 +27,11 @@ def check_shard_option(shard_option): mesh_shape = (2, 2) device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - tracer = ColoTracer() - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(4, 4, 4, 16).to('meta'), - 'others': torch.rand(32, 16).to('meta') - }) + tracer = ColoTracer(bias_addition_split=True) + meta_args = {'input': torch.rand(4, 4, 4, 16).to('meta'), 'others': torch.rand(32, 16).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) linear_func_node = list(graph.nodes)[2] strategies_vector = StrategiesVector(linear_func_node) @@ -112,6 +107,7 @@ def check_shard_option(shard_option): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_shard_option(): # for shard_option in [ShardOption.STANDARD, ShardOption.SHARD, ShardOption.FULL_SHARD, ShardOption.SHARD_LAST_AXIS]: for shard_option in [ShardOption.SHARD_LAST_AXIS]: diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_softmax_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_softmax_handler.py index c43ee292bedf..75ae0416ef98 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_softmax_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_softmax_handler.py @@ -1,21 +1,18 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn import torch.nn.functional as F +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.softmax_handler import SoftmaxHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -31,7 +28,7 @@ def forward(self, input, other): return softmax_node -def check_split_handler(rank, softmax_dim, model_cls, world_size, port): +def check_split_handler(rank, world_size, port, softmax_dim, model_cls): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = model_cls(softmax_dim=softmax_dim).cuda() @@ -54,7 +51,7 @@ def check_split_handler(rank, softmax_dim, model_cls, world_size, port): input_args=[input, other], meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] @@ -62,13 +59,14 @@ def check_split_handler(rank, softmax_dim, model_cls, world_size, port): # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%input_1, %other), kwargs = {bias: None}) # %softmax : [#users=1] = call_method[target=split](args = (%linear,), kwargs = {}) # return split - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 16, 64, 32).to('meta'), + 'other': torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) previous_mod_node = list(graph.nodes)[2] split_node = list(graph.nodes)[3] @@ -173,13 +171,7 @@ def check_split_handler(rank, softmax_dim, model_cls, world_size, port): @parameterize('softmax_dim', [0, 1, 2, 3]) @parameterize('model_cls', [LinearSplitModel]) def test_split_handler(softmax_dim, model_cls): - world_size = 4 - run_func = partial(check_split_handler, - softmax_dim=softmax_dim, - model_cls=model_cls, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_split_handler, 4, softmax_dim=softmax_dim, model_cls=model_cls) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_split_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_split_handler.py index 044aef19d38d..f860c629b0a0 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_split_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_split_handler.py @@ -1,21 +1,18 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import SplitHandler from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -45,7 +42,7 @@ def forward(self, input, other): return split_node -def check_split_handler(rank, split_size, split_dim, model_cls, world_size, port): +def check_split_handler(rank, world_size, port, split_size, split_dim, model_cls): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = model_cls(split_size=split_size, split_dim=split_dim).cuda() @@ -76,7 +73,7 @@ def check_split_handler(rank, split_size, split_dim, model_cls, world_size, port input_args=[input, other], meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) if model_cls.__name__ == 'ConvSplitModel': # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] @@ -84,11 +81,11 @@ def check_split_handler(rank, split_size, split_dim, model_cls, world_size, port # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %other), kwargs = {}) # %split : [#users=1] = call_method[target=split](args = (%conv2d,), kwargs = {}) # return split - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 8, 66, 66).to('meta'), - "other": torch.rand(16, 8, 3, 3).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 8, 66, 66).to('meta'), + 'other': torch.rand(16, 8, 3, 3).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) if model_cls.__name__ == 'LinearSplitModel': # graph(): @@ -97,13 +94,14 @@ def check_split_handler(rank, split_size, split_dim, model_cls, world_size, port # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%input_1, %other), kwargs = {bias: None}) # %split : [#users=1] = call_method[target=split](args = (%linear,), kwargs = {}) # return split - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 16, 64, 32).to('meta'), + 'other': torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) previous_mod_node = list(graph.nodes)[2] split_node = list(graph.nodes)[3] @@ -255,14 +253,7 @@ def check_split_handler(rank, split_size, split_dim, model_cls, world_size, port @parameterize('split_dim', [0, 1, 2]) @parameterize('model_cls', [ConvSplitModel, LinearSplitModel]) def test_split_handler(split_size, split_dim, model_cls): - world_size = 4 - run_func = partial(check_split_handler, - split_size=split_size, - split_dim=split_dim, - model_cls=model_cls, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_split_handler, 4, split_size=split_size, split_dim=split_dim, model_cls=model_cls) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_sum_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_sum_handler.py index 5fda4de1a101..c11291ecac96 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_sum_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_sum_handler.py @@ -1,21 +1,17 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.sum_handler import SumHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, run_on_environment_flag, spawn from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -35,7 +31,7 @@ def forward(self, input, other): return sum_node -def check_sum_handler(rank, sum_dims, keepdim, world_size, port): +def check_sum_handler(rank, world_size, port, sum_dims, keepdim): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') model = LinearSumModel(sum_dims=sum_dims, keepdim=keepdim).cuda() @@ -58,7 +54,7 @@ def check_sum_handler(rank, sum_dims, keepdim, world_size, port): meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] @@ -66,12 +62,13 @@ def check_sum_handler(rank, sum_dims, keepdim, world_size, port): # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%input_1, %other), kwargs = {bias: None}) # %sum_1 : [#users=1] = call_function[target=torch.sum](args = (%linear,), kwargs = {}) # return sum_1 - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + meta_args = { + "input": torch.rand(8, 16, 64, 32).to('meta'), + "other": torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) previous_mod_node = list(graph.nodes)[2] sum_node = list(graph.nodes)[3] @@ -116,107 +113,107 @@ def check_sum_handler(rank, sum_dims, keepdim, world_size, port): # check strategy name if sum_dims == (0, 2) and keepdim == False: - assert '[R, R, R, S1] -> [R, S1]_0' in strategy_name_list - assert '[R, S0, R, S1] -> [S0, S1]_1' in strategy_name_list - assert '[R, R, R, S1] -> [R, S1]_2' in strategy_name_list - assert '[R, R, R, S0] -> [R, S0]_3' in strategy_name_list - assert '[R, S1, R, S0] -> [S1, S0]_4' in strategy_name_list - assert '[R, R, R, S0] -> [R, S0]_5' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_6' in strategy_name_list - assert '[R, S0, R, R] -> [S0, R]_7' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_0' in strategy_name_list + assert '[R, S01, R, R] -> [S01, R]_1' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_2' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_3' in strategy_name_list + assert '[R, R, R, S01] -> [R, S01]_4' in strategy_name_list + assert '[R, R, R, S1] -> [R, S1]_5' in strategy_name_list + assert '[R, R, R, S0] -> [R, S0]_6' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_7' in strategy_name_list assert '[R, R, R, R] -> [R, R]_8' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_9' in strategy_name_list - assert '[R, S1, R, R] -> [S1, R]_10' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_11' in strategy_name_list - assert '[R, R, R, S1] -> [R, S1]_12' in strategy_name_list - assert '[R, R, R, S0] -> [R, S0]_13' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_14' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_15' in strategy_name_list + assert '[R, R, R, S0] -> [R, S0]_9' in strategy_name_list + assert '[R, R, R, S1] -> [R, S1]_10' in strategy_name_list + assert '[R, R, R, S1] -> [R, S1]_11' in strategy_name_list + assert '[R, S0, R, S1] -> [S0, S1]_12' in strategy_name_list + assert '[R, R, R, S1] -> [R, S1]_13' in strategy_name_list + assert '[R, R, R, S0] -> [R, S0]_14' in strategy_name_list + assert '[R, S1, R, S0] -> [S1, S0]_15' in strategy_name_list assert '[R, R, R, S0] -> [R, S0]_16' in strategy_name_list - assert '[R, R, R, S1] -> [R, S1]_17' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_18' in strategy_name_list - assert '[R, S01, R, R] -> [S01, R]_19' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_17' in strategy_name_list + assert '[R, S0, R, R] -> [S0, R]_18' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_19' in strategy_name_list assert '[R, R, R, R] -> [R, R]_20' in strategy_name_list - assert '[R, R, R, R] -> [R, R]_21' in strategy_name_list - assert '[R, R, R, S01] -> [R, S01]_22' in strategy_name_list + assert '[R, S1, R, R] -> [S1, R]_21' in strategy_name_list + assert '[R, R, R, R] -> [R, R]_22' in strategy_name_list assert '[R, R, R, R] -> [R, R]_23' in strategy_name_list if sum_dims == (0, 2) and keepdim == True: - assert '[R, R, R, S1] -> [R, R, R, S1]_0' in strategy_name_list - assert '[R, S0, R, S1] -> [R, S0, R, S1]_1' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, R, S1]_2' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_3' in strategy_name_list - assert '[R, S1, R, S0] -> [R, S1, R, S0]_4' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_5' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_6' in strategy_name_list - assert '[R, S0, R, R] -> [R, S0, R, R]_7' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_0' in strategy_name_list + assert '[R, S01, R, R] -> [R, S01, R, R]_1' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_2' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_3' in strategy_name_list + assert '[R, R, R, S01] -> [R, R, R, S01]_4' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_5' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_6' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_7' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_8' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_9' in strategy_name_list - assert '[R, S1, R, R] -> [R, S1, R, R]_10' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_11' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, R, S1]_12' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_13' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_14' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_15' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_9' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_10' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_11' in strategy_name_list + assert '[R, S0, R, S1] -> [R, S0, R, S1]_12' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_13' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_14' in strategy_name_list + assert '[R, S1, R, S0] -> [R, S1, R, S0]_15' in strategy_name_list assert '[R, R, R, S0] -> [R, R, R, S0]_16' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, R, S1]_17' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_18' in strategy_name_list - assert '[R, S01, R, R] -> [R, S01, R, R]_19' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_17' in strategy_name_list + assert '[R, S0, R, R] -> [R, S0, R, R]_18' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_19' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_20' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_21' in strategy_name_list - assert '[R, R, R, S01] -> [R, R, R, S01]_22' in strategy_name_list + assert '[R, S1, R, R] -> [R, S1, R, R]_21' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_22' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_23' in strategy_name_list if sum_dims == 1 and keepdim == False: - assert '[S0, R, R, S1] -> [S0, R, S1]_0' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, S1]_1' in strategy_name_list - assert '[R, R, S0, S1] -> [R, S0, S1]_2' in strategy_name_list - assert '[S1, R, R, S0] -> [S1, R, S0]_3' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, S0]_4' in strategy_name_list - assert '[R, R, S1, S0] -> [R, S1, S0]_5' in strategy_name_list - assert '[S0, R, R, R] -> [S0, R, R]_6' in strategy_name_list + assert '[S01, R, R, R] -> [S01, R, R]_0' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R]_1' in strategy_name_list + assert '[R, R, S01, R] -> [R, S01, R]_2' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R]_3' in strategy_name_list + assert '[R, R, R, S01] -> [R, R, S01]_4' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, S1]_5' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, S0]_6' in strategy_name_list assert '[R, R, R, R] -> [R, R, R]_7' in strategy_name_list - assert '[R, R, S0, R] -> [R, S0, R]_8' in strategy_name_list - assert '[S1, R, R, R] -> [S1, R, R]_9' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R]_10' in strategy_name_list - assert '[R, R, S1, R] -> [R, S1, R]_11' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R]_8' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, S0]_9' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, S1]_10' in strategy_name_list + assert '[S0, R, R, S1] -> [S0, R, S1]_11' in strategy_name_list assert '[R, R, R, S1] -> [R, R, S1]_12' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, S0]_13' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R]_14' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R]_15' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, S0]_16' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, S1]_17' in strategy_name_list - assert '[S01, R, R, R] -> [S01, R, R]_18' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R]_19' in strategy_name_list - assert '[R, R, S01, R] -> [R, S01, R]_20' in strategy_name_list + assert '[R, R, S0, S1] -> [R, S0, S1]_13' in strategy_name_list + assert '[S1, R, R, S0] -> [S1, R, S0]_14' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, S0]_15' in strategy_name_list + assert '[R, R, S1, S0] -> [R, S1, S0]_16' in strategy_name_list + assert '[S0, R, R, R] -> [S0, R, R]_17' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R]_18' in strategy_name_list + assert '[R, R, S0, R] -> [R, S0, R]_19' in strategy_name_list + assert '[S1, R, R, R] -> [S1, R, R]_20' in strategy_name_list assert '[R, R, R, R] -> [R, R, R]_21' in strategy_name_list - assert '[R, R, R, S01] -> [R, R, S01]_22' in strategy_name_list + assert '[R, R, S1, R] -> [R, S1, R]_22' in strategy_name_list assert '[R, R, R, R] -> [R, R, R]_23' in strategy_name_list if sum_dims == 1 and keepdim == True: - assert '[S0, R, R, S1] -> [S0, R, R, S1]_0' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, R, S1]_1' in strategy_name_list - assert '[R, R, S0, S1] -> [R, R, S0, S1]_2' in strategy_name_list - assert '[S1, R, R, S0] -> [S1, R, R, S0]_3' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_4' in strategy_name_list - assert '[R, R, S1, S0] -> [R, R, S1, S0]_5' in strategy_name_list - assert '[S0, R, R, R] -> [S0, R, R, R]_6' in strategy_name_list + assert '[S01, R, R, R] -> [S01, R, R, R]_0' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_1' in strategy_name_list + assert '[R, R, S01, R] -> [R, R, S01, R]_2' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_3' in strategy_name_list + assert '[R, R, R, S01] -> [R, R, R, S01]_4' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_5' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_6' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_7' in strategy_name_list - assert '[R, R, S0, R] -> [R, R, S0, R]_8' in strategy_name_list - assert '[S1, R, R, R] -> [S1, R, R, R]_9' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_10' in strategy_name_list - assert '[R, R, S1, R] -> [R, R, S1, R]_11' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_8' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_9' in strategy_name_list + assert '[R, R, R, S1] -> [R, R, R, S1]_10' in strategy_name_list + assert '[S0, R, R, S1] -> [S0, R, R, S1]_11' in strategy_name_list assert '[R, R, R, S1] -> [R, R, R, S1]_12' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_13' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_14' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_15' in strategy_name_list - assert '[R, R, R, S0] -> [R, R, R, S0]_16' in strategy_name_list - assert '[R, R, R, S1] -> [R, R, R, S1]_17' in strategy_name_list - assert '[S01, R, R, R] -> [S01, R, R, R]_18' in strategy_name_list - assert '[R, R, R, R] -> [R, R, R, R]_19' in strategy_name_list - assert '[R, R, S01, R] -> [R, R, S01, R]_20' in strategy_name_list + assert '[R, R, S0, S1] -> [R, R, S0, S1]_13' in strategy_name_list + assert '[S1, R, R, S0] -> [S1, R, R, S0]_14' in strategy_name_list + assert '[R, R, R, S0] -> [R, R, R, S0]_15' in strategy_name_list + assert '[R, R, S1, S0] -> [R, R, S1, S0]_16' in strategy_name_list + assert '[S0, R, R, R] -> [S0, R, R, R]_17' in strategy_name_list + assert '[R, R, R, R] -> [R, R, R, R]_18' in strategy_name_list + assert '[R, R, S0, R] -> [R, R, S0, R]_19' in strategy_name_list + assert '[S1, R, R, R] -> [S1, R, R, R]_20' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_21' in strategy_name_list - assert '[R, R, R, S01] -> [R, R, R, S01]_22' in strategy_name_list + assert '[R, R, S1, R] -> [R, R, S1, R]_22' in strategy_name_list assert '[R, R, R, R] -> [R, R, R, R]_23' in strategy_name_list @@ -226,9 +223,7 @@ def check_sum_handler(rank, sum_dims, keepdim, world_size, port): @parameterize('sum_dims', [(0, 2), 1]) @parameterize('keepdim', [False, True]) def test_sum_handler(sum_dims, keepdim): - world_size = 4 - run_func = partial(check_sum_handler, sum_dims=sum_dims, keepdim=keepdim, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_sum_handler, 4, sum_dims=sum_dims, keepdim=keepdim) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_tensor_constructor.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_tensor_constructor.py index de35fe256ac7..5b6ac051a8ef 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_tensor_constructor.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_tensor_constructor.py @@ -1,11 +1,13 @@ import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.tensor_constructor_handler import TensorConstructorHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag class TensorConstructorModel(nn.Module): @@ -20,9 +22,10 @@ def forward(self, x): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_where_handler(): model = TensorConstructorModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %x : torch.Tensor [#users=2] = placeholder[target=x] # %size : [#users=1] = call_method[target=size](args = (%x,), kwargs = {}) @@ -30,10 +33,10 @@ def test_where_handler(): # %arange : [#users=1] = call_function[target=torch.arange](args = (%getitem,), kwargs = {}) # %add : [#users=1] = call_function[target=operator.add](args = (%x, %arange), kwargs = {}) # return add - graph = tracer.trace(model, meta_args={ - "x": torch.rand(10).to('meta'), - }) + meta_args = {'x': torch.rand(10).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_unary_element_wise_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_unary_element_wise_handler.py index a861cb7f57f0..f4e6dafdfd69 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_unary_element_wise_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_unary_element_wise_handler.py @@ -1,13 +1,14 @@ import torch import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.unary_elementwise_handler import UnaryElementwiseHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.fx.tracer.meta_patch.patched_module import linear -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag class ReLuModel(nn.Module): @@ -23,21 +24,23 @@ def forward(self, input, other): @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_elementwise_handler(): model = ReLuModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] # %other : torch.Tensor [#users=1] = placeholder[target=other] # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %other), kwargs = {}) # %act : [#users=1] = call_module[target=act](args = (%conv2d,), kwargs = {}) # return act - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(4, 4, 64, 64).to('meta'), - "other": torch.rand(4, 16, 3, 3).to('meta'), - }) + meta_args = { + 'input': torch.rand(4, 4, 64, 64).to('meta'), + 'other': torch.rand(16, 4, 3, 3).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) @@ -69,13 +72,13 @@ def test_elementwise_handler(): assert mapping['input'].name == "conv2d" assert mapping['input'].data.is_meta - assert mapping['input'].data.shape == torch.Size([4, 4, 62, 62]) + assert mapping['input'].data.shape == torch.Size([4, 16, 62, 62]) assert mapping['input'].type == OperationDataType.ARG - assert mapping['input'].logical_shape == torch.Size([4, 4, 62, 62]) + assert mapping['input'].logical_shape == torch.Size([4, 16, 62, 62]) assert mapping['output'].name == "act" assert mapping['output'].data.is_meta - assert mapping['output'].data.shape == torch.Size([4, 4, 62, 62]) + assert mapping['output'].data.shape == torch.Size([4, 16, 62, 62]) assert mapping['output'].type == OperationDataType.OUTPUT # getitem is a following strategy handler, so the number of strategies is equal to the predecessor node. diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_view_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_view_handler.py index 8a96ac0d66f0..fbb194d8e0b8 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_view_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_view_handler.py @@ -1,21 +1,19 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.node_handler import ViewHandler from colossalai.auto_parallel.tensor_shard.node_handler.conv_handler import ConvFunctionHandler from colossalai.auto_parallel.tensor_shard.node_handler.linear_handler import LinearFunctionHandler from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, parameterize, rerun_if_address_is_in_use +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port from tests.test_auto_parallel.test_tensor_shard.test_node_handler.utils import numerical_test_for_node_strategy @@ -74,7 +72,7 @@ def check_view_handler(rank, tgt_shape, model_cls, world_size, port): input_args=[input, other], meta_arg_names=['input', 'other'], node_type='following') - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) if model_cls.__name__ == 'ConvViewModel': # graph(): # %input_1 : torch.Tensor [#users=1] = placeholder[target=input] @@ -82,11 +80,8 @@ def check_view_handler(rank, tgt_shape, model_cls, world_size, port): # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %other), kwargs = {}) # %view : [#users=1] = call_method[target=view](args = (%conv2d, 2, -1), kwargs = {}) # return view - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 8, 66, 66).to('meta'), - "other": torch.rand(16, 8, 3, 3).to('meta'), - }) + meta_args = {'input': torch.rand(8, 8, 66, 66).to('meta'), 'other': torch.rand(16, 8, 3, 3).to('meta')} + graph = tracer.trace(model, meta_args=meta_args) if model_cls.__name__ == 'LinearViewModel': # graph(): @@ -95,13 +90,14 @@ def check_view_handler(rank, tgt_shape, model_cls, world_size, port): # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%input_1, %other), kwargs = {bias: None}) # %view : [#users=1] = call_method[target=view](args = (%linear, 32, 4, 32, 32, 4), kwargs = {}) # return view - graph = tracer.trace(model, - meta_args={ - "input": torch.rand(8, 16, 64, 32).to('meta'), - "other": torch.rand(64, 32).to('meta'), - }) + meta_args = { + 'input': torch.rand(8, 16, 64, 32).to('meta'), + 'other': torch.rand(64, 32).to('meta'), + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) previous_mod_node = list(graph.nodes)[2] view_node = list(graph.nodes)[3] @@ -255,13 +251,7 @@ def check_view_handler(rank, tgt_shape, model_cls, world_size, port): @parameterize('tgt_shape', [(32, 4, 64, 16, 4), (8, 4, 4, 64, 16, 4)]) @parameterize('model_cls', [ConvViewModel, LinearViewModel]) def test_view_handler(tgt_shape, model_cls): - world_size = 4 - run_func = partial(check_view_handler, - tgt_shape=tgt_shape, - model_cls=model_cls, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_view_handler, 4, tgt_shape=tgt_shape, model_cls=model_cls) if __name__ == '__main__': diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_where_handler.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_where_handler.py index 9838e2eb01c6..bd7635ac1737 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_where_handler.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/test_where_handler.py @@ -1,12 +1,14 @@ +import pytest import torch import torch.nn as nn -from colossalai.auto_parallel.tensor_shard.node_handler.where_handler import \ - WhereHandler -from colossalai.auto_parallel.tensor_shard.sharding_strategy import (OperationData, OperationDataType, StrategiesVector) +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer +from colossalai.auto_parallel.tensor_shard.node_handler.where_handler import WhereHandler +from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, StrategiesVector from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.fx.tracer.meta_patch.patched_module import linear +from colossalai.testing import clear_cache_before_run class ConvModel(nn.Module): @@ -19,22 +21,25 @@ def forward(self, condition, x, y): return output +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') +@clear_cache_before_run() def test_where_handler(): model = ConvModel() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) # graph(): # %condition : torch.Tensor [#users=1] = placeholder[target=condition] # %x : torch.Tensor [#users=1] = placeholder[target=x] # %y : torch.Tensor [#users=1] = placeholder[target=y] # %where : [#users=1] = call_function[target=torch.where](args = (%condition, %x, %y), kwargs = {}) # return where - graph = tracer.trace(model, - meta_args={ - "condition": torch.rand(4, 4, 64, 64).to('meta'), - "x": torch.rand(4, 1, 64, 64).to('meta'), - "y": torch.rand(1, 4, 64, 64).to('meta') - }) + meta_args = { + 'condition': torch.rand(4, 4, 64, 64).to('meta'), + 'x': torch.rand(4, 1, 64, 64).to('meta'), + 'y': torch.rand(1, 4, 64, 64).to('meta') + } + graph = tracer.trace(model, meta_args=meta_args) gm = ColoGraphModule(model, graph) + shape_prop_pass(gm, *meta_args.values()) physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/utils.py b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/utils.py index 0cdfdbc9d0cd..28a8bbd9a4c1 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_node_handler/utils.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_node_handler/utils.py @@ -4,6 +4,9 @@ import torch from torch.fx import GraphModule +from colossalai._analyzer.fx.graph_module import ColoGraphModule +from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass from colossalai.auto_parallel.tensor_shard.options import SolverOptions @@ -11,7 +14,6 @@ from colossalai.auto_parallel.tensor_shard.solver.cost_graph import CostGraph from colossalai.auto_parallel.tensor_shard.solver.solver import Solver from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.tracer.tracer import ColoTracer from colossalai.tensor.shape_consistency import to_global from colossalai.testing.comparison import assert_close @@ -79,14 +81,16 @@ def numerical_test_for_node_strategy(model: torch.nn.Module, model_to_shard, args_to_shard, kwargs_to_shard = _build_model_to_compare(model, input_args, input_kwargs, grad_to_shard_dict) - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) input_sample = {} for input_arg, meta_arg_name in zip(input_args, meta_arg_names): - input_sample[meta_arg_name] = torch.rand(input_arg.shape).to('meta') + input_sample[meta_arg_name] = torch.empty(input_arg.shape, dtype=input_arg.dtype).to('meta') for meta_kwarg_name, input_kwarg in input_kwargs.items(): - input_sample[meta_kwarg_name] = torch.rand(input_kwarg.shape).to('meta') + input_sample[meta_kwarg_name] = torch.empty(input_kwarg.shape, dtype=input_kwarg.dtype).to('meta') graph = tracer.trace(root=model_to_shard, meta_args=input_sample) - gm = GraphModule(model_to_shard, graph, model_to_shard.__class__.__name__) + gm = ColoGraphModule(model_to_shard, graph, model_to_shard.__class__.__name__) + shape_prop_pass(gm, *input_sample.values()) + solver_options = SolverOptions() strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options) strategies_constructor.build_strategies_and_cost() diff --git a/tests/test_auto_parallel/test_tensor_shard/test_param_resharding_cost.py b/tests/test_auto_parallel/test_tensor_shard/test_param_resharding_cost.py deleted file mode 100644 index 92f011ba30d2..000000000000 --- a/tests/test_auto_parallel/test_tensor_shard/test_param_resharding_cost.py +++ /dev/null @@ -1,126 +0,0 @@ -import torch - -from colossalai.auto_parallel.tensor_shard.options import SolverOptions -from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationDataType -from colossalai.auto_parallel.tensor_shard.solver import CostGraph, GraphAnalyser, Solver, StrategiesConstructor -from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx import ColoGraphModule, ColoTracer -from colossalai.testing.pytest_wrapper import run_on_environment_flag - - -def _param_resharding_cost_assertion(node): - for strategy in node.strategies_vector: - for prev_node, resharding_cost in strategy.resharding_costs.items(): - if strategy.get_op_data_by_name(str(prev_node)).type == OperationDataType.PARAM: - for cost in resharding_cost: - assert cost.fwd == 0 - assert cost.bwd == 0 - assert cost.total == 0 - - -class LinearModel(torch.nn.Module): - - def __init__(self, in_features, out_features): - super().__init__() - self.linear = torch.nn.Linear(in_features, out_features) - - def forward(self, x): - x = self.linear(x) - x = x * 2 - - return x - - -class ConvModel(torch.nn.Module): - - def __init__(self, in_channels, out_channels, kernel_size, bias=True): - super().__init__() - self.conv = torch.nn.Conv2d(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - bias=bias) - - def forward(self, x): - x = self.conv(x) - x = x * 2 - - return x - - -@run_on_environment_flag(name='AUTO_PARALLEL') -def test_linear_module(): - model = LinearModel(4, 8) - physical_mesh_id = torch.arange(0, 4) - mesh_shape = (2, 2) - # [[0, 1] - # [2, 3]] - device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - tracer = ColoTracer() - # graph(): - # %x : torch.Tensor [#users=1] = placeholder[target=x] - # %linear_weight : [#users=1] = get_attr[target=linear.weight] - # %linear_bias : [#users=1] = get_attr[target=linear.bias] - # %linear : [#users=1] = call_function[target=torch._C._nn.linear](args = (%x, %linear_weight), kwargs = {}) - # %add : [#users=1] = call_function[target=operator.add](args = (%linear, %linear_bias), kwargs = {}) - # %mul : [#users=1] = call_function[target=operator.mul](args = (%add, 2), kwargs = {}) - # return mul - graph = tracer.trace(root=model, meta_args={'x': torch.rand(4, 4).to('meta')}) - # def forward(self, x : torch.Tensor): - # linear_weight = self.linear.weight - # linear_bias = self.linear.bias - # linear = torch._C._nn.linear(x, linear_weight); x = linear_weight = None - # add = linear + linear_bias; linear = linear_bias = None - # mul = add * 2; add = None - # return mul - gm = ColoGraphModule(model, graph) - gm.recompile() - node_list = list(graph.nodes) - - solver_options = SolverOptions() - strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options) - strategies_constructor.build_strategies_and_cost() - linear_node = node_list[3] - _param_resharding_cost_assertion(linear_node) - - -@run_on_environment_flag(name='AUTO_PARALLEL') -def test_conv_module(): - model = ConvModel(3, 6, 2) - physical_mesh_id = torch.arange(0, 4) - mesh_shape = (2, 2) - # [[0, 1] - # [2, 3]] - device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - tracer = ColoTracer() - # graph(): - # %x : torch.Tensor [#users=1] = placeholder[target=x] - # %conv_weight : [#users=1] = get_attr[target=conv.weight] - # %conv_bias : [#users=1] = get_attr[target=conv.bias] - # %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%x, %conv_weight), kwargs = {}) - # %view : [#users=1] = call_method[target=view](args = (%conv_bias, [1, -1, 1, 1]), kwargs = {}) - # %add : [#users=1] = call_function[target=operator.add](args = (%conv2d, %view), kwargs = {}) - # %mul : [#users=1] = call_function[target=operator.mul](args = (%add, 2), kwargs = {}) - # return mul - graph = tracer.trace(root=model, meta_args={'x': torch.rand(4, 3, 64, 64).to('meta')}) - # def forward(self, x : torch.Tensor): - # conv_weight = self.conv.weight - # conv_bias = self.conv.bias - # conv2d = torch.conv2d(x, conv_weight); x = conv_weight = None - # view = conv_bias.view([1, -1, 1, 1]); conv_bias = None - # add = conv2d + view; conv2d = view = None - # mul = add * 2; add = None - # return mul - gm = ColoGraphModule(model, graph) - - gm.recompile() - node_list = list(graph.nodes) - conv_node = node_list[3] - solver_options = SolverOptions() - strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options) - strategies_constructor.build_strategies_and_cost() - _param_resharding_cost_assertion(conv_node) - - -if __name__ == '__main__': - test_linear_module() - test_conv_module() diff --git a/tests/test_auto_parallel/test_tensor_shard/test_shape_consistency_pass.py b/tests/test_auto_parallel/test_tensor_shard/test_shape_consistency_pass.py deleted file mode 100644 index 24a3ae5b42c3..000000000000 --- a/tests/test_auto_parallel/test_tensor_shard/test_shape_consistency_pass.py +++ /dev/null @@ -1,86 +0,0 @@ -import copy -from functools import partial - -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn - -from colossalai.auto_parallel.tensor_shard.initialize import initialize_model -from colossalai.device.device_mesh import DeviceMesh -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.testing import assert_close, rerun_if_address_is_in_use -from colossalai.testing.pytest_wrapper import run_on_environment_flag -from colossalai.utils import free_port - - -class ConvModel(nn.Module): - - def __init__(self, c_in, c_out): - super().__init__() - self.conv = nn.Conv2d(c_in, c_out, kernel_size=3, padding=1, bias=False) - - def forward(self, x): - x = self.conv(x) - x = torch.flatten(x) - return x - - -def check_apply(rank, world_size, port): - disable_existing_loggers() - launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - input = torch.rand(4, 4, 4, 4).cuda() - test_input = copy.deepcopy(input) - # graph(): - # %x : torch.Tensor [#users=1] = placeholder[target=x] - # %conv : [#users=1] = call_module[target=conv](args = (%mul,), kwargs = {}) - # return conv - model = ConvModel(4, 4).cuda() - test_model = copy.deepcopy(model) - physical_mesh_id = torch.arange(0, 4) - mesh_shape = (2, 2) - # [[0, 1] - # [2, 3]] - device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - meta_args = {'x': torch.rand(4, 4, 4, 4).to('meta')} - gm = initialize_model(model, meta_args, device_mesh) - - output = gm(input) - origin_output = test_model(test_input) - assert output.equal(origin_output) - origin_loss = origin_output.sum() - loss = output.sum() - - origin_loss.backward() - loss.backward() - - grad_0 = test_model.conv.weight.grad.narrow(0, 0, 1) - grad_1 = test_model.conv.weight.grad.narrow(0, 1, 1) - grad_2 = test_model.conv.weight.grad.narrow(0, 2, 1) - grad_3 = test_model.conv.weight.grad.narrow(0, 3, 1) - - if rank == 0: - assert_close(gm.module.conv.weight.grad.data, grad_0.data) - elif rank == 1: - assert_close(gm.module.conv.weight.grad.data, grad_1.data) - elif rank == 2: - assert_close(gm.module.conv.weight.grad.data, grad_2.data) - elif rank == 3: - assert_close(gm.module.conv.weight.grad.data, grad_3.data) - else: - raise ValueError(f'rank {rank} does not exist.') - - -# skip this test due to pulp not installed in CI environment -@run_on_environment_flag(name='AUTO_PARALLEL') -@pytest.mark.dist -@rerun_if_address_is_in_use() -def test_apply(): - world_size = 4 - run_func = partial(check_apply, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_apply() diff --git a/tests/test_auto_parallel/test_tensor_shard/test_solver_with_resnet_v2.py b/tests/test_auto_parallel/test_tensor_shard/test_solver_with_resnet_v2.py index bbfc3e1fcc14..0d93e4e40527 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_solver_with_resnet_v2.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_solver_with_resnet_v2.py @@ -2,16 +2,19 @@ from torch.fx import GraphModule from torchvision.models import resnet50 +from colossalai._analyzer.fx.passes import shape_prop_pass +# from colossalai.fx.tracer.tracer import ColoTracer +from colossalai._analyzer.fx.tracer.tracer import ColoTracer from colossalai.auto_parallel.tensor_shard.constants import BATCHNORM_MODULE_OP from colossalai.auto_parallel.tensor_shard.options import SolverOptions from colossalai.auto_parallel.tensor_shard.solver import CostGraph, GraphAnalyser, Solver, StrategiesConstructor from colossalai.device.device_mesh import DeviceMesh -from colossalai.fx.tracer.tracer import ColoTracer from colossalai.tensor.shape_consistency import ShapeConsistencyManager -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag @run_on_environment_flag(name='AUTO_PARALLEL') +@clear_cache_before_run() def test_cost_graph(): physical_mesh_id = torch.arange(0, 8) mesh_shape = (2, 4) @@ -20,7 +23,7 @@ def test_cost_graph(): device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) shape_consistency_manager = ShapeConsistencyManager() - tracer = ColoTracer() + tracer = ColoTracer(bias_addition_split=True) model = resnet50(num_classes=100000) input_sample = {'x': torch.rand(128, 3, 224, 224).to('meta')} @@ -50,6 +53,7 @@ def test_cost_graph(): # %fc : [#users=1] = call_module[target=fc](args = (%flatten,), kwargs = {}) # return fc gm = GraphModule(model, graph, model.__class__.__name__) + shape_prop_pass(gm, *input_sample.values()) gm.recompile() solver_options = SolverOptions() diff --git a/tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py b/tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py index 9a2240d62de4..d07145e48e1f 100644 --- a/tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py +++ b/tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py @@ -8,7 +8,7 @@ from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port +from colossalai.testing import free_port if AUTOCHUNK_AVAILABLE: from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen diff --git a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_alphafold_utils.py b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_alphafold_utils.py index cb250d6402e2..15610e2b50dc 100644 --- a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_alphafold_utils.py +++ b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_alphafold_utils.py @@ -9,7 +9,7 @@ from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port +from colossalai.testing import free_port if AUTOCHUNK_AVAILABLE: from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen diff --git a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_block.py b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_block.py index 17a5abf4cab8..9e4cb7ee9f95 100644 --- a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_block.py +++ b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_block.py @@ -1,10 +1,8 @@ -from functools import partial from typing import Dict, List, Tuple import pytest import torch import torch.fx -import torch.multiprocessing as mp try: from fastfold.model.nn.evoformer import EvoformerBlock @@ -15,6 +13,7 @@ from test_autochunk_alphafold_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn def get_model(): @@ -66,18 +65,19 @@ def get_chunk_target() -> Dict: not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("max_memory", [None, 20, 24]) -@pytest.mark.parametrize("data_args", [(32, 64)]) # (msa_len, pair_len) +@clear_cache_before_run() +@parameterize("max_memory", [None, 20, 24]) +@parameterize("data_args", [(32, 64)]) def test_evoformer_block(data_args, max_memory): - run_func = partial( + spawn( run_test, + 1, data_args=data_args, max_memory=max_memory, get_model=get_model, get_data=get_data, get_chunk_target=get_chunk_target, ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": diff --git a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_stack.py b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_stack.py index 5210c1c8d48e..6b47033e199f 100644 --- a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_stack.py +++ b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_evoformer_stack.py @@ -1,10 +1,8 @@ -from functools import partial from typing import List, Tuple import pytest import torch import torch.fx -import torch.multiprocessing as mp try: from fastfold.model.nn.evoformer import EvoformerStack @@ -15,6 +13,7 @@ from test_autochunk_alphafold_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn def get_model(): @@ -61,17 +60,18 @@ def get_data(msa_len: int, pair_len: int) -> Tuple[List, List]: not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("max_memory", [None, 20, 24]) -@pytest.mark.parametrize("data_args", [(32, 64)]) # (msa_len, pair_len) +@clear_cache_before_run() +@parameterize("max_memory", [None, 20, 24]) +@parameterize("data_args", [(32, 64)]) # (msa_len, pair_len) def test_evoformer_stack(data_args, max_memory): - run_func = partial( + spawn( run_test, + 1, data_args=data_args, max_memory=max_memory, get_model=get_model, get_data=get_data, ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": diff --git a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_extramsa_block.py b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_extramsa_block.py index ad955479e617..b4c577c18ee6 100644 --- a/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_extramsa_block.py +++ b/tests/test_autochunk/test_autochunk_alphafold/test_autochunk_extramsa_block.py @@ -1,10 +1,8 @@ -from functools import partial from typing import Dict, List, Tuple import pytest import torch import torch.fx -import torch.multiprocessing as mp try: from fastfold.model.nn.evoformer import ExtraMSABlock @@ -14,6 +12,7 @@ from test_autochunk_alphafold_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn def get_model(): @@ -57,17 +56,18 @@ def get_data(msa_len: int, pair_len: int) -> Tuple[List, List]: not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("max_memory", [None, 20, 24]) -@pytest.mark.parametrize("data_args", [(32, 64)]) # (msa_len, pair_len) +@clear_cache_before_run() +@parameterize("max_memory", [None, 20, 24]) +@parameterize("data_args", [(32, 64)]) # (msa_len, pair_len) def test_extramsa_block(data_args, max_memory): - run_func = partial( + spawn( run_test, + 1, data_args=data_args, max_memory=max_memory, get_model=get_model, get_data=get_data, ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": diff --git a/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_diffuser_utils.py b/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_diffuser_utils.py index 529250fe8f51..b6a792f5652c 100644 --- a/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_diffuser_utils.py +++ b/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_diffuser_utils.py @@ -8,7 +8,6 @@ from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port if AUTOCHUNK_AVAILABLE: from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen @@ -93,6 +92,8 @@ def assert_codegen_run( def run_test( rank: int, + world_size: int, + port: int, model: Any, data: tuple, max_memory: int, @@ -106,9 +107,9 @@ def run_test( colossalai.launch( config={}, rank=rank, - world_size=1, + world_size=world_size, host="localhost", - port=free_port(), + port=port, backend="nccl", ) diff --git a/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_unet.py b/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_unet.py index 16c5b10ff4ae..f0cf2a5fcbca 100644 --- a/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_unet.py +++ b/tests/test_autochunk/test_autochunk_diffuser/test_autochunk_unet.py @@ -1,21 +1,23 @@ -from functools import partial from typing import List, Tuple import pytest import torch -import torch.multiprocessing as mp try: - from diffusers import UNet2DModel - MODELS = [UNet2DModel] + import diffusers + MODELS = [diffusers.UNet2DModel] HAS_REPO = True + from packaging import version + SKIP_UNET_TEST = version.parse(diffusers.__version__) > version.parse("0.10.2") except: MODELS = [] HAS_REPO = False + SKIP_UNET_TEST = False from test_autochunk_diffuser_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn BATCH_SIZE = 1 HEIGHT = 448 @@ -33,31 +35,27 @@ def get_data(shape: tuple) -> Tuple[List, List]: return meta_args, concrete_args +@pytest.mark.skipif( + SKIP_UNET_TEST, + reason="diffusers version > 0.10.2", +) @pytest.mark.skipif( not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("shape", [LATENTS_SHAPE]) -@pytest.mark.parametrize("max_memory", [None, 150, 300]) +@clear_cache_before_run() +@parameterize("model", MODELS) +@parameterize("shape", [LATENTS_SHAPE]) +@parameterize("max_memory", [None, 150, 300]) def test_evoformer_block(model, shape, max_memory): - run_func = partial( + spawn( run_test, + 1, max_memory=max_memory, model=model, data=get_data(shape), ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": - run_test( - rank=0, - data=get_data(LATENTS_SHAPE), - max_memory=None, - model=UNet2DModel, - print_code=False, - print_mem=True, - print_est_mem=False, - print_progress=False, - ) + test_evoformer_block() diff --git a/tests/test_autochunk/test_autochunk_transformer/test_autochunk_gpt.py b/tests/test_autochunk/test_autochunk_transformer/test_autochunk_gpt.py index 018a2557a974..82af6c05c6ef 100644 --- a/tests/test_autochunk/test_autochunk_transformer/test_autochunk_gpt.py +++ b/tests/test_autochunk/test_autochunk_transformer/test_autochunk_gpt.py @@ -1,9 +1,7 @@ -from functools import partial from typing import List, Tuple import pytest import torch -import torch.multiprocessing as mp try: from transformers import GPT2Config, GPT2Model @@ -16,6 +14,7 @@ from test_autochunk_transformer_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn BATCH_SIZE = 1 SEQ_LENGTH = 512 @@ -31,22 +30,25 @@ def get_data(shape: tuple) -> Tuple[List, List]: return meta_args, concrete_args, sequence +@pytest.mark.skip("full op is not implemented now") +# FIXME(ver217, oahzxl): implement full op @pytest.mark.skipif( not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("shape", [(BATCH_SIZE, SEQ_LENGTH)]) -@pytest.mark.parametrize("max_memory", [None, 6, 8]) +@clear_cache_before_run() +@parameterize("model", MODELS) +@parameterize("shape", [(BATCH_SIZE, SEQ_LENGTH)]) +@parameterize("max_memory", [None, 6, 8]) def test_autochunk_gpt(model, shape, max_memory): - run_func = partial( + spawn( run_test, + 1, data=get_data(shape), max_memory=max_memory, model=model, config=GPT2Config(n_embd=96, n_positions=shape[1], n_layer=2, n_head=4), ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": diff --git a/tests/test_autochunk/test_autochunk_transformer/test_autochunk_transformer_utils.py b/tests/test_autochunk/test_autochunk_transformer/test_autochunk_transformer_utils.py index bc5eda7edf91..5c863b0df47f 100644 --- a/tests/test_autochunk/test_autochunk_transformer/test_autochunk_transformer_utils.py +++ b/tests/test_autochunk/test_autochunk_transformer/test_autochunk_transformer_utils.py @@ -5,10 +5,8 @@ import colossalai from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE -from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port if AUTOCHUNK_AVAILABLE: from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen @@ -100,6 +98,8 @@ def assert_allclose(out_model: Any, out_gm: Any) -> None: def run_test( rank: int, + world_size: int, + port: int, model: Any, config: Any, data: tuple, @@ -116,9 +116,9 @@ def run_test( colossalai.launch( config={}, rank=rank, - world_size=1, + world_size=world_size, host="localhost", - port=free_port(), + port=port, backend="nccl", ) diff --git a/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit.py b/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit.py index 2b7cbf1390d2..a98aa0e03954 100644 --- a/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit.py +++ b/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit.py @@ -1,9 +1,7 @@ -from functools import partial from typing import List, Tuple import pytest import torch -import torch.multiprocessing as mp try: from timm.models.vision_transformer import vit_large_patch16_384 as vit @@ -16,6 +14,7 @@ from test_autochunk_vit_utils import run_test from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE +from colossalai.testing import clear_cache_before_run, parameterize, spawn def get_data() -> Tuple[List, List]: @@ -28,16 +27,17 @@ def get_data() -> Tuple[List, List]: not (AUTOCHUNK_AVAILABLE and HAS_REPO), reason="torch version is lower than 1.12.0", ) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("max_memory", [None, 32, 40]) +@clear_cache_before_run() +@parameterize("model", MODELS) +@parameterize("max_memory", [None, 32, 40]) def test_evoformer_block(model, max_memory): - run_func = partial( + spawn( run_test, + 1, max_memory=max_memory, model=model, data=get_data(), ) - mp.spawn(run_func, nprocs=1) if __name__ == "__main__": diff --git a/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit_utils.py b/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit_utils.py index 035dd59799b4..3202318fb6d1 100644 --- a/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit_utils.py +++ b/tests/test_autochunk/test_autochunk_vit/test_autochunk_vit_utils.py @@ -8,7 +8,6 @@ from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp -from colossalai.utils import free_port if AUTOCHUNK_AVAILABLE: from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen @@ -85,6 +84,8 @@ def assert_codegen_run( def run_test( rank: int, + world_size: int, + port: int, model: Any, data: tuple, max_memory: int, @@ -98,9 +99,9 @@ def run_test( colossalai.launch( config={}, rank=rank, - world_size=1, + world_size=world_size, host="localhost", - port=free_port(), + port=port, backend="nccl", ) diff --git a/tests/test_booster/test_accelerator.py b/tests/test_booster/test_accelerator.py new file mode 100644 index 000000000000..6f3f66ed41b8 --- /dev/null +++ b/tests/test_booster/test_accelerator.py @@ -0,0 +1,14 @@ +import torch.nn as nn + +from colossalai.booster.accelerator import Accelerator +from colossalai.testing import clear_cache_before_run, parameterize + + +@clear_cache_before_run() +@parameterize('device', ['cpu', 'cuda']) +def test_accelerator(device): + accelerator = Accelerator(device) + model = nn.Linear(8, 8) + model = accelerator.configure_model(model) + assert next(model.parameters()).device.type == device + del model, accelerator diff --git a/tests/test_booster/test_mixed_precision/test_fp16_torch.py b/tests/test_booster/test_mixed_precision/test_fp16_torch.py index c56fcae58a60..26ce00e94869 100644 --- a/tests/test_booster/test_mixed_precision/test_fp16_torch.py +++ b/tests/test_booster/test_mixed_precision/test_fp16_torch.py @@ -1,17 +1,28 @@ import torch from torch.optim import Adam +import colossalai from colossalai.booster.mixed_precision import FP16TorchMixedPrecision +from colossalai.testing import rerun_if_address_is_in_use, spawn from tests.kit.model_zoo import model_zoo -def test_torch_amp(): - for name, (model_fn, data_gen_fn, output_transform_fn, _) in model_zoo.items(): +def run_torch_amp(rank, world_size, port): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + sub_model_zoo = model_zoo.get_sub_registry('timm') + for name, (model_fn, data_gen_fn, output_transform_fn, _, _) in sub_model_zoo.items(): + # dlrm_interactionarch has not parameters, so skip + if name == 'dlrm_interactionarch': + continue + model = model_fn().cuda() optimizer = Adam(model.parameters(), lr=1e-3) criterion = lambda x: x.mean() data = data_gen_fn() - data = {k: v.cuda() if torch.is_tensor(v) else v for k, v in data.items()} + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } mixed_precision = FP16TorchMixedPrecision() model, optimizer, criterion = mixed_precision.configure(model, optimizer, criterion) output = model(**data) @@ -21,3 +32,9 @@ def test_torch_amp(): optimizer.backward(loss) optimizer.clip_grad_by_norm(1.0) optimizer.step() + del model, optimizer, criterion, data, output, mixed_precision + + +@rerun_if_address_is_in_use() +def test_torch_ddp_plugin(): + spawn(run_torch_amp, 1) diff --git a/tests/test_booster/test_plugin/test_dp_plugin_base.py b/tests/test_booster/test_plugin/test_dp_plugin_base.py new file mode 100644 index 000000000000..689b334cae50 --- /dev/null +++ b/tests/test_booster/test_plugin/test_dp_plugin_base.py @@ -0,0 +1,88 @@ +from typing import Callable, Iterator, List, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler as LRScheduler +from torch.utils.data import DataLoader, TensorDataset + +import colossalai +from colossalai.booster.plugin.dp_plugin_base import DPPluginBase +from colossalai.checkpoint_io import CheckpointIO +from colossalai.interface import OptimizerWrapper +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +class DPPluginWrapper(DPPluginBase): + """This is a wrapper class for testing DP plugin initialization and dataloader creation. + """ + + def configure( + self, + model: nn.Module, + optimizer: Optimizer, + criterion: Callable = None, + dataloader: DataLoader = None, + lr_scheduler: LRScheduler = None, + ) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]: + pass + + def control_checkpoint_io(self) -> bool: + pass + + def control_device(self) -> bool: + pass + + def control_precision(self) -> bool: + pass + + def get_checkpoint_io(self) -> CheckpointIO: + pass + + def support_no_sync(self) -> bool: + pass + + def supported_devices(self) -> List[str]: + pass + + def supported_precisions(self) -> List[str]: + pass + + def no_sync(self, model: nn.Module) -> Iterator[None]: + pass + + +def check_dataloader_sharding(): + plugin = DPPluginWrapper() + + # create a custom dataset with 0 to 10 + dataset = TensorDataset(torch.arange(0, 10)) + train_dataloader = plugin.prepare_dataloader(dataset, batch_size=2) + + # get the first batch of data + batch = next(iter(train_dataloader))[0].cuda() + is_rank_0 = dist.get_rank() == 0 + + if is_rank_0: + batch_to_compare = batch.clone() + else: + batch_to_compare = batch + # pass to the rank 1 value to rank 0 + dist.broadcast(batch_to_compare, src=1) + + # compare on rank 0 + if is_rank_0: + assert not torch.equal(batch, + batch_to_compare), 'Same number was found across ranks but expected it to be different' + + +def run_dist(rank, world_size, port): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_dataloader_sharding() + + +@rerun_if_address_is_in_use() +def test_dp_plugin_dataloader(): + spawn(run_dist, 2) diff --git a/tests/test_booster/test_plugin/test_gemini_plugin.py b/tests/test_booster/test_plugin/test_gemini_plugin.py new file mode 100644 index 000000000000..d29c92926066 --- /dev/null +++ b/tests/test_booster/test_plugin/test_gemini_plugin.py @@ -0,0 +1,132 @@ +from contextlib import nullcontext +from typing import Optional + +import torch +import torch.distributed as dist + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin +from colossalai.fx import is_compatible_with_meta +from colossalai.lazy.lazy_init import LazyInitContext +from colossalai.nn.optimizer import HybridAdam +from colossalai.tensor.colo_parameter import ColoParameter +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero import ColoInitContext +from tests.kit.model_zoo import model_zoo + + +def run_fn(init_method, model_fn, data_gen_fn, output_transform_fn) -> Optional[str]: + try: + if init_method == 'colo': + ctx = ColoInitContext() + elif init_method == 'lazy': + ctx = LazyInitContext() + else: + ctx = nullcontext() + plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, max_norm=1.0, initial_scale=2**5) + booster = Booster(plugin=plugin) + with ctx: + model = model_fn() + optimizer = HybridAdam(model.parameters(), lr=1e-3) + criterion = lambda x: x.mean() + data = data_gen_fn() + + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } + + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + for n, p in model.named_parameters(): + assert isinstance(p, ColoParameter), f'{n} is not a ColoParameter' + + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.step() + + except Exception as e: + return repr(e) + + +# TODO(ver217): CI does not support lazy now +# @parameterize('init_method', ['lazy', 'none', 'colo']) + + +@parameterize('init_method', ['none']) +def check_gemini_plugin(init_method: str = 'none', early_stop: bool = True): + """check gemini plugin over model zoo + + Args: + early_stop (bool, optional): Whether to stop when getting the first error. Defaults to True. + """ + is_support_meta = is_compatible_with_meta() + if not is_support_meta and init_method == 'lazy': + return + + passed_models = [] + failed_info = {} # (model_name, error) pair + + for name, (model_fn, data_gen_fn, output_transform_fn, _, _) in model_zoo.items(): + # These models lead to CUDA error + if name in ('diffusers_auto_encoder_kl', 'diffusers_vq_model', 'diffusers_unet2d_model', 'timm_resmlp', + 'timm_gmixer_12_224', 'timm_gmlp_b16_224', 'timm_mixer_b16_224', 'timm_convnext'): + continue + # These models are not compatible with gemini + if name in [ + 'diffusers_clip_vision_model', 'timm_resnet', 'timm_beit', 'timm_beitv2', 'timm_eca_nfnet', + 'timm_efficientformer', 'timm_hrnet_w18_small', 'timm_nf_ecaresnet101', 'timm_nf_regnet_b0', + 'timm_skresnet18', 'timm_wide_resnet50_2', 'timm_convit', 'timm_dm_nfnet', 'timm_swin_transformer', + 'torchaudio_conformer', 'torchaudio_deepspeech', 'torchaudio_wavernn', 'torchaudio_tacotron', + 'deepfm_interactionarch', 'deepfm_simpledeepfmnn', 'dlrm', 'dlrm_interactionarch', + 'torchvision_googlenet', 'torchvision_inception_v3', 'torchvision_mobilenet_v3_small', + 'torchvision_resnet18', 'torchvision_resnext50_32x4d', 'torchvision_wide_resnet50_2', + 'torchvision_vit_b_16', 'torchvision_convnext_base', 'torchvision_swin_s', 'transformers_albert', + 'transformers_albert_for_pretraining', 'transformers_bert', 'transformers_bert_for_pretraining', + 'transformers_gpt_double_heads', 'torchaudio_hubert_base', 'torchaudio_wav2vec2_base', + 'transformers_t5_for_conditional_generation', 'transformers_t5', 'transformers_t5_encoder_model' + ]: + continue + + if init_method == 'lazy' and name in [ + 'timm_convmixer', 'timm_vision_transformer', 'timm_deit', 'timm_deit3', 'timm_inception_v3', + 'timm_tnt_b_patch16_224', 'timm_rexnet', 'torchvision_densenet121', 'torchvision_efficientnet_b0', + 'torchvision_mobilenet_v2', 'torchvision_mnasnet0_5', 'torchvision_regnet_x_16gf', + 'torchvision_shufflenet_v2_x0_5', 'torchvision_efficientnet_v2_s' + ]: + continue + + err = run_fn(init_method, model_fn, data_gen_fn, output_transform_fn) + torch.cuda.empty_cache() + + if err is None: + passed_models.append(name) + else: + failed_info[name] = err + if early_stop: + break + + if dist.get_rank() == 0: + print(f'Init method: {init_method}') + print(f'Passed models({len(passed_models)}): {passed_models}\n\n') + print(f'Failed models({len(failed_info)}): {list(failed_info.keys())}\n\n') + assert len(failed_info) == 0, '\n'.join([f'{k}: {v}' for k, v in failed_info.items()]) + + +def run_dist(rank, world_size, port, early_stop: bool = True): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_gemini_plugin(early_stop=early_stop) + + +@rerun_if_address_is_in_use() +def test_gemini_plugin(early_stop: bool = True): + spawn(run_dist, 4, early_stop=early_stop) + + +if __name__ == '__main__': + test_gemini_plugin(early_stop=False) diff --git a/tests/test_booster/test_plugin/test_low_level_zero_plugin.py b/tests/test_booster/test_plugin/test_low_level_zero_plugin.py new file mode 100644 index 000000000000..eedd8c59a3a8 --- /dev/null +++ b/tests/test_booster/test_plugin/test_low_level_zero_plugin.py @@ -0,0 +1,99 @@ +from typing import Optional + +import torch +import torch.distributed as dist + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import LowLevelZeroPlugin +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from tests.kit.model_zoo import model_zoo + +# These models are not compatible with AMP +_AMP_ERR_MODELS = ['timm_convit', 'dlrm', 'deepfm_interactionarch', 'deepfm_simpledeepfmnn'] +# These models have no parameters +_LOW_LEVEL_ZERO_ERR_MODELS = ['dlrm_interactionarch', 'deepfm_overarch', 'deepfm_sparsearch', 'dlrm_sparsearch'] +# These models will get stuck +_STUCK_MODELS = [ + 'diffusers_vq_model', 'transformers_albert', 'transformers_albert_for_pretraining', 'transformers_bert', + 'transformers_bert_for_pretraining', 'transformers_gpt_double_heads' +] + + +def run_fn(stage, model_fn, data_gen_fn, output_transform_fn) -> Optional[str]: + try: + plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=2**5) + booster = Booster(plugin=plugin) + model = model_fn() + optimizer = HybridAdam(model.parameters(), lr=1e-3) + criterion = lambda x: x.mean() + data = data_gen_fn() + + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } + + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.step() + + except Exception as e: + return repr(e) + + +@parameterize('stage', [2]) +def check_low_level_zero_plugin(stage: int, early_stop: bool = True): + """check low level zero plugin over model zoo + + Args: + stage (int), stage of low level zero plugin + early_stop (bool, optional): Whether to stop when getting the first error. Defaults to True. + """ + passed_models = [] + failed_info = {} # (model_name, error) pair + ignore_models = _AMP_ERR_MODELS + _LOW_LEVEL_ZERO_ERR_MODELS + _STUCK_MODELS + skipped_models = [] + + for name, (model_fn, data_gen_fn, output_transform_fn, _, _) in model_zoo.items(): + # FIXME(ver217): fix these models + if name in ignore_models: + skipped_models.append(name) + continue + err = run_fn(stage, model_fn, data_gen_fn, output_transform_fn) + + torch.cuda.empty_cache() + + if err is None: + passed_models.append(name) + else: + failed_info[name] = err + if early_stop: + break + + if dist.get_rank() == 0: + print(f'Passed models({len(passed_models)}): {passed_models}\n\n') + print(f'Failed models({len(failed_info)}): {list(failed_info.keys())}\n\n') + print(f'Skipped models({len(skipped_models)}): {skipped_models}\n\n') + assert len(failed_info) == 0, '\n'.join([f'{k}: {v}' for k, v in failed_info.items()]) + + +def run_dist(rank, world_size, port, early_stop: bool = True): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_low_level_zero_plugin(early_stop=early_stop) + + +@rerun_if_address_is_in_use() +def test_low_level_zero_plugin(early_stop: bool = True): + spawn(run_dist, 4, early_stop=early_stop) + + +if __name__ == '__main__': + test_low_level_zero_plugin(early_stop=False) diff --git a/tests/test_booster/test_plugin/test_torch_ddp_plugin.py b/tests/test_booster/test_plugin/test_torch_ddp_plugin.py new file mode 100644 index 000000000000..1484273973ae --- /dev/null +++ b/tests/test_booster/test_plugin/test_torch_ddp_plugin.py @@ -0,0 +1,115 @@ +from contextlib import nullcontext + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import SGD + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.interface import OptimizerWrapper +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.kit.model_zoo import model_zoo + + +def run_fn(model_fn, data_gen_fn, output_transform_fn): + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + model = model_fn() + optimizer = SGD(model.parameters(), lr=1e-3) + criterion = lambda x: x.mean() + data = data_gen_fn() + + data = {k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items()} + + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + assert isinstance(model.module, DDP) + assert isinstance(optimizer, OptimizerWrapper) + + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.clip_grad_by_norm(1.0) + optimizer.step() + + +def check_torch_ddp_plugin(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, _) in model_zoo.items(): + if name == 'dlrm_interactionarch': + continue + run_fn(model_fn, data_gen_fn, output_transform_fn) + torch.cuda.empty_cache() + + +class DummyModel(nn.Module): + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.rand(1)) + + def forward(self, x): + return self.weight * x + + +def check_torch_ddp_no_sync(): + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + + model = DummyModel() + criterion = lambda x: x.mean() + optimizer = SGD(model.parameters(), lr=1e-3) + # create a custom dasetset with 0 to 10 + dataset = torch.arange(0, 10) + train_dataloader = plugin.prepare_dataloader(dataset, batch_size=2) + model, optimizer, criterion, train_dataloader, _ = booster.boost(model, + optimizer, + criterion, + dataloader=train_dataloader) + + def fwd_bwd(): + output = model(batch.cuda()) + loss = criterion(output) + booster.backward(loss, optimizer) + + def get_grad_set_over_all_ranks(): + for p in model.parameters(): + # grad shape is (1, ) + assert p.grad.shape == (1,) + grad_list = [torch.empty_like(p.grad) for _ in range(dist.get_world_size())] + dist.all_gather(grad_list, p.grad) + # get grad set of all ranks + grad_set = set([grad.item() for grad in grad_list]) + # as the model only has one parameter, we can return here + return grad_set + + for i, batch in enumerate(train_dataloader): + if i > 1: + # only check the first two batches + break + # no_sync for the first batch, sync for the second batch + ctx = booster.no_sync(model) if i == 0 else nullcontext() + with ctx: + fwd_bwd() + grad_set = get_grad_set_over_all_ranks() + # for the first batch, all ranks should have different grads + # for the second batch, as grad is synchronized,all ranks should have the same grads + target_num_different_grad = dist.get_world_size() if i == 0 else 1 + assert len(grad_set) == target_num_different_grad + + +def run_dist(rank, world_size, port): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_torch_ddp_plugin() + check_torch_ddp_no_sync() + + +@rerun_if_address_is_in_use() +def test_torch_ddp_plugin(): + spawn(run_dist, 2) diff --git a/tests/test_booster/test_plugin/test_torch_fsdp_plugin.py b/tests/test_booster/test_plugin/test_torch_fsdp_plugin.py new file mode 100644 index 000000000000..cbd5d57800db --- /dev/null +++ b/tests/test_booster/test_plugin/test_torch_fsdp_plugin.py @@ -0,0 +1,64 @@ +import pytest +import torch +from packaging import version +from torch.optim import SGD + +import colossalai +from colossalai.booster import Booster + +if version.parse(torch.__version__) >= version.parse('1.12.0'): + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from colossalai.booster.plugin import TorchFSDPPlugin + +from colossalai.interface import OptimizerWrapper +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.kit.model_zoo import model_zoo + + +# test baisc fsdp function +def run_fn(model_fn, data_gen_fn, output_transform_fn): + plugin = TorchFSDPPlugin() + booster = Booster(plugin=plugin) + model = model_fn() + optimizer = SGD(model.parameters(), lr=1e-3) + criterion = lambda x: x.mean() + data = data_gen_fn() + + data = {k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items()} + + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + assert isinstance(model.module, FSDP) + assert isinstance(optimizer, OptimizerWrapper) + + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.clip_grad_by_norm(1.0) + optimizer.step() + + +def check_torch_fsdp_plugin(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, _) in model_zoo.items(): + if any(element in name for element in [ + 'diffusers', 'deepfm_sparsearch', 'dlrm_interactionarch', 'torchvision_googlenet', + 'torchvision_inception_v3' + ]): + continue + run_fn(model_fn, data_gen_fn, output_transform_fn) + torch.cuda.empty_cache() + + +def run_dist(rank, world_size, port): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_torch_fsdp_plugin() + + +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason="requires torch1.12 or higher") +@rerun_if_address_is_in_use() +def test_torch_fsdp_plugin(): + spawn(run_dist, 2) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py new file mode 100644 index 000000000000..7b664419b405 --- /dev/null +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -0,0 +1,123 @@ +import os + +import pytest +import torch +import torch.distributed as dist +from utils import shared_tempdir + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import ( + check_state_dict_equal, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo + + +@clear_cache_before_run() +@parameterize('placement_policy', ['cuda', 'cpu']) +@parameterize('model_name', ['transformers_bert_for_sequence_classification']) +@parameterize('use_safetensors', [False, True]) +def exam_state_dict_with_origin(placement_policy, model_name, use_safetensors: bool): + from transformers import BertForSequenceClassification + (model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values())) + bert_model = model_fn() + + with shared_tempdir() as tempdir: + pretrained_path = os.path.join(tempdir, 'pretrained') + bert_model.config.save_pretrained(save_directory=pretrained_path) + + plugin = GeminiPlugin(placement_policy=placement_policy) + booster = Booster(plugin=plugin) + bert_model, _, _, _, _ = booster.boost(bert_model) + model_size = sum(p.numel() * p.element_size() for p in bert_model.parameters()) / 1024**2 + + booster.save_model(bert_model, + pretrained_path, + True, + True, + '', (model_size / 3), + use_safetensors=use_safetensors) + dist.barrier() + + new_bert_model = BertForSequenceClassification.from_pretrained(pretrained_path) + check_state_dict_equal(bert_model.unwrap().state_dict(only_rank_0=False, dtype=torch.float32), + new_bert_model.state_dict(), False) + + +@clear_cache_before_run() +@parameterize('placement_policy', ['cuda', 'cpu']) +@parameterize('shard', [False, True]) +@parameterize('model_name', ['transformers_gpt']) +@parameterize('size_per_shard', [32]) +def exam_state_dict(placement_policy, shard: bool, model_name: str, size_per_shard: int): + (model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values())) + criterion = lambda x: x.mean() + plugin = GeminiPlugin(placement_policy=placement_policy, precision="fp16", initial_scale=(2**14)) + booster = Booster(plugin=plugin) + + model = model_fn() + new_model = model_fn() + optimizer = HybridAdam(model.parameters(), lr=0.001) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + new_optimizer = HybridAdam(new_model.parameters(), lr=0.001) + new_model, new_optimizer, criterion, _, _ = booster.boost(new_model, new_optimizer, criterion) + + data = data_gen_fn() + data = {k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items()} + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.step() + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optimizer_ckpt_path = f"{tempdir}/optimizer" + booster.save_model(model, model_ckpt_path, shard=shard, size_per_shard=size_per_shard) + + booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard, size_per_shard=size_per_shard) + dist.barrier() + + booster.load_model(new_model, model_ckpt_path) + check_state_dict_equal(model.unwrap().state_dict(only_rank_0=False), + new_model.unwrap().state_dict(only_rank_0=False), False) + + booster.load_optimizer(new_optimizer, optimizer_ckpt_path) + check_state_dict_equal(optimizer.unwrap().state_dict(only_rank_0=False), + new_optimizer.unwrap().state_dict(only_rank_0=False), False) + + # Check the new model/optimizer can successfully run. + data = data_gen_fn() + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } + output = new_model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + booster.backward(loss, new_optimizer) + new_optimizer.step() + booster.save_model(new_model, model_ckpt_path, shard=shard) + booster.save_optimizer(new_optimizer, optimizer_ckpt_path, shard=shard) + + +def run_dist(rank, world_size, port): + config = {} + colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + exam_state_dict() + exam_state_dict_with_origin() + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [2]) +@rerun_if_address_is_in_use() +def test_gemini_ckpIO(world_size): + spawn(run_dist, world_size) diff --git a/tests/test_checkpoint_io/test_gemini_torch_compability.py b/tests/test_checkpoint_io/test_gemini_torch_compability.py new file mode 100644 index 000000000000..464fccb39103 --- /dev/null +++ b/tests/test_checkpoint_io/test_gemini_torch_compability.py @@ -0,0 +1,171 @@ +import pytest +import torch +import torch.distributed as dist +from torch.optim import Adam +from utils import shared_tempdir + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import GeminiPlugin, TorchDDPPlugin +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import ( + check_state_dict_equal, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo + + +@clear_cache_before_run() +@parameterize('shard', [False, True]) +@parameterize('model_name', ['transformers_gpt']) +def exam_torch_load_from_gemini(shard: bool, model_name: str): + + (model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values())) + criterion = lambda x: x.mean() + plugin = GeminiPlugin(precision="fp16", initial_scale=(2**14)) + booster = Booster(plugin=plugin) + + model = model_fn() + optimizer = HybridAdam(model.parameters(), lr=0.001) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + data = data_gen_fn() + data = {k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items()} + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.step() + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optimizer_ckpt_path = f"{tempdir}/optimizer" + + booster.save_model(model, model_ckpt_path, shard=shard) + booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard) + dist.barrier() + + new_model = model_fn() + new_optimizer = Adam(new_model.parameters(), lr=0.001) + new_plugin = TorchDDPPlugin() + new_booster = Booster(plugin=new_plugin) + new_model, new_optimizer, criterion, _, _ = new_booster.boost(new_model, new_optimizer, criterion) + + # Loading HybridAdam states to torch.Adam + new_booster.load_model(new_model, model_ckpt_path, strict=True) + + # Add prefix to get aligned with pytorch parameter names. + check_state_dict_equal( + model.unwrap().state_dict(only_rank_0=False, prefix='module.module.', dtype=torch.float32), + new_model.state_dict(), False) + + new_booster.load_optimizer(new_optimizer, optimizer_ckpt_path) + check_state_dict_equal(optimizer.unwrap().state_dict(only_rank_0=False), new_optimizer.state_dict(), False) + + # Check the new model/optimizer can successfully run. + data = data_gen_fn() + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } + output = new_model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + new_booster.backward(loss, new_optimizer) + new_optimizer.step() + new_booster.save_model(new_model, model_ckpt_path, shard=shard) + new_booster.save_optimizer(new_optimizer, optimizer_ckpt_path, shard=shard) + + +@clear_cache_before_run() +@parameterize('shard', [False, True]) +@parameterize('model_name', ['transformers_gpt']) +def exam_gemini_load_from_torch(shard: bool, model_name: str): + + (model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values())) + criterion = lambda x: x.mean() + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + + model = model_fn() + optimizer = Adam(model.parameters(), lr=0.001) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + data = data_gen_fn() + data = {k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items()} + output = model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + + booster.backward(loss, optimizer) + optimizer.step() + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optimizer_ckpt_path = f"{tempdir}/optimizer" + + booster.save_model(model, model_ckpt_path, shard=shard) + booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard) + dist.barrier() + + new_model = model_fn() + new_optimizer = HybridAdam(new_model.parameters(), lr=0.001) + new_plugin = GeminiPlugin() + new_booster = Booster(plugin=new_plugin) + new_model, new_optimizer, criterion, _, _ = new_booster.boost(new_model, new_optimizer, criterion) + + # Loading torch.Adam states to HybridAdam + new_booster.load_model(new_model, model_ckpt_path, strict=True) + + # Add prefix to get aligned with pytorch parameter names. + check_state_dict_equal( + new_model.unwrap().state_dict(only_rank_0=False, prefix='module.module.', dtype=torch.float32), + model.state_dict(), False) + + new_booster.load_optimizer(new_optimizer, optimizer_ckpt_path) + old_state_dict = optimizer.state_dict() + new_state_dict = new_optimizer.unwrap().state_dict(only_rank_0=False) + + # Comparison of param_groups needs special care here, + # since not all hyperparameters in Adam are used by HybridAdam + hyperparameters_to_examine = ['params', 'lr', 'betas', 'eps', 'weight_decay'] + for old_group, new_group in zip(old_state_dict['param_groups'], new_state_dict['param_groups']): + for k in hyperparameters_to_examine: + assert k in old_group and k in new_group, \ + f"Old group's keys: {list(old_group.keys())}, New group's keys: {list(new_group.keys())}" + assert old_group[k] == new_group[k] + check_state_dict_equal(old_state_dict['state'], new_state_dict['state'], False) + + # Check the new model/optimizer can successfully run. + data = data_gen_fn() + data = { + k: v.to('cuda') if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__ else v for k, v in data.items() + } + output = new_model(**data) + output = output_transform_fn(output) + output_key = list(output.keys())[0] + loss = criterion(output[output_key]) + new_booster.backward(loss, new_optimizer) + new_optimizer.step() + new_booster.save_model(new_model, model_ckpt_path, shard=shard) + new_booster.save_optimizer(new_optimizer, optimizer_ckpt_path, shard=shard) + + +def run_dist(rank, world_size, port): + config = {} + colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + exam_torch_load_from_gemini() + exam_gemini_load_from_torch() + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [2]) +@rerun_if_address_is_in_use() +def test_gemini_ckpIO(world_size): + spawn(run_dist, world_size) diff --git a/tests/test_checkpoint_io/test_general_checkpoint_io.py b/tests/test_checkpoint_io/test_general_checkpoint_io.py new file mode 100644 index 000000000000..0976d4503a61 --- /dev/null +++ b/tests/test_checkpoint_io/test_general_checkpoint_io.py @@ -0,0 +1,210 @@ +import tempfile + +import pytest +import torch +from torch.optim import Adam +from torchvision.models import resnet18 + +from colossalai.booster.plugin.gemini_plugin import GeminiCheckpointIO +from colossalai.checkpoint_io import GeneralCheckpointIO +from colossalai.testing import check_state_dict_equal, clear_cache_before_run, parameterize + +# ======== +# Note: +# 1. due to checkpoint IO can be quite slow if tested with all models, we will only test on resnet for now +# 2. we will test on both sharded and unsharded checkpoints +# 3. implement sharded checkpoint and test it +# ======== + + +@clear_cache_before_run() +@parameterize('use_safetensors', [True, False]) +def test_unsharded_checkpoint(use_safetensors: bool): + # create a model and optimizer + model = resnet18() + optimizer = Adam(model.parameters(), lr=0.001) + + # create test data sample + x = torch.randn(1, 3, 224, 224) + + # run fwd and bwd + y = model(x) + loss = y.sum() + loss.backward() + optimizer.step() + + # create a temp file for checkpoint + if use_safetensors: + suffix = ".safetensors" + else: + suffix = ".bin" + model_ckpt_tempfile = tempfile.NamedTemporaryFile(suffix=suffix) + optimizer_ckpt_tempfile = tempfile.NamedTemporaryFile() + + # save the model and optimizer + ckpt_io = GeneralCheckpointIO() + ckpt_io.save_model(model, model_ckpt_tempfile.name, use_safetensors=use_safetensors) + ckpt_io.save_optimizer(optimizer, optimizer_ckpt_tempfile.name) + + # create new model + new_model = resnet18() + new_optimizer = Adam(new_model.parameters(), lr=0.001) + + # load the model and optimizer + ckpt_io.load_model(new_model, model_ckpt_tempfile.name) + ckpt_io.load_optimizer(new_optimizer, optimizer_ckpt_tempfile.name) + + # check for model and optimizer state dict recursively + check_state_dict_equal(model.state_dict(), new_model.state_dict()) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict()) + + +@pytest.mark.parametrize('use_safetensors', [True, False]) +def test_sharded_model_checkpoint(use_safetensors: bool): + # create a model and optimizer + model = resnet18() + optimizer = Adam(model.parameters(), lr=0.001) + # create test data sample + x = torch.randn(1, 3, 224, 224) + + # run fwd and bwd + y = model(x) + loss = y.sum() + loss.backward() + optimizer.step() + + # create a temp file for checkpoint + if use_safetensors: + suffix = ".safetensors" + SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" + else: + suffix = ".bin" + WEIGHTS_INDEX_NAME = "model.bin.index.json" + + model_ckpt_dir = tempfile.TemporaryDirectory() + optimizer_ckpt_tempfile = tempfile.NamedTemporaryFile() + + # save the model and optimizer + ckpt_io = GeneralCheckpointIO() + + ckpt_io.save_model(model, model_ckpt_dir.name, True, True, "", 10, use_safetensors=use_safetensors) + ckpt_io.save_optimizer(optimizer, optimizer_ckpt_tempfile.name, shard=False) + + # create new model + new_model = resnet18() + new_optimizer = Adam(new_model.parameters(), lr=0.001) + + ckpt_io.load_model(new_model, str(model_ckpt_dir.name), strict=True) + ckpt_io.load_optimizer(new_optimizer, optimizer_ckpt_tempfile.name) + + # check for model and optimizer state dict recursively + check_state_dict_equal(model.state_dict(), new_model.state_dict()) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict()) + + +def test_sharded_optimizer_checkpoint(): + + # create a model and optimizer + model = resnet18() + optimizer = Adam(model.parameters(), lr=0.001) + + # create test data sample + x = torch.randn(1, 3, 224, 224) + + # run fwd and bwd + y = model(x) + loss = y.sum() + loss.backward() + optimizer.step() + + # create temp directories for checkpoint + model_ckpt_dir = tempfile.TemporaryDirectory() + optimizer_ckpt_dir = tempfile.TemporaryDirectory() + + # save the model and optimizer + ckpt_io = GeneralCheckpointIO() + + ckpt_io.save_model(model, model_ckpt_dir.name, True, True, "", 10, use_safetensors=False) + ckpt_io.save_optimizer(optimizer, optimizer_ckpt_dir.name, shard=True, size_per_shard=10) + + # create new model + new_model = resnet18() + new_optimizer = Adam(new_model.parameters(), lr=0.001) + + ckpt_io.load_model(new_model, str(model_ckpt_dir.name), strict=True) + ckpt_io.load_optimizer(new_optimizer, str(optimizer_ckpt_dir.name)) + + # check for model and optimizer state dict recursively + check_state_dict_equal(model.state_dict(), new_model.state_dict()) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict()) + + # continue running fwd and bwd + for _ in range(5): + y = new_model(x) + loss = y.sum() + loss.backward() + new_optimizer.step() + + # save the newly got optimizer + ckpt_io.save_model(new_model, model_ckpt_dir.name, True, True, "", 10, use_safetensors=False) + ckpt_io.save_optimizer(new_optimizer, optimizer_ckpt_dir.name, shard=True, size_per_shard=10) + + # create another new model + new_new_model = resnet18() + new_new_optimizer = Adam(new_new_model.parameters(), lr=0.001) + + ckpt_io.load_model(new_new_model, str(model_ckpt_dir.name), strict=True) + ckpt_io.load_optimizer(new_new_optimizer, str(optimizer_ckpt_dir.name)) + + # check for model and optimizer state dict recursively + check_state_dict_equal(new_model.state_dict(), new_new_model.state_dict()) + check_state_dict_equal(new_optimizer.state_dict(), new_new_optimizer.state_dict()) + + +def test_sharded_optimizer_multiple_param_groups(): + + # create a model and optimizer + model = resnet18() + optimizer = Adam([{ + 'params': model.layer1.parameters() + }, { + 'params': model.layer2.parameters(), + 'lr': 0.002 + }], + lr=0.001) + + # create test data sample + x = torch.randn(1, 3, 224, 224) + + # run fwd and bwd + y = model(x) + loss = y.sum() + loss.backward() + optimizer.step() + + # create temp directories for checkpoint + model_ckpt_dir = tempfile.TemporaryDirectory() + optimizer_ckpt_dir = tempfile.TemporaryDirectory() + + # save the model and optimizer + ckpt_io = GeneralCheckpointIO() + + ckpt_io.save_model(model, model_ckpt_dir.name, True, True, "", 10, use_safetensors=False) + ckpt_io.save_optimizer(optimizer, optimizer_ckpt_dir.name, shard=True, size_per_shard=10) + + # create new model + new_model = resnet18() + new_optimizer = Adam([{ + 'params': new_model.layer1.parameters() + }, { + 'params': new_model.layer2.parameters(), + 'lr': 0.002 + }], + lr=0.001) + + ckpt_io.load_model(new_model, str(model_ckpt_dir.name), strict=True) + ckpt_io.load_optimizer(new_optimizer, str(optimizer_ckpt_dir.name)) + + # check for model and optimizer state dict recursively + check_state_dict_equal(model.state_dict(), new_model.state_dict()) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict()) diff --git a/tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py b/tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py new file mode 100644 index 000000000000..c51b54c82f57 --- /dev/null +++ b/tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py @@ -0,0 +1,64 @@ +import torch +import torch.distributed as dist +from torchvision.models import resnet18 +from utils import shared_tempdir + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import LowLevelZeroPlugin +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import ( + check_state_dict_equal, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) + + +@clear_cache_before_run() +@parameterize('stage', [2]) +@parameterize('shard', [True, False]) +def check_low_level_zero_checkpointIO(stage: int, shard: bool): + plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=32) + booster = Booster(plugin=plugin) + model = resnet18() + criterion = lambda x: x.mean() + optimizer = HybridAdam((model.parameters()), lr=0.001) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + x = torch.randn(4, 3, 224, 224) + x = x.to('cuda') + output = model(x) + loss = criterion(output) + booster.backward(loss, optimizer) + optimizer.step() + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optimizer_ckpt_path = f"{tempdir}/optimizer" + # lr scheduler is tested in test_torch_ddp_checkpoint_io.py and low level zero does not change it, we can skip it here + booster.save_model(model, model_ckpt_path, shard=shard) + if not shard: + # TODO(ver217): optimizer checkpointing is not supported for sharded checkpoint + booster.save_optimizer(optimizer, optimizer_ckpt_path) + dist.barrier() + + new_model = resnet18() + new_optimizer = HybridAdam((new_model.parameters()), lr=0.001) + new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer) + + booster.load_model(new_model, model_ckpt_path) + check_state_dict_equal(model.state_dict(), new_model.state_dict(), False) + if not shard: + booster.load_optimizer(new_optimizer, optimizer_ckpt_path) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict(), False) + + +def run_dist(rank, world_size, port): + colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host='localhost') + check_low_level_zero_checkpointIO() + + +@rerun_if_address_is_in_use() +def test_low_level_zero_checkpointIO(): + spawn(run_dist, 2) diff --git a/tests/test_checkpoint_io/test_torch_ddp_checkpoint_io.py b/tests/test_checkpoint_io/test_torch_ddp_checkpoint_io.py new file mode 100644 index 000000000000..14332b5b3fca --- /dev/null +++ b/tests/test_checkpoint_io/test_torch_ddp_checkpoint_io.py @@ -0,0 +1,70 @@ +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import SGD +from torchvision.models import resnet18 +from utils import shared_tempdir + +import colossalai +from colossalai.booster import Booster +from colossalai.booster.plugin import TorchDDPPlugin +from colossalai.interface import OptimizerWrapper +from colossalai.testing import check_state_dict_equal, parameterize, rerun_if_address_is_in_use, spawn + + +@parameterize('shard', [True, False]) +@parameterize('size_per_shard', [16, 128]) +def check_torch_ddp_checkpointIO(shard: bool, size_per_shard: int): + plugin = TorchDDPPlugin() + booster = Booster(plugin=plugin) + model = resnet18() + criterion = lambda x: x.mean() + optimizer = SGD((model.parameters()), lr=0.001) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1) + model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion, lr_scheduler=scheduler) + + assert isinstance(model.module, DDP) + assert isinstance(optimizer, OptimizerWrapper) + + x = torch.randn(4, 3, 224, 224) + x = x.to('cuda') + output = model(x) + loss = criterion(output) + booster.backward(loss, optimizer) + optimizer.clip_grad_by_norm(1.0) + optimizer.step() + scheduler.step() + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optimizer_ckpt_path = f"{tempdir}/optimizer" + lr_scheduler_ckpt_path = f"{tempdir}/lr_scheduler" + booster.save_model(model, model_ckpt_path, shard=shard, size_per_shard=size_per_shard) + booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard, size_per_shard=size_per_shard) + booster.save_lr_scheduler(scheduler, lr_scheduler_ckpt_path) + dist.barrier() + + new_model = resnet18() + new_optimizer = SGD((new_model.parameters()), lr=0.001) + new_scheduler = torch.optim.lr_scheduler.StepLR(new_optimizer, step_size=1, gamma=0.1) + new_model, new_optimizer, _, _, new_scheduler = booster.boost(new_model, + new_optimizer, + lr_scheduler=new_scheduler) + + booster.load_model(new_model, model_ckpt_path) + check_state_dict_equal(model.state_dict(), new_model.state_dict(), False) + + booster.load_optimizer(new_optimizer, optimizer_ckpt_path) + check_state_dict_equal(optimizer.state_dict(), new_optimizer.state_dict(), False) + booster.load_lr_scheduler(new_scheduler, lr_scheduler_ckpt_path) + check_state_dict_equal(scheduler.state_dict(), new_scheduler.state_dict(), False) + + +def run_dist(rank, world_size, port): + colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host='localhost') + check_torch_ddp_checkpointIO() + + +@rerun_if_address_is_in_use() +def test_torch_ddp_checkpointIO(): + spawn(run_dist, 2) diff --git a/tests/test_checkpoint_io/test_torch_fsdp_checkpoint_io.py b/tests/test_checkpoint_io/test_torch_fsdp_checkpoint_io.py new file mode 100644 index 000000000000..2b6090bb1e29 --- /dev/null +++ b/tests/test_checkpoint_io/test_torch_fsdp_checkpoint_io.py @@ -0,0 +1,113 @@ +import pytest +import torch +from packaging import version +from torch import nn +from torch.optim import SGD +from torchvision.models import resnet18 +from utils import shared_tempdir + +import colossalai +from colossalai.booster import Booster + +if version.parse(torch.__version__) >= version.parse('1.12.0'): + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from colossalai.booster.plugin import TorchFSDPPlugin + +from colossalai.testing import rerun_if_address_is_in_use, spawn, check_state_dict_equal + + +def compare_nested_dict(dict1, dict2): + for key in dict1: + if key in dict2: + if type(dict1[key]) is dict: + assert type(dict2[key]) is dict + diff = compare_nested_dict(dict1[key], dict2[key]) + if not diff: + return diff + elif type(dict1[key]) is list: + assert type(dict2[key]) is list + for i, val in enumerate(dict1[key]): + if isinstance(val, torch.Tensor): + if not torch.equal(dict1[key][i], dict2[key][i]): + return False + elif val != dict2[key][i]: + return False + elif type(dict1[key]) is torch.Tensor: + assert type(dict2[key]) is torch.Tensor + if not torch.equal(dict1[key], dict2[key]): + return False + else: + if dict1[key] != dict2[key]: + return False + else: + return False + return True + + +def check_torch_fsdp_ckpt(): + model = resnet18() + plugin = TorchFSDPPlugin() + booster = Booster(plugin=plugin) + optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9) + criterion = lambda x: x.mean() + fsdp_model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion) + + inputs = torch.randn(4, 3, 224, 224) + outputs = None + + def run_model(): + nonlocal outputs + outputs = fsdp_model(inputs) + optimizer.zero_grad() + criterion(outputs).backward() + optimizer.step() + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + optim_ckpt_path = f"{tempdir}/optimizer" + + run_model() + + booster.save_model(fsdp_model, model_ckpt_path, shard=False) + booster.save_optimizer(optimizer, optim_ckpt_path, shard=False) + + full_msd = fsdp_model.state_dict() + #full_osd = FSDP.full_optim_state_dict(fsdp_model, optimizer) + sharded_osd = optimizer.state_dict() + import copy + sharded_osd = copy.deepcopy(sharded_osd) + + run_model() + + full_msd_updated = fsdp_model.state_dict() + #full_osd_updated = FSDP.full_optim_state_dict(fsdp_model, optimizer, rank0_only=True) + sharded_osd_updated = optimizer.state_dict() + + assert not compare_nested_dict(sharded_osd, sharded_osd_updated) + assert not compare_nested_dict(full_msd_updated, full_msd) + outputs_first = fsdp_model(inputs) + assert criterion(outputs_first) != criterion(outputs) + + booster.load_model(fsdp_model, model_ckpt_path) + booster.load_optimizer(optimizer, optim_ckpt_path) + + full_msd_restore = fsdp_model.state_dict() + #full_osd_restore = FSDP.full_optim_state_dict(fsdp_model, optimizer, rank0_only=True) + sharded_osd_restore = optimizer.state_dict() + + assert compare_nested_dict(sharded_osd, sharded_osd_restore) + assert compare_nested_dict(full_msd_restore, full_msd) + outputs_sec = fsdp_model(inputs) + assert criterion(outputs_sec) == criterion(outputs) + + +def run_dist(rank, world_size, port): + # init dist env + colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost') + check_torch_fsdp_ckpt() + + +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason="requires torch1.12 or higher") +@rerun_if_address_is_in_use() +def test_torch_fsdp_ckpt(): + spawn(run_dist, 2) diff --git a/tests/test_checkpoint_io/utils.py b/tests/test_checkpoint_io/utils.py new file mode 100644 index 000000000000..2d35e157f446 --- /dev/null +++ b/tests/test_checkpoint_io/utils.py @@ -0,0 +1,21 @@ +import tempfile +from contextlib import contextmanager, nullcontext +from typing import Iterator + +import torch.distributed as dist + + +@contextmanager +def shared_tempdir() -> Iterator[str]: + """ + A temporary directory that is shared across all processes. + """ + ctx_fn = tempfile.TemporaryDirectory if dist.get_rank() == 0 else nullcontext + with ctx_fn() as tempdir: + try: + obj = [tempdir] + dist.broadcast_object_list(obj, src=0) + tempdir = obj[0] # use the same directory on all ranks + yield tempdir + finally: + dist.barrier() diff --git a/tests/test_cluster/test_device_mesh_manager.py b/tests/test_cluster/test_device_mesh_manager.py new file mode 100644 index 000000000000..bb818a275879 --- /dev/null +++ b/tests/test_cluster/test_device_mesh_manager.py @@ -0,0 +1,34 @@ +import torch + +from colossalai.cluster.device_mesh_manager import DeviceMeshInfo, DeviceMeshManager +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import spawn + + +def check_device_mesh_manager(rank, world_size, port): + disable_existing_loggers() + launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + device_mesh_manager = DeviceMeshManager() + # TODO(ver217): this test is strictly relies on hardware, temporary skip it + # device_mesh_info_auto = DeviceMeshInfo(physical_ids=[0, 1, 2, 3],) + # device_mesh_auto = device_mesh_manager.create_device_mesh('0', device_mesh_info_auto) + # assert device_mesh_auto.shape == (2, 2) + # assert device_mesh_auto._logical_mesh_id.tolist() == [[0, 1], [2, 3]] + + device_mesh_info_with_shape = DeviceMeshInfo( + physical_ids=[0, 1, 2, 3], + mesh_shape=(2, 2), + ) + device_mesh_with_shape = device_mesh_manager.create_device_mesh('1', device_mesh_info_with_shape) + + assert device_mesh_with_shape.shape == (2, 2) + assert device_mesh_with_shape._logical_mesh_id.tolist() == [[0, 1], [2, 3]] + + +def test_device_mesh_manager(): + spawn(check_device_mesh_manager, 4) + + +if __name__ == '__main__': + test_device_mesh_manager() diff --git a/tests/test_comm/test_boardcast_send_recv_v2.py b/tests/test_comm/test_boardcast_send_recv_v2.py index 1520d6054043..253f6f21cd80 100644 --- a/tests/test_comm/test_boardcast_send_recv_v2.py +++ b/tests/test_comm/test_boardcast_send_recv_v2.py @@ -1,17 +1,12 @@ -from functools import partial -from typing import List - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from colossalai.communication.p2p_v2 import _send_object, _recv_object, init_process_group + +from colossalai.communication.p2p_v2 import _recv_object, _send_object from colossalai.context import ParallelMode from colossalai.core import global_context as gpc from colossalai.initialize import launch -from colossalai.utils import free_port, get_current_device -from colossalai.testing import rerun_if_address_is_in_use from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, spawn disable_existing_loggers() world_size = 4 @@ -45,9 +40,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_object_list_p2p(): - disable_existing_loggers() - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, world_size) if __name__ == '__main__': diff --git a/tests/test_comm/test_comm.py b/tests/test_comm/test_comm.py index 07cb67730d24..747596bd2ded 100644 --- a/tests/test_comm/test_comm.py +++ b/tests/test_comm/test_comm.py @@ -1,15 +1,13 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp + from colossalai.communication import all_gather, all_reduce, reduce_scatter from colossalai.context import ParallelMode from colossalai.core import global_context as gpc from colossalai.initialize import launch -from colossalai.utils import free_port, get_current_device -from colossalai.testing import rerun_if_address_is_in_use +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device CONFIG = dict(parallel=dict(data=8, pipeline=1, tensor=dict(mode=None, size=1))) @@ -66,9 +64,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_comm(): - world_size = 4 - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, 4) if __name__ == '__main__': diff --git a/tests/test_comm/test_object_list_p2p.py b/tests/test_comm/test_object_list_p2p.py index 701e3e8ade79..e9d7630c1543 100644 --- a/tests/test_comm/test_object_list_p2p.py +++ b/tests/test_comm/test_object_list_p2p.py @@ -1,15 +1,18 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from colossalai.communication.p2p import send_forward, recv_forward, send_backward, recv_backward, send_forward_recv_backward, send_backward_recv_forward + +from colossalai.communication.p2p import ( + recv_backward, + recv_forward, + send_backward, + send_backward_recv_forward, + send_forward, + send_forward_recv_backward, +) from colossalai.context import ParallelMode from colossalai.core import global_context as gpc from colossalai.initialize import launch -from colossalai.utils import free_port, get_current_device -from colossalai.testing import rerun_if_address_is_in_use +from colossalai.testing import rerun_if_address_is_in_use, spawn CONFIG = dict(parallel=dict(pipeline=2)) torch.manual_seed(123) @@ -96,9 +99,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_object_list_p2p(): - world_size = 2 - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, 2) if __name__ == '__main__': diff --git a/tests/test_comm/test_object_list_p2p_v2.py b/tests/test_comm/test_object_list_p2p_v2.py index c639ac9f8ef3..cae38385b6e1 100644 --- a/tests/test_comm/test_object_list_p2p_v2.py +++ b/tests/test_comm/test_object_list_p2p_v2.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from colossalai.communication.p2p_v2 import send_forward, recv_forward, send_backward, recv_backward, init_process_group -from colossalai.context import ParallelMode, Initializer_Pipeline + +from colossalai.communication.p2p_v2 import recv_backward, recv_forward, send_backward, send_forward +from colossalai.context import ParallelMode from colossalai.core import global_context as gpc from colossalai.initialize import launch -from colossalai.utils import free_port, get_current_device -from colossalai.testing import rerun_if_address_is_in_use from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, spawn disable_existing_loggers() @@ -121,10 +117,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_object_list_p2p(): - disable_existing_loggers() - run_func = partial(check_layer, world_size=world_size, port=free_port()) - disable_existing_loggers() - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, world_size) if __name__ == '__main__': diff --git a/tests/test_context/test_hybrid_parallel.py b/tests/test_context/test_hybrid_parallel.py index f311b1d2e736..9f26a5af53ce 100644 --- a/tests/test_context/test_hybrid_parallel.py +++ b/tests/test_context/test_hybrid_parallel.py @@ -1,19 +1,17 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial from pathlib import Path + import pytest import torch -import torch.multiprocessing as mp from colossalai import launch +from colossalai.context import reset_seeds from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc -from colossalai.utils import free_port -from colossalai.context import reset_seeds from colossalai.global_variables import tensor_parallel_env as tp_env -from colossalai.testing import rerun_if_address_is_in_use +from colossalai.testing import free_port, rerun_if_address_is_in_use, spawn CONFIG_PATH_LIST = list(Path(__file__).parent.glob('configs/*.py')) @@ -134,9 +132,14 @@ def init_context(config_path, rank, world_size, backend, port, host): torch.cuda.empty_cache() -def run_dist(rank, world_size, backend, port_list, host): - for config_path, port in zip(CONFIG_PATH_LIST, port_list): - init_context(config_path=config_path, rank=rank, world_size=world_size, backend=backend, port=port, host=host) +def run_dist(rank, world_size, port, backend, port_list, host): + for config_path, current_port in zip(CONFIG_PATH_LIST, port_list): + init_context(config_path=config_path, + rank=rank, + world_size=world_size, + backend=backend, + port=current_port, + host=host) reset_seeds() @@ -156,8 +159,7 @@ def test_context(): port_list.append(port) break - test_fn = partial(run_dist, world_size=world_size, backend='gloo', port_list=port_list, host='localhost') - mp.spawn(test_fn, nprocs=world_size) + spawn(run_dist, world_size, backend='gloo', port_list=port_list, host='localhost') if __name__ == '__main__': diff --git a/tests/test_data/test_data_parallel_sampler.py b/tests/test_data/test_data_parallel_sampler.py index 54fa44bdc0c2..2ad3fd696c39 100644 --- a/tests/test_data/test_data_parallel_sampler.py +++ b/tests/test_data/test_data_parallel_sampler.py @@ -2,20 +2,18 @@ # -*- encoding: utf-8 -*- import os -from functools import partial from pathlib import Path import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp +from torchvision import datasets, transforms import colossalai -from torchvision import transforms, datasets -from colossalai.context import ParallelMode, Config +from colossalai.context import Config, ParallelMode from colossalai.core import global_context as gpc -from colossalai.utils import get_dataloader, free_port -from colossalai.testing import rerun_if_address_is_in_use +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_dataloader CONFIG = Config(dict( parallel=dict( @@ -58,9 +56,7 @@ def run_data_sampler(rank, world_size, port): @pytest.mark.cpu @rerun_if_address_is_in_use() def test_data_sampler(): - world_size = 4 - test_func = partial(run_data_sampler, world_size=world_size, port=free_port()) - mp.spawn(test_func, nprocs=world_size) + spawn(run_data_sampler, 4) if __name__ == '__main__': diff --git a/tests/test_data/test_deterministic_dataloader.py b/tests/test_data/test_deterministic_dataloader.py index 4d76e7f137f1..239e79dff7d8 100644 --- a/tests/test_data/test_deterministic_dataloader.py +++ b/tests/test_data/test_deterministic_dataloader.py @@ -2,21 +2,18 @@ # -*- encoding: utf-8 -*- import os -from functools import partial from pathlib import Path import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -from torchvision import transforms, datasets +from torchvision import datasets, transforms import colossalai -from colossalai.context import ParallelMode, Config +from colossalai.context import Config, ParallelMode from colossalai.core import global_context as gpc -from colossalai.utils import get_dataloader, free_port -from colossalai.testing import rerun_if_address_is_in_use -from torchvision import transforms +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_dataloader CONFIG = Config( dict( @@ -70,9 +67,7 @@ def run_data_sampler(rank, world_size, port): @pytest.mark.cpu @rerun_if_address_is_in_use() def test_data_sampler(): - world_size = 4 - test_func = partial(run_data_sampler, world_size=world_size, port=free_port()) - mp.spawn(test_func, nprocs=world_size) + spawn(run_data_sampler, 4) if __name__ == '__main__': diff --git a/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor.py b/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor.py index 3c2390c92837..4992acbd7cc2 100644 --- a/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor.py +++ b/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor.py @@ -1,25 +1,22 @@ import os - -from functools import partial from pathlib import Path -import colossalai import pytest import torch -import torch.multiprocessing as mp +from torchvision import transforms +from torchvision.datasets import CIFAR10 + +import colossalai from colossalai.amp import AMP_TYPE -from colossalai.trainer import Trainer, hooks from colossalai.context import ParallelMode -from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus -from colossalai.utils import free_port from colossalai.core import global_context as gpc from colossalai.logging import get_dist_logger from colossalai.nn import CrossEntropyLoss from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR -from colossalai.utils import get_dataloader from colossalai.pipeline.pipelinable import PipelinableContext -from torchvision.datasets import CIFAR10 -from torchvision import transforms +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn +from colossalai.trainer import Trainer, hooks +from colossalai.utils import get_dataloader BATCH_SIZE = 4 NUM_EPOCHS = 60 @@ -51,7 +48,7 @@ def run_trainer(rank, world_size, port): pipelinable.policy = "uniform" model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE)) - # craete dataloaders + # create dataloaders root = Path(os.environ['DATA']) transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4, pad_if_needed=True), @@ -71,7 +68,7 @@ def run_trainer(rank, world_size, port): # create lr scheduler lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, total_steps=NUM_EPOCHS, warmup_steps=WARMUP_EPOCHS) - # intiailize + # initialize engine, train_dataloader, *_ = colossalai.initialize(model=model, optimizer=optimizer, criterion=criterion, @@ -96,9 +93,7 @@ def run_trainer(rank, world_size, port): @skip_if_not_enough_gpus(min_gpus=8) @rerun_if_address_is_in_use() def test_hybrid_parallel(): - world_size = 8 - run_func = partial(run_trainer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_trainer, 8) if __name__ == '__main__': diff --git a/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor_v2.py b/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor_v2.py index 2bafe0f7e374..62bbb8f50391 100644 --- a/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor_v2.py +++ b/tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor_v2.py @@ -1,111 +1,104 @@ -import os - -from functools import partial -from pathlib import Path - -import colossalai -import pytest -import torch -import torch.multiprocessing as mp -from colossalai.amp import AMP_TYPE -from colossalai.trainer import Trainer, hooks -from colossalai.context import ParallelMode -from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus -from colossalai.utils import free_port -from colossalai.core import global_context as gpc -from colossalai.logging import get_dist_logger -from colossalai.nn import CrossEntropyLoss -from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR -from colossalai.utils import get_dataloader -from colossalai.pipeline.pipelinable import PipelinableContext -from colossalai.logging import disable_existing_loggers -from torchvision.datasets import CIFAR10 -from torchvision import transforms - -from colossalai.engine.schedule._pipeline_schedule_v2 import PipelineScheduleV2 - -disable_existing_loggers() -BATCH_SIZE = 4 -NUM_EPOCHS = 10 -WARMUP_EPOCHS = 5 -CONFIG = dict(NUM_MICRO_BATCHES=2, - parallel=dict(pipeline=2, tensor=dict(size=1, mode='1d')), - fp16=dict(mode=AMP_TYPE.NAIVE), - gradient_accumulation=2) - - -def run_trainer(rank, world_size, port): - disable_existing_loggers() - colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - - disable_existing_loggers() - # get logger - logger = get_dist_logger() - - pipelinable = PipelinableContext() - try: - from titans.model.vit import vit_tiny_patch4_32 - except ImportError: - logger.warning('skip the test_cifar_with_data_pipeline_tensor test because titan is not installed') - logger.warning('please install titan from https://github.com/hpcaitech/Titans') - return - with pipelinable: - model = vit_tiny_patch4_32() - pipelinable.to_layer_list() - pipelinable.policy = "uniform" - model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE)) - - # craete dataloaders - root = Path(os.environ['DATA']) - transform_train = transforms.Compose([ - transforms.RandomCrop(32, padding=4, pad_if_needed=True), - transforms.AutoAugment(policy=transforms.AutoAugmentPolicy.CIFAR10), - transforms.ToTensor(), - transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), - ]) - train_dataset = CIFAR10(root=root, train=True, download=True, transform=transform_train) - train_dataloader = get_dataloader(dataset=train_dataset, shuffle=True, batch_size=BATCH_SIZE, pin_memory=True) - - # create loss function - criterion = CrossEntropyLoss(label_smoothing=0.1) - - # create optimizer - optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0) - - # create lr scheduler - lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, total_steps=NUM_EPOCHS, warmup_steps=WARMUP_EPOCHS) - - # intiailize - engine, train_dataloader, *_ = colossalai.initialize(model=model, - optimizer=optimizer, - criterion=criterion, - train_dataloader=train_dataloader) - - engine._schedule = PipelineScheduleV2(num_microbatches=gpc.config.NUM_MICRO_BATCHES) - - logger = get_dist_logger() - - trainer = Trainer(engine=engine, logger=logger) - - hook_list = [ - hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=False), - ] - - trainer.fit(train_dataloader=train_dataloader, - max_steps=2, - epochs=NUM_EPOCHS, - hooks=hook_list, - display_progress=True) - - -@pytest.mark.dist -@rerun_if_address_is_in_use() -def test_hybrid_parallel(): - world_size = 2 - run_func = partial(run_trainer, world_size=world_size, port=free_port()) - disable_existing_loggers() - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_hybrid_parallel() +import os +from pathlib import Path + +import pytest +import torch +from torchvision import transforms +from torchvision.datasets import CIFAR10 + +import colossalai +from colossalai.amp import AMP_TYPE +from colossalai.context import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.engine.schedule._pipeline_schedule_v2 import PipelineScheduleV2 +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.nn import CrossEntropyLoss +from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR +from colossalai.pipeline.pipelinable import PipelinableContext +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.trainer import Trainer, hooks +from colossalai.utils import get_dataloader + +disable_existing_loggers() +BATCH_SIZE = 4 +NUM_EPOCHS = 10 +WARMUP_EPOCHS = 5 +CONFIG = dict(NUM_MICRO_BATCHES=2, + parallel=dict(pipeline=2, tensor=dict(size=1, mode='1d')), + fp16=dict(mode=AMP_TYPE.NAIVE), + gradient_accumulation=2) + + +def run_trainer(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + disable_existing_loggers() + # get logger + logger = get_dist_logger() + + pipelinable = PipelinableContext() + try: + from titans.model.vit import vit_tiny_patch4_32 + except ImportError: + logger.warning('skip the test_cifar_with_data_pipeline_tensor test because titan is not installed') + logger.warning('please install titan from https://github.com/hpcaitech/Titans') + return + with pipelinable: + model = vit_tiny_patch4_32() + pipelinable.to_layer_list() + pipelinable.policy = "uniform" + model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE)) + + # create dataloaders + root = Path(os.environ['DATA']) + transform_train = transforms.Compose([ + transforms.RandomCrop(32, padding=4, pad_if_needed=True), + transforms.AutoAugment(policy=transforms.AutoAugmentPolicy.CIFAR10), + transforms.ToTensor(), + transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), + ]) + train_dataset = CIFAR10(root=root, train=True, download=True, transform=transform_train) + train_dataloader = get_dataloader(dataset=train_dataset, shuffle=True, batch_size=BATCH_SIZE, pin_memory=True) + + # create loss function + criterion = CrossEntropyLoss(label_smoothing=0.1) + + # create optimizer + optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0) + + # create lr scheduler + lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, total_steps=NUM_EPOCHS, warmup_steps=WARMUP_EPOCHS) + + # initialize + engine, train_dataloader, *_ = colossalai.initialize(model=model, + optimizer=optimizer, + criterion=criterion, + train_dataloader=train_dataloader) + + engine._schedule = PipelineScheduleV2(num_microbatches=gpc.config.NUM_MICRO_BATCHES) + + logger = get_dist_logger() + + trainer = Trainer(engine=engine, logger=logger) + + hook_list = [ + hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=False), + ] + + trainer.fit(train_dataloader=train_dataloader, + max_steps=2, + epochs=NUM_EPOCHS, + hooks=hook_list, + display_progress=True) + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_hybrid_parallel(): + spawn(run_trainer, 2) + disable_existing_loggers() + + +if __name__ == '__main__': + test_hybrid_parallel() diff --git a/tests/test_ddp/test_ddp_ignore_params.py b/tests/test_ddp/test_ddp_ignore_params.py index 679c8b0f6afe..39efcd41a1d4 100644 --- a/tests/test_ddp/test_ddp_ignore_params.py +++ b/tests/test_ddp/test_ddp_ignore_params.py @@ -1,23 +1,20 @@ import os import random -from functools import partial from typing import Callable, Type import numpy as np import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import colossalai -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager -from colossalai.nn.parallel import ColoDDP, ZeroDDP +from colossalai.nn.parallel import ColoDDP from colossalai.tensor import ProcessGroup -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, ZeroDDP +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager def set_seed(seed): @@ -88,8 +85,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [2]) @rerun_if_address_is_in_use() def test_ddp_ignore_params(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ddp/test_ddp_state_dict.py b/tests/test_ddp/test_ddp_state_dict.py index f229364c6eb1..54f89f972765 100644 --- a/tests/test_ddp/test_ddp_state_dict.py +++ b/tests/test_ddp/test_ddp_state_dict.py @@ -1,18 +1,15 @@ -import copy +from collections import OrderedDict import pytest -import colossalai import torch -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use + +import colossalai +from colossalai.nn.parallel import ColoDDP +from colossalai.tensor import ColoParameter, ProcessGroup +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils import free_port -from colossalai.utils.model.colo_init_context import ColoInitContext -from functools import partial +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs -from colossalai.nn.parallel import ColoDDP -from collections import OrderedDict -from colossalai.tensor import ProcessGroup, ColoParameter def check_state_dict_equal(state_dict: OrderedDict, other_state_dict: OrderedDict): @@ -63,8 +60,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_state_dict(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ddp/test_reducer.py b/tests/test_ddp/test_reducer.py index 5b302d99ffb1..e8d3a112c938 100644 --- a/tests/test_ddp/test_reducer.py +++ b/tests/test_ddp/test_reducer.py @@ -1,15 +1,15 @@ +from functools import partial + import pytest -import colossalai import torch -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils.cuda import get_current_device -from colossalai.utils import free_port -from functools import partial -from colossalai.nn.parallel.reducer import Reducer import torch.distributed as dist from torch.distributed.distributed_c10d import _get_default_group +import colossalai +from colossalai.nn.parallel.reducer import Reducer +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.cuda import get_current_device + REDUCE_CNT = 0 @@ -40,8 +40,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_reducer(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_device/test_alpha_beta.py b/tests/test_device/test_alpha_beta.py index 99abacd1342b..ab933ed57d0d 100644 --- a/tests/test_device/test_alpha_beta.py +++ b/tests/test_device/test_alpha_beta.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest -import torch.multiprocessing as mp from colossalai.device import AlphaBetaProfiler from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn -def check_alpha_beta(rank, physical_devices, world_size, port): +def check_alpha_beta(rank, world_size, port, physical_devices): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') profiler = AlphaBetaProfiler(physical_devices) @@ -24,9 +20,7 @@ def check_alpha_beta(rank, physical_devices, world_size, port): @parameterize('physical_devices', [[0, 1, 2, 3], [0, 3]]) @rerun_if_address_is_in_use() def test_profile_alpha_beta(physical_devices): - world_size = 4 - run_func = partial(check_alpha_beta, physical_devices=physical_devices, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_alpha_beta, 4, physical_devices=physical_devices) if __name__ == '__main__': diff --git a/tests/test_device/test_device_mesh.py b/tests/test_device/test_device_mesh.py index 3be057b3a98b..590d6966bff6 100644 --- a/tests/test_device/test_device_mesh.py +++ b/tests/test_device/test_device_mesh.py @@ -1,21 +1,89 @@ -from colossalai.device.device_mesh import DeviceMesh +import pytest import torch +import torch.distributed as dist + +import colossalai +from colossalai.device.device_mesh import DeviceMesh +from colossalai.testing import rerun_if_address_is_in_use, spawn def test_device_mesh(): - physical_mesh_id = torch.arange(0, 16).reshape(2, 8) + physical_mesh_id = torch.arange(0, 16) mesh_shape = (4, 4) # [[0, 1, 2, 3], # [4, 5, 6, 7], # [8, 9, 10,11], # [12,13,14,15]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) - assert device_mesh.convert_map[5] == [1, 1] - assert device_mesh.convert_map[11] == [2, 3] - assert device_mesh.global_rank_to_process_groups_with_logical_rank(0)[0] == [[0, 0], [1, 0], [2, 0], [3, 0]] - assert device_mesh.global_rank_to_process_groups_with_logical_rank(2)[1] == [[0, 0], [0, 1], [0, 2], [0, 3]] - assert device_mesh.global_rank_to_process_groups_with_global_rank(2)[1] == [0, 1, 2, 3] + assert device_mesh.global_rank_to_local_rank(5) == [1, 1] + assert device_mesh.global_rank_to_local_rank(11) == [2, 3] + assert device_mesh.get_ranks_in_process_group(axis=1, global_rank=2) == [0, 1, 2, 3] + + +def check_1d_device_mesh(): + # check for 1D device mesh + process_group = dist.GroupMember.WORLD + device_mesh = DeviceMesh.from_process_group(process_group) + + # checks + assert device_mesh.shape == [4] + assert len(device_mesh.get_process_group_for_all_axes().keys()) == 1, 'Expected 1 axis for the process group dict' + assert device_mesh.get_process_group(axis=0) == process_group, 'Expected world process group' + assert device_mesh.is_initialized + assert device_mesh.num_devices == 4 + assert device_mesh.is_initialized + assert device_mesh.logical_mesh_id is None + assert device_mesh._is_init_from_process_group + + +def check_2d_device_mesh(): + # create process group for 2D device mesh + first_row_ranks = [0, 1] + second_row_ranks = [2, 3] + first_col_ranks = [0, 2] + second_col_ranks = [1, 3] + + first_row_pg = dist.new_group(first_row_ranks, backend='nccl') + second_row_pg = dist.new_group(second_row_ranks, backend='nccl') + first_col_pg = dist.new_group(first_col_ranks, backend='nccl') + second_col_pg = dist.new_group(second_col_ranks, backend='nccl') + + # check for + current_rank = dist.get_rank() + + if current_rank in first_row_ranks: + row_pg = first_row_pg + else: + row_pg = second_row_pg + + if current_rank in first_col_ranks: + col_pg = first_col_pg + else: + col_pg = second_col_pg + + device_mesh = DeviceMesh.from_process_group([col_pg, row_pg]) + + # checks + assert device_mesh.shape == [2, 2] + assert len(device_mesh.get_process_group_for_all_axes().keys()) == 2, 'Expected 2 axes for the process group dict' + assert device_mesh.get_process_group(axis=0) == col_pg, 'Expected column process group' + assert device_mesh.get_process_group(axis=1) == row_pg, 'Expected row process group' + assert device_mesh.num_devices == 4 + assert device_mesh.is_initialized + assert device_mesh.logical_mesh_id is None + assert device_mesh._is_init_from_process_group + + +def check_init_from_process_group(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_device_mesh_from_process_group(): + spawn(check_init_from_process_group, 4) if __name__ == '__main__': test_device_mesh() + test_device_mesh_from_process_group() diff --git a/tests/test_device/test_extract_alpha_beta.py b/tests/test_device/test_extract_alpha_beta.py index e32bebdd908e..52604b9c6a49 100644 --- a/tests/test_device/test_extract_alpha_beta.py +++ b/tests/test_device/test_extract_alpha_beta.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest -import torch.multiprocessing as mp from colossalai.device import AlphaBetaProfiler from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn -def check_extract_alpha_beta(rank, physical_devices, world_size, port): +def check_extract_alpha_beta(rank, world_size, port, physical_devices): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') profiler = AlphaBetaProfiler(physical_devices) @@ -27,12 +23,7 @@ def check_extract_alpha_beta(rank, physical_devices, world_size, port): @parameterize('physical_devices', [[0, 1, 2, 3], [0, 3]]) @rerun_if_address_is_in_use() def test_profile_alpha_beta(physical_devices): - world_size = 4 - run_func = partial(check_extract_alpha_beta, - physical_devices=physical_devices, - world_size=world_size, - port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_extract_alpha_beta, 4, physical_devices=physical_devices) if __name__ == '__main__': diff --git a/tests/test_device/test_init_logical_pg.py b/tests/test_device/test_init_logical_pg.py index 3172897fb5cd..7c6339eff67e 100644 --- a/tests/test_device/test_init_logical_pg.py +++ b/tests/test_device/test_init_logical_pg.py @@ -1,15 +1,12 @@ -import torch -from functools import partial import pytest +import torch import torch.distributed as dist -import torch.multiprocessing as mp from torch.distributed import ReduceOp from colossalai.core import global_context as gpc -from colossalai.initialize import launch -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use from colossalai.device.device_mesh import DeviceMesh +from colossalai.initialize import launch +from colossalai.testing import rerun_if_address_is_in_use, spawn def check_layer(rank, world_size, port): @@ -23,16 +20,12 @@ def check_layer(rank, world_size, port): # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - logical_pg_dict = {0: [[0, 2], [1, 3]], 1: [[0, 1], [2, 3]]} - logical_process_groups = device_mesh.process_groups_dict - for mesh_dim, pgs in logical_pg_dict.items(): - for index, pg in enumerate(pgs): - if rank in pg: - tensor = torch.ones(4).cuda() - group = logical_process_groups[mesh_dim][index][1] - dist.all_reduce(tensor, op=ReduceOp.SUM, group=group) - assert tensor.equal(tensor_to_check) + for axis in range(len(mesh_shape)): + tensor = torch.ones(4).cuda() + pg = device_mesh.get_process_group(axis=axis) + dist.all_reduce(tensor, op=ReduceOp.SUM, group=pg) + assert tensor.equal(tensor_to_check) gpc.destroy() @@ -40,9 +33,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_logical_pg(): - world_size = 4 - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, 4) if __name__ == '__main__': diff --git a/tests/test_device/test_search_logical_device_mesh.py b/tests/test_device/test_search_logical_device_mesh.py index 591eafb2a50d..b22a76eabc2f 100644 --- a/tests/test_device/test_search_logical_device_mesh.py +++ b/tests/test_device/test_search_logical_device_mesh.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest -import torch.multiprocessing as mp from colossalai.device import AlphaBetaProfiler from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn -def check_alpha_beta(rank, physical_devices, world_size, port): +def check_alpha_beta(rank, world_size, port, physical_devices): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') profiler = AlphaBetaProfiler(physical_devices) @@ -27,9 +23,7 @@ def check_alpha_beta(rank, physical_devices, world_size, port): @parameterize('physical_devices', [[0, 1, 2, 3], [0, 3]]) @rerun_if_address_is_in_use() def test_profile_alpha_beta(physical_devices): - world_size = 4 - run_func = partial(check_alpha_beta, physical_devices=physical_devices, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_alpha_beta, 4, physical_devices=physical_devices) if __name__ == '__main__': diff --git a/tests/test_engine/test_engine.py b/tests/test_engine/test_engine.py index fb5bd1e1602e..62493cf3712d 100644 --- a/tests/test_engine/test_engine.py +++ b/tests/test_engine/test_engine.py @@ -1,13 +1,10 @@ -from functools import partial +import pytest import colossalai -import pytest -import torch.multiprocessing as mp from colossalai.amp import AMP_TYPE from colossalai.core import global_context as gpc -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from tests.components_to_test.registry import non_distributed_component_funcs -from colossalai.testing import parameterize, rerun_if_address_is_in_use CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None)), fp16=dict(mode=None), @@ -58,9 +55,7 @@ def run_engine(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_engine(): - world_size = 2 - run_func = partial(run_engine, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_engine, 2) if __name__ == '__main__': diff --git a/tests/test_engine/test_gradient_accumluation.py b/tests/test_engine/test_gradient_accumluation.py index 7f5ee47be8e6..7783827c7c44 100644 --- a/tests/test_engine/test_gradient_accumluation.py +++ b/tests/test_engine/test_gradient_accumluation.py @@ -1,22 +1,20 @@ import os -from functools import partial from pathlib import Path -import colossalai -from colossalai.testing.utils import rerun_if_address_is_in_use import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.core import global_context as gpc -from colossalai.logging import get_dist_logger -from colossalai.utils import free_port, get_dataloader -from colossalai.testing import rerun_if_address_is_in_use from torch.optim import Adam from torchvision import transforms from torchvision.datasets import CIFAR10 from torchvision.models import resnet18 +import colossalai +from colossalai.core import global_context as gpc +from colossalai.logging import get_dist_logger +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_dataloader + # Config BATCH_SIZE = 2 NUM_CLASSES = 10 @@ -90,9 +88,7 @@ def run_no_pipeline(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_engine(): - world_size = 4 - func = partial(run_no_pipeline, world_size=world_size, port=free_port()) - mp.spawn(func, nprocs=world_size) + spawn(run_no_pipeline, 4) if __name__ == '__main__': diff --git a/tests/test_fx/test_codegen/test_activation_checkpoint_codegen.py b/tests/test_fx/test_codegen/test_activation_checkpoint_codegen.py index 83df1bb5e69c..bcac2ec426d9 100644 --- a/tests/test_fx/test_codegen/test_activation_checkpoint_codegen.py +++ b/tests/test_fx/test_codegen/test_activation_checkpoint_codegen.py @@ -1,15 +1,13 @@ import pytest import torch -import torch.multiprocessing as mp import torch.nn.functional as F -from torch.fx import GraphModule from torch.utils.checkpoint import checkpoint import colossalai from colossalai.core import global_context as gpc from colossalai.fx import ColoTracer from colossalai.fx.graph_module import ColoGraphModule -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn try: from colossalai.fx.codegen import ActivationCheckpointCodeGen @@ -65,9 +63,9 @@ def forward(self, x, y): return y1 + y2 + y3 + y4 + y5 + y6 -def _run_act_ckpt_codegen(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_act_ckpt_codegen(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and run forward model = MyModule() @@ -118,13 +116,14 @@ def _run_act_ckpt_codegen(rank): @pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0') +@rerun_if_address_is_in_use() def test_act_ckpt_codegen(): - mp.spawn(_run_act_ckpt_codegen, nprocs=1) + spawn(_run_act_ckpt_codegen, 1) -def _run_act_ckpt_python_code_torch11(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_act_ckpt_python_code_torch11(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and run forward model = MyModule() @@ -174,8 +173,9 @@ def _run_act_ckpt_python_code_torch11(rank): @pytest.mark.skipif(with_codegen, reason='torch version is equal to or higher than 1.12.0') @pytest.mark.skip(reason="currently torch11 ColoGraphModule is not done") +@rerun_if_address_is_in_use() def test_act_ckpt_python_code_torch11(): - mp.spawn(_run_act_ckpt_python_code_torch11, nprocs=1) + spawn(_run_act_ckpt_python_code_torch11, 1) if __name__ == '__main__': diff --git a/tests/test_fx/test_codegen/test_nested_activation_checkpoint_codegen.py b/tests/test_fx/test_codegen/test_nested_activation_checkpoint_codegen.py index 6b3a49d181e1..5b327807a57b 100644 --- a/tests/test_fx/test_codegen/test_nested_activation_checkpoint_codegen.py +++ b/tests/test_fx/test_codegen/test_nested_activation_checkpoint_codegen.py @@ -1,15 +1,11 @@ import pytest import torch -import torch.multiprocessing as mp -import torch.nn.functional as F -from torch.fx import GraphModule -from torch.utils.checkpoint import checkpoint import colossalai from colossalai.core import global_context as gpc from colossalai.fx import ColoTracer from colossalai.fx.graph_module import ColoGraphModule -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn try: from colossalai.fx.codegen import ActivationCheckpointCodeGen @@ -35,9 +31,9 @@ def forward(self, x): return self.linear6(self.linear5(self.linear4(self.linear3(self.linear2(self.linear1(x)))))) -def _run_act_ckpt_codegen(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_act_ckpt_codegen(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and run forward model = MyModule() @@ -89,12 +85,12 @@ def _run_act_ckpt_codegen(rank): @pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0') def test_act_ckpt_codegen(): - mp.spawn(_run_act_ckpt_codegen, nprocs=1) + spawn(_run_act_ckpt_codegen, 1) -def _run_act_ckpt_python_code_torch11(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_act_ckpt_python_code_torch11(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and run forward model = MyModule() @@ -146,8 +142,9 @@ def _run_act_ckpt_python_code_torch11(rank): @pytest.mark.skipif(with_codegen, reason='torch version is equal to or higher than 1.12.0') @pytest.mark.skip(reason="currently torch11 ColoGraphModule is not done") +@rerun_if_address_is_in_use() def test_act_ckpt_python_code_torch11(): - mp.spawn(_run_act_ckpt_python_code_torch11, nprocs=1) + spawn(_run_act_ckpt_python_code_torch11, 1) if __name__ == '__main__': diff --git a/tests/test_fx/test_codegen/test_offload_codegen.py b/tests/test_fx/test_codegen/test_offload_codegen.py index 5d090066c763..c217b96586fe 100644 --- a/tests/test_fx/test_codegen/test_offload_codegen.py +++ b/tests/test_fx/test_codegen/test_offload_codegen.py @@ -2,15 +2,13 @@ import pytest import torch -import torch.multiprocessing as mp -import torch.nn.functional as F from torch.fx import GraphModule import colossalai from colossalai.core import global_context as gpc from colossalai.fx import ColoTracer from colossalai.fx.graph_module import ColoGraphModule -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn try: from colossalai.fx.codegen import ActivationCheckpointCodeGen @@ -58,7 +56,7 @@ def _test_fwd_and_bwd(model: torch.nn.Module, gm: ColoGraphModule, data: torch.T fx_out = gm(data) assert torch.equal(non_fx_out, fx_out), "fx_out doesn't comply with original output" - # test barckward + # test backward loss0 = non_fx_out.sum() loss0.backward() loss1 = fx_out.sum() @@ -66,9 +64,9 @@ def _test_fwd_and_bwd(model: torch.nn.Module, gm: ColoGraphModule, data: torch.T assert _is_all_gradient_close(model, gm), "gm doesn't have the same gradient as original one" -def _run_offload_codegen(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_offload_codegen(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and input model = MyNet().cuda() @@ -116,13 +114,14 @@ def _run_offload_codegen(rank): @pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0') +@rerun_if_address_is_in_use() def test_act_ckpt_codegen(): - mp.spawn(_run_offload_codegen, nprocs=1) + spawn(_run_offload_codegen, 1) -def _run_offload_codegen_torch11(rank): - # launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly - colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl') +def _run_offload_codegen_torch11(rank, world_size, port): + # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') # build model and input model = MyNet().cuda() @@ -171,8 +170,9 @@ def _run_offload_codegen_torch11(rank): @pytest.mark.skip(reason="currently torch11 ColoGraphModule is not implemented") +@rerun_if_address_is_in_use() def test_act_ckpt_python_code_torch11(): - mp.spawn(_run_offload_codegen_torch11, nprocs=1) + spawn(_run_offload_codegen_torch11, 1) if __name__ == "__main__": diff --git a/tests/test_fx/test_coloproxy.py b/tests/test_fx/test_coloproxy.py index 2bb6cf86466c..96cf5198da10 100644 --- a/tests/test_fx/test_coloproxy.py +++ b/tests/test_fx/test_coloproxy.py @@ -1,9 +1,11 @@ +import pytest import torch import torch.nn as nn +from torch.fx import GraphModule + from colossalai.fx.proxy import ColoProxy from colossalai.fx.tracer.tracer import ColoTracer -from torch.fx import GraphModule -import pytest +from colossalai.testing import clear_cache_before_run class Conv1D(nn.Module): @@ -23,6 +25,7 @@ def forward(self, x): return x +@clear_cache_before_run() def test_coloproxy(): tracer = ColoTracer() diff --git a/tests/test_fx/test_comm_size_compute.py b/tests/test_fx/test_comm_size_compute.py index 8825bbb461d6..d3daadd71406 100644 --- a/tests/test_fx/test_comm_size_compute.py +++ b/tests/test_fx/test_comm_size_compute.py @@ -1,13 +1,11 @@ -import colossalai -import colossalai.nn as col_nn -import pytest import torch -import torch.nn as nn +from torch.fx import symbolic_trace + from colossalai.fx._compatibility import is_compatible_with_meta -from colossalai.fx.passes.adding_split_node_pass import (split_with_split_nodes_pass, uniform_split_pass) +from colossalai.fx.passes.adding_split_node_pass import split_with_split_nodes_pass, uniform_split_pass from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.fx.passes.utils import get_comm_size -from torch.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run is_compatible = is_compatible_with_meta() if is_compatible: @@ -35,6 +33,7 @@ def forward(self, x): return x +@clear_cache_before_run() def test_comm_size_compute(): model = MLP(MODEL_DIM) input_sample = torch.rand(BATCH_SIZE, MODEL_DIM, device='meta') diff --git a/tests/test_fx/test_complete_workflow.py b/tests/test_fx/test_complete_workflow.py deleted file mode 100644 index a21a351f8d77..000000000000 --- a/tests/test_fx/test_complete_workflow.py +++ /dev/null @@ -1,87 +0,0 @@ -from functools import partial - -import pytest -import torch -import torch.distributed as dist -import torch.multiprocessing as mp -import torch.nn as nn - -import colossalai -from colossalai.fx import ColoTracer -from colossalai.fx.passes.shard_1d_pass import transformer_mlp_pass -from colossalai.tensor import ProcessGroup -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.utils.model.lazy_init_context import LazyInitContext - - -class MLP(torch.nn.Module): - - def __init__(self, dim: int): - super().__init__() - self.linear1 = torch.nn.Linear(dim, dim) - self.linear2 = torch.nn.Linear(dim, dim) - self.dropout = torch.nn.Dropout(0) - self.relu = torch.nn.ReLU() - - def forward(self, x): - x = self.linear1(x) - x = self.dropout(x) - x = self.relu(x) - x = self.linear2(x) - return x - - -def run_workflow(world_size, dev): - # initailization - with LazyInitContext() as ctx: - model = MLP(16) - - for param in model.parameters(): - assert param.is_meta - - # tracing - tracer = ColoTracer() - graph = tracer.trace(model) - gm = torch.fx.GraphModule(model, graph, model.__class__.__name__) - - # annotate - annotated_gm = transformer_mlp_pass(gm, process_group=ProcessGroup(tp_degree=world_size)) - annotated_gm.recompile() - - # materialization and sharding - ctx.lazy_init_parameters(annotated_gm, device=dev) - for param in model.parameters(): - assert not param.is_meta - - # # check sharding - assert list(model.linear1.weight.shape) == [16 // world_size, 16] - assert list(model.linear1.bias.shape) == [16 // world_size] - assert list(model.linear2.weight.shape) == [16, 16 // world_size] - - # test forward to make sure that IR transform will produce the same results - # like how ColoTensor would do it normally - data = torch.rand(4, 16, device=dev) - non_fx_out = model(data) - fx_out = annotated_gm(data) - assert torch.equal(non_fx_out, fx_out), f'{non_fx_out} vs {fx_out}' - - -def run_dist(rank, world_size, dev, port): - colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - run_workflow(world_size, dev) - - -@pytest.mark.dist -@pytest.mark.parametrize('world_size', [1, 2]) -@pytest.mark.parametrize('dev', ['cuda', 'cpu']) -@rerun_if_address_is_in_use() -def test_complete_workflow(world_size, dev): - if dev == 'cpu' and world_size > 1: - return - run_func = partial(run_dist, world_size=world_size, dev=dev, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_complete_workflow(1, 'cuda') diff --git a/tests/test_fx/test_graph_manipulation.py b/tests/test_fx/test_graph_manipulation.py index fb33e58a778c..175b69dd96fe 100644 --- a/tests/test_fx/test_graph_manipulation.py +++ b/tests/test_fx/test_graph_manipulation.py @@ -1,9 +1,11 @@ -import colossalai import torch -from colossalai.fx.passes.utils import get_leaf, get_top, assign_bfs_level_to_nodes -from colossalai.fx import ColoTracer from torch.fx import GraphModule + +import colossalai +from colossalai.fx import ColoTracer from colossalai.fx.passes.meta_info_prop import MetaInfoProp, TensorMetadata +from colossalai.fx.passes.utils import assign_bfs_level_to_nodes, get_leaf, get_top +from colossalai.testing import clear_cache_before_run class MLP(torch.nn.Module): @@ -25,6 +27,7 @@ def forward(self, x): return l4, l5 +@clear_cache_before_run() def test_graph_manipulation(): model = MLP(4) tracer = ColoTracer() diff --git a/tests/test_fx/test_meta/test_aten.py b/tests/test_fx/test_meta/test_aten.py index 209ded89cfb9..e490522dbf15 100644 --- a/tests/test_fx/test_meta/test_aten.py +++ b/tests/test_fx/test_meta/test_aten.py @@ -3,7 +3,9 @@ import pytest import torch import torch.nn as nn + from colossalai.fx._compatibility import is_compatible_with_meta +from colossalai.testing import clear_cache_before_run if is_compatible_with_meta(): from colossalai.fx.profiler import MetaTensor @@ -71,6 +73,7 @@ def run_and_compare(f: Union[nn.Module, Callable], x: torch.Tensor, requires_bac @pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0') +@clear_cache_before_run() def test_meta_aten(): for (aten_op, requires_backward), v in registered_meta.items(): for f, x in v: diff --git a/tests/test_fx/test_meta/test_backward.py b/tests/test_fx/test_meta/test_backward.py index 351c02c5744a..7aed6fd4597b 100644 --- a/tests/test_fx/test_meta/test_backward.py +++ b/tests/test_fx/test_meta/test_backward.py @@ -2,11 +2,14 @@ import timm.models as tmm import torch import torchvision.models as tm + from colossalai.fx._compatibility import is_compatible_with_meta if is_compatible_with_meta(): from colossalai.fx.profiler import MetaTensor +from colossalai.testing import clear_cache_before_run + tm_models = [ tm.vgg11, tm.resnet18, @@ -28,6 +31,7 @@ @pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0') +@clear_cache_before_run() def test_torchvision_models(): for m in tm_models: model = m() @@ -36,6 +40,7 @@ def test_torchvision_models(): @pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0') +@clear_cache_before_run() def test_timm_models(): for m in tmm_models: model = m() diff --git a/tests/test_fx/test_meta/test_meta_trace.py b/tests/test_fx/test_meta/test_meta_trace.py index 404b6d27d2d4..61614f8a6623 100644 --- a/tests/test_fx/test_meta/test_meta_trace.py +++ b/tests/test_fx/test_meta/test_meta_trace.py @@ -2,11 +2,14 @@ import timm.models as tmm import torch import torchvision.models as tm + from colossalai.fx._compatibility import is_compatible_with_meta if is_compatible_with_meta(): from colossalai.fx import meta_trace +from colossalai.testing import clear_cache_before_run + tm_models = [ tm.vgg11, tm.resnet18, @@ -28,6 +31,7 @@ @pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0') +@clear_cache_before_run() def test_torchvision_models_trace(): for m in tm_models: model = m() @@ -36,6 +40,7 @@ def test_torchvision_models_trace(): @pytest.mark.skipif(not is_compatible_with_meta(), reason='torch version is lower than 1.12.0') +@clear_cache_before_run() def test_timm_models_trace(): for m in tmm_models: model = m() diff --git a/tests/test_fx/test_meta_info_prop.py b/tests/test_fx/test_meta_info_prop.py index 6fac180d8ba2..a12512696a73 100644 --- a/tests/test_fx/test_meta_info_prop.py +++ b/tests/test_fx/test_meta_info_prop.py @@ -1,7 +1,9 @@ import torch +from torch.fx import symbolic_trace + from colossalai.fx._compatibility import is_compatible_with_meta from colossalai.fx.passes.meta_info_prop import MetaInfoProp, TensorMetadata -from torch.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run if is_compatible_with_meta(): from colossalai.fx.profiler import MetaTensor @@ -18,6 +20,7 @@ def meta_check(meta_info_spec: TensorMetadata, orig_tensor: torch.Tensor): assert meta_info_spec.numel == orig_tensor.numel() +@clear_cache_before_run() def test_meta_info_prop(): model = torch.nn.Linear(DIM_IN, DIM_OUT) input_sample = torch.rand(BATCH_SIZE, DIM_IN, device='meta') diff --git a/tests/test_fx/test_parallel_1d.py b/tests/test_fx/test_parallel_1d.py index 8963ba29cb03..1044be7db1f4 100644 --- a/tests/test_fx/test_parallel_1d.py +++ b/tests/test_fx/test_parallel_1d.py @@ -1,18 +1,15 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp -from colossalai.core import global_context as gpc -from colossalai.logging import disable_existing_loggers -from colossalai.initialize import launch -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use from torch.fx import symbolic_trace + +from colossalai.core import global_context as gpc from colossalai.fx.passes import column_shard_linear_pass +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn class MLP(torch.nn.Module): @@ -52,11 +49,10 @@ def check_layer(rank, world_size, port): @pytest.mark.dist +@clear_cache_before_run() @rerun_if_address_is_in_use() def test_1d(): - world_size = 2 - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, 2) if __name__ == '__main__': diff --git a/tests/test_fx/test_pipeline/test_topo/test_topo.py b/tests/test_fx/test_pipeline/test_topo/test_topo.py index 75c74870523c..16da56250dc3 100644 --- a/tests/test_fx/test_pipeline/test_topo/test_topo.py +++ b/tests/test_fx/test_pipeline/test_topo/test_topo.py @@ -1,11 +1,13 @@ import pytest import torch import transformers -from topo_utils import split_model_and_get_DAG, check_topo, MLP +from topo_utils import MLP, check_topo, split_model_and_get_DAG BATCH_SIZE = 1 SEQ_LENGHT = 16 + +@pytest.mark.skip('ShapeProp is not compatible with PyTorch 1.11.0') def test_opt(): MODEL_LIST = [ MLP, @@ -13,7 +15,10 @@ def test_opt(): ] CONFIGS = [ - {'dim': 10, 'layers': 12}, + { + 'dim': 10, + 'layers': 12 + }, transformers.OPTConfig(vocab_size=100, hidden_size=128, num_hidden_layers=4, num_attention_heads=4), ] @@ -21,15 +26,15 @@ def data_gen_MLP(): x = torch.zeros((16, 10)) kwargs = dict(x=x) return kwargs - + def data_gen_OPT(): input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGHT), dtype=torch.int64) attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGHT), dtype=torch.int64) kwargs = dict(input_ids=input_ids, attention_mask=attention_mask) return kwargs - + DATAGEN = [ - data_gen_MLP, + data_gen_MLP, data_gen_OPT, ] @@ -39,5 +44,6 @@ def data_gen_OPT(): # print(f'{top_mod=}\n----\n{topo=}') check_topo(top_mod, topo) + if __name__ == '__main__': - test_opt() \ No newline at end of file + test_opt() diff --git a/tests/test_fx/test_pipeline_passes.py b/tests/test_fx/test_pipeline_passes.py index de8a9402ba56..1078dac9db7c 100644 --- a/tests/test_fx/test_pipeline_passes.py +++ b/tests/test_fx/test_pipeline_passes.py @@ -1,12 +1,17 @@ +import pytest import torch import torch.nn as nn -import colossalai -import colossalai.nn as col_nn from torch.fx import symbolic_trace -from colossalai.fx.passes.adding_split_node_pass import split_with_split_nodes_pass, balanced_split_pass, \ - uniform_split_pass, balanced_split_pass_v2 -import pytest +import colossalai +import colossalai.nn as col_nn +from colossalai.fx.passes.adding_split_node_pass import ( + balanced_split_pass, + balanced_split_pass_v2, + split_with_split_nodes_pass, + uniform_split_pass, +) +from colossalai.testing import clear_cache_before_run MODEL_DIM = 16 BATCH_SIZE = 8 @@ -39,6 +44,7 @@ def pipeline_pass_test_helper(model, data, pass_func): assert output.equal(origin_output) +@clear_cache_before_run() def test_pipeline_passes(): model = MLP(MODEL_DIM) data = torch.rand(BATCH_SIZE, MODEL_DIM) diff --git a/tests/test_fx/test_profiler/test_profiler_meta_info_prop.py b/tests/test_fx/test_profiler/test_profiler_meta_info_prop.py index c717960181ad..b5a6bbe8bf18 100644 --- a/tests/test_fx/test_profiler/test_profiler_meta_info_prop.py +++ b/tests/test_fx/test_profiler/test_profiler_meta_info_prop.py @@ -9,7 +9,7 @@ from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.fx.profiler import calculate_fwd_out, calculate_fwd_tmp, is_compatible_with_meta, parameter_size from colossalai.fx.tracer.tracer import ColoTracer -from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.testing import clear_cache_before_run, run_on_environment_flag if is_compatible_with_meta(): from colossalai.fx.profiler import MetaTensor @@ -126,6 +126,7 @@ def run_gpt_forward(gm: torch.fx.GraphModule): @run_on_environment_flag(name='FX_PROFILER') +@clear_cache_before_run() def test_meta_info_prop(): for m in [ tm.alexnet, tm.resnet18, tm.resnet34, tm.resnet50, tm.resnet101, tm.resnet152, tm.densenet121, @@ -155,6 +156,7 @@ def test_meta_info_prop(): @run_on_environment_flag(name='FX_PROFILER') +@clear_cache_before_run() def test_gpt_meta_info_prop(): for m in [gpt2_medium]: model = m().cuda() diff --git a/tests/test_fx/test_tracer/test_activation_checkpoint_annotation.py b/tests/test_fx/test_tracer/test_activation_checkpoint_annotation.py index a834951bb695..632ab8c09750 100644 --- a/tests/test_fx/test_tracer/test_activation_checkpoint_annotation.py +++ b/tests/test_fx/test_tracer/test_activation_checkpoint_annotation.py @@ -4,6 +4,7 @@ from torch.utils.checkpoint import checkpoint from colossalai.fx import ColoTracer +from colossalai.testing import clear_cache_before_run class MLP(torch.nn.Module): @@ -35,6 +36,7 @@ def forward(self, x): return x +@clear_cache_before_run() def test_activation_checkpoint_annotation(): module = MyModule() diff --git a/tests/test_fx/test_tracer/test_bias_addition_module.py b/tests/test_fx/test_tracer/test_bias_addition_module.py index afa30a217604..2f88d8c784e8 100644 --- a/tests/test_fx/test_tracer/test_bias_addition_module.py +++ b/tests/test_fx/test_tracer/test_bias_addition_module.py @@ -1,6 +1,7 @@ import torch from colossalai.fx import ColoGraphModule, ColoTracer +from colossalai.testing import clear_cache_before_run class LinearModel(torch.nn.Module): @@ -32,6 +33,7 @@ def forward(self, x): return x +@clear_cache_before_run() def test_linear_module(): model = LinearModel(3, 6) tracer = ColoTracer() @@ -68,6 +70,7 @@ def test_linear_module(): assert add_node._meta_data.shape == (3, 6) +@clear_cache_before_run() def test_conv_module(): model = ConvModel(3, 6, 2) tracer = ColoTracer() diff --git a/tests/test_fx/test_tracer/test_control_flow.py b/tests/test_fx/test_tracer/test_control_flow.py index ed842cff2776..820729dadb3e 100644 --- a/tests/test_fx/test_tracer/test_control_flow.py +++ b/tests/test_fx/test_tracer/test_control_flow.py @@ -1,7 +1,9 @@ import torch import torch.nn as nn from torch.fx import GraphModule + from colossalai.fx import ColoTracer as Tracer +from colossalai.testing import clear_cache_before_run class ControlFlowModel(nn.Module): @@ -21,6 +23,7 @@ def forward(self, x, y): return x1 - y1 +@clear_cache_before_run() def test_control_flow(): model = ControlFlowModel() tracer = Tracer() diff --git a/tests/test_fx/test_tracer/test_functional_conv.py b/tests/test_fx/test_tracer/test_functional_conv.py index 95670b85f335..a552e905223d 100644 --- a/tests/test_fx/test_tracer/test_functional_conv.py +++ b/tests/test_fx/test_tracer/test_functional_conv.py @@ -1,8 +1,11 @@ import torch from torch.nn import functional as F + from colossalai.fx.tracer.meta_patch import patched_function +from colossalai.testing import clear_cache_before_run +@clear_cache_before_run() def test_conv(): # test F.conv_1d data_1d = torch.rand(3, 16, 10) diff --git a/tests/test_fx/test_tracer/test_hf_model/hf_tracer_utils.py b/tests/test_fx/test_tracer/test_hf_model/hf_tracer_utils.py index 6d93fe0408d7..58c8132e1490 100644 --- a/tests/test_fx/test_tracer/test_hf_model/hf_tracer_utils.py +++ b/tests/test_fx/test_tracer/test_hf_model/hf_tracer_utils.py @@ -1,24 +1,31 @@ +from typing import List + import torch from numpy import isin from torch.fx import GraphModule from torch.utils._pytree import tree_flatten -from colossalai.fx import symbolic_trace +# from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace -def trace_model_and_compare_output(model, data_gen): +def trace_model_and_compare_output(model, data_gen, ignore_data: List[str] = None): # must turn on eval mode to ensure the output is consistent model.eval() + inputs = data_gen() + + if ignore_data is not None: + # drop the ignore_data key + inputs = {k: v for k, v in inputs.items() if k not in ignore_data} + try: - kwargs = data_gen() - meta_args = {k: v.to('meta') for k, v in kwargs.items()} + meta_args = {k: v.to('meta') for k, v in inputs.items()} gm = symbolic_trace(model, meta_args=meta_args) except Exception as e: raise RuntimeError(f"Failed to trace {model.__class__.__name__}, error: {e}") # run forward - inputs = data_gen() non_fx_out = model(**inputs) fx_out = gm(**inputs) diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py index b1c9c211a9a0..a1470400ad82 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py @@ -1,15 +1,21 @@ +import pytest +import torch from hf_tracer_utils import trace_model_and_compare_output +from packaging import version +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo BATCH_SIZE = 2 SEQ_LENGTH = 16 +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_albert(): sub_registry = model_zoo.get_sub_registry('transformers_albert') - for name, (model_fn, data_gen_fn, _, _) in sub_registry.items(): + for name, (model_fn, data_gen_fn, _, _, _) in sub_registry.items(): model = model_fn() trace_model_and_compare_output(model, data_gen_fn) diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py index 1bf4947c31a0..632ad366ccc4 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py @@ -1,14 +1,20 @@ +import pytest +import torch from hf_tracer_utils import trace_model_and_compare_output +from packaging import version +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_bert(): sub_registry = model_zoo.get_sub_registry('transformers_bert') - for name, (model_fn, data_gen_fn, _, _) in sub_registry.items(): + for name, (model_fn, data_gen_fn, _, _, _) in sub_registry.items(): model = model_fn() - trace_model_and_compare_output(model, data_gen_fn) + trace_model_and_compare_output(model, data_gen_fn, ignore_data=['labels', 'next_sentence_label']) if __name__ == '__main__': diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py index 92ece357bfed..ac87a7fcb13b 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py @@ -2,6 +2,7 @@ import torch from colossalai.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run from colossalai.testing.random import seed_all from tests.kit.model_zoo import model_zoo @@ -40,24 +41,26 @@ def assert_fn(ta, tb): @pytest.mark.skip(reason='cannot pass this test yet') +@clear_cache_before_run() def test_diffusers(): seed_all(9091, cuda_deterministic=True) sub_model_zoo = model_zoo.get_sub_registry('diffusers') - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in sub_model_zoo.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, _, attribute) in sub_model_zoo.items(): data = data_gen_fn() trace_and_compare(model_fn, data, output_transform_fn) torch.cuda.synchronize() print(f"{name:40s} √") +@clear_cache_before_run() def test_torch_diffusers(): seed_all(65535, cuda_deterministic=True) sub_model_zoo = model_zoo.get_sub_registry('diffusers') - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in sub_model_zoo.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, attribute) in sub_model_zoo.items(): data = data_gen_fn() model = model_fn() output = model(**data) diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py index 67a3178fae1b..31bcb7028e25 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py @@ -1,17 +1,27 @@ import pytest +import torch from hf_tracer_utils import trace_model_and_compare_output +from packaging import version +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo -# TODO: remove this skip once we handle the latest gpt model -@pytest.mark.skip +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_gpt(): sub_registry = model_zoo.get_sub_registry('transformers_gpt') - for name, (model_fn, data_gen_fn, _, _) in sub_registry.items(): + for name, (model_fn, data_gen_fn, _, _, _) in sub_registry.items(): model = model_fn() - trace_model_and_compare_output(model, data_gen_fn) + + # TODO: support the following models + # 1. GPT2DoubleHeadsModel + # as they are not supported, let's skip them + if model.__class__.__name__ in ['GPT2DoubleHeadsModel']: + continue + + trace_model_and_compare_output(model, data_gen_fn, ignore_data=['labels']) if __name__ == '__main__': diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py index 740f5a9f0c57..c68b89e82fbe 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py @@ -1,14 +1,19 @@ +import pytest +import torch from hf_tracer_utils import trace_model_and_compare_output +from packaging import version +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_opt(): sub_registry = model_zoo.get_sub_registry('transformers_opt') - - for name, (model_fn, data_gen_fn, _, _) in sub_registry.items(): + for name, (model_fn, data_gen_fn, _, _, _) in sub_registry.items(): model = model_fn() - trace_model_and_compare_output(model, data_gen_fn) + trace_model_and_compare_output(model, data_gen_fn, ignore_data=['labels', 'start_positions', 'end_positions']) if __name__ == '__main__': diff --git a/tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py b/tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py index 7073fd63470b..45e06bc2bbb0 100644 --- a/tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py +++ b/tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py @@ -1,14 +1,25 @@ +import pytest +import torch from hf_tracer_utils import trace_model_and_compare_output +from packaging import version +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_t5(): sub_registry = model_zoo.get_sub_registry('transformers_t5') - for name, (model_fn, data_gen_fn, _, _) in sub_registry.items(): + for name, (model_fn, data_gen_fn, _, _, _) in sub_registry.items(): + if name == "transformers_t5_for_conditional_generation": + # cannot trace for loss function yet + # so we use a data gen which does not produce labels + data_gen_fn = sub_registry.get('transformers_t5')[1] + model = model_fn() - trace_model_and_compare_output(model, data_gen_fn) + trace_model_and_compare_output(model, data_gen_fn, ignore_data=['labels']) if __name__ == '__main__': diff --git a/tests/test_fx/test_tracer/test_patched_module.py b/tests/test_fx/test_tracer/test_patched_module.py index 94a93e16f3c7..ef778e21801a 100644 --- a/tests/test_fx/test_tracer/test_patched_module.py +++ b/tests/test_fx/test_tracer/test_patched_module.py @@ -1,5 +1,7 @@ import torch + from colossalai.fx.tracer.meta_patch import patched_module +from colossalai.testing import clear_cache_before_run def _run(data, module, patch_fn): @@ -31,6 +33,7 @@ def _assert_output_shape(data, module, patch_fn, expect_exception, output_shape) assert output.shape == output_shape +@clear_cache_before_run() def test_linear(): # test linear patch can produce the meta output with correct shape data = torch.rand(2, 4, device='meta') @@ -42,6 +45,7 @@ def test_linear(): _assert_output_shape(data, module, patched_module.torch_nn_linear, True, None) +@clear_cache_before_run() def test_rnn(): # test rnn patch can produce the meta output with correct shape data = (torch.randn(5, 3, 10), torch.randn(2, 3, 20)) @@ -58,6 +62,7 @@ def test_rnn(): _assert_output_shape(meta_data, module, patched_module.torch_nn_rnn, True, None) +@clear_cache_before_run() def test_embedding(): data = torch.rand(2, 4, device='meta') @@ -134,6 +139,7 @@ def test_embedding(): output_shape=None) +@clear_cache_before_run() def test_conv1d(): # test conv 1d data = torch.rand(2, 3, 4) @@ -212,6 +218,7 @@ def test_conv2d(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_conv3d(): # test conv 3d data = torch.rand(2, 3, 4, 4, 4) @@ -253,6 +260,7 @@ def test_conv3d(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_conv_transpose1d(): # test conv transpose1d data = torch.rand(2, 3, 4) @@ -276,6 +284,7 @@ def test_conv_transpose1d(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_conv_transpose2d(): # test conv transpose2d data = torch.rand(2, 3, 4, 4) @@ -299,6 +308,7 @@ def test_conv_transpose2d(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_conv_transpose3d(): # test conv transpose2d data = torch.rand(2, 3, 4, 4, 4) @@ -322,6 +332,7 @@ def test_conv_transpose3d(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_pool1d(): combinations = [[torch.nn.MaxPool1d, patched_module.torch_nn_maxpool1d], [torch.nn.AvgPool1d, patched_module.torch_nn_avgpool1d]] @@ -349,6 +360,7 @@ def test_pool1d(): _assert_output_shape(data=data, module=pooler, patch_fn=patch_func, expect_exception=True, output_shape=None) +@clear_cache_before_run() def test_pool2d(): combinations = [[torch.nn.MaxPool2d, patched_module.torch_nn_maxpool2d], [torch.nn.AvgPool2d, patched_module.torch_nn_avgpool2d]] @@ -379,6 +391,7 @@ def test_pool2d(): _assert_output_shape(data=data, module=pooler, patch_fn=patch_func, expect_exception=True, output_shape=None) +@clear_cache_before_run() def test_pool3d(): combinations = [[torch.nn.MaxPool3d, patched_module.torch_nn_maxpool3d], [torch.nn.AvgPool3d, patched_module.torch_nn_avgpool3d]] @@ -410,6 +423,7 @@ def test_pool3d(): # adapative pooling is different from other pooling, so test it individually +@clear_cache_before_run() def test_adaptive_pooling_1d(): pooler = torch.nn.AdaptiveAvgPool1d(output_size=3) patch_func = patched_module.torch_nn_adapative_pooling_1d @@ -434,6 +448,7 @@ def test_adaptive_pooling_1d(): _assert_output_shape(data=data, module=pooler, patch_fn=patch_func, expect_exception=True, output_shape=None) +@clear_cache_before_run() def test_adaptive_pooling_2d(): pooler = torch.nn.AdaptiveAvgPool2d(output_size=3) patch_func = patched_module.torch_nn_adapative_pooling_2d @@ -458,6 +473,7 @@ def test_adaptive_pooling_2d(): output_shape=output.shape) +@clear_cache_before_run() def test_adaptive_pooling_3d(): pooler = torch.nn.AdaptiveAvgPool3d(output_size=3) patch_func = patched_module.torch_nn_adapative_pooling_3d diff --git a/tests/test_fx/test_tracer/test_patched_op.py b/tests/test_fx/test_tracer/test_patched_op.py index 4406f02db24b..e0c5f560c49e 100644 --- a/tests/test_fx/test_tracer/test_patched_op.py +++ b/tests/test_fx/test_tracer/test_patched_op.py @@ -1,6 +1,9 @@ +from functools import partial + import torch + from colossalai.fx.tracer.meta_patch import patched_function -from functools import partial +from colossalai.testing import clear_cache_before_run def _run(data, patch_fn): @@ -22,6 +25,7 @@ def _assert_output_shape(data, patch_fn, expect_exception, output_shape): assert output.shape == output_shape +@clear_cache_before_run() def test_repeat_interleave(): patch_fn = patched_function.torch_repeat_interleave @@ -63,6 +67,7 @@ def test_repeat_interleave(): output_shape=materialized_output.shape) +@clear_cache_before_run() def test_torch_max(): data = torch.rand(4, 3) out = torch.max(data) diff --git a/tests/test_fx/test_tracer/test_timm_model/test_timm_model.py b/tests/test_fx/test_tracer/test_timm_model/test_timm_model.py index 31baa3e89798..98433b8f7c3b 100644 --- a/tests/test_fx/test_tracer/test_timm_model/test_timm_model.py +++ b/tests/test_fx/test_tracer/test_timm_model/test_timm_model.py @@ -1,8 +1,9 @@ import pytest -import timm.models as tm import torch +from packaging import version -from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo @@ -42,12 +43,20 @@ def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None): f'{model.__class__.__name__} has inconsistent outputs, {fx_output_val} vs {non_fx_output_val}' +# FIXME(ver217): timm/models/convit.py:71: in forward +# if self.rel_indices is None or self.rel_indices.shape[1] != N: +# torch/fx/proxy.py:284: in __bool__ +# return self.tracer.to_bool(self) +# torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow +@pytest.mark.skip("convit is not supported yet") +@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12') +@clear_cache_before_run() def test_timm_models(): torch.backends.cudnn.deterministic = True sub_model_zoo = model_zoo.get_sub_registry('timm') - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in sub_model_zoo.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, _, attribute) in sub_model_zoo.items(): data = data_gen_fn() if attribute is not None and attribute.has_control_flow: meta_args = {k: v.to('meta') for k, v in data.items()} diff --git a/tests/test_fx/test_tracer/test_torchaudio_model/test_torchaudio_model.py b/tests/test_fx/test_tracer/test_torchaudio_model/test_torchaudio_model.py index bf6c7ae551ab..2b7def5bef85 100644 --- a/tests/test_fx/test_tracer/test_torchaudio_model/test_torchaudio_model.py +++ b/tests/test_fx/test_tracer/test_torchaudio_model/test_torchaudio_model.py @@ -1,20 +1,22 @@ -import re - +import pytest import torch +from packaging import version from torchaudio_utils import trace_and_compare +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo +# We cannot handle the tensors constructed with constant during forward, such as ``torch.empty(0).to(device=Proxy.device)`` +# TODO: We could handle this case by hijacking torch.Tensor.to function. +@pytest.mark.skip +@clear_cache_before_run() def test_torchaudio_models(): torch.backends.cudnn.deterministic = True sub_model_zoo = model_zoo.get_sub_registry('torchaudio') - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in sub_model_zoo.items(): - # FIXME(ver217): temporarily skip these models - if re.search(f'(conformer|emformer|tacotron|wav2vec2_base|hubert_base)', name): - continue + for name, (model_fn, data_gen_fn, output_transform_fn, _, _, attribute) in sub_model_zoo.items(): model = model_fn() trace_and_compare(model, data_gen_fn, diff --git a/tests/test_fx/test_tracer/test_torchaudio_model/torchaudio_utils.py b/tests/test_fx/test_tracer/test_torchaudio_model/torchaudio_utils.py index 18d86fc05941..239f38680cec 100644 --- a/tests/test_fx/test_tracer/test_torchaudio_model/torchaudio_utils.py +++ b/tests/test_fx/test_tracer/test_torchaudio_model/torchaudio_utils.py @@ -1,6 +1,6 @@ import torch -from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace def trace_and_compare(model, data_gen, output_transform_fn, need_meta=False, need_concrete=False): diff --git a/tests/test_fx/test_tracer/test_torchrec_model/test_deepfm_model.py b/tests/test_fx/test_tracer/test_torchrec_model/test_deepfm_model.py index 6cbca343d134..f969c8e6c3da 100644 --- a/tests/test_fx/test_tracer/test_torchrec_model/test_deepfm_model.py +++ b/tests/test_fx/test_tracer/test_torchrec_model/test_deepfm_model.py @@ -1,17 +1,13 @@ import pytest import torch -from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo BATCH = 2 SHAPE = 10 -deepfm_models = model_zoo.get_sub_registry('deepfm') -NOT_DFM = False -if not deepfm_models: - NOT_DFM = True - def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None): # trace @@ -52,11 +48,12 @@ def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None): ), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}' -@pytest.mark.skipif(NOT_DFM, reason='torchrec is not installed') -def test_torchrec_deepfm_models(deepfm_models): +@clear_cache_before_run() +def test_torchrec_deepfm_models(): + deepfm_models = model_zoo.get_sub_registry('deepfm') torch.backends.cudnn.deterministic = True - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in deepfm_models.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, attribute) in deepfm_models.items(): data = data_gen_fn() if attribute is not None and attribute.has_control_flow: meta_args = {k: v.to('meta') for k, v in data.items()} @@ -67,4 +64,4 @@ def test_torchrec_deepfm_models(deepfm_models): if __name__ == "__main__": - test_torchrec_deepfm_models(deepfm_models) + test_torchrec_deepfm_models() diff --git a/tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py b/tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py index 7aa868265f15..94fb24f33376 100644 --- a/tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py +++ b/tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py @@ -1,17 +1,13 @@ import pytest import torch -from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo BATCH = 2 SHAPE = 10 -dlrm_models = model_zoo.get_sub_registry('dlrm') -NOT_DLRM = False -if not dlrm_models: - NOT_DLRM = True - def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None): # trace @@ -52,12 +48,19 @@ def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None): ), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}' -@pytest.mark.skipif(NOT_DLRM, reason='torchrec is not installed') -def test_torchrec_dlrm_models(dlrm_models): +@clear_cache_before_run() +def test_torchrec_dlrm_models(): torch.backends.cudnn.deterministic = True + dlrm_models = model_zoo.get_sub_registry('dlrm') - for name, (model_fn, data_gen_fn, output_transform_fn, attribute) in dlrm_models.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, attribute) in dlrm_models.items(): data = data_gen_fn() + + # dlrm_interactionarch is not supported + # TODO(FrankLeeeee): support this model + if name == 'dlrm_interactionarch': + continue + if attribute is not None and attribute.has_control_flow: meta_args = {k: v.to('meta') for k, v in data.items()} else: @@ -67,4 +70,4 @@ def test_torchrec_dlrm_models(dlrm_models): if __name__ == "__main__": - test_torchrec_dlrm_models(dlrm_models) + test_torchrec_dlrm_models() diff --git a/tests/test_fx/test_tracer/test_torchvision_model/test_torchvision_model.py b/tests/test_fx/test_tracer/test_torchvision_model/test_torchvision_model.py index 455638818463..74cb753e2937 100644 --- a/tests/test_fx/test_tracer/test_torchvision_model/test_torchvision_model.py +++ b/tests/test_fx/test_tracer/test_torchvision_model/test_torchvision_model.py @@ -1,14 +1,16 @@ import torch -from colossalai.fx import symbolic_trace +from colossalai._analyzer.fx import symbolic_trace +from colossalai.testing import clear_cache_before_run from tests.kit.model_zoo import model_zoo +@clear_cache_before_run() def test_torchvision_models(): torch.backends.cudnn.deterministic = True tv_sub_registry = model_zoo.get_sub_registry('torchvision') - for name, (model_fn, data_gen_fn, output_transform_fn, model_attribute) in tv_sub_registry.items(): + for name, (model_fn, data_gen_fn, output_transform_fn, _, model_attribute) in tv_sub_registry.items(): data = data_gen_fn() if model_attribute is not None and model_attribute.has_stochastic_depth_prob: diff --git a/tests/test_kernels/test_self_attention.py b/tests/test_kernels/test_self_attention.py new file mode 100644 index 000000000000..b316404a58db --- /dev/null +++ b/tests/test_kernels/test_self_attention.py @@ -0,0 +1,136 @@ +import pytest +from packaging import version +import torch +from torch import nn +import torch.nn.functional as F + +from colossalai.kernel.triton.ops import self_attention_compute_using_triton +from colossalai.kernel.triton.qkv_matmul_kernel import qkv_gemm_4d_kernel + +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + print("please install triton from https://github.com/openai/triton") + +TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.4') + +@pytest.mark.skipif(not TRITON_CUDA_SUPPORT, reason="triton requires cuda version to be higher than 11.4") +def test_qkv_matmul(): + qkv = torch.randn((4, 24, 64*3), device="cuda", dtype=torch.float16) + scale = 1.2 + head_size = 32 + batches = qkv.shape[0] + d_model = qkv.shape[-1] // 3 + num_of_heads = d_model // head_size + + q = qkv[:, :, :d_model] + k = qkv[:, :, d_model:d_model * 2] + + q = q.view(batches, -1, num_of_heads, head_size) + k = k.view(batches, -1, num_of_heads, head_size) + q_copy = q.clone() + k_copy = k.clone() + q = torch.transpose(q, 1, 2).contiguous() + k = torch.transpose(k, 1, 2).contiguous() + k = torch.transpose(k, 2, 3).contiguous() + + torch_ouput = torch.einsum('bnij,bnjk->bnik', q, k) + torch_ouput *= 1.2 + + q, k = q_copy, k_copy + batches, M, H, K = q.shape + N = k.shape[1] + score_output = torch.empty( + (batches, H, M, N), device=q.device, dtype=q.dtype) + + grid = lambda meta: ( + batches, + H, + triton.cdiv(M, meta["BLOCK_SIZE_M"]) * + triton.cdiv(N, meta["BLOCK_SIZE_N"]), + ) + + K = q.shape[3] + qkv_gemm_4d_kernel[grid]( + q, k, score_output, + M, N, K, + q.stride(0), q.stride(2), q.stride(1), q.stride(3), + k.stride(0), k.stride(2), k.stride(3), k.stride(1), + score_output.stride(0), score_output.stride(1), score_output.stride(2), score_output.stride(3), + scale=scale, + # currently manually setting, later on we can use auto-tune config to match best setting + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=32, + BLOCK_SIZE_K=32, + GROUP_SIZE_M=8, + ) + + check = torch.allclose(torch_ouput.cpu(), score_output.cpu(), rtol=1e-3, atol=1e-5) + assert check is True, "the outputs of triton and torch are not matched" + + +def self_attention_compute_using_torch(qkv, + input_mask, + scale, + head_size + ): + + batches = qkv.shape[0] + d_model = qkv.shape[-1] // 3 + num_of_heads = d_model // head_size + + q = qkv[:, :, :d_model] + k = qkv[:, :, d_model:d_model * 2] + v = qkv[:, :, d_model * 2:] + q = q.view(batches, -1, num_of_heads, head_size) + k = k.view(batches, -1, num_of_heads, head_size) + v = v.view(batches, -1, num_of_heads, head_size) + + q = torch.transpose(q, 1, 2).contiguous() + k = torch.transpose(k, 1, 2).contiguous() + v = torch.transpose(v, 1, 2).contiguous() + + k = torch.transpose(k, -1, -2).contiguous() + + score_output = torch.einsum('bnij,bnjk->bnik', q, k) + score_output *= scale + + softmax_output = F.softmax(score_output, dim = -1) + res = torch.einsum('bnij,bnjk->bnik', softmax_output, v) + res = torch.transpose(res, 1, 2) + res = res.contiguous() + + + return res.view(batches, -1, d_model), score_output, softmax_output + +@pytest.mark.skipif(not TRITON_CUDA_SUPPORT, reason="triton requires cuda version to be higher than 11.4") +def test_self_atttention_test(): + + qkv = torch.randn((4, 24, 64*3), device="cuda", dtype=torch.float16) + data_output_torch, score_output_torch, softmax_output_torch = self_attention_compute_using_torch( + qkv.clone(), + input_mask = None, + scale = 1.2, + head_size = 32 + ) + + data_output_triton = self_attention_compute_using_triton( + qkv.clone(), + alibi=None, + head_size=32, + scale=1.2, + input_mask=None, + layer_past=None, + use_flash=False, + triangular=True) + + check = torch.allclose(data_output_triton.cpu(), data_output_torch.cpu(), rtol=1e-4, atol=1e-2) + assert check is True, "the triton output is not matched with torch output" + + +if __name__ == "__main__": + test_qkv_matmul() + test_self_atttention_test() \ No newline at end of file diff --git a/tests/test_kernels/test_softmax.py b/tests/test_kernels/test_softmax.py new file mode 100644 index 000000000000..843d811d019c --- /dev/null +++ b/tests/test_kernels/test_softmax.py @@ -0,0 +1,27 @@ +import pytest +from packaging import version +import torch +from torch import nn + +from colossalai.kernel.triton.ops import softmax + +TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.4') + +@pytest.mark.skipif(not TRITON_CUDA_SUPPORT, reason="triton requires cuda version to be higher than 11.4") +def test_softmax_op(): + data_samples = [ + torch.randn((3, 4, 5, 32), device = "cuda", dtype = torch.float32), + torch.randn((320, 320, 78), device = "cuda", dtype = torch.float32), + torch.randn((2345, 4, 5, 64), device = "cuda", dtype = torch.float16) + ] + + for data in data_samples: + module = nn.Softmax(dim = -1) + data_torch_out = module(data) + data_triton_out = softmax(data) + check = torch.allclose(data_torch_out.cpu(), data_triton_out.cpu(), rtol=1e-3, atol=1e-3) + assert check is True, "softmax outputs from triton and torch are not matched" + + +if __name__ == "__main__": + test_softmax_op() \ No newline at end of file diff --git a/tests/test_layers/test_1d/test_1d.py b/tests/test_layers/test_1d/test_1d.py index 897590f0d9c8..891512542475 100644 --- a/tests/test_layers/test_1d/test_1d.py +++ b/tests/test_layers/test_1d/test_1d.py @@ -1,18 +1,14 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from checks_1d.check_layer_1d import * from colossalai.core import global_context as gpc from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=4, mode='1d')),) @@ -40,9 +36,7 @@ def check_layer(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_1d(): - world_size = 4 - run_func = partial(check_layer, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer, 4) if __name__ == '__main__': diff --git a/tests/test_layers/test_2d/test_2d.py b/tests/test_layers/test_2d/test_2d.py index da235d0cf168..bcea5ce7b25d 100644 --- a/tests/test_layers/test_2d/test_2d.py +++ b/tests/test_layers/test_2d/test_2d.py @@ -1,22 +1,27 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp +from checks_2d.check_layer_2d import ( + check_classifier_given_embed_weight, + check_classifier_no_given_weight, + check_embed, + check_layernorm, + check_linear, + check_loss, + check_patch_embed, + check_vocab_parallel_classifier_given_embed_weight, + check_vocab_parallel_classifier_no_given_weight, + check_vocab_parallel_embed, + check_vocab_parallel_loss, +) +from checks_2d.check_operation_2d import check_AB, check_ABT, check_ATB + from colossalai.core import global_context as gpc from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use -from checks_2d.check_layer_2d import (check_classifier_given_embed_weight, check_classifier_no_given_weight, - check_embed, check_layernorm, check_linear, check_loss, check_patch_embed, - check_vocab_parallel_classifier_given_embed_weight, - check_vocab_parallel_classifier_no_given_weight, check_vocab_parallel_embed, - check_vocab_parallel_loss) -from checks_2d.check_operation_2d import check_AB, check_ABT, check_ATB +from colossalai.testing import rerun_if_address_is_in_use, spawn CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=4, mode='2d')),) @@ -57,9 +62,7 @@ def check_layer_and_operation(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_2d(): - world_size = 4 - run_func = partial(check_layer_and_operation, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer_and_operation, 4) if __name__ == '__main__': diff --git a/tests/test_layers/test_2p5d/test_2p5d.py b/tests/test_layers/test_2p5d/test_2p5d.py index 365e2d934df8..373d834d0032 100644 --- a/tests/test_layers/test_2p5d/test_2p5d.py +++ b/tests/test_layers/test_2p5d/test_2p5d.py @@ -1,15 +1,12 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp +from checks_2p5d.check_layer_2p5d import * +from checks_2p5d.check_operation_2p5d import check_AB, check_ABT, check_ATB + from colossalai.core import global_context as gpc from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use -from checks_2p5d.check_layer_2p5d import * -from checks_2p5d.check_operation_2p5d import check_AB, check_ABT, check_ATB +from colossalai.testing import rerun_if_address_is_in_use, spawn CONFIG = dict(parallel=dict( pipeline=dict(size=1), @@ -53,9 +50,7 @@ def check_layer_and_operation(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_2p5d(): - world_size = 4 - run_func = partial(check_layer_and_operation, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer_and_operation, 4) if __name__ == '__main__': diff --git a/tests/test_layers/test_3d/test_3d.py b/tests/test_layers/test_3d/test_3d.py index 29a8b3aea239..fde71a4a0d26 100644 --- a/tests/test_layers/test_3d/test_3d.py +++ b/tests/test_layers/test_3d/test_3d.py @@ -1,19 +1,24 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp +from checks_3d.check_layer_3d import ( + check_classifier_no_given_weight, + check_embed, + check_layernorm, + check_linear, + check_loss, + check_patch_embed, + check_vocab_parallel_classifier_given_embed_weight, + check_vocab_parallel_classifier_no_given_weight, + check_vocab_parallel_embed, + check_vocab_parallel_loss, +) + from colossalai.core import global_context as gpc from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus -from checks_3d.check_layer_3d import (check_classifier_no_given_weight, check_embed, check_layernorm, check_linear, - check_loss, check_patch_embed, check_vocab_parallel_classifier_given_embed_weight, - check_vocab_parallel_classifier_no_given_weight, check_vocab_parallel_embed, - check_vocab_parallel_loss) +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn CONFIG = dict( parallel=dict( @@ -52,9 +57,7 @@ def check_layer_and_operation(rank, world_size, port): @skip_if_not_enough_gpus(min_gpus=8) @rerun_if_address_is_in_use() def test_3d(): - world_size = 8 - run_func = partial(check_layer_and_operation, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_layer_and_operation, 8) if __name__ == '__main__': diff --git a/tests/test_layers/test_cache_embedding.py b/tests/test_layers/test_cache_embedding.py index cff9072c7a06..22d4f02a48d7 100644 --- a/tests/test_layers/test_cache_embedding.py +++ b/tests/test_layers/test_cache_embedding.py @@ -1,20 +1,21 @@ -import pytest -from functools import partial - -import numpy as np import random +from typing import List +import numpy as np +import pytest import torch -import torch.multiprocessing as mp import colossalai -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.tensor import ColoParameter, ProcessGroup, ShardSpec, ComputePattern, ComputeSpec, \ - ColoTensor, ColoTensorSpec -from colossalai.nn.parallel.layers import CachedParamMgr, CachedEmbeddingBag, ParallelCachedEmbeddingBag, EvictionStrategy, \ - ParallelCachedEmbeddingBagTablewise, TablewiseEmbeddingBagConfig -from typing import List +from colossalai.nn.parallel.layers import ( + CachedEmbeddingBag, + CachedParamMgr, + EvictionStrategy, + ParallelCachedEmbeddingBag, + ParallelCachedEmbeddingBagTablewise, + TablewiseEmbeddingBagConfig, +) +from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec +from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn NUM_EMBED, EMBED_DIM = 10, 8 BATCH_SIZE = 8 @@ -44,6 +45,7 @@ def synthesize_1d_sparse_feature( @pytest.mark.skip +@clear_cache_before_run() def test_cachemgr(): model = torch.nn.EmbeddingBag(10000, 128) # 10 chunks, 5 in cuda @@ -72,6 +74,7 @@ def test_cachemgr(): assert mgr.cuda_available_chunk_num == 5 +@clear_cache_before_run() def test_reorder_with_freq(): num_embed = 100 chunk_size = 1 @@ -102,7 +105,8 @@ def test_reorder_with_freq(): f"offset in chunk: {offset_in_chunk}, mgr: {mgr_offsets}" -@pytest.mark.parametrize('use_LFU', [True, False]) +@clear_cache_before_run() +@parameterize('use_LFU', [True, False]) def test_freq_aware_embed(use_LFU: bool): device = torch.device('cuda', 0) evict_strategy = EvictionStrategy.LFU if use_LFU else EvictionStrategy.DATASET @@ -148,7 +152,8 @@ def test_freq_aware_embed(use_LFU: bool): f"model weight: {model_weight[10:18, :8]}, reference: {ref_weight[10:18, :8]}" -@pytest.mark.parametrize('init_freq', [True, False]) +@clear_cache_before_run() +@parameterize('init_freq', [True, False]) def test_lfu_strategy(init_freq: bool): # minimal test to check behavior Bag = CachedEmbeddingBag(5, @@ -248,7 +253,7 @@ def run_parallel_freq_aware_embed_tablewise(rank, world_size): input0 [1,2,3] [6,7] [] input1 [] [9] [13,15] input2 [1,5] [6,8] [11] - ↑ ↑ ↑ + ↑ ↑ ↑ rank 0 rank 0 rank 1 in KJT format ''' @@ -363,8 +368,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_parallel_freq_aware_embed(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_layers/test_sequence/test_sequence.py b/tests/test_layers/test_sequence/test_sequence.py index 3862c4ccd439..60f2d55f43af 100644 --- a/tests/test_layers/test_sequence/test_sequence.py +++ b/tests/test_layers/test_sequence/test_sequence.py @@ -1,14 +1,11 @@ -import colossalai -import colossalai.nn as col_nn +import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -import pytest -from colossalai.core import global_context as gpc +import colossalai from colossalai.context import ParallelMode -from colossalai.testing import rerun_if_address_is_in_use -from functools import partial +from colossalai.core import global_context as gpc +from colossalai.testing import rerun_if_address_is_in_use, spawn CONFIG = dict(parallel=dict(tensor=dict(size=4, mode='sequence'))) @@ -48,7 +45,7 @@ def check_ring_qk(rank, world_size): ring_qk = colossalai.nn.layer.parallel_sequence.RingQK.apply sub_a = ring_qk(sub_q, sub_k, batch_size, num_heads, sub_seq_length) - # check master and distributed attetion scores + # check master and distributed attention scores sub_master_a = a[:, rank * sub_seq_length:(rank + 1) * sub_seq_length] assert torch.allclose(sub_a, sub_master_a, rtol=1e-5, atol=1e-2) @@ -121,8 +118,8 @@ def check_ring_av(rank, world_size): 'attention output cannot match' -def run_test(rank, world_size): - colossalai.launch(rank=rank, world_size=world_size, config=CONFIG, host='localhost', port=29500) +def run_test(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, config=CONFIG, host='localhost', port=port) # check_ring_qk(rank, world_size) check_ring_av(rank, world_size) @@ -134,9 +131,7 @@ def run_test(rank, world_size): @pytest.mark.dist @rerun_if_address_is_in_use() def test_sequence(): - world_size = 4 - run_func = partial(run_test, world_size=world_size) - mp.spawn(run_func, nprocs=world_size) + spawn(run_test, 4) if __name__ == '__main__': diff --git a/tests/test_lazy/lazy_init_utils.py b/tests/test_lazy/lazy_init_utils.py new file mode 100644 index 000000000000..9d9e9a3a5c76 --- /dev/null +++ b/tests/test_lazy/lazy_init_utils.py @@ -0,0 +1,105 @@ +import random +from copy import deepcopy +from typing import Any, Callable, Optional, Tuple + +import numpy as np +import torch +from packaging import version + +from colossalai.device.device_mesh import DeviceMesh +from colossalai.lazy.lazy_init import LazyInitContext, LazyTensor, _MyTensor +from colossalai.tensor.d_tensor import to_global +from colossalai.tensor.d_tensor.layout import Layout +from tests.kit.model_zoo.registry import ModelAttribute + +SUPPORT_LAZY = version.parse(torch.__version__) >= version.parse('1.12.0') + +# model_fn, data_gen_fn, output_transform_fn, model_attr +TestingEntry = Tuple[Callable[[], torch.nn.Module], Callable[[], dict], Callable[[], dict], Optional[ModelAttribute]] + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def assert_model_equal(m1: torch.nn.Module, m2: torch.nn.Module) -> None: + s1 = m1.state_dict() + s2 = m2.state_dict() + + assert len(s1) == len(s2), f'len {len(s1)} vs {len(s2)}' + + for (n1, t1), (n2, t2) in zip(s1.items(), s2.items()): + assert n1 == n2 + assert torch.equal(t1, t2), f'{n1} {t1} vs {t2}' + + for p1, p2 in zip(m1.parameters(), m2.parameters()): + assert p1.requires_grad == p2.requires_grad + + +def assert_forward_equal(m1: torch.nn.Module, m2: torch.nn.Module, data_gen_fn: Callable[[], dict], + output_transform_fn: Callable[[Any], dict]) -> None: + data = data_gen_fn() + + m1.eval() + m2.eval() + # run forward + with torch.no_grad(): + outputs1 = m1(**data) + outputs2 = m2(**data) + + # compare output + transformed_out1 = output_transform_fn(outputs1) + transformed_out2 = output_transform_fn(outputs2) + + assert len(transformed_out1) == len(transformed_out2) + + for key, out1 in transformed_out1.items(): + out2 = transformed_out2[key] + assert torch.allclose(out1, out2, atol=1e-5), \ + f'{m1.__class__.__name__} has inconsistent outputs, {out1} vs {out2}' + + +def check_lazy_init(entry: TestingEntry, + seed: int = 42, + verbose: bool = False, + check_forward: bool = False, + default_device: str = 'cpu') -> None: + model_fn, data_gen_fn, output_transform_fn, _, model_attr = entry + _MyTensor._pre_op_fn = lambda *args: set_seed(seed) + LazyTensor._pre_op_fn = lambda *args: set_seed(seed) + ctx = LazyInitContext(tensor_cls=_MyTensor, default_device=default_device) + with ctx: + model = model_fn() + ctx = LazyInitContext(default_device=default_device) + with ctx: + deferred_model = model_fn() + copied_deferred_model = deepcopy(deferred_model) + deferred_model = ctx.materialize(deferred_model, verbose=verbose) + copied_deferred_model = ctx.materialize(copied_deferred_model, verbose=verbose) + assert_model_equal(model, deferred_model) + assert_model_equal(deferred_model, copied_deferred_model) + if check_forward: + assert_forward_equal(model, deferred_model, data_gen_fn, output_transform_fn) + assert_forward_equal(deferred_model, copied_deferred_model, data_gen_fn, output_transform_fn) + if verbose: + print(f'{model.__class__.__name__} pass') + + +def assert_dist_model_equal(model: torch.nn.Module, distributed_model: torch.nn.Module, device_mesh: DeviceMesh, + sharding_spec_dict: dict) -> None: + state = model.state_dict() + distributed_state = distributed_model.state_dict() + + assert len(state) == len(distributed_state), f'len {len(state)} vs {len(distributed_state)}' + + for (n1, t1), (n2, t2) in zip(state.items(), distributed_state.items()): + assert n1 == n2 + t1 = t1.cuda() + t2 = t2.cuda() + if n2 in sharding_spec_dict: + layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_dict[n2], global_shape=t1.shape) + t2.dist_layout = layout + t2 = to_global(t2) + assert torch.equal(t1, t2), f'{n1} {t1} vs {t2}' diff --git a/tests/test_lazy/test_distribute.py b/tests/test_lazy/test_distribute.py new file mode 100644 index 000000000000..622d9deb601d --- /dev/null +++ b/tests/test_lazy/test_distribute.py @@ -0,0 +1,102 @@ +from typing import Optional + +import pytest +import torch +import torch.nn as nn + +import colossalai +from colossalai.device.device_mesh import DeviceMesh +from colossalai.tensor.d_tensor.layout import Layout +from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils.common import print_rank_0 + +try: + from colossalai.lazy.lazy_init import LazyInitContext, LazyTensor, _MyTensor +except: + pass +from lazy_init_utils import SUPPORT_LAZY, assert_dist_model_equal, set_seed + +from tests.kit.model_zoo import model_zoo + + +def find_shard_dim(shape: torch.Size) -> Optional[int]: + for dim, size in enumerate(shape): + if size % 2 == 0: + return dim + + +def make_sharding_spec(original_tensor: torch.Tensor) -> Layout: + shard_dim = find_shard_dim(original_tensor.shape) + dim_partition_dict = {shard_dim: [0]} if shard_dim is not None else {} + target_sharding_spec = ShardingSpec(dim_size=original_tensor.dim(), dim_partition_dict=dim_partition_dict) + return target_sharding_spec + + +def _get_current_name(prefix: str, name: str) -> str: + return f'{prefix}.{name}'.lstrip('.') + + +def generate_sharding_spec_dict(model: nn.Module) -> dict: + sharding_spec_dict = {} + + @torch.no_grad() + def generate_recursively(module: nn.Module, prefix: str = ''): + # recursively initialize the module + for name, mod in module.named_children(): + generate_recursively(mod, prefix=_get_current_name(prefix, name)) + + # initialize tensors directly attached to the current module + for name, param in module.named_parameters(recurse=False): + if isinstance(param, LazyTensor): + sharding_spec = make_sharding_spec(param) + sharding_spec_dict[_get_current_name(prefix, name)] = sharding_spec + + for name, buf in module.named_buffers(recurse=False): + if isinstance(buf, LazyTensor): + sharding_spec = make_sharding_spec(buf) + sharding_spec_dict[_get_current_name(prefix, name)] = sharding_spec + + generate_recursively(model) + + return sharding_spec_dict + + +@parameterize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm']) +def run_dist_lazy_init(subset, seed: int = 42): + sub_model_zoo = model_zoo.get_sub_registry(subset) + device_mesh = DeviceMesh(torch.Tensor([0, 1, 2, 3]), (2, 2), init_process_group=True) + _MyTensor._pre_op_fn = lambda *args: set_seed(seed) + LazyTensor._pre_op_fn = lambda *args: set_seed(seed) + + for name, entry in sub_model_zoo.items(): + # TODO(ver217): lazy init does not support weight norm, skip these models + if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base') or name.startswith('transformers_llama'): + continue + print_rank_0(name) + model_fn, data_gen_fn, output_transform_fn, _, model_attr = entry + ctx = LazyInitContext(tensor_cls=_MyTensor) + with ctx: + model = model_fn() + ctx = LazyInitContext() + with ctx: + deferred_model = model_fn() + sharding_spec_dict = generate_sharding_spec_dict(deferred_model) + ctx.distribute(deferred_model, device_mesh, sharding_spec_dict, verbose=True) + assert_dist_model_equal(model, deferred_model, device_mesh, sharding_spec_dict) + + +def run_dist(rank, world_size, port) -> None: + colossalai.launch({}, rank=rank, world_size=world_size, host='localhost', port=port) + run_dist_lazy_init() + + +@pytest.mark.skipif(not SUPPORT_LAZY, reason='torch version should be >= 1.12.0') +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_dist_lazy_init(): + spawn(run_dist, 4) + + +if __name__ == '__main__': + test_dist_lazy_init() diff --git a/tests/test_utils/test_lazy_init/test_models.py b/tests/test_lazy/test_models.py similarity index 58% rename from tests/test_utils/test_lazy_init/test_models.py rename to tests/test_lazy/test_models.py index 9faddecbaca4..e37184125d21 100644 --- a/tests/test_utils/test_lazy_init/test_models.py +++ b/tests/test_lazy/test_models.py @@ -1,22 +1,19 @@ import pytest +from lazy_init_utils import SUPPORT_LAZY, check_lazy_init from tests.kit.model_zoo import model_zoo -# FIXME(ver217): uncomment this line -# from utils import check_lazy_init - -# FIXME(ver217): temporarily skip this test since torch 1.11 does not fully support meta tensor -@pytest.mark.skip +@pytest.mark.skipif(not SUPPORT_LAZY, reason='requires torch >= 1.12.0') @pytest.mark.parametrize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm']) -def test_torchvision_models_lazy_init(subset): +@pytest.mark.parametrize('default_device', ['cpu', 'cuda']) +def test_torchvision_models_lazy_init(subset, default_device): sub_model_zoo = model_zoo.get_sub_registry(subset) for name, entry in sub_model_zoo.items(): # TODO(ver217): lazy init does not support weight norm, skip these models - if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base'): + if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base') or name.startswith('transformers_llama'): continue - # FIXME(ver217): uncomment this line - # check_lazy_init(entry, verbose=True) + check_lazy_init(entry, verbose=True, default_device=default_device) if __name__ == '__main__': diff --git a/tests/test_moe/test_grad_handler.py b/tests/test_moe/test_grad_handler.py index e7b9a55277c6..e7002a75f3f7 100644 --- a/tests/test_moe/test_grad_handler.py +++ b/tests/test_moe/test_grad_handler.py @@ -1,16 +1,15 @@ -from functools import partial import pytest import torch -import torch.nn as nn -import torch.multiprocessing as mp import torch.distributed as dist +import torch.nn as nn + import colossalai -from colossalai.utils import free_port, get_current_device -from colossalai.nn.layer.moe import Top1Router, UniformNoiseGenerator, MoeLayer, Experts from colossalai.context.moe_context import MOE_CONTEXT -from colossalai.utils.moe import sync_moe_model_param from colossalai.engine.gradient_handler import MoeGradientHandler -from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use +from colossalai.nn.layer.moe import Experts, MoeLayer, Top1Router, UniformNoiseGenerator +from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.utils.moe import sync_moe_model_param BATCH_SIZE = 4 DIM = 16 @@ -65,9 +64,7 @@ def run_test(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_grad_handler(): - world_size = 4 - run_func = partial(run_test, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_test, 4) if __name__ == '__main__': diff --git a/tests/test_moe/test_kernel.py b/tests/test_moe/test_kernel.py index 62f9241642b9..39603c158731 100644 --- a/tests/test_moe/test_kernel.py +++ b/tests/test_moe/test_kernel.py @@ -1,15 +1,14 @@ -from functools import partial import pytest import torch import torch.nn as nn -import torch.multiprocessing as mp + import colossalai from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import free_port, get_current_device -from colossalai.nn.layer.moe import Top1Router, Top2Router, MoeLayer, Experts from colossalai.context.moe_context import MOE_CONTEXT -from colossalai.testing import rerun_if_address_is_in_use +from colossalai.core import global_context as gpc +from colossalai.nn.layer.moe import Experts, MoeLayer, Top1Router, Top2Router +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device BATCH_SIZE = 16 NUM_EXPERTS = 4 @@ -42,7 +41,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f if data_type == torch.float16: layer = layer.half() - # use matrix multiplication instead of COL_MOE_KERNL in MOE dispatch and combine + # use matrix multiplication instead of COL_MOE_KERNEL in MOE dispatch and combine layer.use_kernel = False old_out, _ = layer(tokens) ech = old_out.shape @@ -58,7 +57,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f layer.gate_weight.grad.zero_() layer.use_kernel = True - new_out, _ = layer(tokens) # get ouputs through colossal kernel + new_out, _ = layer(tokens) # get outputs through colossal kernel if data_type == torch.float32: check_equal(old_out, new_out) @@ -90,15 +89,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f @pytest.mark.parametrize("router", [Top1Router, Top2Router]) @rerun_if_address_is_in_use() def test_moe_kernel(rs, hidden_size, data_type, router): - world_size = 4 - run_func = partial(run_routing, - world_size=world_size, - port=free_port(), - rs=rs, - hidden_size=hidden_size, - data_type=data_type, - router=router) - mp.spawn(run_func, nprocs=world_size) + spawn(run_routing, 4, rs=rs, hidden_size=hidden_size, data_type=data_type, router=router) if __name__ == '__main__': diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py new file mode 100644 index 000000000000..8a0283ba71fc --- /dev/null +++ b/tests/test_moe/test_moe_checkpoint.py @@ -0,0 +1,50 @@ +import os + +import pytest +import torch +import torch.distributed as dist + +import colossalai +from colossalai.context import MOE_CONTEXT +from colossalai.nn.layer.moe import load_moe_model, save_moe_model +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from tests.test_moe.test_moe_zero_init import MoeModel +from tests.test_zero.test_legacy.common import CONFIG + + +def exam_moe_checkpoint(): + with ColoInitContext(device=get_current_device()): + model = MoeModel(checkpoint=True) + save_moe_model(model, 'temp_path.pth') + + with ColoInitContext(device=get_current_device()): + other_model = MoeModel(checkpoint=True) + load_moe_model(other_model, 'temp_path.pth') + + state_0 = model.state_dict() + state_1 = other_model.state_dict() + for k, v in state_0.items(): + u = state_1.get(k) + assert torch.equal(u.data, v.data) + + if dist.get_rank() == 0: + os.remove('temp_path.pth') + + +def _run_dist(rank, world_size, port): + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + MOE_CONTEXT.setup(seed=42) + exam_moe_checkpoint() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [2, 4]) +@rerun_if_address_is_in_use() +def test_moe_checkpoint(world_size): + spawn(_run_dist) + + +if __name__ == '__main__': + test_moe_checkpoint(world_size=4) diff --git a/tests/test_moe/test_moe_colo_init.py b/tests/test_moe/test_moe_colo_init.py index ae0c1390c129..555338fcf9fc 100644 --- a/tests/test_moe/test_moe_colo_init.py +++ b/tests/test_moe/test_moe_colo_init.py @@ -1,63 +1,56 @@ -from functools import partial - -import colossalai -import pytest -import torch -import torch.multiprocessing as mp -import torch.distributed as dist -from colossalai.testing import parameterize -from colossalai.utils import free_port -from colossalai.context import MOE_CONTEXT -from colossalai.tensor import ColoParameter -from colossalai.utils.model.colo_init_context import ColoInitContext - -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import get_current_device - -from tests.test_zero.common import CONFIG -from tests.test_moe.test_moe_zero_init import MoeModel -from tests.test_tensor.common_utils import debug_print - - -@parameterize("init_device_type", ['cpu', 'cuda']) -def exam_moe_colo_init(init_device_type): - world_size = dist.get_world_size() - - if init_device_type == 'cuda': - init_device = get_current_device() - elif init_device_type == 'cpu': - init_device = torch.device("cpu") - else: - raise NotImplementedError("Unknown device found.") - - with ColoInitContext(device=init_device): - model = MoeModel(checkpoint=True) - - for name, param in model.named_parameters(): - assert isinstance(param, ColoParameter), "parameter `{}` has an init problem".format(name) - - if hasattr(param, "moe_info"): - param.set_process_group(param.moe_info.pg) - - if hasattr(param, "moe_info"): - assert param.process_group.dp_world_size() == param.moe_info.dp_size - else: - assert param.process_group.dp_world_size() == world_size - - -def _run_dist(rank, world_size, port): - colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - MOE_CONTEXT.setup(seed=42) - exam_moe_colo_init() - - -@pytest.mark.dist -@pytest.mark.parametrize("world_size", [4]) -@rerun_if_address_is_in_use() -def test_moe_colo_init(world_size): - run_func = partial(_run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_moe_colo_init(world_size=4) +import pytest +import torch +import torch.distributed as dist + +import colossalai +from colossalai.context import MOE_CONTEXT +from colossalai.tensor import ColoParameter +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from tests.test_moe.test_moe_zero_init import MoeModel +from tests.test_tensor.common_utils import debug_print +from tests.test_zero.test_legacy.common import CONFIG + + +@parameterize("init_device_type", ['cpu', 'cuda']) +def exam_moe_colo_init(init_device_type): + world_size = dist.get_world_size() + + if init_device_type == 'cuda': + init_device = get_current_device() + elif init_device_type == 'cpu': + init_device = torch.device("cpu") + else: + raise NotImplementedError("Unknown device found.") + + with ColoInitContext(device=init_device): + model = MoeModel(checkpoint=True) + + for name, param in model.named_parameters(): + assert isinstance(param, ColoParameter), "parameter `{}` has an init problem".format(name) + + if hasattr(param, "moe_info"): + param.set_process_group(param.moe_info.pg) + + if hasattr(param, "moe_info"): + assert param.process_group.dp_world_size() == param.moe_info.dp_size + else: + assert param.process_group.dp_world_size() == world_size + + +def _run_dist(rank, world_size, port): + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + MOE_CONTEXT.setup(seed=42) + exam_moe_colo_init() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_moe_colo_init(world_size): + spawn(_run_dist, world_size) + + +if __name__ == '__main__': + test_moe_colo_init(world_size=4) diff --git a/tests/test_moe/test_moe_group.py b/tests/test_moe/test_moe_group.py index 3126f59e246e..6dc3f5f18b6d 100644 --- a/tests/test_moe/test_moe_group.py +++ b/tests/test_moe/test_moe_group.py @@ -1,21 +1,20 @@ -from functools import partial import pytest -import torch.nn as nn -import torch.multiprocessing as mp import torch.distributed as dist +import torch.nn as nn + import colossalai -from colossalai.utils import free_port, get_current_device -from colossalai.nn.layer.moe import Experts from colossalai.context.moe_context import MOE_CONTEXT +from colossalai.nn.layer.moe import Experts +from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device from colossalai.utils.moe import sync_moe_model_param -from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use D_MODEL = 4 D_FF = 8 CONFIG = dict() -def run_test(rank, port): +def run_test(rank, world_size, port): world_size = 4 colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') expert_module = nn.Linear @@ -62,9 +61,7 @@ def run_test(rank, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_moe_initialization(): - world_size = 4 - run_func = partial(run_test, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_test, 4) if __name__ == '__main__': diff --git a/tests/test_moe/test_moe_zero_init.py b/tests/test_moe/test_moe_zero_init.py index 04dc9c514dd0..79722f9f4056 100644 --- a/tests/test_moe/test_moe_zero_init.py +++ b/tests/test_moe/test_moe_zero_init.py @@ -1,114 +1,108 @@ -from functools import partial - -import colossalai -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn -from colossalai.nn import CheckpointModule -from colossalai.logging import get_dist_logger -from colossalai.testing import parameterize -from colossalai.utils import free_port -from colossalai.context import MOE_CONTEXT -from colossalai.nn.layer import MoeModule -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy) - -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import get_current_device -from tests.test_zero.common import CONFIG - - -class MoeModel(nn.Module): - - def __init__(self, checkpoint: bool = False): - - class TestSubModule(CheckpointModule): - - def __init__(self): - super().__init__(checkpoint) - expert_cls = nn.Linear - expert_args_dict = dict(in_features=16, out_features=16) - self.moe = MoeModule(dim_model=16, - num_experts=8, - use_residual=True, - expert_cls=expert_cls, - **expert_args_dict) - self.proj = nn.Linear(16, 4) - - def _forward(self, x): - x, y = self.moe(x) - x = self.proj(x) - return x, y - - super().__init__() - self.test_embed = nn.Linear(4, 16) - self.test_transform = TestSubModule() - - def forward(self, x): - MOE_CONTEXT.reset_loss() - - x = self.test_embed(x) - x, y = self.test_transform(x) - - MOE_CONTEXT.add_loss(y) - return x - - -@parameterize("init_device_type", ['cpu', 'cuda']) -@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy]) -def run_moe_zero_init(init_device_type, shard_strategy_class): - logger = get_dist_logger("test_moe_zero_init") - - if init_device_type == 'cuda': - init_device = get_current_device() - elif init_device_type == 'cpu': - init_device = torch.device("cpu") - else: - raise NotImplementedError("Unknown device found.") - - model_numel_tensor = torch.zeros(1, dtype=torch.int) - with ZeroInitContext(target_device=init_device, - shard_strategy=shard_strategy_class(), - shard_param=True, - model_numel_tensor=model_numel_tensor): - model = MoeModel(checkpoint=True) - - for name, param in model.named_parameters(): - assert hasattr(param, 'colo_attr') - - # the parameters in moe experts and its gate should not be sharded - if ('experts' in name) or ('gate' in name) or ('residual_combine' in name): - assert not param.colo_attr.sharded_data_tensor.is_sharded, "`{}` parameter has problem".format(name) - else: - assert param.colo_attr.sharded_data_tensor.is_sharded - - # the parameters in moe experts is not replicated - if 'experts' in name: - assert not param.colo_attr.is_replicated - else: - assert param.colo_attr.is_replicated - - if param.colo_attr.param_is_sharded: - assert param.colo_attr.data_payload.device.type == init_device.type, \ - f'{param.colo_attr.data_payload.device.type} vs. {init_device.type}' - else: - assert param.colo_attr.data_payload.device.type == 'cuda' - - -def _run_dist(rank, world_size, port): - colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - MOE_CONTEXT.setup(seed=42) - run_moe_zero_init() - - -@pytest.mark.dist -@pytest.mark.parametrize("world_size", [2, 4]) -@rerun_if_address_is_in_use() -def test_moe_zero_init(world_size): - run_func = partial(_run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_moe_zero_init(world_size=2) +import pytest +import torch +import torch.nn as nn + +import colossalai +from colossalai.context import MOE_CONTEXT +from colossalai.logging import get_dist_logger +from colossalai.nn import CheckpointModule +from colossalai.nn.layer import MoeModule +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from tests.test_zero.test_legacy.common import CONFIG + + +class MoeModel(nn.Module): + + def __init__(self, checkpoint: bool = False): + + class TestSubModule(CheckpointModule): + + def __init__(self): + super().__init__(checkpoint) + expert_cls = nn.Linear + expert_args_dict = dict(in_features=16, out_features=16) + self.moe = MoeModule(dim_model=16, + num_experts=8, + use_residual=True, + expert_cls=expert_cls, + **expert_args_dict) + self.proj = nn.Linear(16, 4) + + def _forward(self, x): + x, y = self.moe(x) + x = self.proj(x) + return x, y + + super().__init__() + self.test_embed = nn.Linear(4, 16) + self.test_transform = TestSubModule() + + def forward(self, x): + MOE_CONTEXT.reset_loss() + + x = self.test_embed(x) + x, y = self.test_transform(x) + + MOE_CONTEXT.add_loss(y) + return x + + +@parameterize("init_device_type", ['cpu', 'cuda']) +@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy]) +def run_moe_zero_init(init_device_type, shard_strategy_class): + logger = get_dist_logger("test_moe_zero_init") + + if init_device_type == 'cuda': + init_device = get_current_device() + elif init_device_type == 'cpu': + init_device = torch.device("cpu") + else: + raise NotImplementedError("Unknown device found.") + + model_numel_tensor = torch.zeros(1, dtype=torch.int) + with ZeroInitContext(target_device=init_device, + shard_strategy=shard_strategy_class(), + shard_param=True, + model_numel_tensor=model_numel_tensor): + model = MoeModel(checkpoint=True) + + for name, param in model.named_parameters(): + assert hasattr(param, 'colo_attr') + + # the parameters in moe experts and its gate should not be sharded + if ('experts' in name) or ('gate' in name) or ('residual_combine' in name): + assert not param.colo_attr.sharded_data_tensor.is_sharded, "`{}` parameter has problem".format(name) + else: + assert param.colo_attr.sharded_data_tensor.is_sharded + + # the parameters in moe experts is not replicated + if 'experts' in name: + assert not param.colo_attr.is_replicated + else: + assert param.colo_attr.is_replicated + + if param.colo_attr.param_is_sharded: + assert param.colo_attr.data_payload.device.type == init_device.type, \ + f'{param.colo_attr.data_payload.device.type} vs. {init_device.type}' + else: + assert param.colo_attr.data_payload.device.type == 'cuda' + + +def _run_dist(rank, world_size, port): + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + MOE_CONTEXT.setup(seed=42) + run_moe_zero_init() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [2, 4]) +@rerun_if_address_is_in_use() +def test_moe_zero_init(world_size): + spawn(_run_dist, world_size) + + +if __name__ == '__main__': + test_moe_zero_init(world_size=2) diff --git a/tests/test_moe/test_moe_zero_model.py b/tests/test_moe/test_moe_zero_model.py index d608ebf0718e..ec37967f18c5 100644 --- a/tests/test_moe/test_moe_zero_model.py +++ b/tests/test_moe/test_moe_zero_model.py @@ -1,23 +1,19 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.context import MOE_CONTEXT from colossalai.engine.gradient_handler import MoeGradientHandler from colossalai.nn import MoeLoss -from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16 -from colossalai.zero.sharded_model.utils import col_model_deepcopy +from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp16 +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_moe.test_moe_zero_init import MoeModel -from tests.test_zero.common import CONFIG, check_grads_padding, run_fwd_bwd +from tests.test_zero.test_legacy.common import CONFIG, check_grads_padding, run_fwd_bwd @parameterize("enable_autocast", [False]) @@ -67,8 +63,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [2]) @rerun_if_address_is_in_use() def test_moe_zero_model(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_moe/test_moe_zero_optim.py b/tests/test_moe/test_moe_zero_optim.py index 9d9a7bd17390..efc6e9ddae27 100644 --- a/tests/test_moe/test_moe_zero_optim.py +++ b/tests/test_moe/test_moe_zero_optim.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.amp import convert_to_apex_amp @@ -10,17 +7,17 @@ from colossalai.engine.gradient_handler import MoeGradientHandler from colossalai.nn import MoeLoss from colossalai.nn.optimizer import CPUAdam -from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port, get_current_device -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model.utils import col_model_deepcopy -from colossalai.zero.sharded_optim import ShardedOptimizerV2 -from colossalai.zero.sharded_optim._utils import has_inf_or_nan +from colossalai.testing import assert_equal_in_group, parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy +from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 +from colossalai.zero.low_level._utils import has_inf_or_nan from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_moe.test_moe_zero_init import MoeModel -from tests.test_zero.common import CONFIG, check_sharded_model_params +from tests.test_zero.test_legacy.common import CONFIG, check_sharded_model_params def _run_step(model, optimizer, data, label, criterion, grad_handler): @@ -116,8 +113,7 @@ def _run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [2]) @rerun_if_address_is_in_use() def test_moe_zero_optim(world_size): - run_func = partial(_run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(_run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ops/test_addmm_tp.py b/tests/test_ops/test_addmm_tp.py index 5182868b5bbd..ecd3721b902e 100644 --- a/tests/test_ops/test_addmm_tp.py +++ b/tests/test_ops/test_addmm_tp.py @@ -1,14 +1,11 @@ -import colossalai -import torch import pytest +import torch import torch.nn as nn -import torch.multiprocessing as mp -from colossalai.tensor import ColoTensor, ProcessGroup -from colossalai.tensor import ColoTensorSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from functools import partial -from tests.test_tensor.common_utils import tensor_shard_equal, tensor_equal, split_param_row_tp1d, split_param_col_tp1d + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.test_tensor.common_utils import split_param_col_tp1d, split_param_row_tp1d, tensor_equal, tensor_shard_equal class Conv1D(nn.Module): @@ -69,8 +66,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_addmm_1d(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ops/test_embedding_bag_tp.py b/tests/test_ops/test_embedding_bag_tp.py index c7a1604e5455..d3d3dcf7e2c9 100644 --- a/tests/test_ops/test_embedding_bag_tp.py +++ b/tests/test_ops/test_embedding_bag_tp.py @@ -1,14 +1,11 @@ +import pytest +import torch from torch.nn import functional as F -from functools import partial import colossalai -import pytest -import torch -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port from colossalai.tensor import ColoParameter, ColoTensorSpec, ProcessGroup -from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.test_tensor.common_utils import split_param_col_tp1d, tensor_equal, tensor_shard_equal def run_with_spec(spec_init_func): @@ -39,8 +36,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_embedding_bag_1d(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ops/test_embedding_tp.py b/tests/test_ops/test_embedding_tp.py index 541dc5c09324..c0b376e2c92a 100644 --- a/tests/test_ops/test_embedding_tp.py +++ b/tests/test_ops/test_embedding_tp.py @@ -1,14 +1,11 @@ +import pytest +import torch from torch.nn import functional as F -from functools import partial import colossalai -import pytest -import torch -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor -from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.test_tensor.common_utils import split_param_col_tp1d, split_param_row_tp1d, tensor_equal, tensor_shard_equal def run_with_spec(spec_init_func, pg: ProcessGroup): @@ -40,8 +37,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_embedding_1d(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ops/test_linear_tp.py b/tests/test_ops/test_linear_tp.py index 603e98564de8..c88adfdd9a77 100644 --- a/tests/test_ops/test_linear_tp.py +++ b/tests/test_ops/test_linear_tp.py @@ -1,14 +1,11 @@ -from functools import partial - -import colossalai import pytest import torch -import torch.multiprocessing as mp import torch.nn.functional as F -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor -from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup +from colossalai.testing import rerun_if_address_is_in_use, spawn +from tests.test_tensor.common_utils import split_param_col_tp1d, split_param_row_tp1d, tensor_equal, tensor_shard_equal def run_with_spec(spec_init_func, split_bias): @@ -44,8 +41,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_linear_1d(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_ops/test_loss_func.py b/tests/test_ops/test_loss_func.py index 9210242a0a9f..fc55c7f77254 100644 --- a/tests/test_ops/test_loss_func.py +++ b/tests/test_ops/test_loss_func.py @@ -1,52 +1,48 @@ -import torch -import pytest -import colossalai -import torch.nn.functional as F -import torch.multiprocessing as mp -from functools import partial -from colossalai.tensor import ColoTensor, ProcessGroup, ColoTensorSpec -from colossalai.utils import get_current_device -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.tensor import ShardSpec, ComputeSpec, ComputePattern - - -def check_cross_entropy(): - input_t = torch.randn(4, 4, device=get_current_device(), requires_grad=True) - input_ct = torch.randn(4, 4, device=get_current_device(), requires_grad=True) - with torch.no_grad(): - input_ct.copy_(input_t) - - target = torch.randint(4, (4,), dtype=torch.int64, device=get_current_device()) - - world_size = torch.distributed.get_world_size() - pg = ProcessGroup(tp_degree=world_size) - input_t_colo = ColoTensor.from_torch_tensor(tensor=input_ct, spec=ColoTensorSpec(pg)) - input_shard = input_t_colo.redistribute(ShardSpec([-1], [pg.tp_world_size()])) - input_shard.set_tensor_spec(dist_spec=None, compute_spec=ComputeSpec(ComputePattern.TP1D)) - - output = F.cross_entropy(input_t, target) - output_colo = F.cross_entropy(input_shard, target) - assert torch.allclose(output_colo, output) - - output.backward() - output_colo.backward() - - assert torch.allclose(input_t.grad, input_ct.grad) - - -def run_dist(rank, world_size, port): - colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - check_cross_entropy() - - -@pytest.mark.dist -@pytest.mark.parametrize('world_size', [1, 2]) -@rerun_if_address_is_in_use() -def test_loss_func(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_loss_func(1) +import pytest +import torch +import torch.nn.functional as F + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device + + +def check_cross_entropy(): + input_t = torch.randn(4, 4, device=get_current_device(), requires_grad=True) + input_ct = torch.randn(4, 4, device=get_current_device(), requires_grad=True) + with torch.no_grad(): + input_ct.copy_(input_t) + + target = torch.randint(4, (4,), dtype=torch.int64, device=get_current_device()) + + world_size = torch.distributed.get_world_size() + pg = ProcessGroup(tp_degree=world_size) + input_t_colo = ColoTensor.from_torch_tensor(tensor=input_ct, spec=ColoTensorSpec(pg)) + input_shard = input_t_colo.redistribute(ShardSpec([-1], [pg.tp_world_size()])) + input_shard.set_tensor_spec(dist_spec=None, compute_spec=ComputeSpec(ComputePattern.TP1D)) + + output = F.cross_entropy(input_t, target) + output_colo = F.cross_entropy(input_shard, target) + assert torch.allclose(output_colo, output) + + output.backward() + output_colo.backward() + + assert torch.allclose(input_t.grad, input_ct.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_cross_entropy() + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [1, 2]) +@rerun_if_address_is_in_use() +def test_loss_func(world_size): + spawn(run_dist, world_size) + + +if __name__ == '__main__': + test_loss_func(1) diff --git a/tests/test_ops/test_op.py b/tests/test_ops/test_op.py index 8d3cf50ff2aa..4176d3b64d90 100644 --- a/tests/test_ops/test_op.py +++ b/tests/test_ops/test_op.py @@ -1,14 +1,12 @@ -import torch import pytest -import colossalai +import torch import torch.nn.functional as F -import torch.multiprocessing as mp -from functools import partial -from colossalai.tensor import ColoTensor, ProcessGroup, ColoTensorSpec, ShardSpec -from colossalai.utils import get_current_device from torch.nn import Parameter -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup, ShardSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device def _run_layer_norm(): @@ -66,8 +64,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [2]) @rerun_if_address_is_in_use() def test_element_wise_ops(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) def run_dist2(rank, world_size, port): @@ -79,8 +76,7 @@ def run_dist2(rank, world_size, port): @pytest.mark.parametrize('world_size', [1]) @rerun_if_address_is_in_use() def test_ln(world_size): - run_func = partial(run_dist2, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist2, world_size) def check_all(): diff --git a/tests/test_ops/test_view.py b/tests/test_ops/test_view.py index fc6fc2d3c291..a9f2033201c7 100644 --- a/tests/test_ops/test_view.py +++ b/tests/test_ops/test_view.py @@ -1,100 +1,97 @@ -from functools import partial - -import colossalai -import pytest -import torch -import torch.multiprocessing as mp -import torch.distributed as dist -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port, get_current_device -from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor, ShardSpec -from colossalai.tensor.distspec import DistPlacementPattern -from tests.test_tensor.common_utils import split_param_row_tp1d, split_param_col_tp1d, debug_print - - -def exam_view_core(pg): - # the case of replicated ColoTensors - x = torch.randn(4, 4).cuda() - x_colo = ColoTensor(x, ColoTensorSpec(pg)) - - y = x.view(2, -1, 2) - y_colo = x_colo.view(2, -1, 2) - - assert torch.all(y == y_colo) - assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE - # the perfect case of col-sliced ColoTensors - split_param_col_tp1d(x_colo, pg) - - z = x.view(torch.Size((2, 1, 2, -1))) - z_colo = x_colo.view(torch.Size((2, 1, 2, -1))) - if dist.get_rank() == 0: - z = z[:, :, :, 0:2] - else: - z = z[:, :, :, 2:] - assert torch.all(z == z_colo) - assert z_colo.dist_spec == x_colo.dist_spec - # the perfect case of row-sliced ColoTensors - split_param_row_tp1d(x_colo, pg) - - z = x.view(torch.Size((-1, 2, 2))) - z_colo = x_colo.view(torch.Size((-1, 2, 2))) - if dist.get_rank() == 0: - z = z[0:2, :, :] - else: - z = z[2:, :, :] - assert torch.all(z == z_colo) - assert z_colo.dist_spec == x_colo.dist_spec - # the normal case of row-sliced ColoTensors - z = x.view(-1, 2, 2, 2) - z_colo = x_colo.view(-1, 2, 2, 2) - assert torch.all(z == z_colo) - assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE - - -def exam_view_autograd(pg): - x = torch.randn(8, 2, device=get_current_device(), requires_grad=True) - y = torch.randn(8, 2, device=get_current_device(), requires_grad=True) - with torch.no_grad(): - y.copy_(x) - y = ColoTensor(y, ColoTensorSpec(pg)) - y_slice = y.redistribute(ShardSpec([-1], [pg.tp_world_size()])) - - xx = x.view(2, 2, -1) - yy_slice = y_slice.view(2, 2, -1) - yy = yy_slice.to_replicate() - grad = torch.randn(2, 2, 4, device=get_current_device()) - - xx.backward(grad) - yy.backward(grad) - assert torch.all(x.grad == y.grad) - - -def exam_view_errors(pg): - x = torch.randn(8, 2, device=get_current_device()) - x = ColoTensor(x, ColoTensorSpec(pg)) - split_param_row_tp1d(x, pg) - - x.view('a', 'b', 'c') - x.view(8, -1) - x.view([-2, -2, -2]) - x.view((-1, -1, -1)) - - -def run_dist(rank, world_size, port): - colossalai.launch(config=dict(), rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - pg = ProcessGroup(tp_degree=torch.distributed.get_world_size()) - exam_view_core(pg) - exam_view_autograd(pg) - # exam_view_errors(pg) - - -@pytest.mark.dist -@pytest.mark.parametrize('world_size', [2]) -@rerun_if_address_is_in_use() -def test_view(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_view(2) +import pytest +import torch +import torch.distributed as dist + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup, ShardSpec +from colossalai.tensor.distspec import DistPlacementPattern +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from tests.test_tensor.common_utils import debug_print, split_param_col_tp1d, split_param_row_tp1d + + +def exam_view_core(pg): + # the case of replicated ColoTensors + x = torch.randn(4, 4).cuda() + x_colo = ColoTensor(x, ColoTensorSpec(pg)) + + y = x.view(2, -1, 2) + y_colo = x_colo.view(2, -1, 2) + + assert torch.all(y == y_colo) + assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE + # the perfect case of col-sliced ColoTensors + split_param_col_tp1d(x_colo, pg) + + z = x.view(torch.Size((2, 1, 2, -1))) + z_colo = x_colo.view(torch.Size((2, 1, 2, -1))) + if dist.get_rank() == 0: + z = z[:, :, :, 0:2] + else: + z = z[:, :, :, 2:] + assert torch.all(z == z_colo) + assert z_colo.dist_spec == x_colo.dist_spec + # the perfect case of row-sliced ColoTensors + split_param_row_tp1d(x_colo, pg) + + z = x.view(torch.Size((-1, 2, 2))) + z_colo = x_colo.view(torch.Size((-1, 2, 2))) + if dist.get_rank() == 0: + z = z[0:2, :, :] + else: + z = z[2:, :, :] + assert torch.all(z == z_colo) + assert z_colo.dist_spec == x_colo.dist_spec + # the normal case of row-sliced ColoTensors + z = x.view(-1, 2, 2, 2) + z_colo = x_colo.view(-1, 2, 2, 2) + assert torch.all(z == z_colo) + assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE + + +def exam_view_autograd(pg): + x = torch.randn(8, 2, device=get_current_device(), requires_grad=True) + y = torch.randn(8, 2, device=get_current_device(), requires_grad=True) + with torch.no_grad(): + y.copy_(x) + y = ColoTensor(y, ColoTensorSpec(pg)) + y_slice = y.redistribute(ShardSpec([-1], [pg.tp_world_size()])) + + xx = x.view(2, 2, -1) + yy_slice = y_slice.view(2, 2, -1) + yy = yy_slice.to_replicate() + grad = torch.randn(2, 2, 4, device=get_current_device()) + + xx.backward(grad) + yy.backward(grad) + assert torch.all(x.grad == y.grad) + + +def exam_view_errors(pg): + x = torch.randn(8, 2, device=get_current_device()) + x = ColoTensor(x, ColoTensorSpec(pg)) + split_param_row_tp1d(x, pg) + + x.view('a', 'b', 'c') + x.view(8, -1) + x.view([-2, -2, -2]) + x.view((-1, -1, -1)) + + +def run_dist(rank, world_size, port): + colossalai.launch(config=dict(), rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + pg = ProcessGroup(tp_degree=torch.distributed.get_world_size()) + exam_view_core(pg) + exam_view_autograd(pg) + # exam_view_errors(pg) + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [2]) +@rerun_if_address_is_in_use() +def test_view(world_size): + spawn(run_dist, world_size) + + +if __name__ == '__main__': + test_view(2) diff --git a/tests/test_optimizer/test_adam_kernel.py b/tests/test_optimizer/test_adam_kernel.py new file mode 100644 index 000000000000..2186a421fe00 --- /dev/null +++ b/tests/test_optimizer/test_adam_kernel.py @@ -0,0 +1,131 @@ +# This test checks adam kernels +# Baseline is pure fp32 torch adam optimizer +import math +from abc import abstractmethod +from typing import Type + +import pytest +import torch +from torch import Tensor + +from colossalai.utils import get_current_device, multi_tensor_applier + +_FUSED_ALLOWED_P_G_TYPES = [(torch.float, torch.half), (torch.float, torch.float), (torch.half, torch.float), + (torch.half, torch.half), (torch.bfloat16, torch.float), (torch.float, torch.bfloat16), + (torch.bfloat16, torch.bfloat16)] + +_CPU_ALLOWED_P_G_TYPES = [(torch.float, torch.half), (torch.float, torch.float), (torch.half, torch.float), + (torch.half, torch.half)] + + +class AdamKernel: + + def __init__(self, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, use_adamw: bool) -> None: + self.lr = lr + self.beta1 = beta1 + self.beta2 = beta2 + self.eps = eps + self.weight_decay = weight_decay + self.use_adamw = use_adamw + + @abstractmethod + def update(self, step: int, param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor): + pass + + +class TorchAdamKernel(AdamKernel): + + def update(self, step: int, param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor): + bias_correction1 = 1 - self.beta1**step + bias_correction2 = 1 - self.beta2**step + + if self.weight_decay != 0: + if self.use_adamw: + # Perform stepweight decay + param.mul_(1 - self.lr * self.weight_decay) + else: + grad = grad.add(param, alpha=self.weight_decay) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(self.beta1).add_(grad, alpha=1 - self.beta1) + exp_avg_sq.mul_(self.beta2).addcmul_(grad, grad, value=1 - self.beta2) + denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(self.eps) + + step_size = self.lr / bias_correction1 + + param.addcdiv_(exp_avg, denom, value=-step_size) + + +class FusedAdamKernel(AdamKernel): + + def __init__(self, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, use_adamw: bool) -> None: + super().__init__(lr, beta1, beta2, eps, weight_decay, use_adamw) + from colossalai.kernel.op_builder import FusedOptimBuilder + fused_optim = FusedOptimBuilder().load() + self.fused_adam = fused_optim.multi_tensor_adam + self.dummy_overflow_buf = torch.cuda.IntTensor([0]) + + def update(self, step: int, param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor): + multi_tensor_applier(self.fused_adam, self.dummy_overflow_buf, [[grad], [param], [exp_avg], [exp_avg_sq]], + self.lr, self.beta1, self.beta2, self.eps, step, self.use_adamw, True, self.weight_decay, + -1) + + +class CPUAdamKernel(AdamKernel): + + def __init__(self, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, use_adamw: bool) -> None: + super().__init__(lr, beta1, beta2, eps, weight_decay, use_adamw) + from colossalai.kernel.op_builder import CPUAdamBuilder + cpu_optim = CPUAdamBuilder().load() + + self.cpu_adam_op = cpu_optim.CPUAdamOptimizer(lr, beta1, beta2, eps, weight_decay, use_adamw) + + def update(self, step: int, param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor): + self.cpu_adam_op.step(step, self.lr, self.beta1, self.beta2, self.eps, self.weight_decay, True, param.view(-1), + grad.view(-1), exp_avg.view(-1), exp_avg_sq.view(-1), -1) + + +def check_adam_kernel(kernel: Type[AdamKernel], adamw: bool, weight_decay: float, p_dtype: torch.dtype, + g_dtype: torch.dtype, device: torch.device, n_steps: int, rtol: float, atol: float): + lr = 1e-3 + beta1, beta2 = 0.9, 0.999 + eps = 1e-8 + torch_adam = TorchAdamKernel(lr, beta1, beta2, eps, weight_decay, adamw) + adam_kernel = kernel(lr, beta1, beta2, eps, weight_decay, adamw) + master_p = torch.rand(64, device=device) + master_g = torch.rand_like(master_p) + master_exp_avg = torch.zeros_like(master_p) + master_exp_avg_sq = torch.zeros_like(master_p) + p = master_p.clone().to(p_dtype) + g = master_g.clone().to(g_dtype) + exp_avg = master_exp_avg.clone() + exp_avg_sq = master_exp_avg_sq.clone() + + for step in range(1, 1 + n_steps): + torch_adam.update(step, master_p, master_g, master_exp_avg, master_exp_avg_sq) + adam_kernel.update(step, p, g, exp_avg, exp_avg_sq) + # if overflow, the weight won't be updated. so there will be no nan in p + assert not torch.isnan(p).any() + assert torch.allclose(master_p, p.float(), rtol=rtol, atol=atol) + + +@pytest.mark.parametrize('adamw', [False, True]) +@pytest.mark.parametrize('weight_decay', [0.0, 0.1]) +@pytest.mark.parametrize('p_dtype, g_dtype', _FUSED_ALLOWED_P_G_TYPES) +def test_fused_adam_kernel(adamw, weight_decay, p_dtype, g_dtype): + rtol, atol = 1e-5, 1e-8 + if p_dtype is torch.float16 or g_dtype is torch.float16: + rtol, atol = 1e-3, 1e-3 + if p_dtype is torch.bfloat16 or g_dtype is torch.bfloat16: + rtol, atol = 4e-3, 4e-3 + check_adam_kernel(FusedAdamKernel, adamw, weight_decay, p_dtype, g_dtype, get_current_device(), 3, rtol, atol) + + +@pytest.mark.parametrize('adamw', [False, True]) +@pytest.mark.parametrize('weight_decay', [0.0, 0.1]) +@pytest.mark.parametrize('p_dtype, g_dtype', _CPU_ALLOWED_P_G_TYPES) +def test_cpu_adam_kernel(adamw, weight_decay, p_dtype, g_dtype): + rtol, atol = 1e-5, 1e-8 + if p_dtype is torch.float16 or g_dtype is torch.float16: + rtol, atol = 1e-3, 1e-3 + check_adam_kernel(CPUAdamKernel, adamw, weight_decay, p_dtype, g_dtype, torch.device('cpu'), 3, rtol, atol) diff --git a/tests/test_optimizer/test_adam_optim.py b/tests/test_optimizer/test_adam_optim.py new file mode 100644 index 000000000000..0f72bc134809 --- /dev/null +++ b/tests/test_optimizer/test_adam_optim.py @@ -0,0 +1,86 @@ +from copy import deepcopy +from typing import Type, Union + +import pytest +import torch +import torch.nn as nn +from torch.optim import Adam, AdamW + +from colossalai.nn.optimizer import CPUAdam, FusedAdam, HybridAdam +from tests.kit.model_zoo import model_zoo + +_ALLOWED_OPTIM_DEVICES = [ + (FusedAdam, torch.device('cuda:0')), + (CPUAdam, torch.device('cpu')), + (CPUAdam, torch.device('cuda:0')), + (HybridAdam, torch.device('cpu')), + (HybridAdam, torch.device('cuda:0')), +] + +_ALLOWED_P_G_TYPES = [ + (torch.float, torch.float), # pure fp32 + (torch.float, torch.half), # fp16 amp + (torch.float, torch.bfloat16), # bfloat16 amp + # (torch.half, torch.half), # FIXME(ver217): cpu adam kernel does not support pure fp16 + # (torch.bfloat16, torch.bfloat16), # FIXME(ver217): cpu adam kernel does not support pure bfloat16 +] + +N_STEPS = 3 + + +def setup_param_groups(bert_model: nn.Module) -> list: + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [p for n, p in bert_model.named_parameters() if not any(nd in n for nd in no_decay)], + "weight_decay": 0.1, + }, + { + "params": [p for n, p in bert_model.named_parameters() if any(nd in n for nd in no_decay)], + "weight_decay": 0.0, + }, + ] + return optimizer_grouped_parameters + + +def set_grad(model: nn.Module, torch_model: nn.Module, g_dtype: torch.dtype) -> None: + for p, torch_p in zip(model.parameters(), torch_model.parameters()): + torch_p.grad = torch.rand_like(torch_p) + # avoid inconsistent grad and param dtype error + orig_p = p.data + p.data = torch_p.grad.clone().to(g_dtype) + p.grad = p.data + p.data = orig_p + + +@pytest.mark.parametrize('optim_cls, device', _ALLOWED_OPTIM_DEVICES) +@pytest.mark.parametrize('adamw', [False, True]) +@pytest.mark.parametrize('p_dtype, g_dtype', _ALLOWED_P_G_TYPES) +def test_adam_optim_on_bert(optim_cls: Union[Type[FusedAdam], Type[CPUAdam], Type[HybridAdam]], device: torch.device, + adamw: bool, p_dtype: torch.dtype, g_dtype: torch.dtype) -> None: + model_fn, *_ = next(iter(model_zoo.get_sub_registry('transformers_bert_for_sequence_classification').values())) + torch_model = model_fn().to(device) + model = deepcopy(torch_model).to(p_dtype) + lr = 1e-3 + beta1, beta2 = 0.9, 0.999 + eps = 1e-8 + torch_optim_cls = AdamW if adamw else Adam + torch_optim = torch_optim_cls(setup_param_groups(torch_model), lr=lr, betas=(beta1, beta2), eps=eps) + optim = optim_cls(setup_param_groups(model), lr=lr, betas=(beta1, beta2), eps=eps, adamw_mode=adamw) + + rtol, atol = 1e-5, 1e-5 + if p_dtype is torch.float16 or g_dtype is torch.float16: + rtol, atol = 2e-3, 2e-3 + if p_dtype is torch.bfloat16 or g_dtype is torch.bfloat16: + rtol, atol = 4e-3, 4e-3 + + for _ in range(N_STEPS): + set_grad(model, torch_model, g_dtype) + torch_optim.step() + optim.step() + torch_optim.zero_grad() + optim.zero_grad() + for p, torch_p in zip(model.parameters(), torch_model.parameters()): + # if overflow, the weight won't be updated. so there will be no nan in p + assert not torch.isnan(p).any() + assert torch.allclose(p.float(), torch_p, rtol=rtol, atol=atol) diff --git a/tests/test_optimizer/test_cpu_adam.py b/tests/test_optimizer/test_cpu_adam.py deleted file mode 100644 index d317dc2e34ad..000000000000 --- a/tests/test_optimizer/test_cpu_adam.py +++ /dev/null @@ -1,120 +0,0 @@ -import math - -import torch - -from colossalai.testing import parameterize - - -def torch_adam_update( - step, - lr, - beta1, - beta2, - eps, - weight_decay, - param, - grad, - exp_avg, - exp_avg_sq, - use_adamw, -): - bias_correction1 = 1 - beta1**step - bias_correction2 = 1 - beta2**step - - if weight_decay != 0: - if use_adamw: - # Perform stepweight decay - param.mul_(1 - lr * weight_decay) - else: - grad = grad.add(param, alpha=weight_decay) - - # Decay the first and second moment running average coefficient - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) - denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) - - step_size = lr / bias_correction1 - - param.addcdiv_(exp_avg, denom, value=-step_size) - - -def assertLess(data_diff, threshold, msg): - assert data_diff < threshold, msg - - -def assertTrue(condition, msg): - assert condition, msg - - -@parameterize('adamw', [True, False]) -@parameterize('step', [1, 2]) -@parameterize('p_dtype', [torch.float, torch.half]) -@parameterize('g_dtype', [torch.float, torch.half]) -def test_cpu_adam(adamw, step, p_dtype, g_dtype): - lr = 1e-3 - beta1, beta2 = 0.9, 0.999 - eps = 1e-8 - weight_decay = 0 - - for i in range(1024): - p_data = torch.rand(64, dtype=p_dtype) - p_data_copy = p_data.clone().float() - p_grad = torch.rand(64, dtype=g_dtype) - p_grad_copy = p_grad.clone().float() - exp_avg = torch.rand(p_data.shape) - exp_avg_copy = exp_avg.clone() - exp_avg_sq = torch.rand(p_data.shape) - exp_avg_sq_copy = exp_avg_sq.clone() - - from colossalai.kernel.op_builder import CPUAdamBuilder - cpu_optim = CPUAdamBuilder().load() - - cpu_adam_op = cpu_optim.CPUAdamOptimizer(lr, beta1, beta2, eps, weight_decay, adamw) - - cpu_adam_op.step( - step, - lr, - beta1, - beta2, - eps, - weight_decay, - True, - p_data.view(-1), # fp32 data - p_grad.view(-1), # fp32 grad - exp_avg.view(-1), - exp_avg_sq.view(-1), - -1, - ) - - torch_adam_update( - step, - lr, - beta1, - beta2, - eps, - weight_decay, - p_data_copy, # fp32 data - p_grad_copy, # fp32 grad - exp_avg_copy, - exp_avg_sq_copy, - adamw, - ) - var = p_data_copy - p_data - data_diff = torch.max(torch.abs(var)) - threshold = 1e-3 - assertLess( - data_diff, - threshold, - f"p_data diff {data_diff}. failed check, step {step}, lr {lr}, eps " - f"{eps} beta1 {beta1} beta2 {beta2} weight_decay {weight_decay} p_dtype {p_dtype}, g_dtype {g_dtype}", - ) - max_grad_diff = torch.max(torch.abs(p_grad_copy - p_grad)) - assertTrue(max_grad_diff < threshold, f"diff {max_grad_diff}") - max_exp_avg_diff = torch.max(torch.abs(exp_avg_copy - exp_avg)) - assertTrue(max_exp_avg_diff < threshold, f"max_exp_avg_diff {max_exp_avg_diff}") - max_exp_avg_sq_diff = torch.max(torch.abs(exp_avg_sq_copy - exp_avg_sq)) - assertTrue(max_exp_avg_sq_diff < threshold, f"max_exp_avg_sq_diff {max_exp_avg_sq_diff}") - - -if __name__ == '__main__': - test_cpu_adam() diff --git a/tests/test_optimizer/test_fused_adam.py b/tests/test_optimizer/test_fused_adam.py deleted file mode 100644 index f7227c2d57c0..000000000000 --- a/tests/test_optimizer/test_fused_adam.py +++ /dev/null @@ -1,63 +0,0 @@ -import torch -import torch.nn as nn -from torch.optim.adam import Adam -from torch.optim import AdamW - -from colossalai.nn.optimizer.fused_adam import FusedAdam -from colossalai.testing import parameterize - - -class FC(nn.Module): - - def __init__(self) -> None: - super().__init__() - self.fc = nn.Sequential(nn.Linear(64, 64)) - - def forward(self, x): - return self.fc(x) - - -@parameterize('adamw', [False, True]) -@parameterize('p_dtype', [torch.float, torch.half]) -@parameterize('g_dtype', [torch.float, torch.half]) -def test_adam(adamw, p_dtype, g_dtype): - model = FC().cuda().to(p_dtype) - state = model.state_dict() - model_copy = FC().cuda().to(p_dtype) - model_copy.load_state_dict(state.copy()) - - if adamw: - optim = FusedAdam(model.parameters(), lr=1e-3, adamw_mode=True) - torch_optim = AdamW(model_copy.parameters(), lr=1e-3) - else: - optim = FusedAdam(model.parameters(), lr=1e-3) - torch_optim = Adam(model_copy.parameters(), lr=1e-3) - - data = torch.rand(1024, 64).cuda().to(p_dtype) - data_copy = data.clone() - label = torch.rand(1024, 64).cuda().to(p_dtype) - - for d, l in zip(data, label): - y = model(d) - loss = ((l - y)**2).sum() - optim.zero_grad() - loss.backward() - if p_dtype != g_dtype: - for i in range(len(optim.param_groups[0]['params'])): - optim.param_groups[0]['params'][i].grad.data = optim.param_groups[0]['params'][i].grad.data.to(g_dtype) - optim.step() - - for d, l in zip(data_copy, label): - y = model_copy(d) - loss = ((l - y)**2).sum() - torch_optim.zero_grad() - loss.backward() - torch_optim.step() - - assert len(optim.param_groups[0]['params']) == len(torch_optim.param_groups[0]['params']) - - for i in range(len(optim.param_groups[0]['params'])): - if torch.isnan(optim.param_groups[0]['params'][i]).any() \ - or torch.isnan(torch_optim.param_groups[0]['params'][i]).any(): - continue - assert torch.allclose(optim.param_groups[0]['params'][i], torch_optim.param_groups[0]['params'][i], 2e-3, 2e-3) diff --git a/tests/test_optimizer/test_fused_adam_kernel.py b/tests/test_optimizer/test_fused_adam_kernel.py deleted file mode 100644 index 7b9b6e9c48ba..000000000000 --- a/tests/test_optimizer/test_fused_adam_kernel.py +++ /dev/null @@ -1,94 +0,0 @@ -import math - -import torch -import torch.nn as nn -from numpy import dtype - -from colossalai.testing import parameterize -from colossalai.utils import multi_tensor_applier - - -def torch_adam_update( - step, - lr, - beta1, - beta2, - eps, - weight_decay, - param, - grad, - exp_avg, - exp_avg_sq, - use_adamw, -): - bias_correction1 = 1 - beta1**step - bias_correction2 = 1 - beta2**step - - if weight_decay != 0: - if use_adamw: - # Perform stepweight decay - param.mul_(1 - lr * weight_decay) - else: - grad = grad.add(param, alpha=weight_decay) - - # Decay the first and second moment running average coefficient - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) - denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) - - step_size = lr / bias_correction1 - - param.addcdiv_(exp_avg, denom, value=-step_size) - - -@parameterize('adamw', [False, True]) -@parameterize('step', [1, 2]) -@parameterize('p_dtype', [torch.float, torch.half]) -@parameterize('g_dtype', [torch.float, torch.half]) -def test_adam(adamw, step, p_dtype, g_dtype): - from colossalai.kernel.op_builder import FusedOptimBuilder - fused_optim = FusedOptimBuilder().load() - fused_adam = fused_optim.multi_tensor_adam - - dummy_overflow_buf = torch.cuda.IntTensor([0]) - - count = 0 - - for i in range(1024): - p = torch.rand(64, dtype=p_dtype).cuda() - p_copy = p.clone().float() - g = torch.rand(p.shape, dtype=g_dtype).cuda() - g_copy = g.clone().float() - m = torch.rand(p.shape).cuda() - m_copy = m.clone() - v = torch.rand(p.shape).cuda() - v_copy = v.clone() - - lr = 1e-3 - beta1, beta2 = 0.9, 0.999 - eps = 1e-8 - weight_decay = 0 - - multi_tensor_applier(fused_adam, dummy_overflow_buf, [[g], [p], [m], [v]], lr, beta1, beta2, eps, step, adamw, - True, weight_decay, -1) - - torch_adam_update( - step, - lr, - beta1, - beta2, - eps, - weight_decay, - p_copy, # fp32 data - g_copy, # fp32 grad - m_copy, - v_copy, - adamw, - ) - - if torch.isnan(p).any() or torch.isnan(p_copy).any(): - count += 1 - continue - assert count < 200, "too many nans" - assert torch.allclose(p.to(torch.float), p_copy.to(torch.float), 1e-5, - 1e-5), f"failed check, adamw {adamw}, p_dtype {p_dtype}, g_dtype {g_dtype}" diff --git a/tests/test_optimizer/test_hybrid_adam.py b/tests/test_optimizer/test_hybrid_adam.py deleted file mode 100644 index d19192add3fb..000000000000 --- a/tests/test_optimizer/test_hybrid_adam.py +++ /dev/null @@ -1,41 +0,0 @@ -import torch -import torch.nn as nn -from torch.optim.adam import Adam -from torch.optim import AdamW - -from colossalai.nn.optimizer.hybrid_adam import HybridAdam -from colossalai.testing import parameterize - -RE = 1024 - - -@parameterize('adamw', [False, True]) -@parameterize('device', ['cpu', 'cuda:0']) -@parameterize('p_dtype', [torch.float]) -@parameterize('g_dtype', [torch.float, torch.half]) -def test_adam(adamw, device, p_dtype, g_dtype): - rng_state = torch.get_rng_state() - p = nn.Parameter(torch.rand(64).to(device, p_dtype)) - torch.set_rng_state(rng_state) - p_copy = nn.Parameter(torch.rand(64).to(device).float()) - - if adamw: - optim = HybridAdam([p], lr=1e-3, adamw_mode=True) - torch_optim = AdamW([p_copy], lr=1e-3) - else: - optim = HybridAdam([p], lr=1e-3) - torch_optim = Adam([p_copy], lr=1e-3) - - print(f"adaw mode {adamw}, device {device}, p_dtype {p_dtype}, g_dtype {g_dtype}") - for i in range(RE): - p.grad = torch.rand(64).to(device, p_dtype) - p_copy.grad = p.grad.clone().float() - p.grad.data = p.grad.data.to(g_dtype) - - optim.step() - torch_optim.step() - - if torch.isnan(p.data).any() or torch.isnan(p_copy.data).any(): - continue - assert torch.allclose(p.data, p_copy.data, 1e-4, 1e-2), \ - f"adaw mode {adamw}, device {device}, p_dtype {p_dtype}, g_dtype {g_dtype}" diff --git a/tests/test_optimizer/test_nvme.py b/tests/test_optimizer/test_nvme.py index 243f785adaf9..5d794ac2dd1a 100644 --- a/tests/test_optimizer/test_nvme.py +++ b/tests/test_optimizer/test_nvme.py @@ -1,7 +1,9 @@ import pytest import torch -from tests.components_to_test.registry import non_distributed_component_funcs + from colossalai.nn.optimizer import CPUAdam, HybridAdam +from colossalai.testing import clear_cache_before_run, parameterize +from tests.components_to_test.registry import non_distributed_component_funcs def move_some_params_to_cuda(model, torch_model): @@ -16,9 +18,10 @@ def check_params_equal(model, torch_model): assert torch.allclose(p, torch_p, atol=1e-3), f'diff: {torch.abs(p - torch_p)}' -@pytest.mark.parametrize('nvme_offload_fraction', [0.0, 0.5, 1.0]) -@pytest.mark.parametrize('nvme_offload_dir', ['./offload', None]) -@pytest.mark.parametrize('adam_cls', [CPUAdam, HybridAdam]) +@clear_cache_before_run() +@parameterize('nvme_offload_fraction', [0.0, 0.5, 1.0]) +@parameterize('nvme_offload_dir', ['./offload', None]) +@parameterize('adam_cls', [CPUAdam, HybridAdam]) def test_nvme_adam(nvme_offload_fraction, nvme_offload_dir, adam_cls): get_components_func = non_distributed_component_funcs.get_callable('simple_net') model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() diff --git a/tests/test_pipeline/rpc_test_utils.py b/tests/test_pipeline/rpc_test_utils.py index 7ce2cd433b12..dab474a4ee21 100644 --- a/tests/test_pipeline/rpc_test_utils.py +++ b/tests/test_pipeline/rpc_test_utils.py @@ -6,13 +6,14 @@ import torch.distributed as dist import torch.distributed.rpc as rpc import torch.multiprocessing as mp -from colossalai import launch -from colossalai.logging import disable_existing_loggers -from colossalai.pipeline.pipeline_process_group import ppg from torch import nn from torch._C._distributed_rpc import _is_current_rpc_agent_set from torch.optim import SGD, Adam, Optimizer, RMSprop +from colossalai import launch +from colossalai.logging import disable_existing_loggers +from colossalai.pipeline.pipeline_process_group import ppg + rpc_is_initialized = _is_current_rpc_agent_set @@ -20,7 +21,9 @@ def color_debug(text, prefix=' ', color='blue'): color = color.upper() print(getattr(Back, color), prefix, Style.RESET_ALL, text) + class MLP(nn.Module): + def __init__(self, dim: int, layers: int): super().__init__() self.layers = torch.nn.ModuleList() @@ -32,8 +35,10 @@ def forward(self, x): for layer in self.layers: x = layer(x) return x.sum() - + + class DAG_MLP(nn.Module): + def __init__(self, dim: int, layers: int): super().__init__() self.layers = torch.nn.ModuleList() @@ -48,6 +53,7 @@ def forward(self, x, y): y = self.dag_layer(y) return x.sum(), y.sum() + class RpcTestModel(nn.Module): def __init__(self, stage_id, actual_stage_num, feat_num, h) -> None: diff --git a/tests/test_pipeline/test_middleware_1f1b.py b/tests/test_pipeline/test_middleware_1f1b.py index c4dc617b1683..5b3aad703275 100644 --- a/tests/test_pipeline/test_middleware_1f1b.py +++ b/tests/test_pipeline/test_middleware_1f1b.py @@ -1,27 +1,27 @@ -import torch -import pytest import os -import torch.multiprocessing as mp -import torch.distributed.rpc as rpc +from functools import partial -from torch import nn +import pytest +import torch +import torch.distributed.rpc as rpc +from rpc_test_utils import DAG_MLP, MLP from torch._C._distributed_rpc import _is_current_rpc_agent_set + from colossalai import launch +from colossalai.fx import ColoTracer +from colossalai.fx.passes.adding_split_node_pass import balanced_split_pass, split_with_split_nodes_pass from colossalai.logging import disable_existing_loggers +from colossalai.pipeline.middleware.adaptor import get_fx_topology from colossalai.pipeline.pipeline_process_group import ppg from colossalai.pipeline.rpc._pipeline_schedule import OneFOneBPipelineEngine -from colossalai.fx.passes.adding_split_node_pass import split_with_split_nodes_pass, balanced_split_pass -from colossalai.fx import ColoTracer -from colossalai.pipeline.middleware.adaptor import get_fx_topology -from rpc_test_utils import MLP, DAG_MLP -from functools import partial -from colossalai.testing import parameterize, rerun_if_address_is_in_use +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn # global variable for model created batch_size = 16 dim = 10 rpc_is_initialized = _is_current_rpc_agent_set + def create_partition_module(pp_rank: int, stage_num: int, model, data_kwargs): model.eval() tracer = ColoTracer() @@ -34,13 +34,15 @@ def create_partition_module(pp_rank: int, stage_num: int, model, data_kwargs): for submodule in split_submodules: if isinstance(submodule, torch.fx.GraphModule): setattr(submodule, '_topo', topo) - return split_submodules[pp_rank+1] + return split_submodules[pp_rank + 1] + def partition(model, data_kwargs: dict, pp_rank: int, chunk: int, stage_num: int): torch.manual_seed(1024) partition = create_partition_module(pp_rank, stage_num, model, data_kwargs) return partition + def run_master(model_cls, world_size, forward_only): torch.manual_seed(100) @@ -50,23 +52,27 @@ def run_master(model_cls, world_size, forward_only): chunk = 1 num_microbatches = 8 use_checkpoint = 'store_true' - + if model_cls == MLP: + def data_gen(): x = torch.zeros((batch_size, dim)) kwargs = dict(x=x) return kwargs + model = model_cls(dim, stage_num * 3) if forward_only: labels = None else: labels = 1 elif model_cls == DAG_MLP: + def data_gen(): x = torch.zeros((batch_size, dim)) y = torch.zeros((batch_size, dim)) kwargs = dict(x=x, y=y) return kwargs + model = model_cls(dim, stage_num * 3) if forward_only: labels = None @@ -74,15 +80,17 @@ def data_gen(): labels = 1 else: pass - + data_kwargs = data_gen() - - engine = OneFOneBPipelineEngine(partition_fn=partial(partition, model, data_kwargs), - stage_num=stage_num, - num_microbatches=num_microbatches, - device=device, - chunk=chunk, - checkpoint=use_checkpoint,) + + engine = OneFOneBPipelineEngine( + partition_fn=partial(partition, model, data_kwargs), + stage_num=stage_num, + num_microbatches=num_microbatches, + device=device, + chunk=chunk, + checkpoint=use_checkpoint, + ) if not forward_only: engine.initialize_optimizer(getattr(torch.optim, 'SGD'), lr=1e-3) @@ -90,13 +98,14 @@ def data_gen(): input_x = torch.randn((batch_size, dim), device=device) input_y = torch.randn((batch_size, dim), device=device) logits = engine.forward_backward({'x': input_x, 'y': input_y}, labels=labels, forward_only=forward_only) - -def run_worker(rank, model_cls, world_size, forward_only, master_func): + + +def run_worker(rank, world_size, port, model_cls, forward_only, master_func): master_addr = 'localhost' master_port = 29020 os.environ['MASTER_ADDR'] = master_addr os.environ['MASTER_PORT'] = str(master_port) - + disable_existing_loggers() launch(dict(), rank, world_size, master_addr, master_port, 'nccl', verbose=False) @@ -113,7 +122,8 @@ def run_worker(rank, model_cls, world_size, forward_only, master_func): # barrier here if rpc_is_initialized(): rpc.shutdown() - + + @pytest.mark.skip("skip due to CI torch version 1.11") @parameterize('model_cls', [MLP, DAG_MLP]) @parameterize('forward_only', [True, False]) @@ -122,7 +132,14 @@ def run_worker(rank, model_cls, world_size, forward_only, master_func): def test_pp_middleware_fwd(model_cls, forward_only): world_size = 4 master_func = run_master - mp.spawn(run_worker, args=(model_cls, world_size, forward_only, master_func), nprocs=world_size) + spawn( + run_worker, + world_size, + model_cls=model_cls, + forward_only=forward_only, + master_func=master_func, + ) + if __name__ == "__main__": - test_pp_middleware_fwd() \ No newline at end of file + test_pp_middleware_fwd() diff --git a/tests/test_pipeline/test_pipelinable.py b/tests/test_pipeline/test_pipelinable.py index c99a88550b71..627cb5ac6f51 100644 --- a/tests/test_pipeline/test_pipelinable.py +++ b/tests/test_pipeline/test_pipelinable.py @@ -1,9 +1,7 @@ import torch -import torch.multiprocessing as mp from colossalai.pipeline.pipelinable import PipelinableContext - -from colossalai.testing import rerun_on_exception +from colossalai.testing import rerun_if_address_is_in_use, rerun_on_exception, spawn NUM_CHUNKS = 1 PIPELINE_SIZE = 2 @@ -27,7 +25,7 @@ def forward(self, x): return x -def run_pipelinable(rank): +def run_pipelinable(rank, world_size, port): pipelinable = PipelinableContext() with pipelinable: model = MLP() @@ -50,9 +48,9 @@ def run_pipelinable(rank): assert layers_count_in_part_0 + layers_count_in_part_1 == pipelinable.layers_count -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") +@rerun_if_address_is_in_use() def test_pipelinable(): - mp.spawn(run_pipelinable, nprocs=1) + spawn(run_pipelinable, 1) if __name__ == '__main__': diff --git a/tests/test_pipeline/test_pipeline_process_group.py b/tests/test_pipeline/test_pipeline_process_group.py index c67e4175df92..2a00e3ac55b1 100644 --- a/tests/test_pipeline/test_pipeline_process_group.py +++ b/tests/test_pipeline/test_pipeline_process_group.py @@ -1,13 +1,12 @@ import os import torch.distributed.rpc as rpc -import torch.multiprocessing as mp -import pytest +from rpc_test_utils import pg_parse_args, rpc_is_initialized -from colossalai.pipeline.pipeline_process_group import ppg from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from rpc_test_utils import pg_parse_args, rpc_is_initialized +from colossalai.pipeline.pipeline_process_group import ppg +from colossalai.testing import spawn def run_worker(rank, args): @@ -40,4 +39,4 @@ def run_worker(rank, args): if __name__ == "__main__": args = pg_parse_args() world_size = args.world_size - mp.spawn(run_worker, args=(args,), nprocs=world_size) \ No newline at end of file + spawn(run_worker, world_size, args=args) diff --git a/tests/test_shardformer/__init__.py b/tests/test_shardformer/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_shardformer/test_layer/test_dist_crossentropy.py b/tests/test_shardformer/test_layer/test_dist_crossentropy.py new file mode 100644 index 000000000000..72e6e5cf26ed --- /dev/null +++ b/tests/test_shardformer/test_layer/test_dist_crossentropy.py @@ -0,0 +1,42 @@ +import pytest +import torch +import torch.nn.functional as F + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.shardformer.layer import cross_entropy_1d +from colossalai.testing import rerun_if_address_is_in_use, spawn + +CONFIG = dict(parallel=dict(data=1, pipeline=1, tensor=dict(size=2, mode='1d')),) + + +def check_dist_crossentropy(rank, world_size, port, ignore_index): + disable_existing_loggers() + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, port=port, host='localhost', backend='nccl') + + # prepare data + pred = torch.randn(2, 4, 8, requires_grad=True) + labels = torch.randint(8, (2, 4)) + # set some label to -100 to test the ignore index + labels[0, -1] = ignore_index + + org_pred = pred.view(-1, 8) + org_labels = labels.view(-1) + org_loss = F.cross_entropy(org_pred, org_labels) + + dist_pred = pred.chunk(world_size, -1)[rank] + dist_loss = cross_entropy_1d(dist_pred.to('cuda'), labels.to('cuda'), ignore_index=ignore_index) + + assert torch.allclose(org_loss, dist_loss, + atol=1e-5), f"dist cross entropy loss is not equal to orgin loss\n{org_loss}\n{dist_loss}" + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_dist_crossentropy(): + ignore_index = -100 + spawn(check_dist_crossentropy, 2, ignore_index=ignore_index) + + +if __name__ == '__main__': + test_dist_crossentropy() diff --git a/tests/test_shardformer/test_layer/test_dropout.py b/tests/test_shardformer/test_layer/test_dropout.py new file mode 100644 index 000000000000..332e377110a4 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_dropout.py @@ -0,0 +1,70 @@ +import torch +import torch.distributed as dist +import torch.nn as nn + +import colossalai +from colossalai.shardformer.layer import DropoutForParallelInput, DropoutForReplicatedInput +from colossalai.testing import assert_equal, assert_not_equal, rerun_if_address_is_in_use, spawn + + +def check_dropout_parallel_input(): + dropout = nn.Dropout().cuda() + dropout_1d = DropoutForParallelInput.from_native_module(dropout, process_group=None) + + # check computation correctness + x = torch.rand(4, 128).cuda() + + # we set seed so that dropout will generate the same mask + torch.cuda.manual_seed(1024) + out = dropout(x) + + # we set seed to simulate the same scenario + # but expect the dropout mask to be different + # due to the internal randomness control + torch.cuda.manual_seed(1024) + out_1d = dropout_1d(x) + + # ensure out is the same across all ranks + world_size = dist.get_world_size() + out_all = [torch.empty_like(out) for _ in range(world_size)] + dist.all_gather(out_all, out) + + for i in range(world_size): + assert_equal(out_all[i], out_all[0]) + + # ensure out_1d is different across ranks + out_1d_all = [torch.zeros_like(out_1d) for _ in range(world_size)] + dist.all_gather(out_1d_all, out_1d) + for i in range(1, world_size): + assert_not_equal(out_1d_all[i], out_1d_all[0]) + + +def check_dropout_replicated_input(): + dropout = nn.Dropout().cuda() + dropout_replica = DropoutForReplicatedInput.from_native_module(dropout, process_group=None) + + # check computation correctness + x = torch.rand(4, 128).cuda() + out_1d = dropout_replica(x) + + # ensure out_1d is different across ranks + world_size = dist.get_world_size() + out_1d_all = [torch.zeros_like(out_1d) for _ in range(world_size)] + dist.all_gather(out_1d_all, out_1d) + for i in range(1, world_size): + assert_equal(out_1d_all[i], out_1d_all[0]) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_dropout_parallel_input() + check_dropout_replicated_input() + + +@rerun_if_address_is_in_use() +def test_dropout(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_dropout() diff --git a/tests/test_shardformer/test_layer/test_embedding.py b/tests/test_shardformer/test_layer/test_embedding.py new file mode 100644 index 000000000000..8a6aa42a42f2 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_embedding.py @@ -0,0 +1,47 @@ +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.testing import assert_close + +import colossalai +from colossalai.shardformer.layer import Embedding1D +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +def check_embedding_1d(): + embedding = nn.Embedding(32, 128).cuda() + embedding_1d = Embedding1D.from_native_module(embedding, process_group=None) + + assert embedding_1d.weight.shape == torch.Size([32, 64]) + + # ensure state dict is reversibly loadable + embedding.load_state_dict(embedding_1d.state_dict()) + embedding_1d.load_state_dict(embedding.state_dict()) + + # check computation correctness + x = torch.randint(low=0, high=32, size=(4, 32)).cuda() + out = embedding(x) + gather_out = embedding_1d(x) + assert_close(out, gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + rank = dist.get_rank() + target_grad = torch.chunk(embedding.weight.grad, 2, dim=1)[rank] + assert_close(target_grad, embedding_1d.weight.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_embedding_1d() + + +@rerun_if_address_is_in_use() +def test_embedding_1d(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_embedding_1d() diff --git a/tests/test_shardformer/test_layer/test_layernorm.py b/tests/test_shardformer/test_layer/test_layernorm.py new file mode 100644 index 000000000000..a117845545be --- /dev/null +++ b/tests/test_shardformer/test_layer/test_layernorm.py @@ -0,0 +1,44 @@ +import torch +import torch.nn as nn +from torch.testing import assert_close + +import colossalai +from colossalai.shardformer.layer import FusedLayerNorm +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +def check_layernorm(): + norm = nn.LayerNorm(128, 0.00001).cuda() + norm1d = FusedLayerNorm.from_native_module(norm, process_group=None) + + assert norm1d.weight.shape == torch.Size([128]) + + # ensure state dict is reversibly loadable + norm.load_state_dict(norm1d.state_dict()) + norm1d.load_state_dict(norm.state_dict()) + + # check computation correctness + x = torch.rand(4, 128).cuda() + out = norm(x) + gather_out = norm1d(x) + assert_close(out, gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + assert_close(norm.weight.grad, norm1d.weight.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_layernorm() + + +@rerun_if_address_is_in_use() +def test_layernorm(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_layernorm_1d() diff --git a/tests/test_shardformer/test_layer/test_linear_1d.py b/tests/test_shardformer/test_layer/test_linear_1d.py new file mode 100644 index 000000000000..da3bdc1d78d3 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_linear_1d.py @@ -0,0 +1,131 @@ +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.testing import assert_close + +import colossalai +from colossalai.shardformer.layer import Linear1D_Col, Linear1D_Row +from colossalai.tensor.d_tensor import is_distributed_tensor +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +def check_linear_1d_col(): + linear = nn.Linear(32, 128).cuda() + linear_col = Linear1D_Col.from_native_module(linear, process_group=None, gather_output=True) + + # ensure that the parameters are distributed + assert is_distributed_tensor(linear_col.weight) + assert is_distributed_tensor(linear_col.bias) + + # ensure the shape is correct + assert linear_col.weight.shape == torch.Size([64, 32]) + assert linear_col.bias.shape == torch.Size([64]) + + # ensure state dict is reversibly loadable + linear.load_state_dict(linear_col.state_dict()) + linear_col.load_state_dict(linear.state_dict()) + + # check computation correctness + x = torch.rand(4, 32).cuda() + x_for_unshard = x.expand_as(x.clone()) + x_for_unshard.requires_grad_(True) + x_for_shard = x.expand_as(x.clone()) + x_for_shard.requires_grad_(True) + + out = linear(x_for_unshard) + gather_out = linear_col(x_for_shard) + assert_close(out, gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + rank = dist.get_rank() + target_grad = torch.chunk(linear.weight.grad, 2, dim=0)[rank] + assert_close(target_grad, linear_col.weight.grad) + + # check the input gradients + assert x_for_shard.grad is not None + assert x_for_unshard.grad is not None + assert_close(x_for_unshard.grad, x_for_shard.grad) + + +def check_linear_1d_row(): + linear = nn.Linear(32, 128).cuda() + linear_row = Linear1D_Row.from_native_module(linear, process_group=None, parallel_input=False) + + assert linear_row.weight.shape == torch.Size([128, 16]) + assert linear_row.bias.shape == torch.Size([128]) + + # check computation correctness + x = torch.rand(4, 32).cuda() + x_for_unshard = x.expand_as(x.clone()) + x_for_unshard.requires_grad_(True) + x_for_shard = x.expand_as(x.clone()) + x_for_shard.requires_grad_(True) + + # run forward + out = linear(x_for_unshard) + gather_out = linear_row(x_for_shard) + assert_close(out, gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + rank = dist.get_rank() + target_grad = torch.chunk(linear.weight.grad, 2, dim=1)[rank] + assert_close(target_grad, linear_row.weight.grad) + + # check the input gradients + assert x_for_shard.grad is not None + assert x_for_unshard.grad is not None + assert_close(x_for_unshard.grad, x_for_shard.grad) + + +def check_linear_col_plus_row(): + linear_1 = nn.Linear(32, 128).cuda() + linear_2 = nn.Linear(128, 32).cuda() + linear_col = Linear1D_Col.from_native_module(linear_1, process_group=None, gather_output=False) + linear_row = Linear1D_Row.from_native_module(linear_2, process_group=None, parallel_input=True) + + # check computation correctness + x = torch.rand(4, 32).cuda() + x_for_unshard = x.expand_as(x.clone()) + x_for_unshard.requires_grad_(True) + x_for_shard = x.expand_as(x.clone()) + x_for_shard.requires_grad_(True) + + # run forward + unshard_out = linear_2(linear_1(x_for_unshard)) + shard_out = linear_row(linear_col(x_for_shard)) + assert_close(unshard_out, shard_out) + + # check backward correctness + unshard_out.sum().backward() + shard_out.sum().backward() + + rank = dist.get_rank() + target_1_grad = torch.chunk(linear_1.weight.grad, 2, dim=0)[rank] + assert_close(target_1_grad, linear_col.weight.grad) + + # check the input gradients + assert x_for_shard.grad is not None + assert x_for_unshard.grad is not None + assert_close(x_for_unshard.grad, x_for_shard.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_linear_1d_col() + check_linear_1d_row() + check_linear_col_plus_row() + + +@rerun_if_address_is_in_use() +def test_linear(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_linear() diff --git a/tests/test_shardformer/test_layer/test_qkv_fused_linear_1d.py b/tests/test_shardformer/test_layer/test_qkv_fused_linear_1d.py new file mode 100644 index 000000000000..681c4f6dd9f1 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_qkv_fused_linear_1d.py @@ -0,0 +1,120 @@ +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.testing import assert_close + +import colossalai +from colossalai.shardformer.layer import GPT2FusedLinearConv1D_Col, GPT2FusedLinearConv1D_Row +from colossalai.shardformer.layer.qkv_fused_linear import split_fused_qkv_in_gpt2_style +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +# This code is copied from https://github.com/huggingface/transformers +class Conv1D(nn.Module): + """ + 1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2). + + Basically works like a linear layer but the weights are transposed. + + Args: + nf (`int`): The number of output features. + nx (`int`): The number of input features. + """ + + def __init__(self, nf, nx): + super().__init__() + self.nf = nf + self.weight = nn.Parameter(torch.empty(nx, nf)) + self.bias = nn.Parameter(torch.zeros(nf)) + nn.init.normal_(self.weight, std=0.02) + + def forward(self, x): + size_out = x.size()[:-1] + (self.nf,) + x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight) + x = x.view(size_out) + return x + + +def rearrange(tensor: torch.Tensor, dim: int): + tensor = tensor.clone() + world_size = 2 + order = torch.arange(world_size * 3) + new_order = [] + for i in range(world_size): + new_order.append(order[i::world_size]) + new_order = torch.cat(new_order) + + tensor_chunks = torch.chunk(tensor, world_size * 3, dim=dim) + rearanged_tensor_chunks = [tensor_chunks[i] for i in new_order] + rearanged_tensor = torch.cat(rearanged_tensor_chunks, dim=dim) + return rearanged_tensor + + +def check_linear_conv_1d_col(): + linear = Conv1D(192, 48).cuda() + linear_conv_col = GPT2FusedLinearConv1D_Col.from_native_module(linear, + process_group=None, + gather_output=True, + n_fused=3) + + assert linear.weight.shape == torch.Size([48, 192]) + assert linear.bias.shape == torch.Size([192]) + assert linear_conv_col.weight.shape == torch.Size([48, 96]) + assert linear_conv_col.bias.shape == torch.Size([96]) + + # ensure weights are reversibly loadable + linear_conv_col.load_state_dict(linear.state_dict()) + linear.load_state_dict(linear_conv_col.state_dict()) + + # check computation correctness + x = torch.rand(4, 48).cuda() + out = linear(x) + gather_out = linear_conv_col(x) + assert_close(rearrange(out, 1), gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + target_grad = split_fused_qkv_in_gpt2_style(linear.weight.grad, 3, None, True) + assert_close(target_grad, linear_conv_col.weight.grad) + + +def check_linear_conv_1d_row(): + linear = Conv1D(192, 48).cuda() + linear_row = GPT2FusedLinearConv1D_Row.from_native_module(linear, process_group=None, parallel_input=False) + + assert linear.weight.shape == torch.Size([48, 192]) + assert linear_row.weight.shape == torch.Size([24, 192]) + assert linear_row.bias.shape == torch.Size([192]) + + # check computation correctness + x = torch.rand(4, 48).cuda() + out = linear(x) + gather_out = linear_row(x) + assert_close(out, gather_out) + + # check backward correctness + out.sum().backward() + gather_out.sum().backward() + + rank = dist.get_rank() + target_grad = torch.chunk(linear.weight.grad, 2, dim=0)[rank] + assert_close(target_grad, linear_row.weight.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + # test for linear conv + check_linear_conv_1d_col() + check_linear_conv_1d_row() + + +@rerun_if_address_is_in_use() +def test_linearconv(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_linearconv() diff --git a/tests/test_shardformer/test_layer/test_vocab_parallel_embedding_1d.py b/tests/test_shardformer/test_layer/test_vocab_parallel_embedding_1d.py new file mode 100644 index 000000000000..8991d9b304f5 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_vocab_parallel_embedding_1d.py @@ -0,0 +1,49 @@ +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.testing import assert_close + +import colossalai +from colossalai.shardformer.layer import VocabParallelEmbedding1D +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn + + +def check_vocab_embedding_1d(): + embedding = nn.Embedding(128, 32).to('cuda') + dist_embedding_1d = VocabParallelEmbedding1D.from_native_module(embedding, process_group=None) + + assert dist_embedding_1d.weight.shape == torch.Size([64, 32]) + assert dist_embedding_1d.num_embeddings == 64 + assert dist_embedding_1d.embedding_dim == 32 + + # ensure state dict is reversibly loadable + embedding.load_state_dict(dist_embedding_1d.state_dict()) + dist_embedding_1d.load_state_dict(embedding.state_dict()) + + # check embedding correctness + x = torch.randint(0, 128, (4, 32)).to('cuda') + org_out = embedding(x) + dist_out = dist_embedding_1d(x) + assert_close(org_out, dist_out) + + # check backward correctness + org_out.sum().backward() + dist_out.sum().backward() + + rank = dist.get_rank() + target_grad = torch.chunk(embedding.weight.grad, 2, dim=0)[rank] + assert_close(target_grad, dist_embedding_1d.weight.grad) + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + check_vocab_embedding_1d() + + +@rerun_if_address_is_in_use() +def test_vocab_embedding(): + spawn(run_dist, nprocs=2) + + +if __name__ == '__main__': + test_vocab_embedding() diff --git a/tests/test_shardformer/test_model/__init__.py b/tests/test_shardformer/test_model/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_shardformer/test_model/_utils.py b/tests/test_shardformer/test_model/_utils.py new file mode 100644 index 000000000000..d83d9ecd39e0 --- /dev/null +++ b/tests/test_shardformer/test_model/_utils.py @@ -0,0 +1,35 @@ +import copy + +from colossalai.shardformer import ShardConfig, ShardFormer + + +def build_model(model_fn, enable_fused_normalization=True, enable_tensor_parallelism=True): + # create new model + org_model = model_fn().cuda() + + # shard model + shard_config = ShardConfig(enable_fused_normalization=enable_fused_normalization, + enable_tensor_parallelism=enable_tensor_parallelism) + model_copy = copy.deepcopy(org_model) + shard_former = ShardFormer(shard_config=shard_config) + sharded_model = shard_former.optimize(model_copy).cuda() + return org_model, sharded_model + + +def run_forward(original_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # prepare input + data = data_gen_fn() + data = {k: v.cuda() for k, v in data.items()} + + # switch to train mode + original_model.train() + sharded_model.train() + # run forward + org_output = original_model(**data) + org_output = output_transform_fn(org_output) + org_loss = loss_fn(org_output) + + shard_output = sharded_model(**data) + shard_output = output_transform_fn(shard_output) + shard_loss = loss_fn(shard_output) + return org_output, org_loss, shard_output, shard_loss diff --git a/tests/test_shardformer/test_model/test_shard_bert.py b/tests/test_shardformer/test_model/test_shard_bert.py new file mode 100644 index 000000000000..1afedb7079ea --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_bert.py @@ -0,0 +1,95 @@ +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # check forward + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output) + + # do backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + + # check grad + + if org_model.__class__.__name__ == 'BertModel': + bert = org_model + sharded_bert = sharded_model + else: + bert = org_model.bert + sharded_bert = sharded_model.bert + + # compare self attention grad + org_grad = bert.encoder.layer[0].attention.self.query.weight.grad + shard_grad = sharded_bert.encoder.layer[0].attention.self.query.weight.grad + shard_weight = sharded_bert.encoder.layer[0].attention.self.query.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + # compare embedding grad + org_grad = bert.embeddings.word_embeddings.weight.grad + shard_grad = sharded_bert.embeddings.word_embeddings.weight.grad + shard_weight = sharded_bert.embeddings.word_embeddings.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_bert_test(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_bert') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + + torch.cuda.empty_cache() + + +def check_bert(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_bert_test() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_bert(): + spawn(check_bert, 2) + + +if __name__ == "__main__": + test_bert() diff --git a/tests/test_shardformer/test_model/test_shard_bloom.py b/tests/test_shardformer/test_model/test_shard_bloom.py new file mode 100644 index 000000000000..a3389652269c --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_bloom.py @@ -0,0 +1,94 @@ +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # check forward + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output, ignore_keys=['past_key_values']) + + # do backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + + # unwrap model + if org_model.__class__.__name__ == 'BloomModel': + bloom = org_model + sharded_bloom = sharded_model + else: + bloom = org_model.transformer + sharded_bloom = sharded_model.transformer + + # check attention grad + org_grad = bloom.h[0].self_attention.query_key_value.weight.grad + shard_grad = sharded_bloom.h[0].self_attention.query_key_value.weight.grad + shard_weight = sharded_bloom.h[0].self_attention.query_key_value.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + # check embedding weights + org_grad = bloom.word_embeddings.weight.grad + shard_grad = sharded_bloom.word_embeddings.weight.grad + shard_weight = sharded_bloom.word_embeddings.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_bloom_test(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_bloom') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + torch.cuda.empty_cache() + + +def check_bloom(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_bloom_test() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_bloom(): + spawn(check_bloom, 2) + + +if __name__ == "__main__": + test_bloom() diff --git a/tests/test_shardformer/test_model/test_shard_gpt2.py b/tests/test_shardformer/test_model/test_shard_gpt2.py new file mode 100644 index 000000000000..ee7737687d99 --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_gpt2.py @@ -0,0 +1,94 @@ +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # check forward + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output, ignore_keys=['past_key_values']) + + # do backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to origin model loss\n{org_loss}\n{shard_loss}" + + # unwrap model + if org_model.__class__.__name__ == 'GPT2Model': + org_model = org_model + sharded_model = sharded_model + else: + org_model = org_model.transformer + sharded_model = sharded_model.transformer + + # check mlp grad + org_grad = org_model.h[0].mlp.c_fc.weight.grad + shard_grad = sharded_model.h[0].mlp.c_fc.weight.grad + shard_weight = sharded_model.h[0].mlp.c_fc.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=1) + else: + all_shard_grad = shard_grad + assert torch.allclose( + org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to origin model grad\n{org_grad}\n{all_shard_grad}" + + # check embedding weights + org_grad = org_model.wte.weight.grad + shard_grad = sharded_model.wte.weight.grad + shard_weight = sharded_model.wte.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose( + org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to origin model grad\n{org_grad}\n{all_shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_gpt2_test(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_gpt') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + torch.cuda.empty_cache() + + +def check_gpt2(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_gpt2_test() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_gpt2(): + spawn(check_gpt2, 2) + + +if __name__ == "__main__": + test_gpt2() diff --git a/tests/test_shardformer/test_model/test_shard_llama.py b/tests/test_shardformer/test_model/test_shard_llama.py new file mode 100644 index 000000000000..74b5fdd18af8 --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_llama.py @@ -0,0 +1,97 @@ +import os + +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + +os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true' + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + + # forward check + assert_hf_output_close(org_output, shard_output, ignore_keys=['past_key_values'], rtol=1e-4) + + # run backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + + # unwrap model + if hasattr(org_model, 'model'): + llama_model = org_model.model + shard_llama_model = sharded_model.model + else: + llama_model = org_model + shard_llama_model = sharded_model + + # check attention grad + org_grad = llama_model.layers[0].self_attn.q_proj.weight.grad + shard_grad = shard_llama_model.layers[0].self_attn.q_proj.weight.grad + shard_weight = shard_llama_model.layers[0].self_attn.q_proj.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}" + + # check embedding grad + org_grad = llama_model.embed_tokens.weight.grad + shard_grad = shard_llama_model.embed_tokens.weight.grad + shard_weight = shard_llama_model.embed_tokens.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_gpt2_llama(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_llama') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + torch.cuda.empty_cache() + + +def check_llama(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_gpt2_llama() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_llama(): + spawn(check_llama, 4) + + +if __name__ == "__main__": + test_llama() diff --git a/tests/test_shardformer/test_model/test_shard_opt.py b/tests/test_shardformer/test_model/test_shard_opt.py new file mode 100644 index 000000000000..25bccb13b1a8 --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_opt.py @@ -0,0 +1,96 @@ +import copy +import os + +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + +os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true' + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output, ignore_keys=['past_key_values'], rtol=1e-4) + + # run backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + + # unwrap model + if hasattr(org_model, 'model'): + opt_model = org_model.model + shard_opt_model = sharded_model.model + else: + opt_model = org_model + shard_opt_model = sharded_model + + # check attention grad + org_grad = opt_model.decoder.layers[0].self_attn.q_proj.weight.grad + shard_grad = shard_opt_model.decoder.layers[0].self_attn.q_proj.weight.grad + shard_weight = shard_opt_model.decoder.layers[0].self_attn.q_proj.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + # check embedding grad + org_grad = opt_model.decoder.embed_tokens.weight.grad + shard_grad = shard_opt_model.decoder.embed_tokens.weight.grad + shard_weight = shard_opt_model.decoder.embed_tokens.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(4)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_t5_test(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_opt') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + torch.cuda.empty_cache() + + +def check_OPTModel(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_t5_test() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_OPTModel(): + spawn(check_OPTModel, 4) + + +if __name__ == '__main__': + test_OPTModel() diff --git a/tests/test_shardformer/test_model/test_shard_t5.py b/tests/test_shardformer/test_model/test_shard_t5.py new file mode 100644 index 000000000000..0762dc09e5af --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_t5.py @@ -0,0 +1,107 @@ +import os + +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor +from colossalai.testing import ( + assert_hf_output_close, + clear_cache_before_run, + parameterize, + rerun_if_address_is_in_use, + spawn, +) +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # check forward + # the value "past_key_values" is sharded, so we ignore + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output, ignore_keys=['past_key_values']) + + # do backward + org_loss.backward() + shard_loss.backward() + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + + # check attention grad + org_grad = org_model.encoder.block[0].layer[0].SelfAttention.q.weight.grad + shard_grad = sharded_model.encoder.block[0].layer[0].SelfAttention.q.weight.grad + shard_weight = sharded_model.encoder.block[0].layer[0].SelfAttention.q.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{shard_grad}" + + # check self attention embed + org_grad = org_model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias.weight.grad + shard_grad = sharded_model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias.weight.grad + shard_weight = sharded_model.encoder.block[0].layer[0].SelfAttention.relative_attention_bias.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=1) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + # check token embedding grad + org_grad = org_model.shared.weight.grad + + # check weights are tied + if hasattr(org_model, 'lm_head'): + assert org_model.shared.weight.data.data_ptr() == org_model.lm_head.weight.data.data_ptr() + assert sharded_model.shared.weight.data.data_ptr() == sharded_model.lm_head.weight.data.data_ptr() + + shard_grad = sharded_model.shared.weight.grad + shard_weight = sharded_model.shared.weight + + if is_distributed_tensor(shard_weight) or is_customized_distributed_tensor(shard_weight): + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + else: + all_shard_grad = shard_grad + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + +@parameterize('enable_fused_normalization', [True, False]) +@parameterize('enable_tensor_parallelism', [True, False]) +def run_t5_test(enable_fused_normalization, enable_tensor_parallelism): + sub_model_zoo = model_zoo.get_sub_registry('transformers_t5') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + torch.cuda.empty_cache() + + +def check_t5(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + run_t5_test() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_t5(): + spawn(check_t5, 2) + + +if __name__ == "__main__": + test_t5() diff --git a/tests/test_shardformer/test_model/test_shard_vit.py b/tests/test_shardformer/test_model/test_shard_vit.py new file mode 100644 index 000000000000..af1605b6b659 --- /dev/null +++ b/tests/test_shardformer/test_model/test_shard_vit.py @@ -0,0 +1,56 @@ +import pytest +import torch + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.testing import assert_hf_output_close, clear_cache_before_run, rerun_if_address_is_in_use, spawn +from tests.kit.model_zoo import model_zoo +from tests.test_shardformer.test_model._utils import build_model, run_forward + + +def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): + # check forward + org_output, org_loss, shard_output, shard_loss = run_forward(org_model, sharded_model, data_gen_fn, + output_transform_fn, loss_fn) + assert_hf_output_close(org_output, shard_output) + + # do backward + org_loss.backward() + shard_loss.backward() + + # check grad + org_grad = org_model.encoder.layer[0].attention.attention.query.weight.grad + shard_grad = sharded_model.encoder.layer[0].attention.attention.query.weight.grad + + shard_grad_list = [torch.zeros([*shard_grad.shape]).to('cuda') for _ in range(2)] + shard_grad = torch.distributed.all_gather(shard_grad_list, shard_grad) + all_shard_grad = torch.cat(shard_grad_list, dim=0) + + assert torch.allclose(org_loss, shard_loss, + atol=1e-5), f"shard model loss is not equal to orgin model loss\n{org_loss}\n{shard_loss}" + assert torch.allclose(org_grad, all_shard_grad, + atol=1e-5), f"shard model grad is not equal to orgin model grad\n{org_grad}\n{all_shard_grad}" + + +def check_vit(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + sub_model_zoo = model_zoo.get_sub_registry('transformers_vit') + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + org_model, sharded_model = build_model(world_size, model_fn) + check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn) + + torch.cuda.empty_cache() + + +@pytest.mark.dist +@pytest.mark.skip +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_vit(): + spawn(check_vit, 4) + + +if __name__ == "__main__": + test_vit() diff --git a/tests/test_shardformer/test_with_torch_ddp.py b/tests/test_shardformer/test_with_torch_ddp.py new file mode 100644 index 000000000000..9f8a5db6c94f --- /dev/null +++ b/tests/test_shardformer/test_with_torch_ddp.py @@ -0,0 +1,77 @@ +import pytest +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +import colossalai +from colossalai.cluster import DistCoordinator +from colossalai.logging import disable_existing_loggers +from colossalai.shardformer import ShardConfig, ShardFormer +from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn +from tests.kit.model_zoo import model_zoo + + +def check_shardformer_with_ddp(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + sub_model_zoo = model_zoo.get_sub_registry('transformers_gpt') + + # create shardformer + # ranks: [0, 1, 2, 3] + # tp ranks = [0, 1], [2, 3] + # dp ranks = [0, 2], [1, 3] + dp_process_group_1 = dist.new_group([0, 2]) + dp_process_group_2 = dist.new_group([1, 3]) + tp_process_group_1 = dist.new_group([0, 1]) + tp_process_group_2 = dist.new_group([2, 3]) + + coordinator = DistCoordinator() + + if coordinator.rank in [0, 1]: + tp_process_group = tp_process_group_1 + else: + tp_process_group = tp_process_group_2 + + if coordinator.rank in [0, 2]: + dp_process_group = dp_process_group_1 + else: + dp_process_group = dp_process_group_2 + + shard_config = ShardConfig(tensor_parallel_process_group=tp_process_group, enable_fused_normalization=True) + shardformer = ShardFormer(shard_config=shard_config) + + for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): + # create and shard model + model = model_fn().cuda() + sharded_model = shardformer.optimize(model) + + # add ddp + sharded_ddp_model = DDP(sharded_model, process_group=dp_process_group) + + # prepare input + data = data_gen_fn() + data = {k: v.cuda() for k, v in data.items()} + + # switch to train mode + sharded_ddp_model.train() + + # run forward + output = sharded_ddp_model(**data) + loss = loss_fn(output) + + # backward + loss.backward() + torch.cuda.empty_cache() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_gpt2(): + spawn(check_shardformer_with_ddp, 4) + + +if __name__ == "__main__": + test_gpt2() + test_gpt2() diff --git a/tests/test_tensor/core/test_dist_spec_mgr.py b/tests/test_tensor/core/test_dist_spec_mgr.py index e02f4e7977f6..89476a35b63a 100644 --- a/tests/test_tensor/core/test_dist_spec_mgr.py +++ b/tests/test_tensor/core/test_dist_spec_mgr.py @@ -1,13 +1,12 @@ import math + +import pytest import torch import torch.distributed as dist -import pytest + import colossalai -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.tensor import DistSpecManager, ProcessGroup, ShardSpec, ReplicaSpec -from functools import partial +from colossalai.tensor import DistSpecManager, ProcessGroup, ReplicaSpec, ShardSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn def run(): @@ -58,8 +57,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_dist_spec_mgr(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/core/test_tensor.py b/tests/test_tensor/core/test_tensor.py index b48d9e9a2dfa..64d198b350a8 100644 --- a/tests/test_tensor/core/test_tensor.py +++ b/tests/test_tensor/core/test_tensor.py @@ -1,17 +1,11 @@ -import torch import pytest -from colossalai.tensor import ColoTensor +import torch from numpy import allclose import colossalai -from colossalai.utils import free_port -from colossalai.tensor import ColoTensorSpec from colossalai.core import global_context as gpc -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.tensor import distspec, ColoTensor, ProcessGroup, ShardSpec, ReplicaSpec -from functools import partial +from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup, ReplicaSpec, ShardSpec, distspec +from colossalai.testing import rerun_if_address_is_in_use, spawn def _run_tensor_indexing(): @@ -152,8 +146,7 @@ def run_dist_tests(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_dist_cases(world_size): - run_func = partial(run_dist_tests, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist_tests, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/model/test_gpt2.py b/tests/test_tensor/model/test_gpt2.py index ad8ac87b2e1e..337bfa840d5d 100644 --- a/tests/test_tensor/model/test_gpt2.py +++ b/tests/test_tensor/model/test_gpt2.py @@ -1,17 +1,13 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP import colossalai from colossalai.nn.parallel.data_parallel import ColoDDP from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import ( debug_print, @@ -145,8 +141,7 @@ def run_dist(rank, world_size, port, use_ddp): @pytest.mark.parametrize('use_ddp', [False, True]) @rerun_if_address_is_in_use() def test_gpt(world_size, use_ddp): - run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size, use_ddp=use_ddp) if __name__ == '__main__': diff --git a/tests/test_tensor/model/test_model.py b/tests/test_tensor/model/test_model.py index 3f53b94e0642..288bd20e3844 100644 --- a/tests/test_tensor/model/test_model.py +++ b/tests/test_tensor/model/test_model.py @@ -1,17 +1,13 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.nn.optimizer import ColossalaiOptimizer from colossalai.tensor import ColoTensor, ProcessGroup from colossalai.tensor.colo_parameter import ColoParameter -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import free_port, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import ( check_equal, @@ -313,8 +309,7 @@ def run_model_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_model(world_size): - run_func = partial(run_model_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_model_dist, world_size) def run_pretrain_load_dist(rank, world_size, port): @@ -329,12 +324,11 @@ def run_pretrain_load_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_pretrain_load(world_size): - run_func = partial(run_pretrain_load_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_pretrain_load_dist, world_size) if __name__ == '__main__': # test_model_parameters() - # test_colo_optgimizer() + # test_colo_optimizer() test_model(4) # test_pretrain_load(4) diff --git a/tests/test_tensor/model/test_module_spec.py b/tests/test_tensor/model/test_module_spec.py index 997b416f12c3..b50851e5eaf2 100644 --- a/tests/test_tensor/model/test_module_spec.py +++ b/tests/test_tensor/model/test_module_spec.py @@ -1,9 +1,7 @@ from copy import deepcopy -from functools import partial import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.nn.parallel.layers import check_colo_module, init_colo_module @@ -17,10 +15,9 @@ ShardSpec, distspec, ) -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import set_seed, tensor_equal, tensor_shard_equal @@ -207,8 +204,7 @@ def run_dist_check(rank, world_size, port): @pytest.mark.skip("for higher testing speed") @rerun_if_address_is_in_use() def test_module_linear_1d(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) @pytest.mark.dist @@ -216,8 +212,7 @@ def test_module_linear_1d(world_size): @pytest.mark.skip("for higher testing speed") @rerun_if_address_is_in_use() def test_module_model(world_size): - run_func = partial(run_dist_model, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist_model, world_size) @pytest.mark.dist @@ -225,8 +220,7 @@ def test_module_model(world_size): @pytest.mark.skip("for higher testing speed") @rerun_if_address_is_in_use() def test_module_check(world_size): - run_func = partial(run_dist_check, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist_check, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_colo_checkpoint_tools.py b/tests/test_tensor/test_colo_checkpoint_tools.py index aa333d55276c..a53a3f37a664 100644 --- a/tests/test_tensor/test_colo_checkpoint_tools.py +++ b/tests/test_tensor/test_colo_checkpoint_tools.py @@ -1,47 +1,41 @@ -import torch -import pytest -from functools import partial - -import torch.multiprocessing as mp -import torch.distributed as dist - -import colossalai -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils.cuda import get_current_device -from colossalai.utils import free_port -from colossalai.tensor import ComputePattern, ComputeSpec, ColoTensor, ShardSpec, ProcessGroup, ColoTensorSpec -from colossalai.utils.checkpoint.utils import gather_tensor, scatter_tensor -from tests.test_tensor.common_utils import tensor_shard_equal - - -def run_dist(rank, world_size, port, dp_degree, tp_degree): - colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - pg = ProcessGroup(dp_degree=dp_degree, tp_degree=tp_degree) - x = torch.randn(4, 4) - param = ColoTensor(torch.nn.Parameter(x), spec=ColoTensorSpec(pg)) - spec = ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D) - param.set_tensor_spec(*spec) - - gather_tensor(param) - if dist.get_rank() == 0: - assert torch.all(x == param) - else: - assert tensor_shard_equal(x, param.data, pg.tp_local_rank(), pg.tp_world_size()) - dist.barrier() - - scatter_tensor(param, spec[0]) - assert tensor_shard_equal(x, param.data, pg.tp_local_rank(), pg.tp_world_size()) - assert param.requires_grad is True - dist.barrier() - - -@pytest.mark.dist -@pytest.mark.parametrize('world_size', [4]) -@rerun_if_address_is_in_use() -def test_checkpoint(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port(), dp_degree=2, tp_degree=world_size // 2) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_checkpoint(world_size=4) +import pytest +import torch +import torch.distributed as dist + +import colossalai +from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.checkpoint.utils import gather_tensor, scatter_tensor +from tests.test_tensor.common_utils import tensor_shard_equal + + +def run_dist(rank, world_size, port, dp_degree, tp_degree): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + pg = ProcessGroup(dp_degree=dp_degree, tp_degree=tp_degree) + x = torch.randn(4, 4) + param = ColoTensor(torch.nn.Parameter(x), spec=ColoTensorSpec(pg)) + spec = ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D) + param.set_tensor_spec(*spec) + + gather_tensor(param) + if dist.get_rank() == 0: + assert torch.all(x == param) + else: + assert tensor_shard_equal(x, param.data, pg.tp_local_rank(), pg.tp_world_size()) + dist.barrier() + + scatter_tensor(param, spec[0]) + assert tensor_shard_equal(x, param.data, pg.tp_local_rank(), pg.tp_world_size()) + assert param.requires_grad is True + dist.barrier() + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [4]) +@rerun_if_address_is_in_use() +def test_checkpoint(world_size): + spawn(run_dist, world_size, dp_degree=2, tp_degree=world_size // 2) + + +if __name__ == '__main__': + test_checkpoint(world_size=4) diff --git a/tests/test_tensor/test_comm_spec_apply.py b/tests/test_tensor/test_comm_spec_apply.py index 46eee61f1ecf..2c68633aabc8 100644 --- a/tests/test_tensor/test_comm_spec_apply.py +++ b/tests/test_tensor/test_comm_spec_apply.py @@ -1,10 +1,5 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from torch.distributed import ReduceOp from colossalai.core import global_context as gpc from colossalai.device.device_mesh import DeviceMesh @@ -12,8 +7,7 @@ from colossalai.logging import disable_existing_loggers from colossalai.tensor.shape_consistency import CollectiveCommPattern, CommSpec from colossalai.tensor.sharding_spec import ShardingSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn def check_all_gather(device_mesh, rank): @@ -218,8 +212,7 @@ def check_comm(rank, world_size, port): @rerun_if_address_is_in_use() def test_comm_spec(): world_size = 4 - run_func = partial(check_comm, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_comm, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_context.py b/tests/test_tensor/test_context.py index 2f7aebed5bc4..45def034ba8e 100644 --- a/tests/test_tensor/test_context.py +++ b/tests/test_tensor/test_context.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai from colossalai.tensor import ( @@ -14,10 +11,9 @@ ReplicaSpec, ShardSpec, ) -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import set_seed @@ -61,8 +57,7 @@ def run_colo_init_context(rank: int, world_size: int, port: int): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_colo_init_context(world_size): - run_func = partial(run_colo_init_context, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_colo_init_context, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_dtensor/test_comm_spec.py b/tests/test_tensor/test_dtensor/test_comm_spec.py index 547a96b264dc..95fcd2aaf8f3 100644 --- a/tests/test_tensor/test_dtensor/test_comm_spec.py +++ b/tests/test_tensor/test_dtensor/test_comm_spec.py @@ -1,19 +1,12 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from torch.distributed import ReduceOp from colossalai.core import global_context as gpc from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.tensor.d_tensor.comm_spec import CollectiveCommPattern, CommSpec -from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn def check_all_gather(process_groups_dict, rank): @@ -129,23 +122,6 @@ def check_all_reduce_bwd(process_groups_dict, rank): assert tensor_to_comm.equal(tensor_to_check) -def check_all_reduce_in_flatten_device_mesh(process_groups_dict, rank): - # tensor to comm - tensor_to_comm = torch.ones(2, 2).cuda() * rank - - # reduce through logical process axis 0 at flatten device mesh - # tensor to check - # tensor([[6., 6.], - # [6., 6.]]) - tensor_to_check = torch.tensor([[6, 6], [6, 6]], dtype=tensor_to_comm.dtype).cuda() - - # CommSpec:(comm_pattern:all_reduce, logical_process_axis:[0, 1]) - comm_spec = CommSpec(CollectiveCommPattern.ALLREDUCE_FWD_IDENTITY_BWD, process_groups_dict, logical_process_axis=0) - tensor_to_comm = comm_spec.covert_spec_to_action(tensor_to_comm) - - assert tensor_to_comm.equal(tensor_to_check) - - def check_comm(rank, world_size, port): disable_existing_loggers() launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') @@ -157,24 +133,22 @@ def check_comm(rank, world_size, port): # [[0, 1, # [2, 3]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) - process_groups_dict = device_mesh.process_groups_dict + + process_group_dict = device_mesh._process_group_dict[rank] # test all gather - check_all_gather(process_groups_dict, rank) + check_all_gather(process_group_dict, rank) # test shard - check_shard(process_groups_dict, rank) + check_shard(process_group_dict, rank) # test all to all - check_all_to_all(process_groups_dict, rank) + check_all_to_all(process_group_dict, rank) # test all reduce - check_all_reduce_fwd(process_groups_dict, rank) - check_all_reduce_bwd(process_groups_dict, rank) + check_all_reduce_fwd(process_group_dict, rank) + check_all_reduce_bwd(process_group_dict, rank) - flatten_process_groups_dict = device_mesh.flatten_device_mesh.process_groups_dict - # test all reduce in 1D flatten device mesh - check_all_reduce_in_flatten_device_mesh(flatten_process_groups_dict, rank) gpc.destroy() @@ -182,8 +156,7 @@ def check_comm(rank, world_size, port): @rerun_if_address_is_in_use() def test_comm_spec(): world_size = 4 - run_func = partial(check_comm, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_comm, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_dtensor/test_dtensor.py b/tests/test_tensor/test_dtensor/test_dtensor.py index a99ac6e41c5e..5a1aef79f332 100644 --- a/tests/test_tensor/test_dtensor/test_dtensor.py +++ b/tests/test_tensor/test_dtensor/test_dtensor.py @@ -1,15 +1,10 @@ -from functools import partial - import torch -import torch.multiprocessing as mp from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers -from colossalai.tensor.d_tensor.d_tensor import DTensor, distribute_tensor -from colossalai.tensor.d_tensor.layout import Layout -from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec -from colossalai.utils import free_port +from colossalai.tensor.d_tensor import ShardingSpec, distribute_tensor, get_global_shape, redistribute, to_global +from colossalai.testing import rerun_if_address_is_in_use, spawn class TestModel(torch.nn.Module): @@ -34,22 +29,18 @@ def check_dtensor(rank, world_size, port): device_mesh = DeviceMesh(torch.Tensor([0, 1, 2, 3]), (2, 2), init_process_group=True) target_sharding_spec = ShardingSpec(dim_size=original_tensor.dim(), dim_partition_dict={0: [0]}) - layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=target_sharding_spec, - entire_shape=original_tensor.shape) - d_tensor = DTensor(original_tensor, layout) + d_tensor = distribute_tensor(original_tensor, device_mesh, target_sharding_spec) - assert d_tensor.entire_shape == original_tensor.shape - assert d_tensor.data_type == original_tensor.dtype + assert get_global_shape(d_tensor) == original_tensor.shape + assert d_tensor.dtype == original_tensor.dtype if rank in (0, 1): - assert d_tensor.to_local().equal(original_tensor.narrow(0, 0, 2)) + assert d_tensor.equal(original_tensor.narrow(0, 0, 2)) elif rank in (2, 3): - assert d_tensor.to_local().equal(original_tensor.narrow(0, 2, 2)) + assert d_tensor.equal(original_tensor.narrow(0, 2, 2)) else: raise ValueError(f'rank {rank} is not in the device mesh') - assert d_tensor.to_global().equal(original_tensor) + assert to_global(d_tensor).equal(original_tensor) output = test_model(d_tensor) if rank in (0, 1): @@ -60,42 +51,37 @@ def check_dtensor(rank, world_size, port): raise ValueError(f'rank {rank} is not in the device mesh') new_sharding_spec = ShardingSpec(dim_size=original_tensor.dim(), dim_partition_dict={0: [0, 1]}) - new_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=new_sharding_spec, - entire_shape=original_tensor.shape) - - d_tensor.layout_convert(new_layout) + d_tensor = redistribute(d_tensor, device_mesh, new_sharding_spec) if rank == 0: - assert d_tensor.local_tensor.equal(original_tensor.narrow(0, 0, 1)) + assert d_tensor.equal(original_tensor.narrow(0, 0, 1)) elif rank == 1: - assert d_tensor.local_tensor.equal(original_tensor.narrow(0, 1, 1)) + assert d_tensor.equal(original_tensor.narrow(0, 1, 1)) elif rank == 2: - assert d_tensor.local_tensor.equal(original_tensor.narrow(0, 2, 1)) + assert d_tensor.equal(original_tensor.narrow(0, 2, 1)) elif rank == 3: - assert d_tensor.local_tensor.equal(original_tensor.narrow(0, 3, 1)) + assert d_tensor.equal(original_tensor.narrow(0, 3, 1)) else: raise ValueError(f'rank {rank} is not in the device mesh') - dtensor_from_local = distribute_tensor(original_tensor, new_layout) + dtensor_from_local = distribute_tensor(original_tensor, device_mesh, new_sharding_spec) if rank == 0: - assert dtensor_from_local.local_tensor.equal(original_tensor.narrow(0, 0, 1)) + assert dtensor_from_local.equal(original_tensor.narrow(0, 0, 1)) elif rank == 1: - assert dtensor_from_local.local_tensor.equal(original_tensor.narrow(0, 1, 1)) + assert dtensor_from_local.equal(original_tensor.narrow(0, 1, 1)) elif rank == 2: - assert dtensor_from_local.local_tensor.equal(original_tensor.narrow(0, 2, 1)) + assert dtensor_from_local.equal(original_tensor.narrow(0, 2, 1)) elif rank == 3: - assert dtensor_from_local.local_tensor.equal(original_tensor.narrow(0, 3, 1)) + assert dtensor_from_local.equal(original_tensor.narrow(0, 3, 1)) else: raise ValueError(f'rank {rank} is not in the device mesh') +@rerun_if_address_is_in_use() def test_dtensor(): world_size = 4 - run_func = partial(check_dtensor, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_dtensor, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_dtensor/test_layout_converter.py b/tests/test_tensor/test_dtensor/test_layout_converter.py index 70cf8726dbd0..5388fd901e09 100644 --- a/tests/test_tensor/test_dtensor/test_layout_converter.py +++ b/tests/test_tensor/test_dtensor/test_layout_converter.py @@ -1,9 +1,7 @@ import math -from functools import partial import pytest import torch -import torch.multiprocessing as mp from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch @@ -11,13 +9,12 @@ from colossalai.tensor.d_tensor.comm_spec import CollectiveCommPattern from colossalai.tensor.d_tensor.layout import Layout from colossalai.tensor.d_tensor.layout_converter import LayoutConverter -from colossalai.tensor.d_tensor.sharding_spec import DimSpec, ShardingSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn -entire_shape = torch.Size((64, 32, 16)) +global_shape = torch.Size((64, 32, 16)) layout_converter = LayoutConverter() -physical_mesh_id = torch.arange(0, 4).reshape(2, 2) +physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) @@ -33,10 +30,7 @@ def check_one_step_transform(rank, world_size, port): # shard_sequence: S0,S1,R # device_mesh_shape: (2, 2) sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict) - layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec, - entire_shape=entire_shape) + layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec, global_shape=global_shape) rst_dict = layout_converter.all_gather_transform_layouts(layout) @@ -52,10 +46,7 @@ def check_one_step_transform(rank, world_size, port): # shard_sequence: S0,S1,R # device_mesh_shape: (4, 4) sharding_spec_all2all = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict_all2all) - layout_all2all = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_all2all, - entire_shape=entire_shape) + layout_all2all = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_all2all, global_shape=global_shape) rst_dict_all2all = layout_converter.all_to_all_transform_layout(layout_all2all) @@ -74,10 +65,7 @@ def check_one_step_transform(rank, world_size, port): # shard_sequence: S0,R,R # device_mesh_shape: (4, 4) sharding_spec_shard = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_shard) - shard_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_shard, - entire_shape=entire_shape) + shard_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_shard, global_shape=global_shape) rst_dict_shard = layout_converter.shard_transform_layout(shard_layout) @@ -103,19 +91,13 @@ def check_layout_converting(rank, world_size, port): # shard_sequence: R,S01,R # device_mesh_shape: (4, 4) sharding_spec_source = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_source) - source_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_source, - entire_shape=entire_shape) + source_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_source, global_shape=global_shape) # DistSpec: # shard_sequence: S01,R,R # device_mesh_shape: (4, 4) sharding_spec_target = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_target) - target_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_target, - entire_shape=entire_shape) + target_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_target, global_shape=global_shape) transform_path, comm_action_sequence = layout_converter.layout_converting(source_layout, target_layout) @@ -162,21 +144,15 @@ def check_layout_converting_apply(rank, world_size, port): # shard_sequence: R,S01,R # device_mesh_shape: (4, 4) sharding_spec_source = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_source) - source_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_source, - entire_shape=entire_shape) + source_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_source, global_shape=global_shape) # DistSpec: # shard_sequence: S01,R,R # device_mesh_shape: (4, 4) sharding_spec_target = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_target) - target_layout = Layout(device_mesh=device_mesh, - device_type=torch.device('cuda'), - sharding_spec=sharding_spec_target, - entire_shape=entire_shape) + target_layout = Layout(device_mesh=device_mesh, sharding_spec=sharding_spec_target, global_shape=global_shape) - original_tensor = torch.rand(entire_shape).cuda() + original_tensor = torch.rand(global_shape).cuda() # tensor_to_apply: [R, S01, R] tensor_to_apply = original_tensor.narrow(1, rank * 8, 8) @@ -192,14 +168,9 @@ def check_layout_converting_apply(rank, world_size, port): @rerun_if_address_is_in_use() def test_layout_converter(): world_size = 4 - run_func = partial(check_one_step_transform, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - run_func = partial(check_layout_converting, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - run_func = partial(check_layout_converting_apply, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_one_step_transform, world_size) + spawn(check_layout_converting, world_size) + spawn(check_layout_converting_apply, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_mix_gather.py b/tests/test_tensor/test_mix_gather.py index c1ab30601501..9122808eb5a3 100644 --- a/tests/test_tensor/test_mix_gather.py +++ b/tests/test_tensor/test_mix_gather.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from colossalai.core import global_context as gpc from colossalai.device.device_mesh import DeviceMesh @@ -11,7 +8,7 @@ from colossalai.tensor.shape_consistency import CollectiveCommPattern, CommSpec from colossalai.tensor.sharding_spec import ShardingSpec from colossalai.tensor.utils import mix_gather_simulator -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn def check_mix_gather_S0S1(device_mesh, rank): @@ -323,10 +320,10 @@ def check_comm(rank, world_size, port): @pytest.mark.skip(reason="Skip because the check functions assume 8 GPUS but CI only have 4 GPUs") +@rerun_if_address_is_in_use() def test_mix_gather(): world_size = 8 - run_func = partial(check_comm, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_comm, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_parameter.py b/tests/test_tensor/test_parameter.py index 7c3c4b2132e4..9c3f05da1ffa 100644 --- a/tests/test_tensor/test_parameter.py +++ b/tests/test_tensor/test_parameter.py @@ -1,9 +1,10 @@ -from colossalai.tensor import ColoParameter, ColoTensor, ColoTensorSpec, ProcessGroup -import torch import pytest +import torch from common_utils import tensor_equal + import colossalai -from colossalai.utils import free_port +from colossalai.tensor import ColoParameter, ColoTensor, ColoTensorSpec, ProcessGroup +from colossalai.testing import free_port @pytest.mark.skip diff --git a/tests/test_tensor/test_shape_consistency.py b/tests/test_tensor/test_shape_consistency.py index 6fe9ee292cd0..859eef051256 100644 --- a/tests/test_tensor/test_shape_consistency.py +++ b/tests/test_tensor/test_shape_consistency.py @@ -1,9 +1,10 @@ -from colossalai.tensor.shape_consistency import ShapeConsistencyManager, CollectiveCommPattern import torch -from colossalai.tensor.sharding_spec import _DimSpec, ShardingSpec + from colossalai.device.device_mesh import DeviceMesh +from colossalai.tensor.shape_consistency import CollectiveCommPattern, ShapeConsistencyManager +from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec -physical_mesh_id = torch.arange(0, 16).reshape(2, 8) +physical_mesh_id = torch.arange(0, 16) mesh_shape = (4, 4) # [[0, 1, 2, 3], # [4, 5, 6, 7], diff --git a/tests/test_tensor/test_shape_consistency_apply.py b/tests/test_tensor/test_shape_consistency_apply.py index 4c838bc83fad..b57952df401f 100644 --- a/tests/test_tensor/test_shape_consistency_apply.py +++ b/tests/test_tensor/test_shape_consistency_apply.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from colossalai.device.device_mesh import DeviceMesh from colossalai.initialize import launch from colossalai.logging import disable_existing_loggers from colossalai.tensor.shape_consistency import CollectiveCommPattern, ShapeConsistencyManager from colossalai.tensor.sharding_spec import ShardingSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn def check_apply(rank, world_size, port): @@ -73,8 +69,7 @@ def check_apply(rank, world_size, port): @rerun_if_address_is_in_use() def test_apply(): world_size = 4 - run_func = partial(check_apply, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(check_apply, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_sharded_linear.py b/tests/test_tensor/test_sharded_linear.py index 85008c67a9c2..9bd9805e9b8f 100644 --- a/tests/test_tensor/test_sharded_linear.py +++ b/tests/test_tensor/test_sharded_linear.py @@ -1,8 +1,5 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn.functional as F import colossalai @@ -10,8 +7,7 @@ from colossalai.nn._ops._utils import gather_forward_split_backward from colossalai.tensor import ColoParameter, ColoTensor, ProcessGroup from colossalai.tensor.sharding_spec import ShardingSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import rerun_if_address_is_in_use, spawn def run_dist(rank, world_size, port): @@ -30,7 +26,7 @@ def run_dist(rank, world_size, port): # the mesh is in the following topo # [[0, 1], # [2, 3]] - physical_mesh_id = torch.arange(0, 4).reshape(2, 2) + physical_mesh_id = torch.arange(0, 4) mesh_shape = (2, 2) device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) row_id = rank // 2 @@ -229,8 +225,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [4]) @rerun_if_address_is_in_use() def test_sharded_mlp(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_tensor/test_sharding_spec.py b/tests/test_tensor/test_sharding_spec.py index 909c84ef0f0e..5007c4141849 100644 --- a/tests/test_tensor/test_sharding_spec.py +++ b/tests/test_tensor/test_sharding_spec.py @@ -5,7 +5,7 @@ def test_sharding_spec(): - physical_mesh_id = torch.arange(0, 16).reshape(2, 8) + physical_mesh_id = torch.arange(0, 16) mesh_shape = (4, 4) # [[0, 1, 2, 3], # [4, 5, 6, 7], diff --git a/tests/test_tensor/test_tp_with_zero.py b/tests/test_tensor/test_tp_with_zero.py index 1a6d23f6a2eb..539806cb196a 100644 --- a/tests/test_tensor/test_tp_with_zero.py +++ b/tests/test_tensor/test_tp_with_zero.py @@ -1,20 +1,14 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP import colossalai from colossalai.amp import convert_to_apex_amp -from colossalai.gemini.chunk import search_chunk_configuration -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel import GeminiDDP, ZeroDDP from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, GeminiAdamOptimizer, GeminiDDP, ZeroDDP +from colossalai.zero.gemini import search_chunk_configuration from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import set_seed, tensor_shard_equal from tests.test_tensor.model.test_gpt2 import init_megatron_spec @@ -85,7 +79,7 @@ def run_gpt(placement_policy, tp_init_spec_func=None): tp_init_spec_func(model, pg) dp_world_size = pg.dp_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[dp_world_size]['chunk_size'] = 5000 config_dict[dp_world_size]['keep_gathered'] = False if placement_policy != 'cuda': @@ -142,8 +136,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_gpt(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_trainer/test_pipeline/test_p2p.py b/tests/test_trainer/test_pipeline/test_p2p.py index 72820c6a1f0d..8ad366133d18 100644 --- a/tests/test_trainer/test_pipeline/test_p2p.py +++ b/tests/test_trainer/test_pipeline/test_p2p.py @@ -1,21 +1,26 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -from colossalai.communication import (recv_backward, recv_forward, recv_obj_meta, send_backward, - send_backward_recv_forward, send_forward, send_forward_recv_backward, - send_obj_meta) + +from colossalai.communication import ( + recv_backward, + recv_forward, + recv_obj_meta, + send_backward, + send_backward_recv_forward, + send_forward, + send_forward_recv_backward, + send_obj_meta, +) from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.initialize import launch from colossalai.logging import get_dist_logger -from colossalai.utils import free_port, get_current_device -from colossalai.testing import rerun_on_exception +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device BATCH_SIZE = 4 SEQ_LENGTH = 2 @@ -85,7 +90,7 @@ def run_check(rank, world_size, port): prev_rank = gpc.get_prev_global_rank(ParallelMode.PIPELINE) next_rank = gpc.get_next_global_rank(ParallelMode.PIPELINE) logger.info('Rank {0}: prev rank {1}, next rank {2}'.format(rank, prev_rank, next_rank)) - logger.info('Distributed environment is initialzied.') + logger.info('Distributed environment is initialized.') check_comm(world_size, rank, prev_rank, next_rank, logger) gpc.destroy() @@ -93,11 +98,10 @@ def run_check(rank, world_size, port): @pytest.mark.dist -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") +@rerun_if_address_is_in_use() def test_p2p(): world_size = 4 - run_func = partial(run_check, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_check, world_size) if __name__ == '__main__': diff --git a/tests/test_trainer/test_pipeline/test_pipeline_schedule.py b/tests/test_trainer/test_pipeline/test_pipeline_schedule.py index 48f729658134..6d7bf6b3d89f 100644 --- a/tests/test_trainer/test_pipeline/test_pipeline_schedule.py +++ b/tests/test_trainer/test_pipeline/test_pipeline_schedule.py @@ -1,34 +1,26 @@ # referenced from Megatron and used to testify communication import os -import os.path as osp -from functools import partial from pathlib import Path -import colossalai import pytest import torch import torch.nn as nn -import torch.multiprocessing as mp -from colossalai.core import global_context as gpc -from colossalai.context import ParallelMode -from colossalai.initialize import launch -from colossalai.utils import free_port, get_dataloader, print_rank_0 -from colossalai.testing import rerun_on_exception from torchvision import transforms from torchvision.datasets import CIFAR10 from torchvision.models import resnet18 +import colossalai +from colossalai.context import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.initialize import launch +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_dataloader, print_rank_0 BATCH_SIZE = 8 -CONFIG=dict( - NUM_MICRO_BATCHES=2, - parallel = dict( - pipeline=dict(size=2), - tensor=dict(size=1, mode=None) - ) -) +CONFIG = dict(NUM_MICRO_BATCHES=2, parallel=dict(pipeline=dict(size=2), tensor=dict(size=1, mode=None))) + def run_schedule(rank, world_size, port): launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') @@ -85,11 +77,10 @@ def forward(self, x): @pytest.mark.dist -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") +@rerun_if_address_is_in_use() def test_pipeline_schedule(): world_size = 2 - run_func = partial(run_schedule, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_schedule, world_size) if __name__ == '__main__': diff --git a/tests/test_trainer/test_trainer_with_non_pipe_schedule.py b/tests/test_trainer/test_trainer_with_non_pipe_schedule.py index b013433293cd..753f82222f9d 100644 --- a/tests/test_trainer/test_trainer_with_non_pipe_schedule.py +++ b/tests/test_trainer/test_trainer_with_non_pipe_schedule.py @@ -1,15 +1,13 @@ -from functools import partial - -import colossalai import pytest import torch -import torch.multiprocessing as mp + +import colossalai from colossalai.amp.amp_type import AMP_TYPE from colossalai.logging import get_dist_logger +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.trainer import Trainer -from colossalai.utils import MultiTimer, free_port +from colossalai.utils import MultiTimer from tests.components_to_test.registry import non_distributed_component_funcs -from colossalai.testing import parameterize, rerun_if_address_is_in_use BATCH_SIZE = 4 IMG_SIZE = 32 @@ -54,8 +52,7 @@ def run_dist(rank, world_size, port): @rerun_if_address_is_in_use() def test_trainer_no_pipeline(): world_size = 4 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_trainer/test_trainer_with_pipe_schedule.py b/tests/test_trainer/test_trainer_with_pipe_schedule.py index 3698526a8e6c..bb63d51a0b65 100644 --- a/tests/test_trainer/test_trainer_with_pipe_schedule.py +++ b/tests/test_trainer/test_trainer_with_pipe_schedule.py @@ -1,23 +1,21 @@ import os -from functools import partial from pathlib import Path -import colossalai import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.engine.schedule import PipelineSchedule -from colossalai.logging import get_dist_logger -from colossalai.trainer import Trainer -from colossalai.utils import MultiTimer, free_port, get_dataloader from torch.optim import Adam from torchvision import transforms from torchvision.datasets import CIFAR10 from torchvision.models import resnet18 -from colossalai.testing import rerun_if_address_is_in_use + +import colossalai +from colossalai.context.parallel_mode import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.logging import get_dist_logger +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.trainer import Trainer +from colossalai.utils import MultiTimer, get_dataloader BATCH_SIZE = 4 IMG_SIZE = 32 @@ -91,8 +89,7 @@ def forward(self, x): @rerun_if_address_is_in_use() def test_trainer_with_pipeline(): world_size = 4 - run_func = partial(run_trainer_with_pipeline, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_trainer_with_pipeline, world_size) if __name__ == '__main__': diff --git a/tests/test_utils/test_activation_checkpointing.py b/tests/test_utils/test_activation_checkpointing.py index 3ac75fb00c86..2930552cc4e7 100644 --- a/tests/test_utils/test_activation_checkpointing.py +++ b/tests/test_utils/test_activation_checkpointing.py @@ -4,8 +4,10 @@ import pytest import torch import torch.nn.functional as F + from colossalai.context.parallel_mode import ParallelMode -from colossalai.context.random import add_seed, seed, set_mode, reset_seeds +from colossalai.context.random import add_seed, reset_seeds, seed, set_mode +from colossalai.testing import clear_cache_before_run, parameterize from colossalai.utils.activation_checkpoint import checkpoint @@ -39,8 +41,9 @@ def forward_inplace(x, weight): @pytest.mark.gpu -@pytest.mark.parametrize("use_reentrant", [True, False]) -@pytest.mark.parametrize("cpu_offload", [True, False]) +@clear_cache_before_run() +@parameterize("use_reentrant", [True, False]) +@parameterize("cpu_offload", [True, False]) def test_activation_checkpointing(cpu_offload, use_reentrant): # as seed manager is singleton @@ -48,7 +51,7 @@ def test_activation_checkpointing(cpu_offload, use_reentrant): # other tests might affect this test reset_seeds() - # We put initilization here to avoid change cuda rng state below + # We put initialization here to avoid change cuda rng state below inputs = torch.rand(2, 2, requires_grad=True, device='cuda') weight = torch.rand(2, 4, requires_grad=True, device='cuda') diff --git a/tests/test_utils/test_checkpoint/test_checkpoint_1d.py b/tests/test_utils/test_checkpoint/test_checkpoint_1d.py index 8a0fea9ae47a..335be61359ed 100644 --- a/tests/test_utils/test_checkpoint/test_checkpoint_1d.py +++ b/tests/test_utils/test_checkpoint/test_checkpoint_1d.py @@ -1,80 +1,77 @@ -#!/usr/bin/env python -# -*- encoding: utf-8 -*- - -import pprint -from functools import partial - -import colossalai.nn as col_nn -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port, is_using_pp -from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint -from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus - - -def build_pipeline(model): - from colossalai.pipeline.utils import partition_uniform - - pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) - pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) - depth = len(model) - start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] - layers = [] - for i in range(depth): - if start <= i < end: - layers.append(model[i]) - else: - layers.append(nn.Identity()) - return nn.Sequential(*tuple(layers)) - - -def check_equal(A, B): - assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) - - -def check_checkpoint_1d(rank, world_size, port): - config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, mode="1d")),) - - disable_existing_loggers() - launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") - - m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) - sd1 = m1.state_dict() - if gpc.get_global_rank() == 0: - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") - save_checkpoint("test.pt", 0, m1) - - m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) - if is_using_pp(): - m2 = build_pipeline(m2) - - load_checkpoint("test.pt", m2) - sd2 = m2.state_dict() - if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: - sd2 = gather_pipeline_parallel_state_dict(sd2) - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") - - if gpc.get_global_rank() == 0: - for k, v in sd1.items(): - assert k in sd2 - check_equal(v, sd2[k].to(torch.device("cpu"))) - - -@pytest.mark.dist -@pytest.mark.skip("takes too long") -@skip_if_not_enough_gpus(min_gpus=8) -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") -def test_checkpoint_1d(): - world_size = 8 - run_func = partial(check_checkpoint_1d, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == "__main__": - test_checkpoint_1d() +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import pprint + +import pytest +import torch +import torch.nn as nn + +import colossalai.nn as col_nn +from colossalai.context.parallel_mode import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn +from colossalai.utils import is_using_pp +from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint + + +def build_pipeline(model): + from colossalai.pipeline.utils import partition_uniform + + pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) + pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) + depth = len(model) + start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] + layers = [] + for i in range(depth): + if start <= i < end: + layers.append(model[i]) + else: + layers.append(nn.Identity()) + return nn.Sequential(*tuple(layers)) + + +def check_equal(A, B): + assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) + + +def check_checkpoint_1d(rank, world_size, port): + config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, mode="1d")),) + + disable_existing_loggers() + launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + + m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) + sd1 = m1.state_dict() + if gpc.get_global_rank() == 0: + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") + save_checkpoint("test.pt", 0, m1) + + m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) + if is_using_pp(): + m2 = build_pipeline(m2) + + load_checkpoint("test.pt", m2) + sd2 = m2.state_dict() + if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: + sd2 = gather_pipeline_parallel_state_dict(sd2) + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") + + if gpc.get_global_rank() == 0: + for k, v in sd1.items(): + assert k in sd2 + check_equal(v, sd2[k].to(torch.device("cpu"))) + + +@pytest.mark.dist +@pytest.mark.skip("takes too long") +@skip_if_not_enough_gpus(min_gpus=8) +@rerun_if_address_is_in_use() +def test_checkpoint_1d(): + spawn(check_checkpoint_1d, 8) + + +if __name__ == "__main__": + test_checkpoint_1d() diff --git a/tests/test_utils/test_checkpoint/test_checkpoint_2d.py b/tests/test_utils/test_checkpoint/test_checkpoint_2d.py index 26314290d4de..175d9ef6ceb9 100644 --- a/tests/test_utils/test_checkpoint/test_checkpoint_2d.py +++ b/tests/test_utils/test_checkpoint/test_checkpoint_2d.py @@ -1,80 +1,77 @@ -#!/usr/bin/env python -# -*- encoding: utf-8 -*- - -import pprint -from functools import partial - -import colossalai.nn as col_nn -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port, get_current_device, is_using_pp -from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint -from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus - - -def build_pipeline(model): - from colossalai.pipeline.utils import partition_uniform - - pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) - pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) - depth = len(model) - start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] - layers = [] - for i in range(depth): - if start <= i < end: - layers.append(model[i]) - else: - layers.append(nn.Identity()) - return nn.Sequential(*tuple(layers)) - - -def check_equal(A, B): - assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) - - -def check_checkpoint_2d(rank, world_size, port): - config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, mode="2d")),) - - disable_existing_loggers() - launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") - - m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) - sd1 = m1.state_dict() - if gpc.get_global_rank() == 0: - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") - save_checkpoint("test.pt", 0, m1) - - m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) - if is_using_pp(): - m2 = build_pipeline(m2) - - load_checkpoint("test.pt", m2) - sd2 = m2.state_dict() - if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: - sd2 = gather_pipeline_parallel_state_dict(sd2) - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") - - if gpc.get_global_rank() == 0: - for k, v in sd1.items(): - assert k in sd2 - check_equal(v, sd2[k].to(torch.device("cpu"))) - - -@pytest.mark.dist -@pytest.mark.skip("takes too long") -@skip_if_not_enough_gpus(min_gpus=8) -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") -def test_checkpoint_2d(): - world_size = 8 - run_func = partial(check_checkpoint_2d, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == "__main__": - test_checkpoint_2d() +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import pprint + +import pytest +import torch +import torch.nn as nn + +import colossalai.nn as col_nn +from colossalai.context.parallel_mode import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn +from colossalai.utils import is_using_pp +from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint + + +def build_pipeline(model): + from colossalai.pipeline.utils import partition_uniform + + pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) + pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) + depth = len(model) + start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] + layers = [] + for i in range(depth): + if start <= i < end: + layers.append(model[i]) + else: + layers.append(nn.Identity()) + return nn.Sequential(*tuple(layers)) + + +def check_equal(A, B): + assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) + + +def check_checkpoint_2d(rank, world_size, port): + config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, mode="2d")),) + + disable_existing_loggers() + launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + + m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) + sd1 = m1.state_dict() + if gpc.get_global_rank() == 0: + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") + save_checkpoint("test.pt", 0, m1) + + m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) + if is_using_pp(): + m2 = build_pipeline(m2) + + load_checkpoint("test.pt", m2) + sd2 = m2.state_dict() + if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: + sd2 = gather_pipeline_parallel_state_dict(sd2) + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") + + if gpc.get_global_rank() == 0: + for k, v in sd1.items(): + assert k in sd2 + check_equal(v, sd2[k].to(torch.device("cpu"))) + + +@pytest.mark.dist +@pytest.mark.skip("takes too long") +@skip_if_not_enough_gpus(min_gpus=8) +@rerun_if_address_is_in_use() +def test_checkpoint_2d(): + spawn(check_checkpoint_2d, 8) + + +if __name__ == "__main__": + test_checkpoint_2d() diff --git a/tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py b/tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py index 3dbd340fd42d..33cb3a65d184 100644 --- a/tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py +++ b/tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py @@ -1,80 +1,77 @@ -#!/usr/bin/env python -# -*- encoding: utf-8 -*- - -import pprint -from functools import partial - -import colossalai.nn as col_nn -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port, get_current_device, is_using_pp -from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint -from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus - - -def build_pipeline(model): - from colossalai.pipeline.utils import partition_uniform - - pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) - pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) - depth = len(model) - start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] - layers = [] - for i in range(depth): - if start <= i < end: - layers.append(model[i]) - else: - layers.append(nn.Identity()) - return nn.Sequential(*tuple(layers)) - - -def check_equal(A, B): - assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) - - -def check_checkpoint_2p5d(rank, world_size, port): - config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, depth=1, mode="2.5d")),) - - disable_existing_loggers() - launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") - - m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) - sd1 = m1.state_dict() - if gpc.get_global_rank() == 0: - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") - save_checkpoint("test.pt", 0, m1) - - m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) - if is_using_pp(): - m2 = build_pipeline(m2) - - load_checkpoint("test.pt", m2) - sd2 = m2.state_dict() - if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: - sd2 = gather_pipeline_parallel_state_dict(sd2) - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") - - if gpc.get_global_rank() == 0: - for k, v in sd1.items(): - assert k in sd2 - check_equal(v, sd2[k].to(torch.device("cpu"))) - - -@pytest.mark.dist -@pytest.mark.skip("takes too long") -@skip_if_not_enough_gpus(min_gpus=8) -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") -def test_checkpoint_2p5d(): - world_size = 8 - run_func = partial(check_checkpoint_2p5d, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == "__main__": - test_checkpoint_2p5d() +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import pprint + +import pytest +import torch +import torch.nn as nn + +import colossalai.nn as col_nn +from colossalai.context.parallel_mode import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn +from colossalai.utils import is_using_pp +from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint + + +def build_pipeline(model): + from colossalai.pipeline.utils import partition_uniform + + pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) + pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) + depth = len(model) + start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] + layers = [] + for i in range(depth): + if start <= i < end: + layers.append(model[i]) + else: + layers.append(nn.Identity()) + return nn.Sequential(*tuple(layers)) + + +def check_equal(A, B): + assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) + + +def check_checkpoint_2p5d(rank, world_size, port): + config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, depth=1, mode="2.5d")),) + + disable_existing_loggers() + launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + + m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) + sd1 = m1.state_dict() + if gpc.get_global_rank() == 0: + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") + save_checkpoint("test.pt", 0, m1) + + m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) + if is_using_pp(): + m2 = build_pipeline(m2) + + load_checkpoint("test.pt", m2) + sd2 = m2.state_dict() + if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: + sd2 = gather_pipeline_parallel_state_dict(sd2) + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") + + if gpc.get_global_rank() == 0: + for k, v in sd1.items(): + assert k in sd2 + check_equal(v, sd2[k].to(torch.device("cpu"))) + + +@pytest.mark.dist +@pytest.mark.skip("takes too long") +@skip_if_not_enough_gpus(min_gpus=8) +@rerun_if_address_is_in_use() +def test_checkpoint_2p5d(): + spawn(check_checkpoint_2p5d, 8) + + +if __name__ == "__main__": + test_checkpoint_2p5d() diff --git a/tests/test_utils/test_checkpoint/test_checkpoint_3d.py b/tests/test_utils/test_checkpoint/test_checkpoint_3d.py index 38f650547585..73ac2dd5fe18 100644 --- a/tests/test_utils/test_checkpoint/test_checkpoint_3d.py +++ b/tests/test_utils/test_checkpoint/test_checkpoint_3d.py @@ -1,80 +1,77 @@ -#!/usr/bin/env python -# -*- encoding: utf-8 -*- - -import pprint -from functools import partial - -import colossalai.nn as col_nn -import pytest -import torch -import torch.multiprocessing as mp -import torch.nn as nn -from colossalai.context.parallel_mode import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.initialize import launch -from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port, get_current_device, is_using_pp -from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint -from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus - - -def build_pipeline(model): - from colossalai.pipeline.utils import partition_uniform - - pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) - pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) - depth = len(model) - start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] - layers = [] - for i in range(depth): - if start <= i < end: - layers.append(model[i]) - else: - layers.append(nn.Identity()) - return nn.Sequential(*tuple(layers)) - - -def check_equal(A, B): - assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) - - -def check_checkpoint_3d(rank, world_size, port): - config = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=8, mode="3d")),) - - disable_existing_loggers() - launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") - - m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) - sd1 = m1.state_dict() - if gpc.get_global_rank() == 0: - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") - save_checkpoint("test.pt", 0, m1) - - m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) - if is_using_pp(): - m2 = build_pipeline(m2) - - load_checkpoint("test.pt", m2) - sd2 = m2.state_dict() - if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: - sd2 = gather_pipeline_parallel_state_dict(sd2) - print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") - - if gpc.get_global_rank() == 0: - for k, v in sd1.items(): - assert k in sd2 - check_equal(v, sd2[k].to(torch.device("cpu"))) - - -@pytest.mark.dist -@pytest.mark.skip("takes too long") -@skip_if_not_enough_gpus(min_gpus=8) -@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*") -def test_checkpoint_3d(): - world_size = 8 - run_func = partial(check_checkpoint_3d, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == "__main__": - test_checkpoint_3d() +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import pprint + +import pytest +import torch +import torch.nn as nn + +import colossalai.nn as col_nn +from colossalai.context.parallel_mode import ParallelMode +from colossalai.core import global_context as gpc +from colossalai.initialize import launch +from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus, spawn +from colossalai.utils import is_using_pp +from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint + + +def build_pipeline(model): + from colossalai.pipeline.utils import partition_uniform + + pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE) + pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE) + depth = len(model) + start, end = partition_uniform(depth, pipeline_size, 1)[pipeline_rank][0] + layers = [] + for i in range(depth): + if start <= i < end: + layers.append(model[i]) + else: + layers.append(nn.Identity()) + return nn.Sequential(*tuple(layers)) + + +def check_equal(A, B): + assert torch.allclose(A, B, rtol=1e-3, atol=1e-2) + + +def check_checkpoint_3d(rank, world_size, port): + config = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=8, mode="3d")),) + + disable_existing_loggers() + launch(config=config, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + + m1 = nn.Sequential(nn.Linear(4, 8), nn.Linear(8, 4)) + sd1 = m1.state_dict() + if gpc.get_global_rank() == 0: + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd1)}\n") + save_checkpoint("test.pt", 0, m1) + + m2 = nn.Sequential(col_nn.Linear(4, 8), col_nn.Linear(8, 4)) + if is_using_pp(): + m2 = build_pipeline(m2) + + load_checkpoint("test.pt", m2) + sd2 = m2.state_dict() + if is_using_pp() and gpc.get_local_rank(ParallelMode.TENSOR) == 0: + sd2 = gather_pipeline_parallel_state_dict(sd2) + print(f"Rank {gpc.get_global_rank()}:\n{pprint.pformat(sd2)}\n") + + if gpc.get_global_rank() == 0: + for k, v in sd1.items(): + assert k in sd2 + check_equal(v, sd2[k].to(torch.device("cpu"))) + + +@pytest.mark.dist +@pytest.mark.skip("takes too long") +@skip_if_not_enough_gpus(min_gpus=8) +@rerun_if_address_is_in_use() +def test_checkpoint_3d(): + spawn(check_checkpoint_3d, 8) + + +if __name__ == "__main__": + test_checkpoint_3d() diff --git a/tests/test_utils/test_checkpoint_io/test_load.py b/tests/test_utils/test_checkpoint_io/test_load.py index 780c13dc534a..2949c9f0752d 100644 --- a/tests/test_utils/test_checkpoint_io/test_load.py +++ b/tests/test_utils/test_checkpoint_io/test_load.py @@ -3,20 +3,19 @@ from tempfile import TemporaryDirectory from typing import Dict -import colossalai import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.utils.checkpoint_io.io import load, save -from colossalai.utils.checkpoint_io.meta import (ParamDistMeta, ParamRedistMeta, RankRedistMeta, RedistMeta) from torch import Tensor from torch.nn import Module from torch.optim import Adam, Optimizer +import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.checkpoint_io.io import load, save +from colossalai.utils.checkpoint_io.meta import ParamDistMeta, ParamRedistMeta, RankRedistMeta, RedistMeta + def check_model_state_dict(a: Dict[str, Tensor], b: Dict[str, Tensor]) -> None: assert set(a.keys()) == set(b.keys()) @@ -24,7 +23,7 @@ def check_model_state_dict(a: Dict[str, Tensor], b: Dict[str, Tensor]) -> None: assert torch.equal(v, b[k]) -def check_optim_state_dict(a: dict, b: dict, ignore_param_gruops: bool = False) -> None: +def check_optim_state_dict(a: dict, b: dict, ignore_param_groups: bool = False) -> None: assert set(a['state'].keys()) == set(b['state'].keys()) for k, state in a['state'].items(): b_state = b['state'][k] @@ -33,7 +32,7 @@ def check_optim_state_dict(a: dict, b: dict, ignore_param_gruops: bool = False) assert torch.equal(v1, v2) else: assert v1 == v2 - if not ignore_param_gruops: + if not ignore_param_groups: assert a['param_groups'] == b['param_groups'] @@ -120,34 +119,33 @@ def test_save_global_load_global(max_shard_size_gb: float): check_optim_state_dict(optimizer.state_dict(), new_optimizer.state_dict()) -def run_dist(rank, world_size, port, func): +def run_dist(rank, world_size, port, test_fn): colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - func() + test_fn() def launch_dist(fn, world_size: int): - proc_fn = partial(run_dist, world_size=world_size, port=free_port(), func=fn) - mp.spawn(proc_fn, nprocs=world_size) + spawn(run_dist, world_size, test_fn=fn) def save_dist(dir_name: str, zero: bool): - model, optmizer = prepare_model_optim(shard=True, zero=zero) - reset_model_optim(model, optmizer) + model, optimizer = prepare_model_optim(shard=True, zero=zero) + reset_model_optim(model, optimizer) world_size = dist.get_world_size() rank = dist.get_rank() - save(dir_name, model, optmizer, dist_meta=get_dist_metas(world_size, zero)[rank]) + save(dir_name, model, optimizer, dist_meta=get_dist_metas(world_size, zero)[rank]) def load_and_check_dist(dir_name: str): world_size = dist.get_world_size() - model, optmizer = prepare_model_optim(shard=True) - reset_model_optim(model, optmizer) + model, optimizer = prepare_model_optim(shard=True) + reset_model_optim(model, optimizer) model_state_dict = deepcopy(model.state_dict()) - optimizer_state_dict = deepcopy(optmizer.state_dict()) - reset_model_optim(model, optmizer, 1) - load(dir_name, model, optmizer, get_redist_meta(world_size), get_dist_metas(world_size)) + optimizer_state_dict = deepcopy(optimizer.state_dict()) + reset_model_optim(model, optimizer, 1) + load(dir_name, model, optimizer, get_redist_meta(world_size), get_dist_metas(world_size)) check_model_state_dict(model_state_dict, model.state_dict()) - check_optim_state_dict(optimizer_state_dict, optmizer.state_dict()) + check_optim_state_dict(optimizer_state_dict, optimizer.state_dict()) @pytest.mark.dist diff --git a/tests/test_utils/test_checkpoint_io/test_merge.py b/tests/test_utils/test_checkpoint_io/test_merge.py index 04e454dcb713..07d4597f8391 100644 --- a/tests/test_utils/test_checkpoint_io/test_merge.py +++ b/tests/test_utils/test_checkpoint_io/test_merge.py @@ -1,18 +1,18 @@ -from colossalai.utils.checkpoint_io.meta import ParamDistMeta -from colossalai.utils.checkpoint_io.constant import GLOBAL_META_FILE_NAME -from colossalai.utils.checkpoint_io.io import save, merge -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from tempfile import TemporaryDirectory -from torch.optim import Adam -from functools import partial -import torch import os +from functools import partial +from tempfile import TemporaryDirectory + import pytest -import colossalai -import torch.nn as nn +import torch import torch.distributed as dist -import torch.multiprocessing as mp +import torch.nn as nn +from torch.optim import Adam + +import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.checkpoint_io.constant import GLOBAL_META_FILE_NAME +from colossalai.utils.checkpoint_io.io import merge, save +from colossalai.utils.checkpoint_io.meta import ParamDistMeta class DummyModel(nn.Module): @@ -52,7 +52,7 @@ def test_merge_global(): assert len(os.listdir(output_dir)) == 0 -def run_dist(rank, world_size, port, func): +def run_dist(rank, world_size, port, test_fn): colossalai.launch(config={'parallel': { 'tensor': { 'mode': '1d', @@ -64,11 +64,11 @@ def run_dist(rank, world_size, port, func): host='localhost', port=port, backend='nccl') - func() + test_fn() def run_save_dist(dir_name: str, zero: bool): - model, optmizer = prepare_model_optim(shard=True, zero=zero) + model, optimizer = prepare_model_optim(shard=True, zero=zero) rank = dist.get_rank() dp_world_size = dist.get_world_size() // 2 if not zero: @@ -90,7 +90,7 @@ def run_save_dist(dir_name: str, zero: bool): 'fc.bias': ParamDistMeta(rank // 2, dp_world_size, 0, 1, zero_numel=1, zero_orig_shape=[1]) } - save(dir_name, model, optmizer, dist_meta=dist_metas) + save(dir_name, model, optimizer, dist_meta=dist_metas) @pytest.mark.dist @@ -100,8 +100,7 @@ def test_merge_tp_dp(zero: bool): with TemporaryDirectory() as dir_name: fn = partial(run_save_dist, dir_name, zero) world_size = 4 - proc_fn = partial(run_dist, world_size=world_size, port=free_port(), func=fn) - mp.spawn(proc_fn, nprocs=world_size) + spawn(run_dist, world_size, test_fn=fn) with TemporaryDirectory() as output_dir: merge(dir_name, output_dir) assert len(os.listdir(output_dir)) == 5 diff --git a/tests/test_utils/test_checkpoint_io/test_redist.py b/tests/test_utils/test_checkpoint_io/test_redist.py index 6e76f3167e31..fdc849a5ecc0 100644 --- a/tests/test_utils/test_checkpoint_io/test_redist.py +++ b/tests/test_utils/test_checkpoint_io/test_redist.py @@ -2,19 +2,23 @@ from functools import partial from tempfile import TemporaryDirectory -import colossalai import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port +from torch.optim import Adam + +import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.checkpoint_io.constant import GLOBAL_META_FILE_NAME from colossalai.utils.checkpoint_io.io import redist, save -from colossalai.utils.checkpoint_io.meta import (ParamDistMeta, ParamRedistMeta, PipelineRedistMeta, RankRedistMeta, - RedistMeta) -from torch.optim import Adam +from colossalai.utils.checkpoint_io.meta import ( + ParamDistMeta, + ParamRedistMeta, + PipelineRedistMeta, + RankRedistMeta, + RedistMeta, +) class DummyModel(nn.Module): @@ -105,7 +109,7 @@ def test_global_to_dist(): check_checkpoint_shape(output_dir) -def run_dist(rank, world_size, port, func): +def run_dist(rank, world_size, port, test_fn): colossalai.launch(config={'parallel': { 'tensor': { 'mode': '1d', @@ -117,13 +121,13 @@ def run_dist(rank, world_size, port, func): host='localhost', port=port, backend='nccl') - func() + test_fn() def run_save_dist(dir_name: str, zero: bool): - model, optmizer = prepare_model_optim(shard=True, zero=zero) + model, optimizer = prepare_model_optim(shard=True, zero=zero) rank = dist.get_rank() - save(dir_name, model, optmizer, dist_meta=get_dist_metas(4, zero)[rank]) + save(dir_name, model, optimizer, dist_meta=get_dist_metas(4, zero)[rank]) @pytest.mark.dist @@ -133,8 +137,7 @@ def test_dist_to_dist(zero: bool): with TemporaryDirectory() as dir_name: fn = partial(run_save_dist, dir_name, zero) world_size = 4 - proc_fn = partial(run_dist, world_size=world_size, port=free_port(), func=fn) - mp.spawn(proc_fn, nprocs=world_size) + spawn(run_dist, world_size, test_fn=fn) with TemporaryDirectory() as output_dir: redist(dir_name, output_dir, get_redist_meta(4), get_dist_metas(4)) if not zero: diff --git a/tests/test_utils/test_checkpoint_io/test_save.py b/tests/test_utils/test_checkpoint_io/test_save.py index 5ff9d0aa2217..2abdd95a6481 100644 --- a/tests/test_utils/test_checkpoint_io/test_save.py +++ b/tests/test_utils/test_checkpoint_io/test_save.py @@ -3,21 +3,24 @@ from tempfile import TemporaryDirectory from typing import Dict -import colossalai import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.utils.checkpoint_io.constant import (GLOBAL_META_FILE_NAME, META_CKPT_FILE_NAME, MODEL_CKPT_FILE_NAME, - OTHER_CKPT_FILE_NAME) -from colossalai.utils.checkpoint_io.io import save -from colossalai.utils.checkpoint_io.meta import ParamDistMeta from torch import Tensor from torch.optim import Adam +import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.checkpoint_io.constant import ( + GLOBAL_META_FILE_NAME, + META_CKPT_FILE_NAME, + MODEL_CKPT_FILE_NAME, + OTHER_CKPT_FILE_NAME, +) +from colossalai.utils.checkpoint_io.io import save +from colossalai.utils.checkpoint_io.meta import ParamDistMeta + def check_model_state_dict(a: Dict[str, Tensor], b: Dict[str, Tensor]) -> None: assert set(a.keys()) == set(b.keys()) @@ -25,7 +28,7 @@ def check_model_state_dict(a: Dict[str, Tensor], b: Dict[str, Tensor]) -> None: assert torch.equal(v, b[k]) -def check_optim_state_dict(a: dict, b: dict, ignore_param_gruops: bool = False) -> None: +def check_optim_state_dict(a: dict, b: dict, ignore_param_groups: bool = False) -> None: assert set(a['state'].keys()) == set(b['state'].keys()) for k, state in a['state'].items(): b_state = b['state'][k] @@ -34,7 +37,7 @@ def check_optim_state_dict(a: dict, b: dict, ignore_param_gruops: bool = False) assert torch.equal(v1, v2) else: assert v1 == v2 - if not ignore_param_gruops: + if not ignore_param_groups: assert a['param_groups'] == b['param_groups'] @@ -104,18 +107,18 @@ def test_save_global_shard(): }) -def run_dist(rank, world_size, port, func): +def run_dist(rank, world_size, port, test_fn): colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - func() + test_fn() def run_save_dist(dir_name): - model, optmizer = prepare_model_optim() + model, optimizer = prepare_model_optim() dist_metas = { 'fc.weight': ParamDistMeta(dist.get_rank(), dist.get_world_size(), 0, 1), 'fc.bias': ParamDistMeta(dist.get_rank(), dist.get_world_size(), 0, 1) } - save(dir_name, model, optmizer, dist_meta=dist_metas) + save(dir_name, model, optimizer, dist_meta=dist_metas) @pytest.mark.dist @@ -124,8 +127,7 @@ def test_save_dist(): with TemporaryDirectory() as dir_name: fn = partial(run_save_dist, dir_name) world_size = 2 - proc_fn = partial(run_dist, world_size=world_size, port=free_port(), func=fn) - mp.spawn(proc_fn, nprocs=world_size) + spawn(run_dist, world_size, test_fn=fn) assert len(os.listdir(dir_name)) == 8 global_meta = torch.load(os.path.join(dir_name, GLOBAL_META_FILE_NAME)) assert len(global_meta['meta']) == 2 diff --git a/tests/test_utils/test_colo_checkpoint.py b/tests/test_utils/test_colo_checkpoint.py index a5ea75fffc36..89760a5456e7 100644 --- a/tests/test_utils/test_colo_checkpoint.py +++ b/tests/test_utils/test_colo_checkpoint.py @@ -1,25 +1,20 @@ -import os, shutil -import torch -import pytest +import os +import shutil from copy import deepcopy -from functools import partial -import torch.multiprocessing as mp +import pytest +import torch import torch.distributed as dist - -from torch.optim.lr_scheduler import CosineAnnealingLR -from torch.optim.lr_scheduler import MultiplicativeLR -from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR +from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR import colossalai -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils.cuda import get_current_device -from colossalai.utils import free_port -from colossalai.utils.model.colo_init_context import ColoInitContext -from colossalai.tensor import ComputePattern, ComputeSpec, ColoTensor, ShardSpec, ProcessGroup -from colossalai.utils.checkpoint import save_checkpoint, load_checkpoint +from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR from colossalai.nn.optimizer import ColossalaiOptimizer - +from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ProcessGroup, ShardSpec +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils.checkpoint import load_checkpoint, save_checkpoint +from colossalai.utils.cuda import get_current_device +from colossalai.zero import ColoInitContext from tests.components_to_test.registry import non_distributed_component_funcs @@ -204,13 +199,7 @@ def run_dist(rank, world_size, port, use_ddp, use_mp_reload, test_scheduler): # @pytest.mark.parametrize('test_scheduler', ['colossalai_cosine_warmup', 'torch_cosine', 'torch_lambda']) @rerun_if_address_is_in_use() def test_checkpoint(world_size, use_ddp, use_mp_reload, test_scheduler=None): - run_func = partial(run_dist, - world_size=world_size, - port=free_port(), - use_ddp=use_ddp, - use_mp_reload=use_mp_reload, - test_scheduler=test_scheduler) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size, use_ddp=use_ddp, use_mp_reload=use_mp_reload, test_scheduler=test_scheduler) if __name__ == '__main__': diff --git a/tests/test_utils/test_commons.py b/tests/test_utils/test_commons.py index 0ecb7446c788..2633d7da21aa 100644 --- a/tests/test_utils/test_commons.py +++ b/tests/test_utils/test_commons.py @@ -1,16 +1,13 @@ -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.zero.sharded_param import ShardedTensor -from colossalai.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline -import colossalai - import torch -import torch.multiprocessing as mp +import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.gemini.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline +from colossalai.zero.legacy.sharded_param import ShardedTensor -def run_tensor_move(rank): - colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl') +def run_tensor_move(rank, world_size, port): + colossalai.launch(config={}, rank=0, world_size=world_size, host='localhost', port=port, backend='nccl') src_t = torch.ones(2, 3).cuda() tgt_t = torch.zeros(2, 3) @@ -37,7 +34,7 @@ def run_tensor_move(rank): @rerun_if_address_is_in_use() def test_tensor_move(): - mp.spawn(run_tensor_move, nprocs=1) + spawn(run_tensor_move, 1) if __name__ == '__main__': diff --git a/tests/test_utils/test_flash_attention.py b/tests/test_utils/test_flash_attention.py index 58e3b21d97eb..7a28b0157384 100644 --- a/tests/test_utils/test_flash_attention.py +++ b/tests/test_utils/test_flash_attention.py @@ -1,22 +1,14 @@ +import random + import pytest import torch from einops import rearrange -from colossalai.kernel.cuda_native.flash_attention import HAS_FLASH_ATTN, HAS_MEM_EFF_ATTN, HAS_TRITON - -if HAS_FLASH_ATTN: - from colossalai.kernel.cuda_native.flash_attention import ( - MaskedFlashAttention, - flash_attention_q_k_v, - flash_attention_q_kv, - flash_attention_qkv, - ) - -if HAS_TRITON: - from colossalai.kernel.cuda_native.flash_attention import triton_flash_attention +from colossalai.kernel.cuda_native.flash_attention import HAS_MEM_EFF_ATTN +from colossalai.testing import clear_cache_before_run, parameterize if HAS_MEM_EFF_ATTN: - from colossalai.kernel.cuda_native.flash_attention import LowerTriangularMask, MemoryEfficientAttention + from colossalai.kernel.cuda_native.flash_attention import AttnMaskType, ColoAttention def baseline_attention(Z, N_CTX, H, q, k, v, sm_scale): @@ -30,117 +22,92 @@ def baseline_attention(Z, N_CTX, H, q, k, v, sm_scale): return ref_out -@pytest.mark.skipif(HAS_TRITON == False, reason="triton is not available") -@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(3, 4, 2, 16)]) -def test_triton_flash_attention(Z, H, N_CTX, D_HEAD, dtype=torch.float16): - torch.manual_seed(20) - q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - sm_scale = 0.3 - dout = torch.randn_like(q) - - ref_out = baseline_attention(Z, N_CTX, H, q, k, v, sm_scale) - ref_out.backward(dout) - ref_dv, v.grad = v.grad.clone(), None - ref_dk, k.grad = k.grad.clone(), None - ref_dq, q.grad = q.grad.clone(), None - - # triton implementation - tri_out = triton_flash_attention(q, k, v, sm_scale) - tri_out.backward(dout) - tri_dv, v.grad = v.grad.clone(), None - tri_dk, k.grad = k.grad.clone(), None - tri_dq, q.grad = q.grad.clone(), None - # compare - assert torch.allclose(ref_out, tri_out, atol=1e-3) - assert torch.allclose(ref_dv, tri_dv, atol=1e-3) - assert torch.allclose(ref_dk, tri_dk, atol=1e-3) - assert torch.allclose(ref_dq, tri_dq, atol=1e-3) - - -@pytest.mark.skipif(HAS_FLASH_ATTN == False, reason="flash is not available") -@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(3, 4, 2, 16)]) -def test_flash_attention(Z, H, N_CTX, D_HEAD, dtype=torch.float16): - torch.manual_seed(20) - q = torch.randn((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - k = torch.randn((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - v = torch.randn((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - sm_scale = 0.3 - dout = torch.randn_like(q) - - # reference implementation - ref_out = baseline_attention(Z, N_CTX, H, q, k, v, sm_scale) - ref_out.backward(dout) - ref_dv, v.grad = v.grad.clone(), None - ref_dk, k.grad = k.grad.clone(), None - ref_dq, q.grad = q.grad.clone(), None - - # flash implementation - q, k, v = map(lambda x: rearrange(x, 'z h n d -> (z n) h d'), [q, k, v]) - dout = rearrange(dout, 'z h n d -> (z n) h d').detach() - for i in range(3): - if i == 0: - tri_out = flash_attention_q_k_v(q, k, v, sm_scale, Z, N_CTX, N_CTX, causal=True) - elif i == 1: - kv = torch.cat((k.unsqueeze(1), v.unsqueeze(1)), dim=1) - tri_out = flash_attention_q_kv(q, kv, sm_scale, Z, N_CTX, N_CTX, causal=True) - else: - qkv = torch.cat((q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1)), dim=1) - tri_out = flash_attention_qkv(qkv, sm_scale, Z, N_CTX, causal=True) - - tri_out.backward(dout, retain_graph=True) - - if i == 0: - tri_dq, tri_dk, tri_dv, = torch.autograd.grad(tri_out, (q, k, v), dout) - tri_out, tri_dq, tri_dk, tri_dv = map(lambda x: rearrange(x, '(z n) h d -> z h n d', z=Z), - (tri_out, tri_dq, tri_dk, tri_dv)) - elif i == 1: - tri_dq, tri_dkv, = torch.autograd.grad(tri_out, (q, kv), dout) - tri_dk, tri_dv = torch.chunk(tri_dkv, 2, dim=1) - tri_out, tri_dq, tri_dk, tri_dv = map(lambda x: rearrange(x, '(z n) h d -> z h n d', z=Z), - (tri_out, tri_dq, tri_dk.squeeze(1), tri_dv.squeeze(1))) - else: - tri_dqkv, = torch.autograd.grad(tri_out, (qkv), dout) - tri_dq, tri_dk, tri_dv = torch.chunk(tri_dqkv, 3, dim=1) - tri_out, tri_dq, tri_dk, tri_dv = map(lambda x: rearrange(x, '(z n) h d -> z h n d', z=Z), - (tri_out, tri_dq.squeeze(1), tri_dk.squeeze(1), tri_dv.squeeze(1))) - - # compare - assert torch.allclose(ref_out, tri_out, atol=1e-3) - assert torch.allclose(ref_dv, tri_dv, atol=1e-3) - assert torch.allclose(ref_dk, tri_dk, atol=1e-3) - assert torch.allclose(ref_dq, tri_dq, atol=1e-3) - - -@pytest.mark.skipif(HAS_FLASH_ATTN == False, reason="flash is not available") -@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(3, 4, 2, 16)]) -def test_masked_flash_attention(Z, H, N_CTX, D_HEAD, dtype=torch.float16): - attn = MaskedFlashAttention(N_CTX, D_HEAD, 0.1) - - qkv = torch.randn((Z, H, 3 * N_CTX * D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - attention_mask = torch.randint(2, (Z, H)).cuda().bool() - - out = attn(qkv, attention_mask) - - dout = torch.rand_like(out) - out.backward(dout) +@pytest.mark.skipif(HAS_MEM_EFF_ATTN == False, reason="xformers is not available") +@clear_cache_before_run() +@parameterize('B, S, H, D_HEAD', [(6, 8, 4, 16)]) +def test_attention_gpt(B, S, H, D_HEAD, dtype=torch.float16): + D = H * D_HEAD + + c_attn = torch.nn.Linear(D, 3 * D, dtype=dtype, device="cuda") + attn = ColoAttention(D, H, dropout=0.1) + + x = torch.randn((B, S, D), dtype=dtype, device="cuda") + + qkv = c_attn(x) + q, k, v = rearrange(qkv, 'b s (n h d) -> n b s h d', n=3, h=H) + y = attn(q, k, v, attn_mask_type=AttnMaskType.causal) + + assert list(y.shape) == [B, S, D] + + dy = torch.rand_like(y) + y.backward(dy) @pytest.mark.skipif(HAS_MEM_EFF_ATTN == False, reason="xformers is not available") -@pytest.mark.parametrize('Z, H, N_CTX, D_HEAD', [(6, 8, 4, 16)]) -def test_memory_efficient_attention(Z, H, N_CTX, D_HEAD, dtype=torch.float16): - attn = MemoryEfficientAttention(N_CTX * D_HEAD, N_CTX, 0.1) +@clear_cache_before_run() +@parameterize('B, S, H, D_HEAD', [(6, 8, 4, 16)]) +def test_attention_bert(B, S, H, D_HEAD, dtype=torch.float16): + D = H * D_HEAD + + c_attn = torch.nn.Linear(D, 3 * D, dtype=dtype, device="cuda") + attn = ColoAttention(D, H, dropout=0.1) + + x = torch.randn((B, S, D), dtype=dtype, device="cuda") + # attention mask of shape [B, S] with zero padding to max length S + mask = [torch.ones(S - i, dtype=dtype, device="cuda") for i in range(B)] + mask = torch.nn.utils.rnn.pad_sequence(mask, batch_first=True) + + qkv = c_attn(x) + q, k, v = rearrange(qkv, 'b s (n h d) -> b s n h d', n=3, h=H).unbind(dim=2) + y = attn(q, k, v, attn_mask=mask, attn_mask_type=AttnMaskType.padding) + + assert list(y.shape) == [B, S, D] + + dy = torch.rand_like(y) + y.backward(dy) + + +@pytest.mark.skipif(HAS_MEM_EFF_ATTN == False, reason="xformers is not available") +@clear_cache_before_run() +@parameterize('B, S, H, D_HEAD', [(6, 8, 4, 16)]) +def test_attention_no_mask(B, S, H, D_HEAD, dtype=torch.float16): + D = H * D_HEAD + + c_attn = torch.nn.Linear(D, 3 * D, dtype=dtype, device="cuda") + attn = ColoAttention(D, H, dropout=0.1) + + x = torch.randn((B, S, D), dtype=dtype, device="cuda") + qkv = c_attn(x) + q, k, v = rearrange(qkv, 'b s (n h d) -> b s n h d', n=3, h=H).unbind(dim=2) + y = attn(q, k, v) + + assert list(y.shape) == [B, S, D] + + dy = torch.rand_like(y) + y.backward(dy) + + +@pytest.mark.skipif(HAS_MEM_EFF_ATTN == False, reason="xformers is not available") +@clear_cache_before_run() +@parameterize('B, S, T, H, D_HEAD', [(6, 24, 8, 4, 16)]) +def test_cross_attention(B, S, T, H, D_HEAD, dtype=torch.float16): + D = H * D_HEAD + + q_attn = torch.nn.Linear(D, D, dtype=dtype, device="cuda") + kv_attn = torch.nn.Linear(D, 2 * D, dtype=dtype, device="cuda") - q = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - k = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() - v = torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0, std=.5).requires_grad_() + attn = ColoAttention(D, H, dropout=0.1) - out = attn(q, k, v, attention_mask=LowerTriangularMask()) + src = torch.randn((B, S, D), dtype=dtype, device="cuda") + tgt = torch.randn((B, T, D), dtype=dtype, device="cuda") - dout = torch.rand_like(out) - out.backward(dout) + q = q_attn(tgt) + kv = kv_attn(src) + q = rearrange(q, 'b s (h d) -> b s h d', h=H) + k, v = rearrange(kv, 'b s (n h d) -> b s n h d', n=2, h=H).unbind(dim=2) + y = attn(q, k, v, attn_mask_type=AttnMaskType.causal) + assert list(y.shape) == [B, T, D] -if __name__ == '__main__': - test_flash_attention(3, 4, 2, 16) + dy = torch.rand_like(y) + y.backward(dy) diff --git a/tests/test_utils/test_lazy_init/utils.py b/tests/test_utils/test_lazy_init/utils.py deleted file mode 100644 index 47ba534bc434..000000000000 --- a/tests/test_utils/test_lazy_init/utils.py +++ /dev/null @@ -1,69 +0,0 @@ -import random -from typing import Any, Callable, Optional, Tuple - -import numpy as np -import torch - -from colossalai.utils.model.experimental import LazyInitContext, LazyTensor, _MyTensor -from tests.kit.model_zoo.registry import ModelAttribute - -# model_fn, data_gen_fn, output_transform_fn, model_attr -TestingEntry = Tuple[Callable[[], torch.nn.Module], Callable[[], dict], Callable[[], dict], Optional[ModelAttribute]] - - -def set_seed(seed: int) -> None: - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - - -def assert_model_eqaual(m1: torch.nn.Module, m2: torch.nn.Module) -> None: - s1 = m1.state_dict() - s2 = m2.state_dict() - - assert len(s1) == len(s2), f'len {len(s1)} vs {len(s2)}' - - for (n1, t1), (n2, t2) in zip(s1.items(), s2.items()): - assert n1 == n2 - assert torch.equal(t1, t2), f'{n1} {t1} vs {t2}' - - -def assert_forward_equal(m1: torch.nn.Module, m2: torch.nn.Module, data_gen_fn: Callable[[], dict], - output_transform_fn: Callable[[Any], dict]) -> None: - data = data_gen_fn() - - m1.eval() - m2.eval() - # run forward - with torch.no_grad(): - outputs1 = m1(**data) - outputs2 = m2(**data) - - # compare output - transformed_out1 = output_transform_fn(outputs1) - transformed_out2 = output_transform_fn(outputs2) - - assert len(transformed_out1) == len(transformed_out2) - - for key, out1 in transformed_out1.items(): - out2 = transformed_out2[key] - assert torch.allclose(out1, out2, atol=1e-5), \ - f'{m1.__class__.__name__} has inconsistent outputs, {out1} vs {out2}' - - -def check_lazy_init(entry: TestingEntry, seed: int = 42, verbose: bool = False, check_forward: bool = False) -> None: - model_fn, data_gen_fn, output_transform_fn, model_attr = entry - _MyTensor._pre_op_fn = lambda *args: set_seed(seed) - LazyTensor._pre_op_fn = lambda *args: set_seed(seed) - ctx = LazyInitContext(tensor_cls=_MyTensor) - with ctx: - model = model_fn() - ctx = LazyInitContext() - with ctx: - deferred_model = model_fn() - deferred_model = ctx.materialize(deferred_model, verbose=verbose) - assert_model_eqaual(model, deferred_model) - if check_forward: - assert_forward_equal(model, deferred_model, data_gen_fn, output_transform_fn) - if verbose: - print(f'{model.__class__.__name__} pass') diff --git a/tests/test_utils/test_lazy_init_ctx.py b/tests/test_utils/test_lazy_init_ctx.py deleted file mode 100644 index 97efb3367490..000000000000 --- a/tests/test_utils/test_lazy_init_ctx.py +++ /dev/null @@ -1,51 +0,0 @@ -import torch -from colossalai.utils.model.lazy_init_context import LazyInitContext -from torchvision.models import resnet34 -import random -import numpy as np - -MANUAL_SEED = 0 -random.seed(MANUAL_SEED) -np.random.seed(MANUAL_SEED) -torch.manual_seed(MANUAL_SEED) - - -def test_lazy_init_with_meta(): - ctx = LazyInitContext(to_meta=True) - with ctx: - model = resnet34(num_classes=10) - - for param in model.parameters(): - assert param.is_meta - for buffer in model.buffers(): - assert buffer.is_meta - - ctx.lazy_init_parameters(model) - - for name, param in model.named_parameters(): - assert not param.is_meta, name - - for buffer in model.buffers(): - assert not buffer.is_meta - - -def test_lazy_init_without_meta(): - ctx = LazyInitContext(to_meta=False) - with ctx: - model = resnet34(num_classes=10) - - for param in model.parameters(): - assert not param.is_meta - for buffer in model.buffers(): - assert not buffer.is_meta - - conv1_weight_before_init = model.conv1.weight.clone() - ctx.lazy_init_parameters(model) - conv1_weight_after_init = model.conv1.weight.clone() - - assert not torch.allclose(conv1_weight_after_init, conv1_weight_before_init) - - -if __name__ == '__main__': - test_lazy_init_with_meta() - test_lazy_init_without_meta() diff --git a/tests/test_utils/test_memory.py b/tests/test_utils/test_memory.py index 46a5aeba505b..c88c2f8ec3c5 100644 --- a/tests/test_utils/test_memory.py +++ b/tests/test_utils/test_memory.py @@ -1,12 +1,9 @@ import pytest import colossalai +from colossalai.testing import spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.memory import colo_set_process_memory_fraction, colo_device_memory_capacity -from colossalai.utils import free_port - -from functools import partial -import torch.multiprocessing as mp +from colossalai.utils.memory import colo_device_memory_capacity, colo_set_process_memory_fraction def _run_colo_set_process_memory_fraction_and_colo_device_memory_capacity(): @@ -24,8 +21,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist @pytest.mark.parametrize("world_size", [3, 4]) def test_memory_utils(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_utils/test_norm_gradient_clipping.py b/tests/test_utils/test_norm_gradient_clipping.py index 259286663033..c0d678026c5f 100644 --- a/tests/test_utils/test_norm_gradient_clipping.py +++ b/tests/test_utils/test_norm_gradient_clipping.py @@ -1,16 +1,15 @@ -from colossalai.tensor import distspec, ColoTensorSpec, ProcessGroup -from colossalai.tensor.colo_parameter import ColoParameter -import colossalai import pytest import torch -import torch.multiprocessing as mp -from colossalai.logging import disable_existing_loggers -from colossalai.utils import free_port, get_current_device +from torch.nn.parameter import Parameter from torch.nn.utils import clip_grad_norm_ -from functools import partial -from colossalai.testing import parameterize, rerun_if_address_is_in_use + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.tensor import ColoTensorSpec, ProcessGroup, distspec +from colossalai.tensor.colo_parameter import ColoParameter +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device from colossalai.utils.common import clip_grad_norm -from torch.nn.parameter import Parameter def close(num: float, other: float, rtol: float = 1e-5, atol: float = 1e-8): @@ -71,8 +70,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_zero_clip_grad(world_size: int): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_utils/test_zero_gradient_clippling.py b/tests/test_utils/test_zero_gradient_clippling.py index 8bdae88464b1..e99cf388e929 100644 --- a/tests/test_utils/test_zero_gradient_clippling.py +++ b/tests/test_utils/test_zero_gradient_clippling.py @@ -1,22 +1,21 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -import copy +from functools import partial -import colossalai -from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import torch.nn as nn -from colossalai.logging import disable_existing_loggers -from colossalai.utils import checkpoint, clip_grad_norm_fp32, free_port from torch.nn.parallel import DistributedDataParallel as DDP from torch.nn.utils import clip_grad_norm_ -from colossalai.zero.shard_utils.tensor_shard_strategy import TensorShardStrategy -from functools import partial -from colossalai.testing import parameterize, rerun_if_address_is_in_use + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import checkpoint, clip_grad_norm_fp32 +from colossalai.zero.legacy.shard_utils.tensor_shard_strategy import TensorShardStrategy +from colossalai.zero.legacy.sharded_model.sharded_model_v2 import ShardedModelV2 def checkpoint_wrapper(module, enable=True): @@ -105,8 +104,7 @@ def run_dist(rank, world_size, port): @rerun_if_address_is_in_use() def test_zero_clip_grad(): world_size = 4 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_chunk_mgrv2.py b/tests/test_zero/test_gemini/test_chunk_mgrv2.py similarity index 89% rename from tests/test_gemini/update/test_chunk_mgrv2.py rename to tests/test_zero/test_gemini/test_chunk_mgrv2.py index 7d192fc631a6..7ea063877b5c 100644 --- a/tests/test_gemini/update/test_chunk_mgrv2.py +++ b/tests/test_zero/test_gemini/test_chunk_mgrv2.py @@ -1,14 +1,10 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai -from colossalai.gemini.chunk import ChunkManager from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero.gemini.chunk import ChunkManager from tests.test_tensor.common_utils import debug_print CUDA_MEM_0 = {False: 512, True: 1024} @@ -64,8 +60,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [2]) @rerun_if_address_is_in_use() def test_chunk_manager(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_chunkv2.py b/tests/test_zero/test_gemini/test_chunkv2.py similarity index 89% rename from tests/test_gemini/update/test_chunkv2.py rename to tests/test_zero/test_gemini/test_chunkv2.py index 96855410bea6..1cb31b260a99 100644 --- a/tests/test_gemini/update/test_chunkv2.py +++ b/tests/test_zero/test_gemini/test_chunkv2.py @@ -1,17 +1,14 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import colossalai -from colossalai.gemini import TensorState -from colossalai.gemini.chunk import Chunk from colossalai.tensor import ColoParameter from colossalai.tensor import ProcessGroup as ColoProcessGroup -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port, get_current_device +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero.gemini import TensorState +from colossalai.zero.gemini.chunk import Chunk def dist_sum(x): @@ -26,7 +23,7 @@ def add_param(param_list, param_cp_list, *args, **kwargs): param_cp_list.append(param.clone()) -def check_euqal(param, param_cp): +def check_equal(param, param_cp): if param.device != param_cp.device: temp = param.data.to(param_cp.device) else: @@ -60,7 +57,7 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory): my_chunk.append_tensor(param) assert my_chunk.utilized_size == 597 for param, param_cp in zip(param_list, param_cp_list): - check_euqal(param, param_cp) + check_equal(param, param_cp) my_chunk.close_chunk() if keep_gathered is False: @@ -80,7 +77,7 @@ def exam_chunk_basic(init_device, keep_gathered, pin_memory): my_chunk.access_chunk() assert my_chunk.device_type == 'cuda' for param, param_cp in zip(param_list, param_cp_list): - check_euqal(param, param_cp) + check_equal(param, param_cp) assert my_chunk.tensor_state_cnter[TensorState.HOLD] == 4 my_chunk.tensor_trans_state(param_list[0], TensorState.COMPUTE) @@ -117,8 +114,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2, 4]) @rerun_if_address_is_in_use() def test_chunk_function(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_fwd_bwd.py b/tests/test_zero/test_gemini/test_fwd_bwd.py similarity index 56% rename from tests/test_gemini/update/test_fwd_bwd.py rename to tests/test_zero/test_gemini/test_fwd_bwd.py index 0d35ba83d2e9..9c5455b8371b 100644 --- a/tests/test_gemini/update/test_fwd_bwd.py +++ b/tests/test_zero/test_gemini/test_fwd_bwd.py @@ -1,24 +1,18 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai from colossalai.amp import convert_to_apex_amp -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP from colossalai.tensor import ProcessGroup -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext -from tests.components_to_test import run_fwd_bwd +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager +from tests.components_to_test import run_fwd, run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import set_seed @@ -34,17 +28,17 @@ def check_grad(model: ZeroDDP, torch_model: torch.nn.Module): assert_close(p0, p1.grad, rtol=1e-3, atol=5e-5) -@parameterize('init_device', [get_current_device()]) @parameterize('placement_policy', ['cuda', 'cpu', 'auto', 'const']) @parameterize('keep_gather', [False, True]) @parameterize('model_name', ['gpt2', 'bert', 'albert']) @parameterize('use_grad_checkpoint', [False, True]) -def exam_gpt_fwd_bwd(placement_policy, - keep_gather, - model_name: str, - use_grad_checkpoint: bool = False, - init_device=get_current_device()): - +def exam_gpt_fwd_bwd( + placement_policy, + keep_gather, + model_name: str, + use_grad_checkpoint: bool = False, +): + init_device = get_current_device() get_components_func = non_distributed_component_funcs.get_callable(model_name) model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() @@ -58,7 +52,7 @@ def exam_gpt_fwd_bwd(placement_policy, torch_p.data.copy_(p.data) world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = keep_gather chunk_manager = ChunkManager(config_dict) @@ -95,18 +89,72 @@ def exam_gpt_fwd_bwd(placement_policy, check_grad(model, torch_model) +@parameterize('placement_policy', ['cuda', 'cpu']) +@parameterize('keep_gather', [False, True]) +@parameterize('model_name', ['gpt2', 'bert', 'albert']) +@parameterize('scatter_after_inference', [False, True]) +def exam_gpt_inference( + placement_policy, + keep_gather, + model_name: str, + scatter_after_inference: bool = False, +): + init_device = get_current_device() + get_components_func = non_distributed_component_funcs.get_callable(model_name) + model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() + + set_seed(42) + with ColoInitContext(device=init_device): + model = model_builder() + + set_seed(42) + torch_model = model_builder().cuda() + for torch_p, p in zip(torch_model.parameters(), model.parameters()): + torch_p.data.copy_(p.data) + + world_size = torch.distributed.get_world_size() + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) + config_dict[world_size]['chunk_size'] = 5000 + config_dict[world_size]['keep_gathered'] = keep_gather + chunk_manager = ChunkManager(config_dict) + gemini_manager = GeminiManager(placement_policy, chunk_manager) + model = ZeroDDP(model, gemini_manager, pin_memory=True, scatter_after_inference=scatter_after_inference) + + pg = ProcessGroup() + amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False, loss_scale=1) + torch_optim = torch.optim.Adam(torch_model.parameters(), lr=1e-3) + torch_model, torch_optim = convert_to_apex_amp(torch_model, torch_optim, amp_config) + torch_model = DDP(torch_model, device_ids=[pg.rank()], process_group=pg.dp_process_group()) + + set_seed(pg.dp_local_rank()) + model.eval() + torch_model.eval() + for i, (input_ids, label) in enumerate(train_dataloader): + # you can only test a single fwd + bwd. + # after bwd param is grad for Gemini, due to the chunk reuse optimization. + if i > 0: + break + with torch.no_grad(): + input_ids, label = input_ids.cuda(), label.cuda() + + torch_loss = run_fwd(torch_model, input_ids, label, criterion) + loss = run_fwd(model, input_ids, label, criterion) + + assert torch.equal(torch_loss, loss) + + def run_dist(rank, world_size, port): config = {} colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') exam_gpt_fwd_bwd() + exam_gpt_inference() @pytest.mark.dist @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_gpt(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_gemini_use_rmt.py b/tests/test_zero/test_gemini/test_gemini_use_rmt.py similarity index 84% rename from tests/test_gemini/update/test_gemini_use_rmt.py rename to tests/test_zero/test_gemini/test_gemini_use_rmt.py index 8cf17a0a726e..00e712050b32 100644 --- a/tests/test_gemini/update/test_gemini_use_rmt.py +++ b/tests/test_zero/test_gemini/test_gemini_use_rmt.py @@ -1,19 +1,13 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager -from colossalai.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer -from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer -from colossalai.nn.parallel import GeminiDDP, ZeroDDP from colossalai.tensor import ProcessGroup -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero import ColoInitContext, ZeroDDP +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager +from colossalai.zero.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer from tests.components_to_test import run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import set_seed @@ -62,7 +56,7 @@ def run_gemini_use_rmt(placement_policy, keep_gather, model_name: str, use_grad_ assert len(step_list) == 4 world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = keep_gather chunk_manager = ChunkManager(config_dict) @@ -100,8 +94,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_gemini_use_rmt(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_get_torch_model.py b/tests/test_zero/test_gemini/test_get_torch_model.py similarity index 81% rename from tests/test_gemini/update/test_get_torch_model.py rename to tests/test_zero/test_gemini/test_get_torch_model.py index e6d586b37041..b3e3b2b22fc3 100644 --- a/tests/test_gemini/update/test_get_torch_model.py +++ b/tests/test_zero/test_gemini/test_get_torch_model.py @@ -1,18 +1,12 @@ -import os -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import colossalai -from colossalai.nn.parallel import GeminiDDP -from colossalai.nn.parallel.utils import get_static_torch_model from colossalai.tensor import ColoParameter -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, GeminiDDP +from colossalai.zero.gemini.utils import get_static_torch_model from tests.components_to_test.registry import non_distributed_component_funcs @@ -51,8 +45,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_convert_torch_module(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_grad_clip.py b/tests/test_zero/test_gemini/test_grad_clip.py similarity index 85% rename from tests/test_gemini/update/test_grad_clip.py rename to tests/test_zero/test_gemini/test_grad_clip.py index d97ba94399c0..ac19a27f4a37 100644 --- a/tests/test_gemini/update/test_grad_clip.py +++ b/tests/test_zero/test_gemini/test_grad_clip.py @@ -1,27 +1,20 @@ -from functools import partial -from time import time - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai from colossalai.amp import convert_to_apex_amp -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager from tests.components_to_test import run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs -from tests.test_tensor.common_utils import debug_print, set_seed +from tests.test_tensor.common_utils import set_seed def check_param(model: ZeroDDP, torch_model: torch.nn.Module): @@ -58,7 +51,7 @@ def exam_grad_clipping(placement_policy, model_name: str): p.data.copy_(torch_p.data) world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = False if placement_policy != 'cuda': @@ -107,8 +100,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_grad_clip(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_inference.py b/tests/test_zero/test_gemini/test_inference.py similarity index 88% rename from tests/test_gemini/update/test_inference.py rename to tests/test_zero/test_gemini/test_inference.py index b057448ad378..fb2018f7b477 100644 --- a/tests/test_gemini/update/test_inference.py +++ b/tests/test_zero/test_gemini/test_inference.py @@ -1,24 +1,19 @@ -from functools import partial from typing import Callable import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai from colossalai.amp import convert_to_apex_amp -from colossalai.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer, post_process_colo_init_ctx, zero_model_wrapper +from colossalai.zero.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager from tests.components_to_test import run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import debug_print, set_seed @@ -39,7 +34,7 @@ def check_param(model: ZeroDDP, torch_model: torch.nn.Module): def multi_chunk_init(model: torch.nn.Module, placement_policy: str): world_size = dist.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = False if placement_policy != 'cuda': @@ -130,8 +125,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_inference(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_optim.py b/tests/test_zero/test_gemini/test_optim.py similarity index 71% rename from tests/test_gemini/update/test_optim.py rename to tests/test_zero/test_gemini/test_optim.py index cd3aa6051d78..a9ee67368e9d 100644 --- a/tests/test_gemini/update/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -1,24 +1,17 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai from colossalai.amp import convert_to_apex_amp -from colossalai.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP -from colossalai.tensor import ColoParameter, ColoTensor -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer, post_process_colo_init_ctx +from colossalai.zero.gemini.chunk import ChunkManager, init_chunk_manager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager from tests.components_to_test import run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import debug_print, set_seed @@ -28,23 +21,40 @@ # these models are too small, all parameters in these models are compacted into one chunk EXAMPLE_MODELS = ['albert', 'beit', 'bert', 'hanging_param_model', 'nested_model', 'repeated_computed_layers'] +# bfloat16 cannot represent them exactly +BF16_IGNORED_KEYS = [ + 'albert.embeddings.word_embeddings.weight', + 'albert.embeddings.position_embeddings.weight', + 'masked_bias', +] + -def check_param(model: ZeroDDP, torch_model: torch.nn.Module): - zero_dict = model.state_dict(only_rank_0=False) +def check_param(model: ZeroDDP, torch_model: torch.nn.Module, dtype: torch.dtype): + zero_dict = model.state_dict(only_rank_0=False, dtype=dtype) torch_dict = torch_model.state_dict() for key, value in torch_dict.items(): # key is 'module.model.PARAMETER', so we truncate it key = key[7:] assert key in zero_dict, "{} not in ZeRO dictionary.".format(key) - temp_zero_value = zero_dict[key].to(device=value.device, dtype=value.dtype) + temp_zero_value = zero_dict[key].to(device=value.device) + if dtype is torch.bfloat16 and any(k in key for k in BF16_IGNORED_KEYS): + continue + rtol, atol = 1e-3, 4e-3 + if dtype is torch.bfloat16: + rtol, atol = 4e-3, 8e-3 # debug_print([0], "max range: ", key, torch.max(torch.abs(value - temp_zero_value))) - assert_close(value, temp_zero_value, rtol=1e-3, atol=4e-3) + assert_close(value.float(), + temp_zero_value.float(), + rtol=rtol, + atol=atol, + msg=lambda s: s + f'\n{key}\n{temp_zero_value.dtype}') @parameterize('placement_policy', ['cuda', 'cpu', 'auto', 'const']) @parameterize('model_name', TEST_MODELS) -def exam_model_step(placement_policy, model_name: str): +@parameterize('mixed_precision', [torch.half, torch.bfloat16]) +def exam_model_step(placement_policy, model_name: str, mixed_precision: torch.dtype): set_seed(42) get_components_func = non_distributed_component_funcs.get_callable(model_name) model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() @@ -63,7 +73,7 @@ def exam_model_step(placement_policy, model_name: str): p.data.copy_(torch_p.data) world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = False if placement_policy != 'cuda': @@ -72,7 +82,7 @@ def exam_model_step(placement_policy, model_name: str): init_device = None chunk_manager = ChunkManager(config_dict, init_device=init_device) gemini_manager = GeminiManager(placement_policy, chunk_manager) - model = ZeroDDP(model, gemini_manager, pin_memory=True) + model = ZeroDDP(model, gemini_manager, pin_memory=True, mixed_precision=mixed_precision) optimizer = HybridAdam(model.parameters(), lr=1e-3) zero_optim = ZeroOptimizer(optimizer, model, initial_scale=128) @@ -81,6 +91,7 @@ def exam_model_step(placement_policy, model_name: str): torch_model.eval() set_seed(dist.get_rank() * 3 + 128) + rtol, atol = 1e-4, 1e-5 for i, (input_ids, label) in enumerate(train_dataloader): if i > 2: break @@ -90,17 +101,18 @@ def exam_model_step(placement_policy, model_name: str): torch_loss = run_fwd_bwd(torch_model, input_ids, label, criterion, torch_optim) loss = run_fwd_bwd(model, input_ids, label, criterion, zero_optim) - assert_close(torch_loss, loss) + assert_close(torch_loss, loss, rtol=rtol, atol=atol) zero_optim.step() torch_optim.step() - check_param(model, torch_model) + check_param(model, torch_model, mixed_precision) @parameterize('placement_policy', ['cuda', 'cpu', 'auto', 'const']) @parameterize('model_name', EXAMPLE_MODELS) -def exam_tiny_example(placement_policy, model_name: str): +@parameterize('mixed_precision', [torch.half, torch.bfloat16]) +def exam_tiny_example(placement_policy, model_name: str, mixed_precision: torch.dtype): set_seed(2008) get_components_func = non_distributed_component_funcs.get_callable(model_name) model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() @@ -118,9 +130,9 @@ def exam_tiny_example(placement_policy, model_name: str): for torch_p, p in zip(torch_model.parameters(), model.parameters()): p.data.copy_(torch_p.data) - chunk_manager = init_chunk_manager(model=model, init_device=get_current_device(), search_range_mb=1) + chunk_manager = init_chunk_manager(model=model, init_device=get_current_device(), search_range_m=1) gemini_manager = GeminiManager(placement_policy, chunk_manager) - model = ZeroDDP(model, gemini_manager, pin_memory=True) + model = ZeroDDP(model, gemini_manager, pin_memory=True, mixed_precision=mixed_precision) optimizer = HybridAdam(model.parameters(), lr=1e-3) zero_optim = ZeroOptimizer(optimizer, model, initial_scale=2) @@ -128,6 +140,9 @@ def exam_tiny_example(placement_policy, model_name: str): torch_model.eval() set_seed(dist.get_rank() * 3 + 128) + rtol, atol = 1.5e-6, 2e-5 + if mixed_precision is torch.bfloat16: + rtol, atol = 2e-3, 2e-3 for i, (input_ids, label) in enumerate(train_dataloader): if i > 2: break @@ -140,12 +155,12 @@ def exam_tiny_example(placement_policy, model_name: str): torch_loss = run_fwd_bwd(torch_model, input_ids, label, criterion, torch_optim) loss = run_fwd_bwd(model, input_ids, label, criterion, zero_optim) - assert_close(torch_loss, loss, rtol=1.5e-6, atol=2e-5) # atol should be 2e-5 for torch lower than 1.12 + assert_close(torch_loss, loss, rtol=rtol, atol=atol) # atol should be 2e-5 for torch lower than 1.12 zero_optim.step() torch_optim.step() - check_param(model, torch_model) + check_param(model, torch_model, mixed_precision) def run_dist(rank, world_size, port): @@ -159,8 +174,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_optim(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/test_runtime_mem_tracer.py b/tests/test_zero/test_gemini/test_runtime_mem_tracer.py similarity index 89% rename from tests/test_gemini/test_runtime_mem_tracer.py rename to tests/test_zero/test_gemini/test_runtime_mem_tracer.py index 294868458c47..0e6f283aa5d2 100644 --- a/tests/test_gemini/test_runtime_mem_tracer.py +++ b/tests/test_zero/test_gemini/test_runtime_mem_tracer.py @@ -3,12 +3,14 @@ import numpy as np import torch -from colossalai.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.testing import clear_cache_before_run +from colossalai.zero import ColoInitContext +from colossalai.zero.gemini.memory_tracer.runtime_mem_tracer import RuntimeMemTracer from tests.components_to_test import run_fwd_bwd from tests.components_to_test.registry import non_distributed_component_funcs +@clear_cache_before_run() def test_runtime_mem_tracer(): test_models = ['gpt2', 'bert', 'simple_net', 'repeated_computed_layers', 'nested_model', 'albert'] diff --git a/tests/test_gemini/update/test_search.py b/tests/test_zero/test_gemini/test_search.py similarity index 85% rename from tests/test_gemini/update/test_search.py rename to tests/test_zero/test_gemini/test_search.py index 2fcdd5380906..51dd84aace5b 100644 --- a/tests/test_gemini/update/test_search.py +++ b/tests/test_zero/test_gemini/test_search.py @@ -1,16 +1,12 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp import colossalai -from colossalai.gemini.chunk import init_chunk_manager, search_chunk_configuration from colossalai.tensor import ComputePattern, ComputeSpec, ProcessGroup, ShardSpec -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port, get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext +from colossalai.zero.gemini.chunk import init_chunk_manager, search_chunk_configuration from tests.components_to_test.registry import non_distributed_component_funcs @@ -34,9 +30,9 @@ def exam_search_chunk_size(): model = model_builder() init_1d_row_spec(model, pg_tp) config_dict, *_ = search_chunk_configuration(model, - search_range_mb=1, - search_interval_byte=16, - min_chunk_size_mb=0, + search_range_m=1, + search_interval=16, + min_chunk_size_m=0, filter_exlarge_params=True) for key in config_dict: @@ -58,9 +54,9 @@ def exam_search_strict_ddp(): with ColoInitContext(device=get_current_device()): ddp_model = model_builder() re_dict, re_total, re_wasted = search_chunk_configuration(ddp_model, - search_range_mb=1, - search_interval_byte=16, - min_chunk_size_mb=0, + search_range_m=1, + search_interval=16, + min_chunk_size_m=0, filter_exlarge_params=True, strict_ddp_flag=False) # get the chunk configuration over sharded ddp models @@ -68,9 +64,9 @@ def exam_search_strict_ddp(): default_dist_spec=default_shard_spec): sharded_ddp_model = model_builder() sh_dict, sh_total, sh_wasted = search_chunk_configuration(sharded_ddp_model, - search_range_mb=1, - search_interval_byte=16, - min_chunk_size_mb=0, + search_range_m=1, + search_interval=16, + min_chunk_size_m=0, filter_exlarge_params=True, strict_ddp_flag=True) assert re_dict == sh_dict @@ -95,8 +91,8 @@ def exam_chunk_manager(): chunk_manager = init_chunk_manager(sharded_ddp_model, get_current_device(), hidden_dim=16, - search_range_mb=1, - min_chunk_size_mb=0, + search_range_m=1, + min_chunk_size_m=0, filter_exlarge_params=True, strict_ddp_flag=True) config_dict = chunk_manager.dp_degree_chunk_size_dict @@ -115,8 +111,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_search(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/update/test_zeroddp_state_dict.py b/tests/test_zero/test_gemini/test_zeroddp_state_dict.py similarity index 87% rename from tests/test_gemini/update/test_zeroddp_state_dict.py rename to tests/test_zero/test_gemini/test_zeroddp_state_dict.py index 00d835842f79..2a5a4ab83029 100644 --- a/tests/test_gemini/update/test_zeroddp_state_dict.py +++ b/tests/test_zero/test_gemini/test_zeroddp_state_dict.py @@ -1,19 +1,13 @@ -from functools import partial - import pytest import torch -import torch.distributed as dist -import torch.multiprocessing as mp from torch.testing import assert_close import colossalai -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager -from colossalai.nn.parallel import ZeroDDP -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, ZeroDDP +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import debug_print, set_seed @@ -41,7 +35,7 @@ def exam_state_dict(placement_policy, keep_gathered, model_name: str): torch_p.data.copy_(p.data) world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = keep_gathered chunk_manager = ChunkManager(config_dict) @@ -73,7 +67,7 @@ def exam_load_state_dict(placement_policy, keep_gathered, model_name: str): torch_model = model_builder() # get a different model world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = keep_gathered @@ -106,8 +100,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_zero_ddp(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_gemini/test_zeroddp_state_dict_shard.py b/tests/test_zero/test_gemini/test_zeroddp_state_dict_shard.py new file mode 100644 index 000000000000..d16bfb7d1622 --- /dev/null +++ b/tests/test_zero/test_gemini/test_zeroddp_state_dict_shard.py @@ -0,0 +1,56 @@ +import pytest +import torch +from torch.testing import assert_close + +import colossalai +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils.cuda import get_current_device +from colossalai.zero import ColoInitContext, ZeroDDP +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager +from tests.components_to_test.registry import non_distributed_component_funcs + + +@parameterize('placement_policy', ['cuda', 'cpu']) +@parameterize('model_name', ['gpt2', 'bert']) +def exam_state_dict(placement_policy, model_name: str): + get_components_func = non_distributed_component_funcs.get_callable(model_name) + model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func() + + with ColoInitContext(device=get_current_device()): + model = model_builder() + + model_size = sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2 + + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) + chunk_manager = ChunkManager(config_dict) + gemini_manager = GeminiManager(placement_policy, chunk_manager) + model = ZeroDDP(model, gemini_manager) + model.train() + + zero_dict = model.state_dict(only_rank_0=False) + accumulated_keys = set() + # ensure number of shards > 1 + for shard, _ in model.state_dict_shard(max_shard_size=(model_size / 3), only_rank_0=False): + for key, value in shard.items(): + assert key not in accumulated_keys, f"key `{key}` is duplicated." + accumulated_keys.add(key) + assert key in zero_dict, f"{key} not in ZeRO dictionary." + assert torch.equal(value, zero_dict[key]), f"{key} not equal." + + +def run_dist(rank, world_size, port): + config = {} + colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + exam_state_dict() + + +@pytest.mark.dist +@pytest.mark.parametrize('world_size', [1, 4]) +@rerun_if_address_is_in_use() +def test_zero_ddp_state_dict_shard(world_size): + spawn(run_dist, world_size) + + +if __name__ == '__main__': + test_zero_ddp_state_dict_shard(1) diff --git a/tests/test_gemini/update/test_zerooptim_state_dict.py b/tests/test_zero/test_gemini/test_zerooptim_state_dict.py similarity index 82% rename from tests/test_gemini/update/test_zerooptim_state_dict.py rename to tests/test_zero/test_gemini/test_zerooptim_state_dict.py index fd13af6b2b0a..ba016d6528dc 100644 --- a/tests/test_gemini/update/test_zerooptim_state_dict.py +++ b/tests/test_zero/test_gemini/test_zerooptim_state_dict.py @@ -1,20 +1,14 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import colossalai -from colossalai.gemini.chunk import ChunkManager, search_chunk_configuration -from colossalai.gemini.gemini_mgr import GeminiManager from colossalai.nn.optimizer import HybridAdam -from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer -from colossalai.nn.parallel import ZeroDDP -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.zero import ColoInitContext, ZeroDDP, ZeroOptimizer +from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration +from colossalai.zero.gemini.gemini_mgr import GeminiManager from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_tensor.common_utils import debug_print, set_seed @@ -33,7 +27,7 @@ def exam_zero_optim_state_dict(placement_policy, keep_gathered): torch_model = model_builder() # get a different model world_size = torch.distributed.get_world_size() - config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100) + config_dict, *_ = search_chunk_configuration(model, search_range_m=1, search_interval=100) config_dict[world_size]['chunk_size'] = 5000 config_dict[world_size]['keep_gathered'] = keep_gathered @@ -85,8 +79,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 4]) @rerun_if_address_is_in_use() def test_zero_optim(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/common.py b/tests/test_zero/test_legacy/common.py similarity index 97% rename from tests/test_zero/common.py rename to tests/test_zero/test_legacy/common.py index bc6cd75a6a60..2c3d122c79af 100644 --- a/tests/test_zero/common.py +++ b/tests/test_zero/test_legacy/common.py @@ -2,10 +2,11 @@ import torch import torch.distributed as dist + from colossalai.logging import get_dist_logger from colossalai.utils import checkpoint -from colossalai.zero.shard_utils import TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.shard_utils import TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 LOGGER = get_dist_logger('zero_test') diff --git a/tests/test_zero/test_found_inf.py b/tests/test_zero/test_legacy/test_found_inf.py similarity index 78% rename from tests/test_zero/test_found_inf.py rename to tests/test_zero/test_legacy/test_found_inf.py index 34283f5015e1..e90158e0a43b 100644 --- a/tests/test_zero/test_found_inf.py +++ b/tests/test_zero/test_legacy/test_found_inf.py @@ -1,72 +1,67 @@ -from functools import partial - -import colossalai -from colossalai.utils.cuda import get_current_device -import pytest -import torch -import torch.multiprocessing as mp -from colossalai.nn.optimizer import HybridAdam -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_optim import ShardedOptimizerV2 -from colossalai.zero.sharded_optim._utils import has_inf_or_nan -from tests.components_to_test.registry import non_distributed_component_funcs -from tests.test_zero.test_sharded_optim_v2 import _run_step - -from common import CONFIG - - -@parameterize("cpu_offload", [True, False]) -@parameterize("shard_strategy_class", [BucketTensorShardStrategy]) -@parameterize("gpu_margin_mem_ratio", [0.0, 0.7]) -def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio): - test_models = ['repeated_computed_layers'] - shard_strategy = shard_strategy_class() - - for model_name in test_models: - get_components_func = non_distributed_component_funcs.get_callable(model_name) - model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func() - - with ZeroInitContext(target_device=torch.device(f'cpu:0') if cpu_offload else get_current_device(), - shard_strategy=shard_strategy, - shard_param=True): - zero_model = model_builder(checkpoint=True) - zero_model = ShardedModelV2( - zero_model, - shard_strategy, - tensor_placement_policy='cpu' if cpu_offload else 'cuda', - reuse_fp16_shard=True, - ) - - sharded_optim = HybridAdam(zero_model.parameters(), lr=1e-3) - sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, gpu_margin_mem_ratio=gpu_margin_mem_ratio) - - for i, (data, label) in enumerate(train_dataloader): - if i > 1: - break - assert zero_model.overflow_counter == 0 - data, label = data.cuda(), label.cuda() - _run_step(zero_model, sharded_optim, data, label, criterion, False) - for param in zero_model.parameters(): - assert not has_inf_or_nan(param.colo_attr.data_payload) - - -def _run_dist(rank, world_size, port): - colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - _run_test_found_inf() - - -# use_cpuadam = True can be used with cpu_offload = False -@pytest.mark.dist -@pytest.mark.parametrize("world_size", [1, 2]) -@rerun_if_address_is_in_use() -def test_found_inf(world_size): - run_func = partial(_run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) - - -if __name__ == '__main__': - test_found_inf(world_size=2) +import pytest +import torch +from common import CONFIG +from test_sharded_optim_v2 import _run_step + +import colossalai +from colossalai.nn.optimizer import HybridAdam +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils.cuda import get_current_device +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 +from colossalai.zero.low_level._utils import has_inf_or_nan +from tests.components_to_test.registry import non_distributed_component_funcs + + +@parameterize("cpu_offload", [True, False]) +@parameterize("shard_strategy_class", [BucketTensorShardStrategy]) +@parameterize("gpu_margin_mem_ratio", [0.0, 0.7]) +def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio): + test_models = ['repeated_computed_layers'] + shard_strategy = shard_strategy_class() + + for model_name in test_models: + get_components_func = non_distributed_component_funcs.get_callable(model_name) + model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func() + + with ZeroInitContext(target_device=torch.device(f'cpu:0') if cpu_offload else get_current_device(), + shard_strategy=shard_strategy, + shard_param=True): + zero_model = model_builder(checkpoint=True) + zero_model = ShardedModelV2( + zero_model, + shard_strategy, + tensor_placement_policy='cpu' if cpu_offload else 'cuda', + reuse_fp16_shard=True, + ) + + sharded_optim = HybridAdam(zero_model.parameters(), lr=1e-3) + sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, gpu_margin_mem_ratio=gpu_margin_mem_ratio) + + for i, (data, label) in enumerate(train_dataloader): + if i > 1: + break + assert zero_model.overflow_counter == 0 + data, label = data.cuda(), label.cuda() + _run_step(zero_model, sharded_optim, data, label, criterion, False) + for param in zero_model.parameters(): + assert not has_inf_or_nan(param.colo_attr.data_payload) + + +def _run_dist(rank, world_size, port): + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + _run_test_found_inf() + + +# use_cpuadam = True can be used with cpu_offload = False +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [1, 2]) +@rerun_if_address_is_in_use() +def test_found_inf(world_size): + spawn(_run_dist, world_size) + + +if __name__ == '__main__': + test_found_inf(world_size=2) diff --git a/tests/test_gemini/test_gemini_manager.py b/tests/test_zero/test_legacy/test_gemini_manager.py similarity index 94% rename from tests/test_gemini/test_gemini_manager.py rename to tests/test_zero/test_legacy/test_gemini_manager.py index 0c138f101f75..0e956f7cc617 100644 --- a/tests/test_gemini/test_gemini_manager.py +++ b/tests/test_zero/test_legacy/test_gemini_manager.py @@ -1,73 +1,75 @@ -import pytest -import torch - -from colossalai.gemini.stateful_tensor import TensorState, StatefulTensor - - -@pytest.mark.dist -def test_gemini_manager(): - # reset the manager, in case that there exists memory information left - manager = StatefulTensor.GST_MGR - manager.reset() - - # occupation 8 - st1 = StatefulTensor(torch.empty(2, 2, dtype=torch.float16, device='cuda')) - # occupation 60 - st2 = StatefulTensor(torch.empty(3, 5, dtype=torch.float32, device='cpu')) - - # occupation 28 - t1 = torch.empty(7, device='cuda') - # occupation 12 - t2 = torch.empty(3, device='cpu') - st3 = StatefulTensor(t1, TensorState.HOLD_AFTER_FWD) - st4 = StatefulTensor(None, TensorState.FREE) - - assert manager.total_number == 4 - assert manager.total_mem['cpu'] == 60 - assert manager.total_mem['cuda'] == 36 - assert manager.state_mem['cpu'][TensorState.HOLD] == 60 - assert manager.state_mem['cuda'][TensorState.HOLD] == 8 - assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 28 - - st4.payload_reset(t2) - st3.payload_reset(t2) - - assert manager.total_number == 4 - assert manager.total_mem['cpu'] == 84 - assert manager.total_mem['cuda'] == 8 - assert manager.state_mem['cpu'][TensorState.HOLD] == 72 - assert manager.state_mem['cuda'][TensorState.HOLD] == 8 - assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 12 - assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 0 - - st1.move_to(torch.device('cpu')) - st2.move_to(torch.device('cpu')) - st3.move_to(torch.device('cuda', 0)) - - assert manager.total_number == 4 - assert manager.total_mem['cpu'] == 80 - assert manager.total_mem['cuda'] == 12 - assert manager.state_mem['cpu'][TensorState.HOLD] == 80 - assert manager.state_mem['cuda'][TensorState.HOLD] == 0 - assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 0 - assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 12 - - st1.trans_state(TensorState.COMPUTE) - st2.trans_state(TensorState.COMPUTE) - st2.trans_state(TensorState.HOLD_AFTER_BWD) - - assert manager.total_number == 4 - assert manager.total_mem['cpu'] == 80 - assert manager.total_mem['cuda'] == 12 - assert manager.state_mem['cpu'][TensorState.HOLD] == 12 - assert manager.state_mem['cuda'][TensorState.HOLD] == 0 - assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 0 - assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 12 - assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_BWD] == 60 - assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_BWD] == 0 - assert manager.state_mem['cpu'][TensorState.COMPUTE] == 8 - assert manager.state_mem['cuda'][TensorState.COMPUTE] == 0 - - -if __name__ == '__main__': - test_gemini_manager() +import pytest +import torch + +from colossalai.testing import clear_cache_before_run +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor, TensorState + + +@pytest.mark.dist +@clear_cache_before_run() +def test_gemini_manager(): + # reset the manager, in case that there exists memory information left + manager = StatefulTensor.GST_MGR + manager.reset() + + # occupation 8 + st1 = StatefulTensor(torch.empty(2, 2, dtype=torch.float16, device='cuda')) + # occupation 60 + st2 = StatefulTensor(torch.empty(3, 5, dtype=torch.float32, device='cpu')) + + # occupation 28 + t1 = torch.empty(7, device='cuda') + # occupation 12 + t2 = torch.empty(3, device='cpu') + st3 = StatefulTensor(t1, TensorState.HOLD_AFTER_FWD) + st4 = StatefulTensor(None, TensorState.FREE) + + assert manager.total_number == 4 + assert manager.total_mem['cpu'] == 60 + assert manager.total_mem['cuda'] == 36 + assert manager.state_mem['cpu'][TensorState.HOLD] == 60 + assert manager.state_mem['cuda'][TensorState.HOLD] == 8 + assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 28 + + st4.payload_reset(t2) + st3.payload_reset(t2) + + assert manager.total_number == 4 + assert manager.total_mem['cpu'] == 84 + assert manager.total_mem['cuda'] == 8 + assert manager.state_mem['cpu'][TensorState.HOLD] == 72 + assert manager.state_mem['cuda'][TensorState.HOLD] == 8 + assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 12 + assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 0 + + st1.move_to(torch.device('cpu')) + st2.move_to(torch.device('cpu')) + st3.move_to(torch.device('cuda', 0)) + + assert manager.total_number == 4 + assert manager.total_mem['cpu'] == 80 + assert manager.total_mem['cuda'] == 12 + assert manager.state_mem['cpu'][TensorState.HOLD] == 80 + assert manager.state_mem['cuda'][TensorState.HOLD] == 0 + assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 0 + assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 12 + + st1.trans_state(TensorState.COMPUTE) + st2.trans_state(TensorState.COMPUTE) + st2.trans_state(TensorState.HOLD_AFTER_BWD) + + assert manager.total_number == 4 + assert manager.total_mem['cpu'] == 80 + assert manager.total_mem['cuda'] == 12 + assert manager.state_mem['cpu'][TensorState.HOLD] == 12 + assert manager.state_mem['cuda'][TensorState.HOLD] == 0 + assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_FWD] == 0 + assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_FWD] == 12 + assert manager.state_mem['cpu'][TensorState.HOLD_AFTER_BWD] == 60 + assert manager.state_mem['cuda'][TensorState.HOLD_AFTER_BWD] == 0 + assert manager.state_mem['cpu'][TensorState.COMPUTE] == 8 + assert manager.state_mem['cuda'][TensorState.COMPUTE] == 0 + + +if __name__ == '__main__': + test_gemini_manager() diff --git a/tests/test_zero/test_init_context.py b/tests/test_zero/test_legacy/test_init_context.py similarity index 86% rename from tests/test_zero/test_init_context.py rename to tests/test_zero/test_legacy/test_init_context.py index 0cba7a492380..84493827193e 100644 --- a/tests/test_zero/test_init_context.py +++ b/tests/test_zero/test_legacy/test_init_context.py @@ -1,22 +1,18 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from common import CONFIG import colossalai -from colossalai.gemini.memory_tracer.utils import colo_model_mem_usage from colossalai.logging import get_dist_logger -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device from colossalai.utils.memory import colo_device_memory_used -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.gemini.memory_tracer.utils import colo_model_mem_usage +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy from tests.components_to_test.registry import non_distributed_component_funcs @@ -70,8 +66,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 4]) @rerun_if_address_is_in_use() def test_zero_init_context(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_gemini/test_param_op.py b/tests/test_zero/test_legacy/test_param_op.py similarity index 94% rename from tests/test_gemini/test_param_op.py rename to tests/test_zero/test_legacy/test_param_op.py index daf386d6d6af..b91371b98922 100644 --- a/tests/test_gemini/test_param_op.py +++ b/tests/test_zero/test_legacy/test_param_op.py @@ -2,7 +2,8 @@ import torch -from colossalai.gemini.paramhooks import BaseParamHookMgr +from colossalai.testing import clear_cache_before_run +from colossalai.zero.legacy.gemini.paramhooks import BaseParamHookMgr from tests.components_to_test.registry import non_distributed_component_funcs @@ -49,6 +50,7 @@ def hook(param, grad) -> torch.Tensor or None: return hookwrapper.hook_triggered_times +@clear_cache_before_run() def test_base_param_hook(): test_models = ['repeated_computed_layers', 'resnet18', 'hanging_param_model', 'inline_op_model'] # test_models = ['bert'] diff --git a/tests/test_zero/test_shard_model_v2.py b/tests/test_zero/test_legacy/test_shard_model_v2.py similarity index 79% rename from tests/test_zero/test_shard_model_v2.py rename to tests/test_zero/test_legacy/test_shard_model_v2.py index 95a9dee38acf..93d624aa2bbd 100644 --- a/tests/test_zero/test_shard_model_v2.py +++ b/tests/test_zero/test_legacy/test_shard_model_v2.py @@ -1,22 +1,18 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - import pytest import torch -import torch.multiprocessing as mp from common import CONFIG, check_grads_padding, run_fwd_bwd from torch.nn.parallel import DistributedDataParallel as DDP import colossalai -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16 -from colossalai.zero.sharded_model.utils import col_model_deepcopy +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp16 +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy from tests.components_to_test.registry import non_distributed_component_funcs @@ -61,8 +57,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 2]) @rerun_if_address_is_in_use() def test_shard_model_v2(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_shard_param.py b/tests/test_zero/test_legacy/test_shard_param.py similarity index 80% rename from tests/test_zero/test_shard_param.py rename to tests/test_zero/test_legacy/test_shard_param.py index 8db2b7e79604..4ba43edceb5d 100644 --- a/tests/test_zero/test_shard_param.py +++ b/tests/test_zero/test_legacy/test_shard_param.py @@ -1,17 +1,15 @@ from copy import deepcopy -from functools import partial -import colossalai import pytest import torch -import torch.multiprocessing as mp -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy) -from colossalai.zero.sharded_param import ShardedTensor -from colossalai.zero.sharded_param.sharded_param import ShardedParamV2 -from tests.test_zero.common import CONFIG, allclose -from colossalai.gemini.stateful_tensor import StatefulTensor +from common import CONFIG, allclose + +import colossalai +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.legacy.sharded_param import ShardedTensor +from colossalai.zero.legacy.sharded_param.sharded_param import ShardedParamV2 @parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy]) @@ -38,8 +36,7 @@ def _run_shard_tensor(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 2]) @rerun_if_address_is_in_use() def test_shard_tensor(world_size): - run_func = partial(_run_shard_tensor, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(_run_shard_tensor, world_size) def _run_shard_param_v2(rank, world_size, port): @@ -86,8 +83,7 @@ def _run_shard_param_v2(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 2]) @rerun_if_address_is_in_use() def test_shard_param_v2(world_size): - run_func = partial(_run_shard_param_v2, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(_run_shard_param_v2, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_sharded_optim_state_dict.py b/tests/test_zero/test_legacy/test_sharded_optim_state_dict.py similarity index 83% rename from tests/test_zero/test_sharded_optim_state_dict.py rename to tests/test_zero/test_legacy/test_sharded_optim_state_dict.py index f8c42930b281..1ca144662722 100644 --- a/tests/test_zero/test_sharded_optim_state_dict.py +++ b/tests/test_zero/test_legacy/test_sharded_optim_state_dict.py @@ -1,20 +1,17 @@ import pytest -import colossalai import torch -import torch.multiprocessing as mp -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils.cuda import get_current_device -from colossalai.utils import free_port -from functools import partial -from tests.test_tensor.common_utils import set_seed -from tests.components_to_test.registry import non_distributed_component_funcs -from colossalai.testing import parameterize + +import colossalai from colossalai.nn.optimizer import HybridAdam -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_optim import ShardedOptimizerV2 from colossalai.tensor import ProcessGroup +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils.cuda import get_current_device +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 +from tests.components_to_test.registry import non_distributed_component_funcs +from tests.test_tensor.common_utils import set_seed def init_zero(model_builder, placement_policy): @@ -85,8 +82,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize('world_size', [1, 2]) @rerun_if_address_is_in_use() def test_sharded_optim_state_dist(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_sharded_optim_v2.py b/tests/test_zero/test_legacy/test_sharded_optim_v2.py similarity index 86% rename from tests/test_zero/test_sharded_optim_v2.py rename to tests/test_zero/test_legacy/test_sharded_optim_v2.py index 8fe7eb639eab..c6f77995ebcd 100644 --- a/tests/test_zero/test_sharded_optim_v2.py +++ b/tests/test_zero/test_legacy/test_sharded_optim_v2.py @@ -1,24 +1,20 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp from common import CONFIG, check_sharded_model_params from torch.nn.parallel import DistributedDataParallel as DDP import colossalai from colossalai.amp import convert_to_apex_amp from colossalai.nn.optimizer import CPUAdam -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model.utils import col_model_deepcopy -from colossalai.zero.sharded_optim import ShardedOptimizerV2 -from colossalai.zero.sharded_optim._utils import has_inf_or_nan +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy +from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 +from colossalai.zero.low_level._utils import has_inf_or_nan from tests.components_to_test.registry import non_distributed_component_funcs @@ -107,8 +103,7 @@ def _run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 2]) @rerun_if_address_is_in_use() def test_sharded_optim_v2(world_size): - run_func = partial(_run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(_run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_sharded_optim_with_sync_bn.py b/tests/test_zero/test_legacy/test_sharded_optim_with_sync_bn.py similarity index 86% rename from tests/test_zero/test_sharded_optim_with_sync_bn.py rename to tests/test_zero/test_legacy/test_sharded_optim_with_sync_bn.py index ea5b315188a3..0223f18c29d6 100644 --- a/tests/test_zero/test_sharded_optim_with_sync_bn.py +++ b/tests/test_zero/test_legacy/test_sharded_optim_with_sync_bn.py @@ -1,20 +1,17 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - -import colossalai import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp +from torchvision.models import resnet50 + +import colossalai from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import TensorShardStrategy -from torchvision.models import resnet50 +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import TensorShardStrategy def run_dist(rank, world_size, port): @@ -23,7 +20,7 @@ def run_dist(rank, world_size, port): # need to configure cudnn deterministic so that # randomness of convolution layers will be disabled zero_config = dict(model_config=dict(shard_strategy=TensorShardStrategy())) - colossalai.launch(config=dict(zero=zero_config, cudnn_determinstic=True, cudnn_benchmark=False), + colossalai.launch(config=dict(zero=zero_config, cudnn_deterministic=True, cudnn_benchmark=False), rank=rank, world_size=world_size, host='localhost', @@ -83,9 +80,7 @@ def test_sharded_optim_with_sync_bn(): wanted if we are doing predictions. """ - world_size = 2 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 2) if __name__ == '__main__': diff --git a/tests/test_zero/test_state_dict.py b/tests/test_zero/test_legacy/test_state_dict.py similarity index 78% rename from tests/test_zero/test_state_dict.py rename to tests/test_zero/test_legacy/test_state_dict.py index 7ac9b151e4d6..5f76fff3e5c3 100644 --- a/tests/test_zero/test_state_dict.py +++ b/tests/test_zero/test_legacy/test_state_dict.py @@ -1,23 +1,20 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from copy import deepcopy from functools import partial -import colossalai import pytest import torch -import torch.multiprocessing as mp -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy) -from colossalai.zero.sharded_model import ShardedModelV2 -from colossalai.zero.sharded_model.utils import col_model_deepcopy -from tests.components_to_test.registry import non_distributed_component_funcs - from common import CONFIG +import colossalai +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy +from colossalai.zero.legacy.sharded_model import ShardedModelV2 +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy +from tests.components_to_test.registry import non_distributed_component_funcs + @parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy]) def run_zero_state_dict(shard_strategy_class): @@ -51,8 +48,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [1, 2]) @rerun_if_address_is_in_use() def test_zero_state_dict(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_tensor_utils.py b/tests/test_zero/test_legacy/test_tensor_utils.py similarity index 82% rename from tests/test_zero/test_tensor_utils.py rename to tests/test_zero/test_legacy/test_tensor_utils.py index 81855ff5e10a..238bc3fe1a98 100644 --- a/tests/test_zero/test_tensor_utils.py +++ b/tests/test_zero/test_legacy/test_tensor_utils.py @@ -1,18 +1,17 @@ import pytest +import torch import colossalai +from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.utils.cuda import get_current_device -from colossalai.gemini.tensor_utils import (colo_tensor_mem_usage, colo_model_data_tensor_move, - colo_model_data_tensor_move_inline, colo_model_data_move_to_cpu, - colo_model_tensor_clone) -from colossalai.gemini.stateful_tensor import StatefulTensor -from colossalai.utils import free_port -from colossalai.testing import rerun_if_address_is_in_use - -import torch - -from functools import partial -import torch.multiprocessing as mp +from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor +from colossalai.zero.legacy.gemini.tensor_utils import ( + colo_model_data_move_to_cpu, + colo_model_data_tensor_move, + colo_model_data_tensor_move_inline, + colo_model_tensor_clone, + colo_tensor_mem_usage, +) def _run_colo_tensor_mem_usage(): @@ -88,8 +87,7 @@ def run_dist(rank, world_size, port): @pytest.mark.parametrize("world_size", [2, 4]) @rerun_if_address_is_in_use() def test_zero_tensor_utils(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size) if __name__ == '__main__': diff --git a/tests/test_zero/test_zero_engine.py b/tests/test_zero/test_legacy/test_zero_engine.py similarity index 74% rename from tests/test_zero/test_zero_engine.py rename to tests/test_zero/test_legacy/test_zero_engine.py index 80ded65d634c..826a543db861 100644 --- a/tests/test_zero/test_zero_engine.py +++ b/tests/test_zero/test_legacy/test_zero_engine.py @@ -1,26 +1,26 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from functools import partial - -import colossalai import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp -from colossalai.core import global_context as gpc -from colossalai.testing import rerun_if_address_is_in_use -from colossalai.utils import free_port -from colossalai.zero.init_ctx import ZeroInitContext -from colossalai.zero.sharded_model.utils import col_model_deepcopy -from colossalai.zero.sharded_optim._utils import has_inf_or_nan -from tests.components_to_test.registry import non_distributed_component_funcs +from common import MP_PARALLEL_CONFIG, ZERO_PARALLEL_CONFIG, check_params, check_sharded_model_params from torch.nn.parallel import DistributedDataParallel as DDP -from common import (MP_PARALLEL_CONFIG, ZERO_PARALLEL_CONFIG, check_params, check_sharded_model_params) +import colossalai +from colossalai.core import global_context as gpc +from colossalai.testing import rerun_if_address_is_in_use, spawn +from colossalai.zero.legacy.init_ctx import ZeroInitContext +from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy +from colossalai.zero.low_level._utils import has_inf_or_nan +from tests.components_to_test.registry import non_distributed_component_funcs -def run_dist(rank, world_size, port, parallel_config): +def run_dist(rank, world_size, port, parallel_config, bf16): + is_mp_config = parallel_config == MP_PARALLEL_CONFIG + is_zero_config = parallel_config == ZERO_PARALLEL_CONFIG + if bf16: + parallel_config['zero']['model_config']['bf16'] = True colossalai.launch(config=parallel_config, rank=rank, world_size=world_size, @@ -34,7 +34,8 @@ def run_dist(rank, world_size, port, parallel_config): model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func() with ZeroInitContext(target_device=torch.cuda.current_device(), shard_strategy=gpc.config.zero.model_config.shard_strategy, - shard_param=True): + shard_param=True, + bf16=bf16): colo_model = model_builder(checkpoint=True) colo_optimizer = optimizer_class(colo_model.parameters(), lr=1e-3) @@ -42,7 +43,8 @@ def run_dist(rank, world_size, port, parallel_config): optimizer=colo_optimizer, criterion=criterion, train_dataloader=train_dataloader) - torch_model = model_builder(checkpoint=True).half() + dtype = torch.bfloat16 if bf16 else torch.float16 + torch_model = model_builder(checkpoint=True).to(dtype) col_model_deepcopy(engine.model, torch_model) torch_model = torch_model.cuda().float() @@ -84,9 +86,9 @@ def run_dist(rank, world_size, port, parallel_config): torch_optimizer.step() i += 1 - if parallel_config == MP_PARALLEL_CONFIG: + if is_mp_config: check_params(torch_model, colo_model, loose=True) - elif parallel_config == ZERO_PARALLEL_CONFIG: + elif is_zero_config: check_sharded_model_params(torch_model, colo_model, loose=True) @@ -96,16 +98,15 @@ def run_dist(rank, world_size, port, parallel_config): @pytest.mark.parametrize("world_size", [2, 4]) @rerun_if_address_is_in_use() def test_mp_engine(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port(), parallel_config=MP_PARALLEL_CONFIG) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, world_size, parallel_config=MP_PARALLEL_CONFIG) @pytest.mark.dist @pytest.mark.parametrize("world_size", [1, 2]) +@pytest.mark.parametrize("bf16", [True, False]) @rerun_if_address_is_in_use() -def test_zero_engine(world_size): - run_func = partial(run_dist, world_size=world_size, port=free_port(), parallel_config=ZERO_PARALLEL_CONFIG) - mp.spawn(run_func, nprocs=world_size) +def test_zero_engine(world_size, bf16): + spawn(run_dist, world_size, parallel_config=ZERO_PARALLEL_CONFIG, bf16=bf16) if __name__ == '__main__': diff --git a/tests/test_zero/low_level_zero/test_grad_acc.py b/tests/test_zero/test_low_level/test_grad_acc.py similarity index 92% rename from tests/test_zero/low_level_zero/test_grad_acc.py rename to tests/test_zero/test_low_level/test_grad_acc.py index 504df202e168..c264a8077d2a 100644 --- a/tests/test_zero/low_level_zero/test_grad_acc.py +++ b/tests/test_zero/test_low_level/test_grad_acc.py @@ -1,16 +1,14 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai +from colossalai.testing import spawn from colossalai.testing.random import seed_all -from colossalai.utils import free_port from colossalai.zero import LowLevelZeroOptimizer @@ -84,7 +82,6 @@ def fwd_bwd_func(number, cur_data): def exam_zero_1_grad_acc(): local_rank = torch.distributed.get_rank() - grad_scale = 32 seed_all(2008) # create models @@ -103,7 +100,6 @@ def exam_zero_1_grad_acc(): # level 1 and 2 will produce exactly the same results zero_optimizer = LowLevelZeroOptimizer(zero_optimizer, overlap_communication=False, - initial_scale=grad_scale, reduce_bucket_size=262144, clip_grad_norm=1.0) @@ -130,9 +126,8 @@ def fwd_bwd_func(number, cur_data, check_flag): if check_flag: # check grad for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()): - unscale_grad = z1p.grad / grad_scale # print(n, p.shape, torch.max(torch.abs(p.grad - unscale_grad))) - assert torch.equal(p.grad, unscale_grad) + assert torch.equal(p.grad, z1p.grad) zero_optimizer._sync_grad() @@ -158,9 +153,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist def test_grad_accumulation(): - world_size = 2 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 2) if __name__ == '__main__': diff --git a/tests/test_zero/low_level_zero/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py similarity index 84% rename from tests/test_zero/low_level_zero/test_zero1_2.py rename to tests/test_zero/test_low_level/test_zero1_2.py index 930b6129174e..8e2206fe6c8d 100644 --- a/tests/test_zero/low_level_zero/test_zero1_2.py +++ b/tests/test_zero/test_low_level/test_zero1_2.py @@ -1,16 +1,14 @@ import copy -from functools import partial import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.random import seed_all -from colossalai.utils import free_port from colossalai.zero import LowLevelZeroOptimizer @@ -27,15 +25,18 @@ def forward(self, x): return x -def half_close(a, b, loose=False): +def loose_close(a, b, dtype: torch.dtype = torch.float32): rtol = None atol = None - if loose: + if dtype is torch.float16: rtol = 5e-2 atol = 5e-4 + elif dtype is torch.bfloat16: + rtol = 4e-3 + atol = 4e-3 - a = a.detach().half() - b = b.detach().half() + a = a.detach().to(dtype) + b = b.detach().to(dtype) assert_close(a, b, rtol=rtol, atol=atol) @@ -98,7 +99,8 @@ def exam_zero_1_2(): assert torch.equal(z1p.data, z2p.data) -def exam_zero_1_torch_ddp(): +@parameterize('dtype', [torch.float16, torch.bfloat16]) +def exam_zero_1_torch_ddp(dtype: torch.dtype): """ In this test, two pairs of model and optimizers are created. 1. zero: use sharded optimizer and fp16 parameters @@ -111,15 +113,10 @@ def exam_zero_1_torch_ddp(): seed_all(1453) # create models - zero_model = MlpModel() - torch_model = copy.deepcopy(zero_model) + torch_model = MlpModel().cuda() + zero_model = copy.deepcopy(torch_model).to(dtype) - zero_model = zero_model.cuda().half() - torch_model = DDP(torch_model.cuda(), bucket_cap_mb=0) - torch_model = torch_model.cuda() - - # for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()): - # half_close(p.data, z1p.data) + torch_model = DDP(torch_model.cuda(), bucket_cap_mb=0).cuda() # create optimizer zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1) @@ -139,11 +136,11 @@ def exam_zero_1_torch_ddp(): input_data = torch.rand(32, 128).cuda() # zero-dp forward - zero_output = zero_model(input_data.half()) + zero_output = zero_model(input_data.to(dtype)) # torch-ddp forward torch_output = torch_model(input_data) - half_close(zero_output, torch_output, loose=True) + loose_close(zero_output, torch_output, dtype=dtype) # zero-dp backward zero_optimizer.backward(zero_output.mean().float(), sync_grad=False) @@ -153,7 +150,7 @@ def exam_zero_1_torch_ddp(): # check grad for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()): - half_close(p.grad, z1p.grad, loose=True) + loose_close(p.grad, z1p.grad, dtype=dtype) # zero-dp step zero_optimizer._sync_grad() @@ -165,7 +162,7 @@ def exam_zero_1_torch_ddp(): # check updated param for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()): # print(n, torch.max(torch.abs(p.data - z1p.data))) - half_close(p.data, z1p.data, loose=True) + loose_close(p.data, z1p.data, dtype=dtype) def run_dist(rank, world_size, port): @@ -176,10 +173,9 @@ def run_dist(rank, world_size, port): @pytest.mark.dist +@rerun_if_address_is_in_use() def test_zero_1_2(): - world_size = 2 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 2) if __name__ == '__main__': diff --git a/tests/test_zero/low_level_zero/test_zero_init.py b/tests/test_zero/test_low_level/test_zero_init.py similarity index 80% rename from tests/test_zero/low_level_zero/test_zero_init.py rename to tests/test_zero/test_low_level/test_zero_init.py index 1305da5df9c5..aeeaff5b5cb9 100644 --- a/tests/test_zero/low_level_zero/test_zero_init.py +++ b/tests/test_zero/test_low_level/test_zero_init.py @@ -1,16 +1,13 @@ -from functools import partial - import pytest import torch import torch.distributed as dist -import torch.multiprocessing as mp import torch.nn as nn import colossalai from colossalai.tensor import ProcessGroup -from colossalai.utils import free_port, get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext -from colossalai.zero import LowLevelZeroOptimizer +from colossalai.testing import spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer class MlpModel(nn.Module): @@ -52,9 +49,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist def test_zero_init(): - world_size = 4 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 4) if __name__ == '__main__': diff --git a/tests/test_zero/low_level_zero/test_zero_tp.py b/tests/test_zero/test_low_level/test_zero_tp.py similarity index 88% rename from tests/test_zero/low_level_zero/test_zero_tp.py rename to tests/test_zero/test_low_level/test_zero_tp.py index 15d3530ff90a..f0804f4bb5ba 100644 --- a/tests/test_zero/low_level_zero/test_zero_tp.py +++ b/tests/test_zero/test_low_level/test_zero_tp.py @@ -1,18 +1,14 @@ -from functools import partial - import pytest import torch -import torch.multiprocessing as mp import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP from torch.testing import assert_close import colossalai from colossalai.tensor import ProcessGroup -from colossalai.testing import parameterize, rerun_if_address_is_in_use -from colossalai.utils import free_port, get_current_device -from colossalai.utils.model.colo_init_context import ColoInitContext -from colossalai.zero import LowLevelZeroOptimizer +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.utils import get_current_device +from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer from tests.test_tensor.common_utils import set_seed, split_param_col_tp1d, split_param_row_tp1d, tensor_shard_equal @@ -90,9 +86,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist @rerun_if_address_is_in_use() def test_zero_with_tp(): - world_size = 4 - run_func = partial(run_dist, world_size=world_size, port=free_port()) - mp.spawn(run_func, nprocs=world_size) + spawn(run_dist, 4) if __name__ == '__main__': diff --git a/version.txt b/version.txt index b0032849c80b..0d91a54c7d43 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.2.7 +0.3.0