Update aistudio-sdk requirement

.github/actions/check-bypass/action.yml

@@ -1,41 +0,0 @@
-name: "Check bypass"
-description: "A custom action to encapsulate PFCCLab/ci-bypass"
-inputs:
-  github-token:
-    description: "GitHub token"
-    required: true
-  workflow-name:
-    description: "Workflow name"
-    required: true
-outputs:
-  can-skip:
-    description: "Whether the workflow can be skipped."
-    value: ${{ steps.check-bypass.outputs.can-skip }}
-
-runs:
-  using: "composite"
-  steps:
-    - id: check-bypass
-      name: Check Bypass
-      env:
-        CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1", "XieYunshen","luotao1","From00"]'
-      uses: PFCCLab/ci-bypass@v1
-      with:
-        github-token: ${{ inputs.github-token }}
-        non-pull-request-event-strategy: 'never-skipped'
-        type: 'composite'
-        composite-rule: |
-          {
-            "any": [
-              {
-                "type": "labeled",
-                "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
-                "username": ${{ env.CI_TEAM_MEMBERS }}
-              },
-              {
-                "type": "commented",
-                "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
-                "username": ${{ env.CI_TEAM_MEMBERS }}
-              }
-            ]
-          }

.github/actions/rerun-workflow/action.yml

@@ -1,30 +0,0 @@
-name: 'Rerun Workflow'
-description: 'Re-run GitHub Actions workflow for a given Pull Request'
-inputs:
-  GITHUB_TOKEN:
-    description: 'GitHub token with repo scope'
-    required: true
-  OWNER:
-    description: 'Repository owner'
-    required: true
-  REPO:
-    description: 'Repository name'
-    required: true
-  PR_ID:
-    description: 'Pull Request ID'
-    required: true
-  JOB_NAME:
-    description: 'Job name to rerun'
-    required: true
-
-runs:
-  using: 'composite'
-  steps:
-    - run: bash ./.github/actions/rerun-workflow/rerun.sh
-      shell: bash
-      env:
-        GITHUB_TOKEN: ${{ inputs.GITHUB_TOKEN }}
-        OWNER: ${{ inputs.OWNER }}
-        REPO: ${{ inputs.REPO }}
-        PR_ID: ${{ inputs.PR_ID }}
-        JOB_NAME: ${{ inputs.JOB_NAME }}

.github/actions/rerun-workflow/rerun.sh

@@ -1,77 +0,0 @@
-# Copyright (c) 2025 PaddleNLP Authors. 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.
-
-set -e
-
-COMMIT_SHA=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
-  "https://api.github.com/repos/$OWNER/$REPO/pulls/$PR_ID" | jq -r '.head.sha')
-
-echo "Commit SHA: $COMMIT_SHA"
-
-response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
-  "https://api.github.com/repos/$OWNER/$REPO/actions/runs?head_sha=$COMMIT_SHA&per_page=100")
-
-echo "Response: $response"
-
-run_ids=$(echo "$response" | jq -r '.workflow_runs[].id')
-
-if [ -n "$run_ids" ]; then
-  echo "Found run_ids for commit $COMMIT_SHA: $run_ids"
-
-  for run_id in $run_ids; do
-    if [ "$JOB_NAME" = "all-failed" ]; then
-      echo "Rerunning all failed jobs for run_id: $run_id"
-
-      rerun_response=$(curl -X POST -s -w "%{http_code}" -o /dev/null \
-        -H "Accept: application/vnd.github.v3+json" \
-        -H "Authorization: Bearer $GITHUB_TOKEN" \
-        "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/rerun-failed-jobs")
-      if [ "$rerun_response" -eq 201 ]; then
-        echo "Successfully requested rerun for all blocked jobs in run_id: $run_id"
-      else
-        echo "Failed to request rerun for run_id: $run_id with status code $rerun_response"
-      fi
-
-    else
-      jobs_response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
-      "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/jobs")
-
-      echo "Jobs Response for run_id $run_id: $jobs_response"
-
-      # if [[ "$JOB_NAME" == *"bypass"* ]]; then
-        block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
-        '.jobs[] | select(.name == $job_name) | .id')
-      # else
-      #   block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
-      #   '.jobs[] | select(.name == $job_name and .conclusion != "success") | .id')
-      # fi
-
-      if [ -n "$block_jobs" ]; then
-        echo "Found block jobs for run_id $run_id: $block_jobs"
-
-        for job_id in $block_jobs; do
-          echo "Rerunning job_id: $job_id"
-          curl -X POST -H "Accept: application/vnd.github.v3+json" \
-            -H "Authorization: token $GITHUB_TOKEN" \
-            "https://api.github.com/repos/$OWNER/$REPO/actions/jobs/$job_id/rerun"
-        done
-      else
-        echo "No block jobs found for run_id $run_id with name $JOB_NAME."
-      fi
-    fi
-  done
-else
-  echo "No matching workflow runs found for commit $COMMIT_SHA."
-  exit 1
-fi

.github/workflows/PaddleFormers-PR-Check.yml

@@ -1,47 +0,0 @@
-name: PaddleFormers PR Check
-
-on:
-  pull_request:
-    types: [opened, synchronize]
-
-env:
-  BRANCH: ${{ github.base_ref }}
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  token: ${{ vars.ACTION_GITHUB_TOKEN }}
-
-jobs:
-  check-approvers:
-    name: Check approval
-    runs-on: 
-      group: APPROVAL
-    steps:
-      - name: Cleanup
-        run: |
-          rm -rf * .[^.]*
-      - name: Update paddle
-        run: |
-          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
-          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
-          rm -rf PaddleNLP-develop.tar.gz >/dev/null
-          git fetch origin pull/${PR_ID}/head
-          git checkout -b origin_pr FETCH_HEAD
-          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-          git fetch upstream
-          git checkout -b test_pr upstream/${BRANCH}
-          git merge --no-edit origin_pr
-          git log --pretty=oneline -10
-
-      - name: Check bypass
-        id: check-bypass
-        uses: ./.github/actions/check-bypass
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          workflow-name: approval
-        
-      - name: Display Required Approvers
-        if: steps.check-bypass.outputs.can-skip != 'true'
-        run: |
-          cd scripts/ci_approval
-          bash run_ci_approval_cherry-pick.sh

.github/workflows/approval.yml

@@ -1,47 +0,0 @@
-name: Approval
-
-on:
-  pull_request:
-    types: [opened, synchronize]
-
-env:
-  BRANCH: ${{ github.base_ref }}
-  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  token: ${{ vars.ACTION_GITHUB_TOKEN }}
-
-jobs:
-  check-approvers:
-    name: Check approval
-    runs-on: 
-      group: APPROVAL
-    steps:
-      - name: Cleanup
-        run: |
-          rm -rf * .[^.]*
-      - name: Update paddle
-        run: |
-          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
-          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
-          rm -rf PaddleNLP-develop.tar.gz >/dev/null
-          git fetch origin pull/${PR_ID}/head
-          git checkout -b origin_pr FETCH_HEAD
-          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-          git fetch upstream
-          git checkout -b test_pr upstream/${BRANCH}
-          git merge --no-edit origin_pr
-          git log --pretty=oneline -10
-
-      - name: Check bypass
-        id: check-bypass
-        uses: ./.github/actions/check-bypass
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          workflow-name: approval
-        
-      - name: Display Required Approvers
-        if: steps.check-bypass.outputs.can-skip != 'true'
-        run: |
-          cd scripts/ci_approval
-          bash -x run_ci_approval.sh

.github/workflows/check-bypass.yml

@@ -1,52 +0,0 @@
-on:
-  workflow_call:
-    inputs:
-      workflow-name:
-        required: true
-        type: string
-    secrets:
-      github-token:
-        required: true
-    outputs:
-      can-skip:
-        description: "Whether the workflow can be skipped."
-        value: ${{ jobs.check-bypass.outputs.can-skip }}
-
-jobs:
-  check-bypass:
-    name: Check bypass
-    runs-on:
-      group: APPROVAL
-    permissions:
-      contents: read
-    env:
-      CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1" , "XieYunshen","luotao1","From00"]'
-    outputs:
-      can-skip: ${{ steps.check-bypass.outputs.can-skip }}
-    steps:
-      - name: Cleanup
-        run: |
-          rm -rf * .[^.]*
-
-      - id: check-bypass
-        name: Check Bypass
-        uses: PFCCLab/ci-bypass@v1
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          non-pull-request-event-strategy: 'never-skipped'
-          type: 'composite'
-          composite-rule: |
-            {
-              "any": [
-                {
-                  "type": "labeled",
-                  "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
-                  "username": ${{ env.CI_TEAM_MEMBERS }}
-                },
-                {
-                  "type": "commented",
-                  "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
-                  "username": ${{ env.CI_TEAM_MEMBERS }}
-                }
-              ]
-            }

.github/workflows/lint.yml

@@ -1,80 +1,17 @@
-name: Codestyle Check
+name: Lint
 
 on: [push, pull_request]
 
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
-  cancel-in-progress: true
-
-env:
-  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  BRANCH: ${{ github.event.pull_request.base.ref }}
-  TASK: PaddleNLP-CI-Lint-${{ github.event.pull_request.number }}
-
 jobs:
-  check-bypass:
-    name: Check bypass
-    uses: ./.github/workflows/check-bypass.yml
-    with:
-      workflow-name: 'lint'
-    secrets:
-      github-token: ${{ secrets.GITHUB_TOKEN }}
-
   Lint:
     name: Lint
-    needs: check-bypass
-    if: ${{ needs.check-bypass.outputs.can-skip != 'true' }}
-    runs-on: [self-hosted, ernie-cpu]
+    runs-on: ubuntu-24.04
     steps:
-      - name: Run Container
-        env:
-          work_dir: ${{ github.workspace }}
-          python_version: "3.10"
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: checkout develop
         run: |
-          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
-          echo "container_name=${container_name}" >> ${{ github.env }}
-          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
-          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
-            -v $work_dir/../../..:$work_dir/../../.. \
-            -v $work_dir:/workspace \
-            -v /home/.cache/pip:/home/.cache/pip \
-            -e BRANCH \
-            -e PR_ID \
-            -e COMMIT_ID \
-            -e work_dir \
-            -e no_proxy \
-            -e python_version \
-            -w /workspace ${docker_image}
-
-      - name: Download Code
-        env:
-          work_dir: ${{ github.workspace }}
-        run: |
-          docker exec -t ${container_name} /bin/bash -c '
-          rm -rf * .[^.]*
-          echo "Downloading PaddleNLP.tar"
-          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
-          echo "Extracting PaddleNLP.tar"
-          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
-          source $work_dir/../../../proxy
-          cd PaddleNLP
-          git config --global user.name "PaddleCI"
-          git config --global user.email "paddle_ci@example.com"
-          git pull
-          git submodule update --init --recursive --force
-          if [ -n "${PR_ID}" ]; then
-            git fetch origin pull/${PR_ID}/head
-            git checkout -b PR_${PR_ID} FETCH_HEAD
-            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-            git fetch upstream ${BRANCH}
-            git merge ${BRANCH} --no-edit
-            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
-          else
-            echo "Not in a pull_request event. Skipping PR-specific operations."
-          fi
-          git log --pretty=oneline -10
-
           if ! git show-ref --quiet refs/heads/develop; then \
               echo "local develop branch is missing, creating local develop branch that tracks remote develop branch"
               git fetch origin develop
@@ -82,32 +19,14 @@ jobs:
           else
             echo "local develop branch exist, skipping"
           fi
-          '
-          
-      - name: Setup Environment
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip' # caching pip dependencies
+      - name: Install dependencies
         run: |
-          docker exec -t $container_name /bin/bash -c '
-          unlink /usr/local/bin/python
-          ln -sf $(which python${python_version}) /usr/local/bin/python
-          set -e
-          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
-          python -m pip config set global.cache-dir "/home/.cache/pip"
-          source $work_dir/../../../proxy
           python -m pip install --upgrade pip
-          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
           make install
-          '
-
-      - name: Test
-        run: |
-          docker exec -t $container_name /bin/bash -c '
-          set -e
-          source $work_dir/../../../proxy
-          cd /workspace/PaddleNLP
-          make lint
-          '
-      
-      - name: Terminate And Delete the Container
-        if: always()
-        run: |
-          docker rm -f $container_name 2>/dev/null || true
+      - name: run the command
+        run: make lint
+    

.github/workflows/llm.yml

@@ -1,207 +0,0 @@
-# Disabled LLM CI workflow - manually enabled when needed
-name: LLM CI (DISABLED)
-
-#on:
-#  pull_request:
-#    types: [opened, synchronize, reopened]
-#    branches: [develop]
-#    schedule:
-#    - cron: "2 0 * * *"
-#  workflow_call:
-#    inputs:
-#      run_downstream:
-#        required: true
-#        type: string
-#      image_name:
-#        required: true
-#        type: string
-
-#concurrency:
-#  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
-#  cancel-in-progress: true
-
-#env:
-#  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-llm
-  ci_scripts: /workspace/PaddleNLP/scripts/regression
-  BRANCH: ${{ github.event.pull_request.base.ref }}
-  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
-  CI_name: llm-ci
-  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
-  HF_ENDPOINT: https://hf-mirror.com
-  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
-  PPNLP_HOME: /ssd1/paddlenlp
-  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
-  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
-  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
-  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
-
-#defaults:
-#  run:
-#    shell: bash
-
-#jobs:
-#  llm-ci:
-#    name: llm-ci (DISABLED)
-    runs-on: [self-hosted, ernie-8gpu]
-    steps:
-      - name: Determine Image Name
-        env:
-          IMAGE_NAME: ${{ inputs.image_name }}
-        run: |
-          if [[ -n "${IMAGE_NAME}" ]]; then
-            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
-          else
-            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
-          fi
-
-      - name: Run Container
-        env:
-          work_dir: ${{ github.workspace }}
-          CACHE_DIR: /home/data/cfs/.cache
-          FLAGS_dynamic_static_unified_comm: "True"
-          python_version: "3.10"
-          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
-        run: |
-          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
-          echo "container_name=${container_name}" >> "$GITHUB_ENV"
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
-              -v $work_dir/../../..:$work_dir/../../.. \
-              -v $work_dir:/workspace \
-              -v /home/.cache/pip:/home/.cache/pip \
-              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
-              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
-              -e BRANCH \
-              -e AGILE_COMPILE_BRANCH \
-              -e PR_ID \
-              -e COMMIT_ID \
-              -e work_dir \
-              -e ci_scripts \
-              -e no_proxy \
-              -e CI_name \
-              -e paddle_whl \
-              -e HF_ENDPOINT \
-              -e STUDIO_GIT_HOST \
-              -e PPNLP_HOME \
-              -e HF_DATASETS_CACHE \
-              -e TRANSFORMERS_CACHE \
-              -e CACHE_DIR \
-              -e FLAGS_dynamic_static_unified_comm \
-              -e python_version \
-              -w /workspace --runtime=nvidia $IMAGE_NAME
-          fi
-
-      - name: Download Code
-        env:
-          work_dir: ${{ github.workspace }}
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping.."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            rm -rf * .[^.]*
-            echo "Downloading PaddleNLP.tar.gz"
-            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
-            echo "Extracting PaddleNLP.tar.gz"
-            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
-            source $work_dir/../../../proxy
-            cd PaddleNLP
-            git config --global user.name "PaddleCI"
-            git config --global user.email "paddle_ci@example.com"
-            git pull
-            git submodule update --init --recursive --force
-            if [ -n "${PR_ID}" ]; then
-              git fetch origin pull/${PR_ID}/head
-              git checkout -b PR_${PR_ID} FETCH_HEAD
-              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-              git fetch upstream ${BRANCH}
-              git merge ${BRANCH} --no-edit
-              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
-            else
-              echo "Not in a pull_request event. Skipping PR-specific operations."
-            fi
-            git log --pretty=oneline -10
-            '
-          fi
-
-      - name: Skip For Bug
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            cd /workspace/PaddleNLP
-            git revert f2477c07272d04244cd3287d1f21c70482a4a85f --no-edit  # 套件PR#10413引入bug-待修复
-            git revert d74c950e15a35b7a100d1688c89318195cc83bca --no-edit  # 框架升级后配合修改，当前因框架bug固定commit适配
-            '
-          fi
-      
-      - name: Test
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            ldconfig
-            unlink /usr/bin/python3
-            ln -sf $(which python${python_version}) /usr/bin/python3
-            pip config set global.cache-dir "/home/.cache/pip"
-            set -e
-            source $work_dir/../../../proxy
-            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
-            echo "Paddle PR#73283  PR#74484 import bug, set paddle_commit=8ae742 to skip "
-            export paddle_whl=https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/8ae7423e99b2ea96e410968a0ebb3f1795e37205/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
-            timeout 2h bash scripts/regression/run_ci.sh python${python_version} ${paddle_whl}
-            '
-          fi
-
-      - name: Upload Allure-reports & Logs
-        if: always()
-        env:
-          home_path: ${{ github.workspace }}/../../..
-          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
-          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            unset http_proxy && unset https_proxy
-            if [ ! -f "${{ env.bos_file }}" ]; then
-              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
-              mkdir ${{ env.home_path }}/bos
-              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
-            fi
-            if [ ! -f "${{ env.allure_file }}" ]; then
-              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
-              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
-            fi
-            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
-              bos_prefix="${PR_ID}/${COMMIT_ID}"
-            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
-              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
-            else
-              bos_prefix="schedule/$(date +%Y%m%d)"
-            fi
-            cd /workspace/PaddleNLP/model_logs
-            for FILE in /workspace/PaddleNLP/model_logs/*; do
-              file=$(basename "$FILE")
-              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
-              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/$file"
-            done
-            cd /workspace/PaddleNLP/
-            ${{ env.allure_file }} generate result -o report
-            tar -czf products.tar.gz report model_logs
-            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
-            echo "products: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/products.tar.gz"
-            '
-          fi
-
-      - name: Terminate And Delete the Container
-        if: always()
-        run: |
-          docker rm -f ${{ env.container_name }} 2>/dev/null || true

.github/workflows/rerun.yml

@@ -1,73 +0,0 @@
-name: Re-run
-
-on:
-  issue_comment:
-    types: [created]
-
-jobs:
-  re-run:
-    if: ${{ github.event.issue.pull_request && contains(github.event.comment.body, '/re-run')  && github.event.comment.user.login == github.event.issue.user.login }}
-    runs-on: [self-hosted, ernie-cpu]
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-
-      - name: Rerun all failed jobs
-        if: ${{ contains(github.event.comment.body, 'all-failed') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'all-failed'
-
-      - name: Rerun LLM
-        if: ${{ contains(github.event.comment.body, 'LLM') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'LLM CI / llm-ci'
-
-      - name: Unittest GPU
-        if: ${{ contains(github.event.comment.body, 'Unittest GPU') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'Unittest GPU CI / unittest-gpu-ci'
-
-      - name: Rerun Unittest CPU
-        if: ${{ contains(github.event.comment.body, 'Unittest CPU') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'Unittest CPU CI / unit-cpu-ci'
-
-      - name: Rerun Distribute (V100)
-        if: ${{ contains(github.event.comment.body, 'Distribute (V100)') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'Distribute CI (V100)/ distribute-v100-ci'
-
-      - name: Rerun Codestyle Check
-        if: ${{ contains(github.event.comment.body, 'Codestyle Check') }}
-        uses: ./.github/actions/rerun-workflow
-        with:
-          PR_ID: ${{ github.event.issue.number }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          OWNER: ${{ github.repository_owner }}
-          REPO: ${{ github.event.repository.name }}
-          JOB_NAME: 'Codestyle Check / Lint'

.github/workflows/tests.yml

@@ -0,0 +1,29 @@
+name: Test
+
+on: [push, pull_request]
+
+jobs:
+  Test:
+    name: Test
+    runs-on: ubuntu-24.04
+    permissions:
+      pull-requests: write
+      contents: read
+      id-token: write
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip' # caching pip dependencies
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r tests/requirements.txt
+          make install
+      - name: run the command
+        run: make test
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v4
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/unittest-cpu.yml

@@ -1,120 +0,0 @@
-# Disabled unittest CPU workflow - manually enabled when needed
-name: Unittest CPU CI (DISABLED)
-
-#on: [pull_request]
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
-  cancel-in-progress: true
-
-env:
-  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  BRANCH: ${{ github.event.pull_request.base.ref }}
-  TASK: PaddleNLP-CI-${{ github.event.pull_request.number }}-unittest-cpu
-  HF_ENDPOINT: https://hf-mirror.com
-  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
-  PPNLP_HOME: /home/disk1/cache
-  HF_DATASETS_CACHE: /home/disk1/cache/huggingface/datasets
-  TRANSFORMERS_CACHE: /home/disk1/cache/huggingface
-
-#jobs:
-#  unittest-cpu-ci:
-#    name: unittest-cpu-ci (DISABLED)
-    runs-on: [self-hosted, ernie-cpu]
-    permissions:
-      pull-requests: write
-      contents: read
-      id-token: write
-    steps:
-      - name: Run Container
-        env:
-          work_dir: ${{ github.workspace }}
-          python_version: "3.10"
-        run: |
-          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
-          echo "container_name=${container_name}" >> ${{ github.env }}
-          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
-          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
-            -v $work_dir/../../..:$work_dir/../../.. \
-            -v $work_dir:/workspace \
-            -v /home/.cache/pip:/home/.cache/pip \
-            -v /home/disk1/cache:/home/disk1/cache \
-            -e BRANCH \
-            -e PR_ID \
-            -e COMMIT_ID \
-            -e work_dir \
-            -e no_proxy \
-            -e python_version \
-            -e HF_ENDPOINT \
-            -e STUDIO_GIT_HOST \
-            -e PPNLP_HOME \
-            -e HF_DATASETS_CACHE \
-            -e TRANSFORMERS_CACHE \
-            -w /workspace ${docker_image}
-
-      - name: Download Code
-        env:
-          work_dir: ${{ github.workspace }}
-        run: |
-          docker exec -t ${container_name} /bin/bash -c '
-          rm -rf * .[^.]*
-          echo "Downloading PaddleNLP.tar"
-          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
-          echo "Extracting PaddleNLP.tar"
-          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
-          source $work_dir/../../../proxy
-          cd PaddleNLP
-          git config --global user.name "PaddleCI"
-          git config --global user.email "paddle_ci@example.com"
-          git pull
-          git submodule update --init --recursive --force
-          if [ -n "${PR_ID}" ]; then
-            git fetch origin pull/${PR_ID}/head
-            git checkout -b PR_${PR_ID} FETCH_HEAD
-            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-            git fetch upstream ${BRANCH}
-            git merge ${BRANCH} --no-edit
-            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
-          else
-            echo "Not in a pull_request event. Skipping PR-specific operations."
-          fi
-          git log --pretty=oneline -10
-          '
-
-      - name: Setup Environment
-        run: |
-          docker exec -t $container_name /bin/bash -c '
-          unlink /usr/bin/python3
-          ln -sf $(which python${python_version}) /usr/bin/python3
-          set -e
-          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
-          python -m pip config set global.cache-dir "/home/.cache/pip"
-          source $work_dir/../../../proxy
-          python -m pip install --upgrade pip
-          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
-          pip install -r tests/requirements.txt
-          make install
-          '
-
-      - name: Test
-        run: |
-          docker exec -t $container_name /bin/bash -c '
-          source $work_dir/../../../proxy
-          cd /workspace/PaddleNLP
-          set -e
-          make test
-          '
-
-      # - name: Upload Coverage To Codecov
-      #   uses: codecov/codecov-action@v3
-      #   env:
-      #     CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
-      #   with:
-      #     token: ${{ secrets.CODECOV_TOKEN }}
-      #     files: ./PaddleNLP/coverage.xml
-      
-      - name: Terminate And Delete the Container
-        if: always()
-        run: |
-          docker rm -f $container_name 2>/dev/null || true

.github/workflows/unittest-gpu.yml

@@ -1,204 +0,0 @@
-# Disabled unittest GPU workflow - manually enabled when needed
-name: Unittest GPU CI (DISABLED)
-
-#on:
-#  pull_request:
-#    types: [opened, synchronize, reopened]
-#    branches: [develop]
-#  schedule:
-#    - cron: "3 0 * * *"
-#  workflow_call:
-#    inputs:
-#      run_downstream:
-#        required: true
-#        type: string
-#      image_name:
-#        required: true
-#        type: string
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
-  cancel-in-progress: true
-
-env:
-  PR_ID: ${{ github.event.pull_request.number }}
-  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
-  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-unit-gpu
-  ci_scripts: /workspace/PaddleNLP/scripts/unit_test
-  BRANCH: ${{ github.event.pull_request.base.ref }}
-  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
-  CI_name: unittest-gpu-ci
-  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
-  HF_ENDPOINT: https://hf-mirror.com
-  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
-  PPNLP_HOME: /ssd1/paddlenlp
-  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
-  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
-  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
-  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
-
-defaults:
-  run:
-    shell: bash
-
-#jobs:
-#  unittest-gpu-ci:
-#    name: unittest-gpu-ci (DISABLED)
-    runs-on: [self-hosted, ernie-8gpu]
-    steps:
-      - name: Determine Image Name
-        env:
-          IMAGE_NAME: ${{ inputs.image_name }}
-        run: |
-          if [[ -n "${IMAGE_NAME}" ]]; then
-            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
-          else
-            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
-          fi
-      
-      - name: Run Container
-        env:
-          work_dir: ${{ github.workspace }}
-          CACHE_DIR: /home/data/cfs/.cache
-          FLAGS_dynamic_static_unified_comm: "True"
-          python_version: "3.10"
-          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
-        run: |
-          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
-          echo "container_name=${container_name}" >> "$GITHUB_ENV"
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
-              -v $work_dir/../../..:$work_dir/../../.. \
-              -v $work_dir:/workspace \
-              -v /home/.cache/pip:/home/.cache/pip \
-              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
-              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
-              -e BRANCH \
-              -e AGILE_COMPILE_BRANCH \
-              -e PR_ID \
-              -e COMMIT_ID \
-              -e work_dir \
-              -e ci_scripts \
-              -e no_proxy \
-              -e CI_name \
-              -e paddle_whl \
-              -e HF_ENDPOINT \
-              -e STUDIO_GIT_HOST \
-              -e PPNLP_HOME \
-              -e HF_DATASETS_CACHE \
-              -e TRANSFORMERS_CACHE \
-              -e CACHE_DIR \
-              -e FLAGS_dynamic_static_unified_comm \
-              -e python_version \
-              -w /workspace --runtime=nvidia $IMAGE_NAME
-          fi
-          
-      - name: Download Code
-        env:
-          work_dir: ${{ github.workspace }}
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping.."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            rm -rf * .[^.]*
-            echo "Downloading PaddleNLP.tar.gz"
-            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
-            echo "Extracting PaddleNLP.tar.gz"
-            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
-            source $work_dir/../../../proxy
-            cd PaddleNLP
-            git config --global user.name "PaddleCI"
-            git config --global user.email "paddle_ci@example.com"
-            git pull
-            git submodule update --init --recursive --force
-            if [ -n "${PR_ID}" ]; then
-              git fetch origin pull/${PR_ID}/head
-              git checkout -b PR_${PR_ID} FETCH_HEAD
-              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
-              git fetch upstream ${BRANCH}
-              git merge ${BRANCH} --no-edit
-              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
-            else
-              echo "Not in a pull_request event. Skipping PR-specific operations."
-            fi
-            git log --pretty=oneline -10
-            '
-          fi
-      
-      - name: Skip For Bug
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            cd /workspace/PaddleNLP
-            echo "no skip for bug"
-            '
-          fi
-      
-      - name: Test
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            ldconfig
-            unlink /usr/bin/python3
-            ln -sf $(which python${python_version}) /usr/bin/python3
-            pip config set global.cache-dir "/home/.cache/pip"
-            source $work_dir/../../../proxy
-            set -e
-            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
-            timeout 50m bash scripts/unit_test/ci_unit.sh ${paddle_whl}
-            '
-          fi
-
-      - name: Upload Allure-reports & Logs
-        if: always()
-        env:
-          home_path: ${{ github.workspace }}/../../..
-          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
-          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
-        run: |
-          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
-            echo "Not in a pull_request or test_build event. Skipping..."
-          else
-            docker exec -t $container_name /bin/bash -c '
-            unset http_proxy && unset https_proxy
-            if [ ! -f "${{ env.bos_file }}" ]; then
-              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
-              mkdir ${{ env.home_path }}/bos
-              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
-            fi
-            if [ ! -f "${{ env.allure_file }}" ]; then
-              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
-              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
-            fi
-            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
-              bos_prefix="${PR_ID}/${COMMIT_ID}"
-            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
-              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
-            else
-              bos_prefix="schedule/$(date +%Y%m%d)"
-            fi
-            cd /workspace/PaddleNLP/unittest_logs
-            for FILE in /workspace/PaddleNLP/unittest_logs/*; do
-              file=$(basename "$FILE")
-              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
-              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/$file"
-            done
-            cd /workspace/PaddleNLP/
-            ${{ env.allure_file }} generate result -o report
-            tar -czf products.tar.gz report unittest_logs
-            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
-            echo "report: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/products.tar.gz"
-            '
-          fi
-
-      - name: Terminate And Delete the Container
-        if: always()
-        run: |
-          docker rm -f $container_name 2>/dev/null || true

.gitignore

@@ -140,4 +140,4 @@ autogen/
 #fp8
 ops/csrc/fp8/deep_gemm/include/cutlass
 ops/csrc/fp8/deep_gemm/include/cute
-.ccls-cache
+.ccls-cache

.pre-commit-config.yaml

@@ -62,12 +62,3 @@ repos:
         language: python
         files: \.(md|markdown|rst)$
         pass_filenames: true
-
--   repo: local
-    hooks:
-    -   id: clang-format
-        name: clang-format
-        description: Format files with ClangFormat.
-        entry: bash ./tools/codestyle/clang_format.sh -i
-        language: system
-        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|xpu|kps)$

csrc/gpu/moe/preprocess_for_moe.cu

@@ -140,14 +140,12 @@ std::vector<paddle::Tensor> preprocess_for_moe_kernel(const paddle::Tensor& topk
     
     if (num_experts == 8) {
       run_align_kernel(8);
-    } else if (num_experts == 64) {
-      run_align_kernel(64);
     } else if (num_experts == 256) {
       run_align_kernel(256);
     } else if (num_experts == 2) {
       run_align_kernel(2);
     } else {
-      printf("【ERROR】, please add the current number of experts in the case");
+      printf("errors");
     }
 
     const int block_threads = 256;

csrc/gpu/unittest/test_get_padding_offset_v2.py

@@ -64,5 +64,6 @@ class GetPaddingOffsetV2Test(unittest.TestCase):
         assert sum(ref_cu_seqlens_q - cu_seqlens_q) == 0, "Check cu_seqlens_q failed."
         assert sum(ref_cu_seqlens_k - cu_seqlens_k) == 0, "Check cu_seqlens_k failed."
 
+
 if __name__ == "__main__":
     unittest.main()

docs/zh/llm/alignment/ppo/README.md

@@ -0,0 +1 @@
+../../../../../llm/alignment/ppo/README.md

llm/alignment/rl/gsm8k_processor.py

@@ -1,74 +0,0 @@
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
-#
-# 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.
-"""
-Preprocess the GSM8k dataset to parquet format
-"""
-
-import argparse
-import os
-import re
-
-import datasets
-
-
-def extract_solution(solution_str):
-    solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
-    assert solution is not None
-    final_solution = solution.group(0)
-    final_solution = final_solution.split("#### ")[1].replace(",", "")
-    return final_solution
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--local_dir", default="./gsm8k")
-
-    args = parser.parse_args()
-
-    data_source = "openai/gsm8k"
-
-    dataset = datasets.load_dataset(data_source, "main")
-
-    train_dataset = dataset["train"]
-    test_dataset = dataset["test"]
-
-    instruction_following = 'Let\'s think step by step and output the final answer after "####".'
-
-    # add a row to each data item that represents a unique id
-    def make_map_fn(split):
-        def process_fn(example, idx):
-            question_raw = "<|im_start|>user\n" + example.pop("question")
-
-            system_raw = (
-                "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n"
-            )
-            question = system_raw + question_raw + " " + instruction_following + "<|im_end|>\n<|im_start|>assistant\n"
-
-            answer_raw = example.pop("answer")
-            solution = extract_solution(answer_raw)
-            data = {
-                "src": question,
-                "tgt": solution,
-            }
-            return data
-
-        return process_fn
-
-    train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True)
-    test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True)
-
-    local_dir = args.local_dir
-
-    train_dataset.to_json(os.path.join(local_dir, "train.jsonl"), orient="records", lines=True)
-    test_dataset.to_json(os.path.join(local_dir, "test.jsonl"), orient="records", lines=True)

llm/alignment/rl/run_rl.py

@@ -42,7 +42,6 @@ from paddlenlp.trainer import (
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
-    AutoModelForTokenClassification,
     AutoTokenizer,
     PretrainedConfig,
 )
@@ -135,6 +134,7 @@ def create_actor_models(
         )
         if not training_args.autotuner_benchmark:
             reference_model.set_state_dict(actor_model.state_dict())
+
     actor_tokenizer = AutoTokenizer.from_pretrained(
         model_args.actor_model_name_or_path,
         model_max_length=data_args.max_length,
@@ -210,43 +210,46 @@ def create_critic_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
+    reward_model,
 ):
     with timers_scope_runtimer("Critic model loading time"):
-        critic_model_config = AutoConfig.from_pretrained(
-            model_args.critic_model_name_or_path,
-            tensor_parallel_output=training_args.tensor_parallel_output,
-            tensor_parallel_degree=training_args.tensor_parallel_degree,
-            tensor_parallel_rank=training_args.tensor_parallel_rank,
-            dtype=training_args.model_dtype,
-            recompute=training_args.critic_recompute,
-            recompute_granularity=model_args.critic_recompute_granularity,
-            recompute_use_reentrant=training_args.recompute_use_reentrant,
-            **common_config,
-        )
-        LlmMetaConfig.set_llm_config(critic_model_config, training_args)
-
-        critic_model_config.max_position_embeddings = data_args.max_length
-        critic_model_config.use_sparse_head_and_loss_fn = False
-        critic_model_config.num_labels = 1
-        critic_model_config.classifier_dropout = 0.0
-        critic_model_config.hidden_dropout = 0.0
-        logger.info(f"Loading Critic model with config:\n\t{critic_model_config}\n")
-
-        if not training_args.autotuner_benchmark:
-            critic_model = AutoModelForTokenClassification.from_pretrained(
-                model_args.critic_model_name_or_path,
-                config=critic_model_config,
+        reward_model_config = reward_model.config
+        if model_args.critic_model_name_or_path is None:
+            model_args.critic_model_name_or_path = model_args.reward_model_name_or_path
+            critic_model = AutoModelForScore.from_config(
+                reward_model_config,
+                dtype=training_args.model_dtype,
+                score_type="critic",
+                do_normalize=False,
+                clip_range_value=training_args.clip_range_value,
+                **common_config,
             )
+            if not training_args.autotuner_benchmark:
+                critic_model.set_state_dict(reward_model.state_dict())
         else:
-            critic_model = AutoModelForTokenClassification.from_config(
-                critic_model_config,
-            )
+            if not training_args.autotuner_benchmark:
+                critic_model = AutoModelForScore.from_pretrained(
+                    model_args.critic_model_name_or_path,
+                    config=reward_model_config,
+                    score_type="critic",
+                    do_normalize=False,
+                    clip_range_value=training_args.clip_range_value,
+                    **common_config,
+                )
+            else:
+                critic_model = AutoModelForScore.from_config(
+                    reward_model_config,
+                    score_type="critic",
+                    do_normalize=False,
+                    clip_range_value=training_args.clip_range_value,
+                    **common_config,
+                )
 
     critic_tokenizer = AutoTokenizer.from_pretrained(
         model_args.critic_model_name_or_path,
         model_max_length=data_args.max_length,
         padding_side="left",
-        tokenizer_alpha=model_args.critic_tokenizer_alpha,
+        tokenizer_alpha=model_args.reward_critic_tokenizer_alpha,
         use_fast=True,
     )
     if critic_tokenizer.pad_token_id is None:
@@ -258,8 +261,8 @@ def create_critic_models(
         if training_args.eval_mode == "single":
             config.tensor_parallel_degree = -1
             config.tensor_parallel_rank = 0
-        with timers_scope_runtimer("Critic eval model loading time"):
-            critic_eval_model = AutoModelForTokenClassification.from_config(config)
+        with timers_scope_runtimer("Reward critic eval model loading time"):
+            critic_eval_model = AutoModelForScore.from_config(config)
     else:
         critic_eval_model = None
 
@@ -267,7 +270,7 @@ def create_critic_models(
 
 
 def create_rl_dataset(data_args, training_args, tokenizer):
-    requires_label = True if training_args.use_rm_server or training_args.use_rule_reward else False
+    requires_label = True if training_args.use_rm_server else False
     train_ds = RLHFDataset(
         dataset_name_or_path=data_args.train_datasets,
         tokenizer=tokenizer,
@@ -330,16 +333,15 @@ def main():
     actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
         model_args, data_args, training_args, common_config, reshard_controller
     )
-    if training_args.use_rule_reward:
-        reward_model, reward_tokenizer = None, actor_tokenizer
-    elif not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
+
+    if not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
         reward_model, reward_tokenizer = create_reward_models(model_args, data_args, training_args, common_config)
     else:
         reward_model, reward_tokenizer = model_args.reward_server, actor_tokenizer
 
     if training_args.rl_algorithm == "ppo":
         critic_model, critic_eval_model, critic_tokenizer = create_critic_models(
-            model_args, data_args, training_args, common_config
+            model_args, data_args, training_args, common_config, reward_model
         )
     else:
         critic_model, critic_eval_model, critic_tokenizer = None, None, None
@@ -353,8 +355,7 @@ def main():
         offload_tensor_to_cpu((reference_model, "freeze_model"))
 
         if training_args.rl_algorithm == "ppo":
-            if not training_args.use_rm_server and not training_args.use_rule_reward:
-                offload_tensor_to_cpu((reward_model, "freeze_model"))
+            offload_tensor_to_cpu((reward_model, "freeze_model"))
             if critic_eval_model is not None:
                 offload_tensor_to_cpu((critic_eval_model, "freeze_model"))
 
@@ -362,14 +363,7 @@ def main():
         paddle.device.cuda.empty_cache()
 
     def compute_metrics(eval_preds):
-        '''
-        If "use_rm_server" is TRUE, the score ranges from -3 to 3, with 3 being the only correct score (format + result).
-        If using the "Regularized Matching Function (use_rule_reward=True)" (currently only implemented for the gsm8k dataset), the score ranges from 0 to 1.
-        '''
-        if training_args.use_rule_reward:
-            accuracy = (eval_preds.predictions == 1).astype("float32").mean().item()
-        else:
-            accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
+        accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
         return {"accuracy": accuracy}
 
     try:
@@ -395,7 +389,7 @@ def main():
         data_collator=partial(
             collate_fn,
             pad_token_id=actor_tokenizer.pad_token_id,
-            requires_label=True if training_args.use_rm_server or training_args.use_rule_reward else False,
+            requires_label=True if training_args.use_rm_server else False,
             max_prompt_len=data_args.max_prompt_len if training_args.balance_batch else None,
         ),  # NOTE: enforce prompt padding to max_prompt_len when using balance_batch
         compute_metrics=compute_metrics,  # TODO: only used for grpo (kk datasets)

llm/auto_parallel/deepseek-v3/run_pretrain_auto.py

@@ -460,25 +460,10 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    do_enable_linear_fused_grad_add = training_args.enable_linear_fused_grad_add
-    do_enable_mp_async_allreduce = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and "enable_mp_async_allreduce" in training_args.tensor_parallel_config
-        and not training_args.sequence_parallel
-    )
-    do_enable_sp_async_reduce_scatter = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and training_args.sequence_parallel
-        and "enable_sp_async_reduce_scatter" in training_args.tensor_parallel_config
-    )
-    if (
-        do_enable_linear_fused_grad_add or do_enable_mp_async_allreduce or do_enable_sp_async_reduce_scatter
-    ) and not training_args.to_static:
-        from llm.utils.fused_layers import mock_layers
+    if training_args.enable_linear_fused_grad_add:
+        from fused_layers import mock_layers
 
-        mock_layers(do_enable_linear_fused_grad_add, do_enable_mp_async_allreduce, do_enable_sp_async_reduce_scatter)
+        mock_layers()
 
     if model_args.tokenizer_name_or_path is None:
         model_args.tokenizer_name_or_path = model_args.model_name_or_path

llm/auto_parallel/gpt-3/run_pretrain_auto.py

@@ -38,7 +38,6 @@ from paddlenlp.transformers import (
     CosineAnnealingWithWarmupDecay,
     GPTConfig,
     GPTForCausalLMAuto,
-    GPTForCausalLMAutoPP,
     GPTForCausalLMNet,
     GPTPretrainingCriterionAuto,
     GPTPretrainingCriterionNet,
@@ -49,7 +48,6 @@ from paddlenlp.utils.tools import get_env_device
 
 MODEL_CLASSES = {
     "gpt": (GPTConfig, GPTForCausalLMAuto, GPTPretrainingCriterionAuto),
-    "gpt_pp": (GPTConfig, GPTForCausalLMAutoPP, GPTPretrainingCriterionAuto),
     "gpt_network": (GPTConfig, GPTForCausalLMNet, GPTPretrainingCriterionNet),
 }
 
@@ -74,12 +72,6 @@ class PreTrainingArguments(AutoTrainingArguments):
             "help": "The steps use to control the learing rate. If the step > decay_steps, will use the min_learning_rate."
         },
     )
-    enable_linear_fused_grad_add: bool = field(
-        default=False,
-        metadata={
-            "help": "Enable fused linear grad add strategy, which will reduce elementwise add for grad accumulation in the backward of nn.Linear ."
-        },
-    )
     job_schedule_profiler_start: int = field(
         default=-1,
         metadata={"help": "The step to start job_schedule_profiler."},
@@ -93,18 +85,13 @@ class PreTrainingArguments(AutoTrainingArguments):
     )
     sr: Optional[int] = field(default=0, metadata={"help": "The count of chunks without recompute."})
     virtual_pipeline_seg_method: str = field(
-        default="LlamaDecoderLayerAuto",
-        metadata={"help": "The seg method of splitting pp layer for virtual pipeline."},
+        default="LlamaDecoderLayerAuto", metadata={"help": "The seg method of splitting pp layer for virtual pipeline."}
     )
     # NOTE(gongenlei): new add autotuner_benchmark
     autotuner_benchmark: bool = field(
         default=False,
         metadata={"help": "Weather to run benchmark by autotuner. True for from_scratch and pad_max_length."},
     )
-    n_microbatches: int = field(
-        default=1,
-        metadata={"help": "Control the num of microbatches in one pp step."},
-    )
     pre_alloc_memory: float = field(
         default=0.0,
         metadata={
@@ -444,26 +431,6 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    do_enable_linear_fused_grad_add = training_args.enable_linear_fused_grad_add
-    do_enable_mp_async_allreduce = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and "enable_mp_async_allreduce" in training_args.tensor_parallel_config
-        and not training_args.sequence_parallel
-    )
-    do_enable_sp_async_reduce_scatter = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and training_args.sequence_parallel
-        and "enable_sp_async_reduce_scatter" in training_args.tensor_parallel_config
-    )
-    if (
-        do_enable_linear_fused_grad_add or do_enable_mp_async_allreduce or do_enable_sp_async_reduce_scatter
-    ) and not training_args.to_static:
-        from llm.utils.fused_layers import mock_layers
-
-        mock_layers(do_enable_linear_fused_grad_add, do_enable_mp_async_allreduce, do_enable_sp_async_reduce_scatter)
-
     if model_args.tokenizer_name_or_path is None:
         model_args.tokenizer_name_or_path = model_args.model_name_or_path
 
@@ -558,7 +525,6 @@ def main():
     config.use_recompute = training_args.recompute
     config.tensor_parallel_degree = training_args.tensor_parallel_degree
     config.tensor_parallel_rank = training_args.tensor_parallel_rank
-    config.to_static = training_args.to_static
 
     if training_args.strategy.pipeline.enable and config.virtual_pp_degree > 1:
         pipeline = training_args.strategy.pipeline
@@ -568,14 +534,6 @@ def main():
     config.hidden_dropout_prob = model_args.hidden_dropout_prob
     config.attention_probs_dropout_prob = model_args.attention_probs_dropout_prob
     print("Final pre-training config:", config)
-    if (
-        "replace_with_parallel_cross_entropy" in training_args.tensor_parallel_config
-        and config.tensor_parallel_degree > 1
-        and config.to_static is False
-    ):
-        from llm.utils.replace_ops import replace_cross_entropy
-
-        replace_cross_entropy()
 
     # Set the dtype for loading model
     dtype = "float32"
@@ -631,7 +589,6 @@ def main():
 
     trainer = PretrainingTrainer(
         model=model,
-        model_type=model_args.model_type,
         criterion=criterion,
         args=training_args,
         data_collator=data_collator,

llm/auto_parallel/llama/run_pretrain_auto.py

@@ -41,7 +41,6 @@ from paddlenlp.transformers import (
     LinearAnnealingWithWarmupDecay,
     LlamaConfig,
     LlamaForCausalLM3DAuto,
-    LlamaForCausalLM3DAutoPP,
     LlamaForCausalLMNet,
     LlamaPretrainingCriterion3DAuto,
     LlamaPretrainingCriterionNet,
@@ -50,7 +49,6 @@ from paddlenlp.utils.log import logger
 
 MODEL_CLASSES = {
     "llama": (LlamaConfig, LlamaForCausalLM3DAuto, LlamaPretrainingCriterion3DAuto),
-    "llama_pp": (LlamaConfig, LlamaForCausalLM3DAutoPP, LlamaPretrainingCriterion3DAuto),
     "llama_network": (LlamaConfig, LlamaForCausalLMNet, LlamaPretrainingCriterionNet),
 }
 
@@ -96,10 +94,6 @@ class PreTrainingArguments(AutoTrainingArguments):
         default=False,
         metadata={"help": "Weather to run benchmark by autotuner. True for from_scratch and pad_max_length."},
     )
-    n_microbatches: int = field(
-        default=1,
-        metadata={"help": "Control the num of microbatches in one pp step."},
-    )
 
     def __post_init__(self):
         super().__post_init__()
@@ -410,20 +404,12 @@ def init_seed(seed: int = 1234, args=None):
                 order = ["pp", "dp", "sharding", "mp", "sep"]
             elif args.hybrid_parallel_topo_order == "sharding_first":
                 order = ["dp", "sharding", "pp", "mp", "sep"]
-            if args.context_parallel_degree is not None and args.context_parallel_degree > 1:
-                sep_degree = args.context_parallel_degree 
-            elif args.sep_parallel_degree is not None and args.sep_parallel_degree > 1:
-                sep_degree = args.sep_parallel_degree
-            else:
-                sep_degree = 1
-            sep_degree=args.sep_parallel_degree if args.sep_parallel_degree > 1 else args.context_parallel_degree
             topo = Topology(
                 dist.get_rank(),
                 dist.get_world_size(),
                 dp_degree=args.dataset_world_size,
                 pp_degree=args.pipeline_parallel_degree,
                 mp_degree=args.tensor_parallel_degree,
-                sep_degree=sep_degree,
                 sharding_degree=1,  # auto_parallel's sharding is not orthogonal with dp, mp and pp
                 order=order,
             )
@@ -464,25 +450,16 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    do_enable_linear_fused_grad_add = training_args.enable_linear_fused_grad_add
-    do_enable_mp_async_allreduce = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and "enable_mp_async_allreduce" in training_args.tensor_parallel_config
-        and not training_args.sequence_parallel
-    )
     do_enable_sp_async_reduce_scatter = (
         training_args.enable_auto_parallel
         and training_args.tensor_parallel_degree > 1
         and training_args.sequence_parallel
         and "enable_sp_async_reduce_scatter" in training_args.tensor_parallel_config
     )
-    if (
-        do_enable_linear_fused_grad_add or do_enable_mp_async_allreduce or do_enable_sp_async_reduce_scatter
-    ) and not training_args.to_static:
+    if training_args.enable_linear_fused_grad_add and not do_enable_sp_async_reduce_scatter:
         from llm.utils.fused_layers import mock_layers
 
-        mock_layers(do_enable_linear_fused_grad_add, do_enable_mp_async_allreduce, do_enable_sp_async_reduce_scatter)
+        mock_layers()
 
     if model_args.tokenizer_name_or_path is None:
         model_args.tokenizer_name_or_path = model_args.model_name_or_path
@@ -570,8 +547,6 @@ def main():
     config.tensor_parallel_rank = training_args.tensor_parallel_rank
     config.sharding_parallel_degree = training_args.sharding_parallel_degree
     config.to_static = training_args.to_static
-    config.sep_parallel_degree = training_args.sep_parallel_degree
-    config.context_parallel_degree = training_args.context_parallel_degree
 
     if training_args.strategy.pipeline.enable and config.virtual_pp_degree > 1:
         pipeline = training_args.strategy.pipeline
@@ -618,6 +593,13 @@ def main():
 
         model.apply(fn)
 
+    if do_enable_sp_async_reduce_scatter:
+        from llm.utils.sp_async_reduce_scatter import (
+            mock_layers_sp_async_reduce_scatter,
+        )
+
+        mock_layers_sp_async_reduce_scatter(model)
+
     # Create the learning_rate scheduler and optimizer
     if training_args.decay_steps is None:
         training_args.decay_steps = training_args.max_steps
@@ -655,7 +637,6 @@ def main():
     )
     trainer = PretrainingTrainer(
         model=model,
-        model_type=model_args.model_type,
         criterion=criterion,
         args=training_args,
         data_collator=data_collator,

llm/auto_parallel/qwen/run_pretrain_auto.py

@@ -433,25 +433,10 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    do_enable_linear_fused_grad_add = training_args.enable_linear_fused_grad_add
-    do_enable_mp_async_allreduce = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and "enable_mp_async_allreduce" in training_args.tensor_parallel_config
-        and not training_args.sequence_parallel
-    )
-    do_enable_sp_async_reduce_scatter = (
-        training_args.enable_auto_parallel
-        and training_args.tensor_parallel_degree > 1
-        and training_args.sequence_parallel
-        and "enable_sp_async_reduce_scatter" in training_args.tensor_parallel_config
-    )
-    if (
-        do_enable_linear_fused_grad_add or do_enable_mp_async_allreduce or do_enable_sp_async_reduce_scatter
-    ) and not training_args.to_static:
-        from llm.utils.fused_layers import mock_layers
+    if training_args.enable_linear_fused_grad_add:
+        from fused_layers import mock_layers
 
-        mock_layers(do_enable_linear_fused_grad_add, do_enable_mp_async_allreduce, do_enable_sp_async_reduce_scatter)
+        mock_layers()
 
     if model_args.tokenizer_name_or_path is None:
         model_args.tokenizer_name_or_path = model_args.model_name_or_path

llm/config/llama/dislora_argument.json

@@ -1,37 +0,0 @@
-{
-    "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
-    "dataset_name_or_path": "/home/bjh/bjh/Dislora/cs_5_lite",
-    "output_dir": "./checkpoints/dislora_ckpts_3",
-    "dislora": true,
-    "per_device_train_batch_size": 1,
-    "gradient_accumulation_steps": 5,
-    "num_train_epochs": 1,
-    "learning_rate": 2e-05,
-    "lr_scheduler_type": "linear",
-    "warmup_steps": 30,
-    "logging_steps": 1,
-    "evaluation_strategy": "no",
-    "save_strategy": "steps",
-    "save_steps": 500,
-    "src_length": 256,
-    "max_length": 512,
-    "bf16": true,
-    "do_train": true,
-    "do_eval": false,
-    "disable_tqdm": false,
-    "load_best_model_at_end": false,
-    "eval_with_do_generation": false,
-    "recompute": false,
-    "save_total_limit": 5,
-    "fp16_opt_level": "O2",
-    "sharding": "stage3",
-    "zero_padding": false,
-    "use_flash_attention": false,
-    "unified_checkpoint": false,
-    "dislora_rank": 8,
-    "dislora_dropout": 0.05,
-    "target_modules": [".*q_proj.*", ".*v_proj.*", ".*k_proj.*", ".*o_proj.*"],
-    "s_tsd": 8,
-    "ortho_lambda": 1.0,
-    "prefer_small_sigma": true
-}

llm/config/qwen/dislora_argument.json

@@ -1,36 +0,0 @@
-{
-    "model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
-    "dataset_name_or_path": "/home/bjh/bjh/Dislora/cs_5_lite",
-    "output_dir": "./checkpoints/dislora_ckpts",
-    "dislora": true,
-    "per_device_train_batch_size": 1,
-    "gradient_accumulation_steps": 1,
-    "num_train_epochs": 1,
-    "learning_rate": 2e-05,
-    "lr_scheduler_type": "linear",
-    "warmup_steps": 30,
-    "logging_steps": 1,
-    "evaluation_strategy": "no",
-    "save_strategy": "steps",
-    "save_steps": 500,
-    "src_length": 256,
-    "max_length": 512,
-    "bf16": true,
-    "do_train": true,
-    "do_eval": false,
-    "disable_tqdm": false,
-    "load_best_model_at_end": false,
-    "eval_with_do_generation": false,
-    "recompute": false,
-    "save_total_limit": 5,
-    "fp16_opt_level": "O2",
-    "sharding": "stage3",
-    "zero_padding": false,
-    "use_flash_attention": false,
-    "unified_checkpoint": false,
-    "dislora_rank": 8,
-    "dislora_dropout": 0.05,
-    "s_tsd": 8,
-    "ortho_lambda": 1.0,
-    "prefer_small_sigma": true
-}

llm/config/qwen/ppo_argument.yaml

@@ -1,131 +0,0 @@
-# RL algorithms
-rl_algorithm: "ppo" # The reinforcement learning algorithm used, supported: "ppo", "grpo", "reinforce_plus_plus"
-
-# models
-actor_model_name_or_path: "Qwen/Qwen2.5-1.5B-Instruct" # The name or path of the actor model
-reward_model_name_or_path: "" # The name or path of the reward model
-critic_model_name_or_path: "Qwen/Qwen2.5-1.5B-Instruct" # The name or path of the critic model
-use_rm_server: false # Whether to use the reward model server
-reward_server: "http://127.0.0.1:8731" # The address of the reward model server
-use_rule_reward: True # The reward for gsm8k dataset. If use_rule_reward: use_rm_server = false
-
-# logging
-logging_dir: ppo-logs # Directory for logging
-logging_steps: 1 # Number of steps between logging
-output_dir: "qwen2.5-1.5b-gsm8k-ppo/checkpoints" # Directory for output ckpts
-report_to: "wandb" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
-wandb_http_proxy: "http://agent.baidu.com:8188" # HTTP proxy for wandb
-run_name: "qwen2.5-1.5b-gsm8k-ppo" # Name of the run
-
-# data
-train_datasets: "gsm8k/train.jsonl" # Path to the training dataset
-eval_datasets: "gsm8k/test.jsonl" # Path to the evaluation dataset
-prompt_key: "src" # Key for the prompt in the dataset
-response_key: "tgt" # Key for the response in the dataset
-dataloader_drop_last: true # Whether to drop the last incomplete batch in the DataLoader
-balance_batch: true # Whether to balance batch size across dataset_world_size
-use_remove_padding: true # Whether to remove padding tokens in the input
-
-# distributed training args
-tensor_parallel_degree: 2 # Degree of tensor parallelism
-sequence_parallel: true # Whether to enable sequence parallelism
-sharding_parallel_degree: -1 # Degree of sharding parallelism
-sharding: "stage1" # Sharding strategy, e.g., "stage1" or "stage2"
-sharding_parallel_config: "enable_release_grads" # Configuration for sharding parallelism
-pipeline_parallel_degree: 1 # Degree of pipeline parallelism
-virtual_pp_degree: 1 # Degree of virtual pipeline parallelism
-
-# rollout args
-max_prompt_len: 1024 # Maximum length of the prompt, exceeding which will be automatically truncated
-max_dec_len: 512 # Maximum length of the response
-min_dec_len: 32 # Minimum length of the response
-top_p: 1.0 # Top-p sampling parameter
-temperature: 1.0 # Temperature parameter for sampling
-repetition_penalty: 1.0 # Repetition penalty parameter
-rollout_max_num_seqs: 1024 # The maximum number of sequences that can be processed in a single inference
-rollout_quant_type: "" # Quantization type, e.g., "weight_only_int8"
-
-# training args
-do_train: true # Whether to perform training
-seed: 42 # Random seed for reproducibility
-global_batch_size: 256 # Global batch size for training (rollouts = rollout_n * global_batch_size)
-global_gen_batch_size: -1 # Global generation batch size for dynamic sampling
-global_mini_batch_size: 64 # Mini-batch size for training, default = global_batch_size
-rollout_n: 1 # Number of rollouts
-update_iters: 1 # Number of training iterations for rollout samples
-per_device_logprob_batch_size: 4 # Log probability batch size per device
-per_device_reward_batch_size: 2 # Reward batch size per device
-per_device_value_batch_size: 2 # Value batch size per device
-per_device_train_batch_size: 2 # Training micro batch size per device
-# gradient_accumulation_steps: 4 # Gradient accumulation steps (auto-calculated): global_bz * rollout_n * 
-num_train_epochs: 5 # Number of training epochs
-max_length: 2048 # Maximum length for training, should be larger than max_prompt_len + max_dec_len
-adam_beta1: 0.9 # AdamW optimizer beta1
-adam_beta2: 0.999 # AdamW optimizer beta2
-adam_epsilon: 1e-8 # AdamW optimizer epsilon
-max_grad_norm: 1.0 # Maximum gradient norm for clipping
-max_steps: -1 # Maximum number of training steps
-save_steps: 300 # Number of steps between model saves
-save_strategy: "steps" # Strategy for saving models
-ignore_save_lr_and_optim: true # Whether to ignore saving learning rate and optimizer state (leave empty if not specified)
-disable_tqdm: true # Whether to disable tqdm progress bar
-
-# actor training args
-learning_rate: 1e-6 # Learning rate for training
-min_learning_rate: 1e-6 # Minimum learning rate
-lr_scheduler_type: "constant" # Learning rate scheduler type
-weight_decay: 1e-2 # Weight decay for the AdamW optimizer
-warmup_ratio: 0.0 # Number of warmup steps
-
-# critic training args
-critic_learning_rate: 1e-5 # Learning rate for critic model
-critic_min_learning_rate: 1e-5 # Minimum learning rate for critic model
-critic_lr_scheduler_type: "constant" # Learning rate scheduler type for critic model
-critic_weight_decay: 1e-2 # Weight decay for the AdamW optimizer of critic model
-critic_warmup_ratio: 0.0 # Number of warmup steps for critic model
-
-# RL args
-kl_coeff: 0.0 # KL coefficient
-kl_loss_coeff: 0.001 # KL loss coefficient
-pg_loss_coeff: 1.0 # Policy gradient loss coefficient
-entropy_coeff: 0.001 # Entropy coefficient
-clip_range_ratio: 0.2 # The clipping range for ratio between the old and new policy. (PPO algorithm)
-clip_range_ratio_low: 0.2 # The clipping range for ratio between the old and new policy. (PPO algorithm)
-clip_range_ratio_high: 0.2 # The clipping range for ratio between the old and new policy. (PPO algorithm)
-clip_range_score: 10.0 # The clipping range for the output of the score model. The reward is clipped into [-clip_range_score, clip_range_score].
-enable_overlong_reward_buffer: false # Whether to enable overlong reward buffer
-overlong_reward_buffer: 256 # The length of the overlong reward buffer
-overlong_penalty_factor: 1.0 # The penalty factor for overlong reward buffer
-clip_range_value: 0.5 # The clipping range for the output of the value model. The value is clipped into [-clip_range_value, clip_range_value].
-normalize_reward: false # Whether to normalize reward
-normalize_advantage: false # Whether to normalize advantage
-dynamic_sampling: false # Whether to use dynamic sampling, which is introcuded in DAPO algorithm https://arxiv.org/abs/2503.14476
-max_gen_batches: 2 # Maximum number of generation batches for dynamic sampling
-use_fp32_compute: true # Whether to use fp32 to compute xx_log_prob,rewards, advantages and loss
-
-# eval args
-do_eval: true # Whether to perform evaluation
-per_device_eval_batch_size: 1319 # Evaluation batch size per device
-evaluation_strategy: "steps" # Evaluation strategy, e.g., "steps"
-eval_steps: 10 # Number of steps between evaluations
-
-# device memory optimization args
-use_flash_attention: true # Whether to use fused attention operations
-use_fused_rms_norm: true # Whether to use fused RMS norm operations, which needs to install fused_ln in slm/model_zoo/gpt-3/external_ops
-use_fused_rope: false # Whether to use fused rope operations
-use_fused_head_and_loss_fn: true # Whether to use fused head and loss function
-use_fused_linear: true # Whether to use fused linear operations. 像是一个没有用的参数
-recompute: false # Whether to enable gradient checkpointing for memory optimization
-recompute_use_reentrant: false # Whether to use reentrant recompute
-recompute_granularity: "full" # Granularity of recompute
-bf16: true # Whether to use mixed precision with bfloat16
-fp16_opt_level: "O2" # Optimization level for fp16 and bf16 training
-amp_master_grad: false # Whether to use float32 weight gradients for master weights in amp opt level=’O2’
-amp_custom_black_list: ["reduce_sum", "softmax_with_cross_entropy", "c_softmax_with_cross_entropy", "elementwise_div", "sin", "cos"] # Custom black list for amp
-amp_custom_white_list: ["lookup_table", "lookup_table_v2", "flash_attn", "matmul", "matmul_v2", "fused_gemm_epilogue"] # Custom white list for amp
-offload_level: "freeze_model" # Level of model offloading to pinned memory, supported values: freeze_model, train_model, optimizer
-release_grads: true # Whether to release gradients
-offload_optim: false # Whether to offload optimizer to pinned memory
-
-# benchmark args
-skip_profile_timer: false # Whether to skip profiling time

llm/docs/finetune.md

@@ -94,7 +94,7 @@ python  -u  -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7"  run_finetune.
 3. 可以通过设置`weight_quantize_algo`将主干模型量化低比特，例如'weight_only_int4','weight_only_int8'，'nf4'或'fp4'。具体参考精调参数介绍
 4. 设置`use_flash_attention`为 True 使用 FlashAttention。在 FlashAttention 打开的基础上设置`flash_mask`为 True 使用 FlashMask。
 5. LoRA API 支持4D 并行策略，可以通过控制`tensor_parallel_degree`、`pipeline_parallel_degree`、 `sharding`、`sharding_parallel_degree`调整并行训练策略，可拓展至**单机 LoRA 微调千亿模型**。
-6. 可配置`rslora`、`lora_plus_scale`、`pissa`、`lora_use_mixer`、`mixer_num`、`use_mora`等参数，使用 rsLoRA、LoRa+、PiSSA、MosLoRA（暂不支持张量模型并行）、LinChain（暂不支持张量模型并行）、MoRA（暂不支持张量模型并行） 等算法。
+6. 可配置`rslora`、`lora_plus_scale`、`pissa`、`lora_use_mixer`、`use_mora`等参数，使用 rsLoRA、LoRa+、PiSSA、MosLoRA（暂不支持张量模型并行）、MoRA（暂不支持张量模型并行） 等算法。
 
 为了后续的**压缩**和**静态图推理**方便，我们提供 LoRA 参数合并脚本，可以将 LoRA 参数合并到主干模型并保存相应的权重。
 ```
@@ -177,33 +177,7 @@ python merge_lokr_params.py \
 - `device`: 运行环境，默认为 gpu。
 </div>
 
-#### 3.4.5 DisLoRA
-```
-# 单卡DisLoRA
-python  run_finetune.py ./config/llama/dislora_argument.json
-
-# 多卡DisLoRA（暂不支持张量模型并行）
-python  -u  -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7"  run_finetune.py ./config/llama/dislora_argument.json
-```
-为了后续的**压缩**和**静态图推理**方便，我们提供 DisLoRA 参数合并脚本，可以将 DisLoRA 参数合并到主干模型并保存相应的权重。
-```
-python merge_dislora_params.py \
-    --model_name_or_path ./base_model \
-    --dislora_path ./checkpoints/dislora_ckpts \
-    --merge_dislora_model_path ./checkpoints/dislora_merge \
-    --device "gpu" \
-    --low_gpu_mem True
-```
-
-<summary>&emsp; 脚本参数介绍</summary><div>
-
-- `dislora_path`: DisLoRA 参数和配置路径，对 DisLoRA 参数进行初始化，默认为 None。
-- `model_name_or_path`: 必须，主干模型参数路径，默认为 None。
-- `merge_dislora_model_path`: 必须，合并参数后保存路径，默认为 None。
-- `device`: 运行环境，默认为 gpu。
-</div>
-
-#### 3.4.6 ReFT
+#### 3.4.4 ReFT
 ```
 # 单卡ReFT
 python  run_finetune.py ./config/llama/reft_argument.json
@@ -254,8 +228,6 @@ python ./predict/reft_predictor.py \
 - `vera_rank`: VeRA 算法中 rank（秩）的值，默认为8。
 - `lokr`: 是否开启 [LoKr](https://arxiv.org/abs/2309.14859) 微调策略，默认为 False。
 - `lokr_rank`: LoKr 算法中 rank（秩）的值，默认为8。
-- `dislora`: 是否开启 [DisLoRA] 微调策略，默认为 False。
-- `dislora_rank`: DisLoRA 算法中 rank（秩）的值，默认为8。
 - `use_long_sequence_strategies`: 是否使用长序列扩展策略，默认为 False。
 - `reft`: 是否开启 [ReFT](https://arxiv.org/abs/2404.03592) 微调策略，默认为 False。
 - `use_mora`: 是否开启 [MoRA](https://arxiv.org/abs/2405.12130) 微调策略，默认为 False。

llm/predict/predictor.py

@@ -40,13 +40,7 @@ except:
     pass
 
 from paddlenlp.generation import GenerationConfig, TextIteratorStreamer
-from paddlenlp.peft import (
-    LoRAConfig,
-    LoRAModel,
-    PrefixConfig,
-    PrefixModelForCausalLM,
-    TAREModel,
-)
+from paddlenlp.peft import LoRAConfig, LoRAModel, PrefixConfig, PrefixModelForCausalLM
 from paddlenlp.taskflow.utils import static_mode_guard
 from paddlenlp.trainer import PdArgumentParser
 from paddlenlp.transformers import (
@@ -91,9 +85,6 @@ class PredictorArgument:
     device: str = field(default="gpu", metadata={"help": "Device"})
     dtype: str = field(default=None, metadata={"help": "Model dtype"})
     lora_path: str = field(default=None, metadata={"help": "The directory of LoRA parameters. Default to None"})
-    tare_path: str = field(default=None, metadata={"help": "The directory of TARE parameters. Default to None"})
-    tare_n: int = field(default=8, metadata={"help": "The num of TARE editors. Default to 8."})
-    tare_k: int = field(default=7, metadata={"help": "The num of TARE selected editors. Default to 7."})
     export_precache: bool = field(default=False, metadata={"help": "whether use prefix weight to do infer"})
     prefix_path: str = field(
         default=None, metadata={"help": "The directory of Prefix Tuning parameters. Default to None"}
@@ -205,7 +196,6 @@ class PredictorArgument:
     )
     dynamic_insert: bool = field(default=False, metadata={"help": "whether use dynamic insert"})
     total_request_num: int = field(default=None, metadata={"help": "The total number of request data"})
-    kv_cache_reuse: int = field(default=0)
 
     def __post_init__(self):
         if self.speculate_method is not None:
@@ -364,11 +354,6 @@ class DygraphPredictor(BasePredictor):
                 prefix_path=config.prefix_path,
                 postprocess_past_key_value=prefix_tuning_params["postprocess_past_key_value"],
             )
-
-        if config.tare_path is not None:
-            self.model = TAREModel(base_model=self.model, n=config.tare_n, k=config.tare_k)
-            self.model.load_model(os.path.join(config.tare_path, "delta_vector.pth"))
-
         self.model.eval()
 
     @paddle.no_grad()
@@ -1170,11 +1155,6 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
             for cache_k_shape, cache_v_shape in zip(self.cache_k_shapes, self.cache_v_shapes):
                 self.cache_kvs.append(paddle.zeros(cache_k_shape, dtype=cachekv_dtype))
                 self.cache_kvs.append(paddle.zeros(cache_v_shape, dtype=cachekv_dtype))
-                if self.config.kv_cache_reuse:
-                    logger.warning(
-                        f"self.config.kv_cache_reuse = {self.config.kv_cache_reuse}, break, len(self.cache_kvs) = {len(self.cache_kvs)}"
-                    )
-                    break
         else:
             # for mla's absorption
             assert self.cache_v_shapes is None

llm/run_finetune.py

@@ -31,15 +31,12 @@ from paddlenlp.datasets import (
 )
 from paddlenlp.metrics import BLEU, Rouge1, Rouge2, RougeL
 from paddlenlp.peft import (
-    DisLoRAConfig,
-    DisLoRAModel,
     LoKrConfig,
     LoKrModel,
     LoRAConfig,
     LoRAModel,
     PrefixConfig,
     PrefixModelForCausalLM,
-    TAREModel,
     VeRAConfig,
     VeRAModel,
 )
@@ -71,7 +68,7 @@ from paddlenlp.transformers import (
 )
 from paddlenlp.transformers.configuration_utils import LlmMetaConfig
 from paddlenlp.transformers.longlora import replace_llama_attn, set_group_size
-from paddlenlp.trl import DataConfig, DisLoRATrainer, ModelConfig, SFTConfig, SFTTrainer
+from paddlenlp.trl import DataConfig, ModelConfig, SFTConfig, SFTTrainer
 from paddlenlp.trl.llm_utils import (
     ZeroPaddingIterDatasetCallback,
     compute_metrics,
@@ -169,13 +166,6 @@ def main():
         qlora_weight_blocksize=model_args.qlora_weight_blocksize,
         qlora_weight_double_quant=model_args.qlora_weight_double_quant,
         qlora_weight_double_quant_block_size=model_args.qlora_weight_double_quant_block_size,
-        apply_hadamard=model_args.apply_hadamard,
-        hadamard_block_size=model_args.hadamard_block_size,
-        quant_input_grad=model_args.quant_input_grad,
-        quant_weight_grad=model_args.quant_weight_grad,
-        apply_online_actscale_step=model_args.apply_online_actscale_step,
-        actscale_moving_rate=model_args.actscale_moving_rate,
-        fp8_format_type=model_args.fp8_format_type,
     )
 
     model_config = AutoConfig.from_pretrained(
@@ -314,15 +304,6 @@ def main():
         tokenizer.pad_token_id = tokenizer.eos_token_id
 
     train_ds, dev_ds, test_ds = create_dataset(data_args, training_args)
-
-    train_dataset_size = None
-    if train_ds is not None and model_args.dislora:
-        train_dataset_size = get_dataset_size(train_ds)
-        if train_dataset_size is not None:
-            logger.info(f"Original training dataset size: {train_dataset_size}")
-        else:
-            logger.warning("Unable to determine training dataset size for dynamic dash_flag calculation")
-
     # TODO(ZHUI & sijunhe): Temporary implementation. Generalize this logic and move to Trainer later.
     if training_args.resume_from_checkpoint is not None and data_args.lazy:
         logger.info(
@@ -389,9 +370,7 @@ def main():
         if eval_zero_padding and test_ds is not None:
             test_ds = intoken_dataset(test_ds, tokenizer=tokenizer, max_length=data_args.max_length)
 
-    model = create_peft_model(
-        model_args, reft_args, training_args, dtype, model_config, model, reft_layers, train_dataset_size
-    )
+    model = create_peft_model(model_args, reft_args, training_args, dtype, model_config, model, reft_layers)
 
     def compute_metrics_do_generation(eval_preds):
         rouge1 = Rouge1()
@@ -455,33 +434,20 @@ def main():
         return_attention_mask=not model_args.flash_mask,
         pad_to_multiple_of=data_args.pad_to_multiple_of,
     )
-
-    if model_args.dislora and hasattr(model_args, "ortho_lambda"):
-        training_args.dislora_ortho_lambda = model_args.ortho_lambda
-
-    trainer_kwargs = {
-        "model": model,
-        "args": training_args,
-        "train_dataset": train_ds,
-        "eval_dataset": dev_ds,
-        "tokenizer": tokenizer,
-        "compute_metrics": metrics,
-        "data_collator": data_collator_fn if not model_args.reft else ReftDataCollator(data_collator=data_collator_fn),
-        "do_generation": data_args.eval_with_do_generation,
-        "callbacks": [ZeroPaddingIterDatasetCallback()] if isinstance(train_ds, ZeroPaddingIterableDataset) else None,
-        "gen_args": gen_args,
-        "data_args": data_args,
-    }
-
-    if model_args.dislora:
-        logger.info("Using DisLoRATrainer for training.")
-        trainer = DisLoRATrainer(**trainer_kwargs)
-    else:
-        trainer = SFTTrainer(**trainer_kwargs)
-
-    trainable_parameters = [
-        p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name)
-    ]
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=dev_ds,
+        tokenizer=tokenizer,
+        compute_metrics=metrics,
+        data_collator=data_collator_fn if not model_args.reft else ReftDataCollator(data_collator=data_collator_fn),
+        do_generation=data_args.eval_with_do_generation,
+        callbacks=[ZeroPaddingIterDatasetCallback()] if isinstance(train_ds, ZeroPaddingIterableDataset) else None,
+        gen_args=gen_args,
+        data_args=data_args,
+    )
+    trainable_parameters = [p for p in model.parameters() if not p.stop_gradient]
     trainer.set_optimizer_grouped_parameters(trainable_parameters)
     if model_args.lorapro:
         optimizer = AdamWLoRAPro(
@@ -501,10 +467,6 @@ def main():
         elif last_checkpoint is not None:
             checkpoint = last_checkpoint
         train_result = trainer.train(resume_from_checkpoint=checkpoint)
-
-        if model_args.tare:
-            model.save_model(os.path.join(training_args.output_dir, "delta_vector.pth"))
-
         if model_args.neftune:
             neft_post_hook_handle.remove()
         if training_args.benchmark:
@@ -560,9 +522,7 @@ def save_to_aistudio(model_args, training_args, trainer):
     )
 
 
-def create_peft_model(
-    model_args, reft_args, training_args, dtype, model_config, model, reft_layers, train_dataset_size
-):
+def create_peft_model(model_args, reft_args, training_args, dtype, model_config, model, reft_layers):
     if model_args.prefix_tuning:
         if training_args.pipeline_parallel_degree > 1:
             raise NotImplementedError("Prefix tuning is not implemented for pipeline parallelism.")
@@ -611,9 +571,6 @@ def create_peft_model(
                 use_quick_lora=model_args.use_quick_lora,
                 lora_use_mixer=model_args.lora_use_mixer,
                 use_mora=model_args.use_mora,
-                nola=model_args.nola,
-                nola_basis_num=model_args.nola_basis_num,
-                mixer_num=model_args.mixer_num,
                 lorapro=model_args.lorapro,
             )
             if model_args.lorapro:
@@ -640,53 +597,6 @@ def create_peft_model(
         else:
             model = LoKrModel.from_pretrained(model=model, lokr_path=model_args.lokr_path)
 
-    if model_args.dislora:
-        # Calculate dynamic dash_flag based on training configuration
-        if train_dataset_size is not None and training_args.do_train:
-            # Calculate warmup steps: len(train_data) * num_epochs // (batch_size * gradient_accumulation_steps * 3)
-            effective_batch_size = (
-                training_args.per_device_train_batch_size
-                * training_args.gradient_accumulation_steps
-                * training_args.dataset_world_size  # Consider data parallel
-            )
-            calculated_dash_flag = (train_dataset_size * training_args.num_train_epochs) // (effective_batch_size * 3)
-
-            # Use calculated value if it's reasonable, otherwise fall back to model_args
-            if calculated_dash_flag > 0:
-                dash_flag = calculated_dash_flag
-                logger.info(
-                    f"Calculated dynamic dash_flag: {dash_flag} based on dataset size: {train_dataset_size}, "
-                    f"epochs: {training_args.num_train_epochs}, effective batch size: {effective_batch_size}"
-                )
-            else:
-                dash_flag = model_args.dash_flag
-                logger.warning(
-                    f"Calculated dash_flag was {calculated_dash_flag}, using model_args.dash_flag: {dash_flag}"
-                )
-        else:
-            dash_flag = getattr(model_args, "dash_flag", 50)
-            if train_dataset_size is None:
-                logger.info(
-                    f"Unable to calculate dynamic dash_flag (dataset size unknown), using configured dash_flag: {dash_flag}"
-                )
-            else:
-                logger.info(f"Not in training mode, using configured dash_flag: {dash_flag}")
-        if model_args.dislora_path is None:
-            dislora_config = DisLoRAConfig(
-                target_modules=model_args.target_modules
-                if model_args.target_modules
-                else get_lora_target_modules(model),
-                r=model_args.dislora_rank,
-                dislora_alpha=1.5 * model_args.dislora_rank,
-                dislora_dropout=model_args.dislora_dropout,
-                dtype=dtype,
-                base_model_name_or_path=model_args.model_name_or_path,
-                s_tsd=model_args.s_tsd,
-                dash_flag=dash_flag,  # Use calculated dash_flag
-                ortho_lambda=model_args.ortho_lambda,
-            )
-        model = DisLoRAModel(model, dislora_config)
-
     if model_args.reft:
         intervention_dtype = dtype
         intervention_params = {
@@ -730,10 +640,6 @@ def create_peft_model(
         model.mark_only_vera_as_trainable(notfreezeB=True)
         model.print_trainable_parameters()
 
-    if model_args.tare:
-        model = TAREModel(base_model=model, n=model_args.tare_n, k=model_args.tare_k)
-        model.print_trainable_parameters()
-
     return model
 
 
@@ -830,24 +736,5 @@ def create_dataset(data_args, training_args):
     return train_ds, dev_ds, test_ds
 
 
-def get_dataset_size(dataset):
-    """Get the size of a dataset, handling both lazy and regular datasets"""
-    if dataset is None:
-        return None
-
-    try:
-        if hasattr(dataset, "__len__"):
-            return len(dataset)
-        elif hasattr(dataset, "_length"):
-            return dataset._length
-        else:
-            # For lazy datasets, we might need to iterate once to count
-            logger.warning("Unable to determine dataset size directly for lazy loading dataset")
-            return None
-    except Exception as e:
-        logger.warning(f"Error getting dataset size: {e}")
-        return None
-
-
 if __name__ == "__main__":
     main()

llm/tools/merge_dislora_params.py

@@ -1,290 +0,0 @@
-# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-
-import paddle
-
-from paddlenlp.peft import DisLoRAConfig, DisLoRAModel
-from paddlenlp.transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
-from paddlenlp.utils.env import CONFIG_NAME
-
-
-def parse_arguments():
-    """解析命令行参数"""
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name_or_path", default=None, help="The directory of pretrained model.")
-    parser.add_argument("--dislora_path", default="", help="The directory of dislora parameters. Default to None")
-    parser.add_argument(
-        "--merge_dislora_model_path",
-        default="",
-        help="The directory of merged parameters. Default to None",
-    )
-    parser.add_argument("--device", type=str, default="gpu", help="Device")
-    parser.add_argument(
-        "--low_gpu_mem", type=bool, default=True, help="Whether to use low gpu memory. Default to False"
-    )
-    return parser.parse_args()
-
-
-def weight_process(name, dislora_config, state_dict):
-    """
-    Based on the DisLoRA algorithm for processing weight merging:
-    The final weight = W_prin + W_res + W_TSD
-    However, here we do not directly add the adapter to the base model; instead, we reconstruct the entire weight matrix.
-    Args:
-        name: Layer name (e.g. "model.layers.0.self_attn.q_proj")
-        dislora_config: DisLoRA configuration
-        state_dict: Model state dictionary
-    # Define the weight_process function to handle the DisLoRA weight merging. The parameters include the layer name, DisLoRA configuration, and the model state dictionary.
-    """
-
-    weight_key = name + ".weight"
-
-    if weight_key not in state_dict:
-        print(f"Warning: {weight_key} not found in state_dict")
-        return
-
-    w_prin = state_dict[weight_key]
-    print(f"Processing layer: {name}")
-    print(f"  W_prin shape: {w_prin.shape}")
-
-    scaling = dislora_config.dislora_alpha / dislora_config.r
-
-    final_weight = w_prin.clone()
-
-    ur_key = name + ".Direc_Ur.weight"
-    sr_key = name + ".Direc_Sr"
-    vhr_key = name + ".Direc_Vhr.weight"
-
-    w_res_added = False
-
-    if all(key in state_dict for key in [ur_key, sr_key, vhr_key]):
-
-        direc_ur = state_dict[ur_key]  # [r, out_features]
-        direc_sr = state_dict[sr_key]  # [r]
-        direc_vhr = state_dict[vhr_key]  # [in_features, r]
-
-        s_diag = paddle.diag(direc_sr)  # [r, r]
-
-        w_res = direc_vhr @ s_diag @ direc_ur * scaling  # [in_features, out_features]
-
-        if w_res.shape != w_prin.shape:
-            print(f"  Error: W_res shape {w_res.shape} doesn't match W_prin shape {w_prin.shape}")
-            return
-
-        final_weight += w_res
-        w_res_added = True
-        print(f"  ✓ Added W_res with scaling factor: {scaling}")
-    else:
-        print(f"  ⚠ W_res components not found for {name}")
-
-    utsd_key = name + ".Direc_Utsd.weight"
-    stsd_key = name + ".Direc_Stsd"
-    vhtsd_key = name + ".Direc_Vhtsd.weight"
-
-    w_tsd_added = False
-    if all(key in state_dict for key in [utsd_key, stsd_key, vhtsd_key]):
-
-        direc_utsd = state_dict[utsd_key]  # [s_tsd, out_features]
-        direc_stsd = state_dict[stsd_key]  # [s_tsd]
-        direc_vhtsd = state_dict[vhtsd_key]  # [in_features, s_tsd]
-
-        if not paddle.all(direc_stsd == 0.0):
-
-            s_diag_tsd = paddle.diag(direc_stsd)  # [s_tsd, s_tsd]
-
-            w_tsd = direc_vhtsd @ s_diag_tsd @ direc_utsd * scaling  # [in_features, out_features]
-
-            if w_tsd.shape != w_prin.shape:
-                print(f"  Error: W_TSD shape {w_tsd.shape} doesn't match W_prin shape {w_prin.shape}")
-                return
-
-            final_weight += w_tsd
-            w_tsd_added = True
-            print(f"  ✓ Added W_TSD with scaling factor: {scaling}")
-        else:
-            print(f"  ⚠ W_TSD parameters are uninitialized (all zeros) for {name}")
-    else:
-        print(f"  ⚠ W_TSD components not found for {name}")
-
-    state_dict[weight_key] = final_weight
-
-    keys_to_remove = []
-    for key in state_dict.keys():
-        if key.startswith(name + ".Direc_") or key == name + ".step":
-            keys_to_remove.append(key)
-
-    for key in keys_to_remove:
-        removed_param = state_dict.pop(key)
-        print(f"  ✓ Removed DisLoRA parameter: {key} (shape: {removed_param.shape})")
-
-    components = []
-    if w_res_added:
-        components.append("W_res")
-    if w_tsd_added:
-        components.append("W_TSD")
-
-    if components:
-        print(f"  ✓ Successfully merged: W_prin + {' + '.join(components)}")
-    else:
-        print("  ✓ Kept original W_prin (no adaptations found)")
-    print()
-
-
-def merge():
-
-    args = parse_arguments()
-    paddle.set_device(args.device)
-
-    print("Loading DisLoRA configuration...")
-    dislora_config = DisLoRAConfig.from_pretrained(args.dislora_path)
-    if dislora_config.base_model_name_or_path is None:
-        if args.model_name_or_path is None:
-            raise ValueError("We can not find a valid model_name_or_path.")
-        else:
-            dislora_config.base_model_name_or_path = args.model_name_or_path
-
-    print("Loading model configuration...")
-    if os.path.isfile(os.path.join(args.dislora_path, CONFIG_NAME)):
-        config = AutoConfig.from_pretrained(args.dislora_path)
-    elif args.model_name_or_path is not None:
-        config = AutoConfig.from_pretrained(args.model_name_or_path)
-    else:
-        raise ValueError(
-            f"We can not find config.json in dislora_path: {args.dislora_path} or find a valid model_name_or_path."
-        )
-
-    config.dtype = dislora_config.dtype
-
-    if (
-        dislora_config.dtype == "bfloat16"
-        or (
-            hasattr(config, "quantization_config")
-            and hasattr(config.quantization_config, "weight_quantize_algo")
-            and config.quantization_config.weight_quantize_algo in ["nf4", "fp4"]
-        )
-    ) and args.device == "cpu":
-        raise ValueError("We can not apply bfloat16 or nf4/fp4 dislora merge on cpu.")
-
-    print("Loading base model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        dislora_config.base_model_name_or_path,
-        config=config,
-        low_cpu_mem_usage=args.low_gpu_mem,
-    )
-
-    print("Loading DisLoRA model...")
-    model = DisLoRAModel.from_pretrained(model=model, dislora_path=args.dislora_path, dislora_config=dislora_config)
-
-    model.eval()
-    model_state_dict = model.model.state_dict()
-
-    print(f"Total parameters in state_dict: {len(model_state_dict)}")
-
-    step_keys = [key for key in model_state_dict.keys() if key.endswith(".step")]
-    if step_keys:
-        print(f"Found {len(step_keys)} step parameters in loaded model:")
-        for key in step_keys[:5]:
-            print(f"  {key}")
-        if len(step_keys) > 5:
-            print(f"  ... and {len(step_keys) - 5} more")
-    else:
-        print("No step parameters found in loaded model")
-    print()
-
-    print("Identifying DisLoRA layers...")
-    dislora_name_set = set()
-    for key in model_state_dict.keys():
-        if any(
-            dislora_param in key
-            for dislora_param in ["Direc_Ur", "Direc_Sr", "Direc_Vhr", "Direc_Utsd", "Direc_Stsd", "Direc_Vhtsd"]
-        ):
-
-            for param_type in ["Direc_Ur", "Direc_Sr", "Direc_Vhr", "Direc_Utsd", "Direc_Stsd", "Direc_Vhtsd"]:
-                if f".{param_type}" in key:
-                    layer_name = key.split(f".{param_type}")[0]
-                    dislora_name_set.add(layer_name)
-                    break
-
-    dislora_name_list = sorted(list(dislora_name_set))
-
-    print(f"Found {len(dislora_name_list)} DisLoRA layers:")
-    for i, name in enumerate(dislora_name_list, 1):
-        print(f"  {i:2d}. {name}")
-    print()
-
-    print("Merging DisLoRA parameters...")
-
-    for i, name in enumerate(dislora_name_list, 1):
-        print(f"[{i}/{len(dislora_name_list)}] Processing: {name}")
-        weight_process(name, dislora_config, model_state_dict)
-
-    print("Cleaning up remaining step parameters...")
-    step_keys_to_remove = [key for key in model_state_dict.keys() if key.endswith(".step")]
-    for key in step_keys_to_remove:
-        removed_param = model_state_dict.pop(key)
-        print(f"  ✓ Removed step parameter: {key} (shape: {removed_param.shape})")
-
-    if step_keys_to_remove:
-        print(f"✓ Removed {len(step_keys_to_remove)} step parameters")
-    else:
-        print("✓ No step parameters found")
-    print()
-
-    print("Verifying parameter cleanup...")
-    remaining_dislora_params = []
-    remaining_step_params = []
-    for key in model_state_dict.keys():
-        if any(
-            dislora_param in key
-            for dislora_param in ["Direc_Ur", "Direc_Sr", "Direc_Vhr", "Direc_Utsd", "Direc_Stsd", "Direc_Vhtsd"]
-        ):
-            remaining_dislora_params.append(key)
-        if key.endswith(".step"):
-            remaining_step_params.append(key)
-
-    if remaining_dislora_params:
-        print(f"Warning: {len(remaining_dislora_params)} DisLoRA parameters still remain:")
-        for param in remaining_dislora_params:
-            print(f"  - {param}")
-    else:
-        print("✓ All DisLoRA parameters successfully removed")
-
-    if remaining_step_params:
-        print(f"Warning: {len(remaining_step_params)} step parameters still remain:")
-        for param in remaining_step_params:
-            print(f"  - {param}")
-    else:
-        print("✓ All step parameters successfully removed")
-    print()
-
-    print("Saving merged model...")
-    os.makedirs(args.merge_dislora_model_path, exist_ok=True)
-    model.model.save_pretrained(args.merge_dislora_model_path, state_dict=model_state_dict)
-
-    print("Saving tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(dislora_config.base_model_name_or_path)
-    tokenizer.save_pretrained(args.merge_dislora_model_path)
-
-    print("=" * 80)
-    print("✓ DisLoRA merge completed successfully!")
-    print(f"✓ Merged model saved to: {args.merge_dislora_model_path}")
-    print(f"✓ Processed {len(dislora_name_list)} DisLoRA layers")
-    print("=" * 80)
-
-
-if __name__ == "__main__":
-    merge()

llm/tools/merge_lora_params.py

@@ -78,15 +78,7 @@ def weight_process(name, quant_config, lora_config, state_dict, device):
         raise ValueError(f"quant_config.weight_quantize_algo {quant_config.weight_quantize_algo} is not supported.")
 
 
-def get_mixer(mixer, mixer_num, index=0):
-    if index == mixer_num - 1:
-        return mixer[index]
-    else:
-        return mixer[index] @ get_mixer(mixer, mixer_num, index + 1)
-
-
 def lora_process(name, layer, lora_config, state_dict, device, lora_state_dict=None):
-
     target_device = device if device == "cpu" else device + ":0"
 
     if (name + ".weight") not in state_dict.keys():
@@ -94,44 +86,35 @@ def lora_process(name, layer, lora_config, state_dict, device, lora_state_dict=N
 
     weight = state_dict.pop(name + ".weight")
     lora_use_mixer = lora_config.lora_use_mixer
-
-    mixer_num = lora_config.mixer_num
-    mixer = {}
     use_mora = lora_config.use_mora
-
     if lora_state_dict is None:
         lora_A = state_dict.pop(name + ".lora_A")
         if not use_mora:
             lora_B = state_dict.pop(name + ".lora_B")
         if lora_use_mixer:
-            for i in range(mixer_num):
-                mixer[i] = state_dict.pop(name + ".lora_mixer_" + str(i))
+            lora_AB = state_dict.pop(name + ".lora_AB")
     else:
         lora_A = lora_state_dict.pop(name + ".lora_A")
         if not use_mora:
             lora_B = lora_state_dict.pop(name + ".lora_B")
         if lora_use_mixer:
-            for i in range(mixer_num):
-                mixer[i] = state_dict.pop(name + ".lora_mixer_" + str(i))
+            lora_AB = lora_state_dict.pop(name + ".lora_AB")
     if device != "cpu":
         weight = weight.to(target_device)
         lora_A = lora_A.to(target_device)
         if not use_mora:
             lora_B = lora_B.to(target_device)
         if lora_use_mixer:
-            for key in mixer.keys():
-                mixer[key] = mixer[key].to(target_device)
+            lora_AB = lora_AB.to(target_device)
 
     if device == "cpu" and weight.dtype.name == "BF16":
         weight = weight.astype("float32")
         lora_A = lora_A.astype("float32")
         if not use_mora:
             lora_B = lora_B.astype("float32")
-
         if lora_use_mixer:
-            for key in mixer.keys():
-                mixer[key] = mixer[key].astype(lora_config.dtype)
-            delta_weight = layer.get_delta_weight(lora_A, lora_B, get_mixer(mixer, mixer_num))
+            lora_AB = lora_AB.astype(lora_config.dtype)
+            delta_weight = layer.get_delta_weight(lora_A, lora_B, lora_AB)
         elif use_mora:
             delta_weight = layer.get_delta_weight(lora_A)
         else:
@@ -139,7 +122,7 @@ def lora_process(name, layer, lora_config, state_dict, device, lora_state_dict=N
         out = (weight + delta_weight).astype(lora_config.dtype)
     else:
         if lora_use_mixer:
-            delta_weight = layer.get_delta_weight(lora_A, lora_B, get_mixer(mixer, mixer_num))
+            delta_weight = layer.get_delta_weight(lora_A, lora_B, lora_AB)
         elif use_mora:
             delta_weight = layer.get_delta_weight(lora_A)
         else:

llm/tools/preprocess/create_pretraining_data.py

@@ -176,7 +176,7 @@ def get_whole_word_mask_tokens(tokens, words, max_word_length=6):
             i += 1
             continue
 
-        # add "##" mark on the middle tokens of Chinese words
+        # add "##" mark on the middel tokens of Chinese words
         # such as ["通过", "利用"] -> ["通", "##过"， "利", "##用"]
         has_add = False
         for length in range(max_word_length, 0, -1):

llm/utils/fused_layers.py

@@ -11,31 +11,10 @@
 # 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 os
-
 import paddle
-import paddle.distributed as dist
 from paddle import _C_ops
-from paddle.distributed import fleet
-from paddle.distributed.fleet.utils.log_util import logger
 from paddle.framework import core
 
-_raise_cuda_env_unset_warning = True
-_mp_async_allreduce = False
-_sp_async_reduce_scatter = False
-
-
-def _check_environment_for_overlap():
-    if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
-        global _raise_cuda_env_unset_warning
-        if _raise_cuda_env_unset_warning:
-            logger.warning(
-                "You set enable_mp_async_allreduce or enable_sp_async_reduce_scatter, but you forget to set environment "
-                "variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
-                "loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
-            )
-            _raise_cuda_env_unset_warning = False
-
 
 def is_fused_matmul_bias_supported():
     if paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm() or paddle.is_compiled_with_xpu():
@@ -50,125 +29,6 @@ else:
     origin_linear = paddle.nn.functional.linear
 
 
-def mp_async_allreduce(x_grad):
-    if _mp_async_allreduce and x_grad.process_mesh is not None:
-        _check_environment_for_overlap()
-        mp_placement_index = x_grad.process_mesh.dim_names.index("mp")
-        if mp_placement_index != -1 and x_grad.placements[mp_placement_index].is_partial():
-            hcg = fleet.get_hybrid_communicate_group()
-            model_parallel_group = hcg.get_model_parallel_group()
-            task = dist.stream.all_reduce(
-                x_grad._local_value(),
-                group=model_parallel_group,
-                sync_op=False,
-            )
-            return task
-        else:
-            return None
-    else:
-        return None
-
-
-def sp_async_reducesctter(x_grad):
-    if _sp_async_reduce_scatter and x_grad.process_mesh is not None:
-        _check_environment_for_overlap()
-        mp_placement_index = x_grad.process_mesh.dim_names.index("mp")
-        if mp_placement_index != -1 and x_grad.placements[mp_placement_index].is_partial():
-            hcg = fleet.get_hybrid_communicate_group()
-            model_parallel_group = hcg.get_model_parallel_group()
-            parallelism = model_parallel_group.nranks
-
-            assert (
-                x_grad.shape[0] % parallelism == 0
-            ), f"Input sequence length {x_grad.shape[0]} can't be divided exactly by sequence parallelism {parallelism}"
-
-            # reduce-scatter dx
-            x_grad_global_shape = x_grad.shape
-            x_grad_global_shape[0] = x_grad_global_shape[0] // parallelism
-            x_grad_local = x_grad._local_value()
-            x_grad_local_shape = x_grad_local.shape
-            x_grad_local_shape[0] = x_grad_local_shape[0] // parallelism
-            dx_local = paddle.empty(shape=x_grad_local_shape, dtype=x_grad.dtype)
-            task = dist.stream.reduce_scatter(
-                dx_local,
-                x_grad_local,
-                op=dist.ReduceOp.SUM,
-                group=model_parallel_group,
-                sync_op=False,
-            )
-            return task, dx_local, x_grad_global_shape
-        else:
-            return None
-    else:
-        return None
-
-
-def sync_mp_allreduce(task, dist_tensor):
-    mp_placement_index = dist_tensor.process_mesh.dim_names.index("mp")
-    new_placments = list()
-    for idx, placement in enumerate(dist_tensor.placements):
-        if idx == mp_placement_index:
-            new_placments.append(dist.Replicate())
-        else:
-            new_placments.append(placement)
-    place = paddle.framework._current_expected_place()
-    place = paddle.framework._get_paddle_place(place)
-
-    task.wait()
-
-    return paddle.Tensor(
-        dist_tensor._local_value(),
-        dims=dist_tensor.shape,
-        process_mesh=dist_tensor.process_mesh,
-        placements=new_placments,
-        place=place,
-    )
-
-
-def sync_sp_reducescatter(task, dist_tensor):
-    task, dx_local, x_grad_global_shape = task
-    placements = [dist.Shard(1), dist.Shard(0)]
-    place = paddle.framework._current_expected_place()
-    place = paddle.framework._get_paddle_place(place)
-
-    task.wait()
-
-    return paddle.Tensor(
-        dx_local,
-        dims=x_grad_global_shape,
-        process_mesh=dist_tensor.process_mesh,
-        placements=placements,
-        place=place,
-    )
-
-
-# modify from Paddle/python/paddle/distributed/auto_parallel/moe_utils.py
-def _dist_reshape(dist_tensor):
-    local_tensor = dist_tensor._local_value()
-    tgt_global_shape = [dist_tensor.shape[0] * dist_tensor.shape[1], dist_tensor.shape[2]]
-    tgt_local_shape = [local_tensor.shape[0] * local_tensor.shape[1], local_tensor.shape[2]]
-
-    place = paddle.framework._current_expected_place()
-    place = paddle.framework._get_paddle_place(place)
-
-    local_tensor = local_tensor.reshape(tgt_local_shape)
-
-    if dist_tensor.placements[1].is_shard():
-        new_placements = [dist.Shard(0), dist.Shard(1)]
-    else:
-        new_placements = [dist.Shard(0), dist.Replicate()]
-
-    out = paddle.Tensor(
-        local_tensor,
-        dims=tgt_global_shape,
-        process_mesh=dist_tensor.process_mesh,
-        placements=new_placements,
-        place=place,
-    )
-    out.stop_gradient = dist_tensor.stop_gradient
-    return out
-
-
 class FusedLinearWithGradAdd(paddle.autograd.PyLayer):
     @staticmethod
     def forward(ctx, x, weight, bias=None, name=None):
@@ -181,131 +41,39 @@ class FusedLinearWithGradAdd(paddle.autograd.PyLayer):
         x, weight, bias = ctx.saved_tensor()
         x_grad = paddle.matmul(y_grad, weight, transpose_y=True)
 
-        if _sp_async_reduce_scatter:
-            task = sp_async_reducesctter(x_grad)
-        else:
-            task = mp_async_allreduce(x_grad)
-
-        # _C_ops.fused_linear_param_grad_add(x, y_grad, weight_grad, bias_grad, multi precision, has bias)
+        # _C_ops.fused_linear_param_grad_add(x, y_grad, dw, db, multi precision, has bias)
         if bias is None:
             if hasattr(weight, "main_grad"):
                 weight.main_grad, _ = _C_ops.fused_linear_param_grad_add(
                     x, y_grad, weight.main_grad, None, True, False
                 )
-                if task is not None:
-                    if _sp_async_reduce_scatter:
-                        x_grad = sync_sp_reducescatter(task, x_grad)
-                    else:
-                        x_grad = sync_mp_allreduce(task, x_grad)
                 return x_grad, None
             else:
                 if weight.grad is not None:
-                    weight.grad, _ = _C_ops.fused_linear_param_grad_add(
-                        x, y_grad, weight.grad, None, False if weight.grad.dtype != paddle.float32 else True, False
-                    )
-                    if task is not None:
-                        if _sp_async_reduce_scatter:
-                            x_grad = sync_sp_reducescatter(task, x_grad)
-                        else:
-                            x_grad = sync_mp_allreduce(task, x_grad)
+                    weight.grad, _ = _C_ops.fused_linear_param_grad_add(x, y_grad, weight.grad, None, False, False)
                     return x_grad, None
                 else:
                     weight_grad, _ = _C_ops.fused_linear_param_grad_add(x, y_grad, None, None, False, False)
-                    if task is not None:
-                        if _sp_async_reduce_scatter:
-                            x_grad = sync_sp_reducescatter(task, x_grad)
-                        else:
-                            x_grad = sync_mp_allreduce(task, x_grad)
                     return x_grad, weight_grad
 
         if hasattr(weight, "main_grad") and hasattr(bias, "main_grad"):
             weight.main_grad, bias.main_grad = _C_ops.fused_linear_param_grad_add(
                 x, y_grad, weight.main_grad, bias.main_grad, True, True
             )
-            if task is not None:
-                if _sp_async_reduce_scatter:
-                    x_grad = sync_sp_reducescatter(task, x_grad)
-                else:
-                    x_grad = sync_mp_allreduce(task, x_grad)
             return x_grad, None, None
         else:
             if weight.grad is not None:
                 assert bias.grad is not None
                 weight.grad, bias.grad = _C_ops.fused_linear_param_grad_add(
-                    x, y_grad, weight.grad, bias.grad, False if weight.grad.dtype != paddle.float32 else True, True
+                    x, y_grad, weight.grad, bias.grad, False, True
                 )
-                if task is not None:
-                    if _sp_async_reduce_scatter:
-                        x_grad = sync_sp_reducescatter(task, x_grad)
-                    else:
-                        x_grad = sync_mp_allreduce(task, x_grad)
                 return x_grad, None, None
             else:
                 weight_grad, bias_grad = _C_ops.fused_linear_param_grad_add(x, y_grad, None, None, False, True)
-                if task is not None:
-                    if _sp_async_reduce_scatter:
-                        x_grad = sync_sp_reducescatter(task, x_grad)
-                    else:
-                        x_grad = sync_mp_allreduce(task, x_grad)
                 return x_grad, weight_grad, bias_grad
 
 
-class OverlapLinear(paddle.autograd.PyLayer):
-    @staticmethod
-    def forward(ctx, x, weight, bias=None, name=None):
-        y = origin_linear(x, weight, bias)
-        ctx.save_for_backward(x, weight, bias)
-        return y
-
-    @staticmethod
-    def backward(ctx, y_grad):
-        x, weight, bias = ctx.saved_tensor()
-        x_grad = paddle.matmul(y_grad, weight, transpose_y=True)
-
-        if _sp_async_reduce_scatter:
-            task = sp_async_reducesctter(x_grad)
-        else:
-            task = mp_async_allreduce(x_grad)
-
-        if _sp_async_reduce_scatter:
-            y_grad = _dist_reshape(y_grad)
-        else:
-            y_grad = y_grad.reshape([-1, y_grad.shape[-1]])
-        weight_grad = paddle.matmul(
-            _dist_reshape(x) if _sp_async_reduce_scatter else x.reshape([-1, x.shape[-1]]),
-            y_grad,
-            transpose_x=True,
-        )
-        if bias is None:
-            if task is not None:
-                if _sp_async_reduce_scatter:
-                    x_grad = sync_sp_reducescatter(task, x_grad)
-                else:
-                    x_grad = sync_mp_allreduce(task, x_grad)
-            return x_grad, weight_grad
-        else:
-            bias_grad = paddle.sum(y_grad, axis=0)
-            if task is not None:
-                if _sp_async_reduce_scatter:
-                    x_grad = sync_sp_reducescatter(task, x_grad)
-                else:
-                    x_grad = sync_mp_allreduce(task, x_grad)
-            return x_grad, weight_grad, bias_grad
-
-
-def mock_layers(
-    enable_fused_linear_grad_add=True, enable_mp_async_allreduce=False, enable_sp_async_reduce_scatter=False
-):
-    global _mp_async_allreduce
-    global _sp_async_reduce_scatter
-    _mp_async_allreduce = enable_mp_async_allreduce
-    _sp_async_reduce_scatter = enable_sp_async_reduce_scatter
-
-    if enable_fused_linear_grad_add:
-        paddle.nn.functional.linear = FusedLinearWithGradAdd.apply
-        if is_fused_matmul_bias_supported():
-            paddle.incubate.nn.functional.fused_linear = FusedLinearWithGradAdd.apply
-    else:
-        paddle.nn.functional.linear = OverlapLinear.apply
-        if is_fused_matmul_bias_supported():
-            paddle.incubate.nn.functional.fused_linear = OverlapLinear.apply
+def mock_layers():
+    paddle.nn.functional.linear = FusedLinearWithGradAdd.apply
+    if is_fused_matmul_bias_supported():
+        paddle.incubate.nn.functional.fused_linear = FusedLinearWithGradAdd.apply

llm/utils/replace_ops.py

@@ -106,7 +106,7 @@ def parallel_cross_entropy(
             "1. soft_label=False is set for parallel computation (current value: {}) \n"
             "2. Input tensor is properly sharded (current sharding status: {}) \n".format(
                 soft_label, 
-                input.placements,
+                input_placement,
             )
         )
 

llm/utils/sp_async_reduce_scatter.py

@@ -172,7 +172,7 @@ def forward_pre_hook(layer, input):
     ipp = id2ipp[id(layer)]
 
 
-def forward_post_hook(layer, input, output):
+def forward_post_hook(layer, input, ouput):
     paddle.nn.functional.linear = paddle_nn_functional_linear
     if is_fused_matmul_bias_supported():
         paddle.incubate.nn.functional.fused_linear = paddle_incubate_nn_functional_fused_linear

ops/csrc/fp8/deep_gemm/include/deep_gemm/fp8_gemm.cuh

@@ -39,34 +39,17 @@ enum class Layout {
     ColMajor
 };
 
-__device__ __host__ constexpr int get_num_math_warpgroups(int block_m) {
-    return block_m == 64 ? 1 : 2;
-}
-
 template <uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup>
 __device__ __host__ constexpr int get_num_threads_per_sm(int block_m) {
     DG_STATIC_ASSERT(kNumMathThreadsPerGroup == 128, "Only support 128 threads per math group");
-    return get_num_math_warpgroups(block_m) * kNumMathThreadsPerGroup + kNumTMAThreads;
-}
-
-template <int kNumFormerIters, int kGap, int kEnd>
-__device__ __host__ void outer_launch_k_iterations(const auto& inner_launch_k_iterations, const auto& func, int num_former_iters) {
-    if (num_former_iters == kNumFormerIters) {
-        inner_launch_k_iterations(func, cute::Int<kNumFormerIters>{});
-        return;
-    }
-
-    if constexpr (kNumFormerIters + kGap <= kEnd)
-        outer_launch_k_iterations<kNumFormerIters + kGap, kGap, kEnd>(inner_launch_k_iterations, func, num_former_iters);
+    return (block_m == 64 ? 1 : 2) * kNumMathThreadsPerGroup + kNumTMAThreads;
 }
 
 template <uint32_t SHAPE_N, uint32_t SHAPE_K,
           uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
-          uint32_t BLOCK_N_PADDING,
-          uint32_t kSwizzleDMode,
           uint32_t kNumGroups, uint32_t kNumStages,
           uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup,
-          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
+          uint32_t kNumTMAMulticast,
           GemmType kGemmType>
 __global__ void __launch_bounds__(get_num_threads_per_sm<kNumTMAThreads, kNumMathThreadsPerGroup>(BLOCK_M), 1)
 fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
@@ -78,16 +61,15 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
     // Scaling checks
     DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
-    DG_STATIC_ASSERT(ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
+    DG_STATIC_ASSERT(ceil_div(BLOCK_N, BLOCK_K) == 1, "Too much B scales in a single block");
 
     // Types
     using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
     using Barrier = cutlass::arch::ClusterTransactionBarrier;
-    DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
 
     // Shared memory
     static constexpr int kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
-    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * (BLOCK_N + BLOCK_N_PADDING) * sizeof(__nv_bfloat16);
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
     static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
     static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
     static constexpr uint32_t SMEM_SCALES_A_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
@@ -107,11 +89,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
         cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_a));
         cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_b));
         cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_scales_a));
-
-        // `tensor_map_d` is only used in swizzling mode
-        // For the `kSwizzleDMode == 0 and BLOCK_N_PADDING == 0` case, it will be treated as padding mode
-        if constexpr (kSwizzleDMode > 0)
-            cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_d));
+        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_d));
     }
     __syncwarp();
 
@@ -150,8 +128,6 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     // Initialize barriers
     DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
     if (threadIdx.x == kNumMathThreads) {
-        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
-        // even with TMA multicast disabled, we want to make the behavior aligned
         #pragma unroll
         for (int i = 0; i < kNumStages; ++ i) {
             full_barriers[i]->init(1);
@@ -169,23 +145,15 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     // For pipeline unrolling
     struct DivisibleK {};
     struct NotDivisibleK {};
-    auto launch_k_iterations = [](const auto& func, int num_former_iters) {
-        constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
-        constexpr int kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
-        constexpr int kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
-
-        // NOTES: for too-many branches (> 5), we disable this optimization
-        // Otherwise, the compiler must know the dynamic variable `num_former_iters`'s real value
-        outer_launch_k_iterations<0, kGap, kEnd>([](const auto& func, auto num_former_iters_type) {
-            if constexpr (SHAPE_K % kFullKOfAllStages == 0) {
-                for (int k_iter = 0; k_iter < kNumIterations; ++ k_iter)
-                    func(k_iter, DivisibleK{}, num_former_iters_type);
-            } else {
-                for (int k_iter = 0; k_iter < kNumIterations - 1; ++ k_iter)
-                    func(k_iter, DivisibleK{}, num_former_iters_type);
-                func(kNumIterations - 1, NotDivisibleK{}, num_former_iters_type);
-            }
-        }, func, kShouldOptimize ? num_former_iters : 0);
+    auto launch_k_iterations = [](const auto& func) {
+        if constexpr (SHAPE_K % kFullKOfAllStages == 0) {
+            for (int k_iter = 0; k_iter < kNumIterations; ++ k_iter)
+                func(k_iter, DivisibleK{});
+        } else {
+            for (int k_iter = 0; k_iter < kNumIterations - 1; ++ k_iter)
+                func(k_iter, DivisibleK{});
+            func(kNumIterations - 1, NotDivisibleK{});
+        }
     };
 
     // Register reconfigurations
@@ -194,7 +162,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 
     // Block scheduler
     uint32_t m_block_idx, n_block_idx;
-    auto scheduler = Scheduler<kGemmType, SHAPE_N, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA>(shape_m, grouped_layout);
+    auto scheduler = Scheduler<kGemmType, SHAPE_N, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast>(shape_m, grouped_layout);
 
     if (threadIdx.x >= kNumMathThreads) {
         // TMA warp-group for loading data
@@ -204,34 +172,28 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
         if (threadIdx.x == kNumMathThreads) {
             // Persistently schedule over blocks
             while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
-                launch_k_iterations([&](int k_iter, auto type, auto _) {
+                launch_k_iterations([&](int k_iter, auto type) {
                     constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
                     constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
                     DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
 
-                    // Assign TMA multicast number into A and B
-                    constexpr int kNumTMAMulticastOnA = kIsTMAMulticastOnA ? kNumTMAMulticast : 1;
-                    constexpr int kNumTMAMulticastOnB = kIsTMAMulticastOnA ? 1 : kNumTMAMulticast;
-
-                    // NOTES: unrolling and `kNumInnerStages` are vital for performance, NVCC will try to eliminate all
-                    // shared memory pointers, e.g. `full_barriers` registers, if all the access indices are constant
                     #pragma unroll
                     for (uint32_t s = 0; s < kNumInnerStages; ++ s) {
                         // Wait consumer release
                         empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
 
-                        // Issue TMA A
+                        // Issue TMA A with broadcasting
                         auto& full_barrier = *full_barriers[s];
                         int k_idx = k_iter * kFullKOfAllStages + s * BLOCK_K;
-                        tma_copy<kNumTMAMulticastOnA>(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
-                                                      smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
-                        tma_copy<kNumTMAMulticastOnA>(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
-                                                      smem_scales_a[s], m_block_idx * BLOCK_M,
-                                                      scheduler.get_global_idx(SHAPE_K_SCALES, 1, k_idx / BLOCK_K));
+                        tma_copy<kNumTMAMulticast>(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                                   smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
+                        tma_copy<kNumTMAMulticast>(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                                   smem_scales_a[s], m_block_idx * BLOCK_M,
+                                                   scheduler.get_global_idx(SHAPE_K_SCALES, 1, k_idx / BLOCK_K));
 
-                        // Issue TMA B
-                        tma_copy<kNumTMAMulticastOnB>(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
-                                                      smem_b[s], k_idx, scheduler.get_global_idx<false>(SHAPE_N, BLOCK_N, n_block_idx, m_block_idx));
+                        // Issue TMA B without broadcasting
+                        tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_b[s], k_idx, scheduler.get_global_idx<false>(SHAPE_N, BLOCK_N, n_block_idx, m_block_idx));
                         full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE);
                     }
 
@@ -241,7 +203,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                         empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
                         full_barriers[s]->arrive();
                     }
-                }, 0);
+                });
             }
 
             // To safely deconstruct distributed shared barriers, we need another round of empty waits
@@ -282,9 +244,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Accumulation for WGMMA or CUDA promotion
-            constexpr int WAVE_BLOCK_M = WGMMA::M * get_num_math_warpgroups(BLOCK_M);
-            DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
-            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
+            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
 
             // Empty barrier arrival
             auto empty_barrier_arrive = [&](int s) {
@@ -296,7 +256,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
             };
 
             // Launch MMAs
-            launch_k_iterations([&](int k_iter, auto type, auto num_former_iters_type) {
+            launch_k_iterations([&](int k_iter, auto type) {
                 constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
                 constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
                 DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
@@ -312,54 +272,42 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                     // Wait TMA arrivals
                     full_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter) & 1);
 
-                    // TODO: remove some useless computation for unaligned Ms
+                    // Read A scales
+                    // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
+                    auto scale_a_0 = ld_shared(smem_scales_a[s] + r_0), scale_a_1 = ld_shared(smem_scales_a[s] + r_1);
+
+                    // Commit WGMMA instructions
                     #pragma unroll
-                    for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
-                      	auto m_offset = local_idx * WAVE_BLOCK_M;
+                    for (int i = 0; i < WGMMA::kNumAccum; ++ i)
+                        warpgroup_fence_operand(accum[i]);
+                    warpgroup_arrive();
+                    #pragma unroll
+                    for (int k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
+                        auto desc_a = make_smem_desc(smem_a[s] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
+                        auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
+                        WGMMA::wgmma(desc_a, desc_b, accum, k);
+                    }
+                    warpgroup_commit_batch();
+                    #pragma unroll
+                    for (int i = 0; i < WGMMA::kNumAccum; ++ i)
+                        warpgroup_fence_operand(accum[i]);
+                    warpgroup_wait<0>();
 
-                    	// Read A scales
-                    	// NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
-                    	auto scale_a_0 = ld_shared(smem_scales_a[s] + r_0 + m_offset);
-                        auto scale_a_1 = ld_shared(smem_scales_a[s] + r_1 + m_offset);
+                    // Notify barrier arrival
+                    empty_barrier_arrive(s);
 
-                    	// Commit WGMMA instructions
-                    	#pragma unroll
-                    	for (int i = 0; i < WGMMA::kNumAccum; ++ i)
-                            warpgroup_fence_operand(accum[i]);
-                    	warpgroup_arrive();
-                    	#pragma unroll
-                    	for (int k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
-                            auto desc_a = make_smem_desc(smem_a[s] + (math_wg_idx * WGMMA::M + m_offset) * BLOCK_K + k * WGMMA::K, 1);
-                            auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
-                            WGMMA::wgmma(desc_a, desc_b, accum, k);
-                    	}
-                    	warpgroup_commit_batch();
-                    	#pragma unroll
-                    	for (int i = 0; i < WGMMA::kNumAccum; ++ i)
-                            warpgroup_fence_operand(accum[i]);
-                    	warpgroup_wait<0>();
-
-                    	// Notify barrier arrival at the last warpgroup wave
-                        if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
-                    	    empty_barrier_arrive(s);
-
-                    	// Promote with scales
-                    	// NOTES: making it as predicates is very important for performance, comparing to two loops
-                    	float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
-                    	float scale_0_1, scale_1_1;
-                    	if constexpr (not kMustUseUniformedScaleB)
-                            scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
-
-                        auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
-                    	#pragma unroll
-                    	for (int i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
-                            // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
-                            bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
-                            shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
-                            shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
-                            shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
-                            shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
-                    	}
+                    // Promote with scales
+                    float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
+                    float scale_0_1, scale_1_1;
+                    if constexpr (not kMustUseUniformedScaleB)
+                        scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
+                    #pragma unroll
+                    for (int i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                        bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
+                        final_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
+                        final_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
+                        final_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
+                        final_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
                     }
                 }
 
@@ -369,81 +317,34 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                     full_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter) & 1);
                     empty_barrier_arrive(s);
                 }
-            }, num_former_iters);
+            });
 
-            // TMA checks
-            constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
-            constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
-            constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
-            DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
-            DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
-                            "Unaligned TMA store or too many TMA store instructions");
-            DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
-            DG_STATIC_ASSERT(static_cast<int>(kSwizzleDMode > 0) + static_cast<int>(BLOCK_N_PADDING > 0) <= 1,
-                            "Swizzling and padding are not compatible");
-
-            // Write back to shared memory using STSM and issue TMA stores
+            // Write back to shared memory using STSM
             DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
             #pragma unroll
-            for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
-                auto m_offset = local_idx * WAVE_BLOCK_M;
-                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
-                #pragma unroll
-                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
-                    // Swizzle or padding into the correct address
-                    uint8_t* smem_ptr = nullptr;
-                    if constexpr (kSwizzleDMode > 0) {
-                        // Calculate the swizzling atom offset and in-atom offset
-                        constexpr int kNumBankGroupBytes = 16;
-                        auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
-
-                        // Calculate the index of the bank group to be written in the atom
-                        auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
-
-                        // Reshape the atom in another view and swizzle
-                        //  - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
-                        //  - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
-                        constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
-                        auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
-                        auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
-                        col ^= row % (kSwizzleDMode / 16);
-
-                        // Add back into the base pointer
-                        // NOTES: think twice before modifying this, as changes may affect the number of instructions
-                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d) +                // Base pointer
-                            warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) +            // Warp offset
-                            m_offset * kSwizzleDMode +                                 // Wave offset
-                            atom_offset * BLOCK_M * kSwizzleDMode +                    // Swizzle atom offset (constants)
-                            row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
-                    } else {
-                        // No swizzling, just padding
-                        // NOTES: padding must be zero for BF16 output
-                        DG_STATIC_ASSERT(BLOCK_N_PADDING == 0, "Padding must be zero for BF16 output");
-                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * (BLOCK_N + BLOCK_N_PADDING) + i * 8);
-                    }
-
-                    // NOTES: only 16 lanes' addresses are used
-                    SM90_U32x2_STSM_N<nv_bfloat162>::copy(
-                        __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
-                        __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
-                        smem_ptr
-                    );
-                }
+            for (auto i = 0; i < WGMMA::kNumAccum / 8; ++ i) {
+                SM90_U32x4_STSM_N<nv_bfloat162>::copy(
+                    __float22bfloat162_rn({final_accum[i * 8 + 0], final_accum[i * 8 + 1]}),
+                    __float22bfloat162_rn({final_accum[i * 8 + 2], final_accum[i * 8 + 3]}),
+                    __float22bfloat162_rn({final_accum[i * 8 + 4], final_accum[i * 8 + 5]}),
+                    __float22bfloat162_rn({final_accum[i * 8 + 6], final_accum[i * 8 + 7]}),
+                    smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + i * 16 + 8 * (lane_idx / 16)
+                );
+            }
+            if constexpr (WGMMA::kNumAccum % 8 != 0) {
+                SM90_U32x2_STSM_N<nv_bfloat162>::copy(
+                    __float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 0], final_accum[WGMMA::kNumAccum / 8 * 8 + 1]}),
+                    __float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 2], final_accum[WGMMA::kNumAccum / 8 * 8 + 3]}),
+                    smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + WGMMA::kNumAccum / 8 * 16
+                );
             }
             cute::tma_store_fence();
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Use TMA store to write back to global memory
-            // TODO: compatible with FP32 output
-            DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
-            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
-                auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
-                auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
-                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
-                                              n_block_idx * BLOCK_N + in_block_n_offset,
+            if (threadIdx.x == 0) {
+                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_d, n_block_idx * BLOCK_N,
                                               scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
-
-                // Wait TMA to be finished
                 cute::tma_store_arrive();
                 cute::tma_store_wait<0>();
             }
@@ -458,10 +359,8 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 
 template <uint32_t SHAPE_N, uint32_t SHAPE_K,
           uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
-          uint32_t BLOCK_N_PADDING,
-          uint32_t kSwizzleDMode,
           uint32_t kNumGroups, uint32_t kNumStages,
-          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
+          uint32_t kNumTMAMulticast,
           GemmType kGemmType>
 class Gemm {
 private:
@@ -481,13 +380,9 @@ public:
         // NOTES: we must use 4 warps to do TMA, because `setmaxnreg.aligned` requires 4 warps
         constexpr uint32_t kNumTMAThreads = 128;
         constexpr uint32_t kNumMathThreadsPerGroup = 128;
-        auto kernel = fp8_gemm_kernel<SHAPE_N, SHAPE_K,
-                                      BLOCK_M, BLOCK_N, BLOCK_K,
-                                      BLOCK_N_PADDING,
-                                      kSwizzleDMode,
-                                      kNumGroups, kNumStages,
-                                      kNumTMAThreads, kNumMathThreadsPerGroup,
-                                      kNumTMAMulticast, kIsTMAMulticastOnA, kGemmType>;
+        auto kernel = fp8_gemm_kernel<SHAPE_N, SHAPE_K, BLOCK_M, BLOCK_N, BLOCK_K,
+                                      kNumGroups, kNumStages, kNumTMAThreads, kNumMathThreadsPerGroup,
+                                      kNumTMAMulticast, kGemmType>;
         DG_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess);
 
         // Cluster launch
@@ -527,17 +422,10 @@ public:
 
     template <typename T>
     static CUtensorMap make_2d_tma_d_desc(T* global_address, uint32_t shape_m) {
-        auto swizzle_mode = CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE;
-        if constexpr (kSwizzleDMode == 32)  swizzle_mode = CU_TENSOR_MAP_SWIZZLE_32B;
-        if constexpr (kSwizzleDMode == 64)  swizzle_mode = CU_TENSOR_MAP_SWIZZLE_64B;
-        if constexpr (kSwizzleDMode == 128) swizzle_mode = CU_TENSOR_MAP_SWIZZLE_128B;
-
-        // Swizzling requires the inner box dim less or equal than `kSwizzleDMode` bytes
-        // So `BLOCK_N * sizeof(T) / kSwizzleDMode` TMA stores are required
         return make_2d_tma_desc(global_address, Layout::RowMajor,
                                 shape_m * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), SHAPE_N,
-                                BLOCK_M, kSwizzleDMode == 0 ? BLOCK_N : kSwizzleDMode / sizeof(T),
-                                swizzle_mode);
+                                min(BLOCK_M, shape_m), BLOCK_N,
+                                CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
     }
 
     template <typename T>

ops/csrc/fp8/deep_gemm/include/deep_gemm/mma_utils.cuh

@@ -20,13 +20,742 @@
 
 #include <cuda.h>
 
-#include <cute/arch/mma_sm90_gmma.hpp>
-#include <cute/arch/mma_sm90_gmma_ext.hpp>
-
 #include "utils.cuh"
 
 namespace deep_gemm {
 
+struct SM90_64x16x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %10, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n16k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7},"
+                     " %8,"
+                     " %9,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07)
+                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 16;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x24x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %14, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n24k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
+                     " %8,   %9,   %10,  %11},"
+                     " %12,"
+                     " %13,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11)
+                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 24;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x32x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %18, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15},"
+                     " %16,"
+                     " %17,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15)
+                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 32;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x40x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %22, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n40k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15, "
+                     " %16,  %17,  %18,  %19},"
+                     " %20,"
+                     " %21,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 40;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x48x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %26, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n48k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15, "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23},"
+                     " %24,"
+                     " %25,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 48;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x56x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %30, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n56k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27}, "
+                     " %28,"
+                     " %29,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                       "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                       "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27)
+                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 56;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x64x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %34, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31}, "
+                     " %32,"
+                     " %33,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 64;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x72x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %38, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n72k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35}, "
+                     " %36,"
+                     " %37,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 72;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x80x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %42, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n80k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39}, "
+                     " %40,"
+                     " %41,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 80;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x88x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %46, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n88k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43}, "
+                     " %44,"
+                     " %45,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 88;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x96x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %50, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47}, "
+                     " %48,"
+                     " %49,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 96;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x104x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 float& d48, float& d49, float& d50, float& d51,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %54, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n104k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
+                     " %48,  %49,  %50,  %51}, "
+                     " %52,"
+                     " %53,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
+                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              d[48], d[49], d[50], d[51],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 104;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x112x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %58, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n112k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
+                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55}, "
+                     " %56,"
+                     " %57,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
+                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 112;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x120x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
+                                 float& d56, float& d57, float& d58, float& d59,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %62, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n120k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
+                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
+                     " %56,  %57,  %58,  %59}, "
+                     " %60,"
+                     " %61,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
+                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
+                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
+              d[56], d[57], d[58], d[59],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 120;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x128x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
+                                 float& d56, float& d57, float& d58, float& d59, float& d60, float& d61, float& d62, float& d63,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %66, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
+                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
+                     " %56,  %57,  %58,  %59,  %60,  %61,  %62,  %63}, "
+                     " %64,"
+                     " %65,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
+                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
+                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
+              d[56], d[57], d[58], d[59], d[60], d[61], d[62], d[63],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 128;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
+struct SM90_64x192x32_F32E4M3E4M3_SS {
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
+                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
+                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
+                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
+                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
+                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
+                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
+                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
+                                 float& d56, float& d57, float& d58, float& d59, float& d60, float& d61, float& d62, float& d63,
+                                 float& d64, float& d65, float& d66, float& d67, float& d68, float& d69, float& d70, float& d71,
+                                 float& d72, float& d73, float& d74, float& d75, float& d76, float& d77, float& d78, float& d79,
+                                 float& d80, float& d81, float& d82, float& d83, float& d84, float& d85, float& d86, float& d87,
+                                 float& d88, float& d89, float& d90, float& d91, float& d92, float& d93, float& d94, float& d95,
+                                 bool scale_d) {
+        asm volatile("{\n"
+                     ".reg .pred p;\n"
+                     "setp.ne.b32 p, %98, 0;\n"
+                     "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3"
+                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
+                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
+                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
+                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
+                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
+                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
+                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
+                     " %56,  %57,  %58,  %59,  %60,  %61,  %62,  %63,  "
+                     " %64,  %65,  %66,  %67,  %68,  %69,  %70,  %71,  "
+                     " %72,  %73,  %74,  %75,  %76,  %77,  %78,  %79,  "
+                     " %80,  %81,  %82,  %83,  %84,  %85,  %86,  %87,  "
+                     " %88,  %89,  %90,  %91,  %92,  %93,  %94,  %95},  "
+                     " %96,"
+                     " %97,"
+                     " p   , 1,    1;\n"
+                     "}\n"
+                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
+                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
+                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
+                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
+                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
+                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
+                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
+                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63),
+                      "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71),
+                      "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79),
+                      "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87),
+                      "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95)
+                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
+    }
+
+    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        wgmma(desc_a, desc_b,
+              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
+              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
+              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
+              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
+              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
+              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
+              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
+              d[56], d[57], d[58], d[59], d[60], d[61], d[62], d[63],
+              d[64], d[65], d[66], d[67], d[68], d[69], d[70], d[71],
+              d[72], d[73], d[74], d[75], d[76], d[77], d[78], d[79],
+              d[80], d[81], d[82], d[83], d[84], d[85], d[86], d[87],
+              d[88], d[89], d[90], d[91], d[92], d[93], d[94], d[95],
+              scale_d);
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = 192;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
 template <typename dtype_t>
 struct SM90_U32x2_STSM_N {
     __device__ __forceinline__ static void
@@ -48,15 +777,15 @@ struct SM90_U32x4_STSM_N {
     }
 };
 
-__forceinline__ __device__ void warpgroup_arrive() {
+__device__ void warpgroup_arrive() {
     asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
 }
 
-__forceinline__ __device__ void warpgroup_commit_batch() {
+__device__ void warpgroup_commit_batch() {
     asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
 }
 
-__forceinline__ __device__ void warpgroup_fence_operand(float& reg) {
+__device__ void warpgroup_fence_operand(float& reg) {
     asm volatile("" : "+f"(reg) :: "memory");
 }
 
@@ -147,53 +876,25 @@ __device__ GmmaDescriptor make_smem_desc(PointerType smem_ptr, int layout_type,
     return desc;
 }
 
-template <int N_, typename MMA>
-struct FP8MMA {
-    
-    template <size_t ...Idx>
-    __forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, std::index_sequence<Idx...>) {
-        using namespace cute::SM90::GMMA;
-        MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
-    }
-
-    __forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        call_fma_impl(desc_a, desc_b, d, scale_d, std::make_index_sequence<N_/2>{});
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = N_;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
 template <int N>
 struct FP8MMASelector {
-
-    static constexpr auto select_mma() {
-        using namespace cute::SM90::GMMA;
-        if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_SS_TN();
-        if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_SS_TN();
-    }
-
     static constexpr auto select_type() {
-        return FP8MMA<N, decltype(select_mma())>();
+        if constexpr (N == 16) return SM90_64x16x32_F32E4M3E4M3_SS();
+        if constexpr (N == 24) return SM90_64x24x32_F32E4M3E4M3_SS();
+        if constexpr (N == 32) return SM90_64x32x32_F32E4M3E4M3_SS();
+        if constexpr (N == 40) return SM90_64x40x32_F32E4M3E4M3_SS();
+        if constexpr (N == 48) return SM90_64x48x32_F32E4M3E4M3_SS();
+        if constexpr (N == 56) return SM90_64x56x32_F32E4M3E4M3_SS();
+        if constexpr (N == 64) return SM90_64x64x32_F32E4M3E4M3_SS();
+        if constexpr (N == 72) return SM90_64x72x32_F32E4M3E4M3_SS();
+        if constexpr (N == 80) return SM90_64x80x32_F32E4M3E4M3_SS();
+        if constexpr (N == 88) return SM90_64x88x32_F32E4M3E4M3_SS();
+        if constexpr (N == 96) return SM90_64x96x32_F32E4M3E4M3_SS();
+        if constexpr (N == 104) return SM90_64x104x32_F32E4M3E4M3_SS();
+        if constexpr (N == 112) return SM90_64x112x32_F32E4M3E4M3_SS();
+        if constexpr (N == 120) return SM90_64x120x32_F32E4M3E4M3_SS();
+        if constexpr (N == 128) return SM90_64x128x32_F32E4M3E4M3_SS();
+        if constexpr (N == 192) return SM90_64x192x32_F32E4M3E4M3_SS();
     }
 
     using type = decltype(select_type());

ops/csrc/fp8/deep_gemm/include/deep_gemm/scheduler.cuh

@@ -30,10 +30,9 @@ enum class GemmType {
 #pragma ide diagnostic ignored "cppcoreguidelines-pro-type-member-init"
 template <GemmType kGemmType,
           uint32_t SHAPE_N, uint32_t BLOCK_M, uint32_t BLOCK_N,
-          uint32_t kNumGroups,
-          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
+          uint32_t kNumGroups, uint32_t kNumTMAMulticast,
           uint32_t kNumNBlocks = ceil_div(SHAPE_N, BLOCK_N),
-          uint32_t kNum1DBlocksPerGroup = 16>
+          uint32_t kNumNBlocksPerGroup = 16>
 struct Scheduler {
     int current_iter = -1;
     uint32_t num_aligned_m_blocks;
@@ -62,27 +61,16 @@ struct Scheduler {
     }
 
     __device__ __forceinline__ void get_swizzled_block_idx(const uint32_t num_m_blocks, int block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
-        DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
+        DG_STATIC_ASSERT(kNumNBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
 
         // Swizzle for better L2 usages
-        // TODO: unify these 2 branches
-        if constexpr (kIsTMAMulticastOnA) {
-            auto num_blocks_per_group = num_m_blocks * kNum1DBlocksPerGroup;
-            auto group_idx = block_idx / num_blocks_per_group;
-            auto first_n_block_idx = group_idx * kNum1DBlocksPerGroup;
-            auto num_n_blocks_in_group = min(kNum1DBlocksPerGroup, kNumNBlocks - first_n_block_idx);
-            auto in_group_idx = block_idx % num_blocks_per_group;
-            m_block_idx = in_group_idx / num_n_blocks_in_group;
-            n_block_idx = first_n_block_idx + in_group_idx % num_n_blocks_in_group;
-        } else {
-            auto num_blocks_per_group = kNumNBlocks * kNum1DBlocksPerGroup;
-            auto group_idx = block_idx / num_blocks_per_group;
-            auto first_m_block_idx = group_idx * kNum1DBlocksPerGroup;
-            auto num_m_blocks_in_group = min(kNum1DBlocksPerGroup, num_m_blocks - first_m_block_idx);
-            auto in_group_idx = block_idx % num_blocks_per_group;
-            m_block_idx = first_m_block_idx + in_group_idx % num_m_blocks_in_group;
-            n_block_idx = in_group_idx / num_m_blocks_in_group;
-        }
+        auto num_blocks_per_group = num_m_blocks * kNumNBlocksPerGroup;
+        auto group_idx = block_idx / num_blocks_per_group;
+        auto first_n_block_idx = group_idx * kNumNBlocksPerGroup;
+        auto num_n_blocks_in_group = min(kNumNBlocksPerGroup, kNumNBlocks - first_n_block_idx);
+        auto in_group_idx = block_idx % num_blocks_per_group;
+        m_block_idx = in_group_idx / num_n_blocks_in_group;
+        n_block_idx = first_n_block_idx + in_group_idx % num_n_blocks_in_group;
     }
 
     template <bool kIgnoreGroupedForGroupedContiguous=true>
@@ -128,7 +116,6 @@ struct Scheduler {
         return true;
     }
 };
-
 #pragma clang diagnostic pop
 
 } // namespace deep_gemm

ops/csrc/fp8/deep_gemm/include/deep_gemm/tma_utils.cuh

@@ -58,18 +58,20 @@ constexpr CUtensorMapDataType get_CUtensorMapDataType() {
     }
 }
 
-inline PFN_cuTensorMapEncodeTiled get_cuTensorMapEncodeTiled() {
+PFN_cuTensorMapEncodeTiled get_cuTensorMapEncodeTiled() {
     // Get pointer to `cuTensorMapEncodeTiled`
     cudaDriverEntryPointQueryResult driver_status;
     void* cuTensorMapEncodeTiled_ptr = nullptr;
 
+/*
 #if CUDA_VERSION >= 12050
     cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr, 12000,
                                      cudaEnableDefault, &driver_status);
 #else
+*/
     cudaGetDriverEntryPoint("cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr,
                             cudaEnableDefault, &driver_status);
-#endif
+//#endif
 
     if (driver_status != cudaDriverEntryPointSuccess)
         throw std::runtime_error("driver_status != cudaDriverEntryPointSuccess");
@@ -81,15 +83,16 @@ CUtensorMap make_2d_tma_copy_desc(T* global_address, uint64_t gmem_dim[2],
                                   uint64_t stride_in_bytes, uint32_t smem_dim[2],
                                   CUtensorMapSwizzle swizzle_type,
                                   PFN_cuTensorMapEncodeTiled encode_func = nullptr) {
-    CUtensorMap tensor_map = {};
-    uint64_t global_stride[1] = {stride_in_bytes};
-    uint32_t elem_strides[2] = {1, 1};
+    CUtensorMap tensor_map{};
+    constexpr uint32_t rank = 2;
+    uint64_t global_stride[rank - 1] = {stride_in_bytes};
+    uint32_t elem_strides[rank] = {1, 1};
 
     if (encode_func == nullptr)
         encode_func = get_cuTensorMapEncodeTiled();
 
     auto result = encode_func(
-            &tensor_map, get_CUtensorMapDataType<std::remove_cv_t<T>>(), 2,
+            &tensor_map, get_CUtensorMapDataType<typename std::remove_cv<T>::type>(), rank,
             global_address, gmem_dim, global_stride, smem_dim, elem_strides,
             CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle_type,
             CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,

ops/csrc/fp8/deep_gemm/include/deep_gemm/utils.cuh

@@ -63,9 +63,4 @@ do {
 template <typename T>
 __device__ __host__ constexpr T ceil_div(T a, T b) {
     return (a + b - 1) / b;
-}
-
-template <typename T>
-__device__ __host__ constexpr T constexpr_gcd(T a, T b) {
-    return b == 0 ? a : constexpr_gcd(b, a % b);
 }

ops/csrc/fp8/deep_gemm/jit/__init__.py

@@ -16,6 +16,6 @@
 # Copyright (c) 2025 DeepSeek
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
-from .compiler import get_nvcc_compiler, build
-from .template import cpp_format, generate
+from .compiler import build, get_nvcc_compiler
 from .runtime import Runtime
+from .template import cpp_format, generate

ops/csrc/fp8/deep_gemm/jit/compiler.py

@@ -16,8 +16,8 @@
 # Copyright (c) 2025 DeepSeek
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
-import hashlib
 import functools
+import hashlib
 import os
 import re
 import subprocess
@@ -75,7 +75,9 @@ def get_nvcc_compiler() -> Tuple[str, str]:
             match = version_pattern.search(os.popen(f"{path} --version").read())
             version = match.group(1)
             assert match, f"Cannot get the version of NVCC compiler {path}"
-            assert version >= least_version_required, f"NVCC {path} version {version} is lower than {least_version_required}"
+            assert (
+                version >= least_version_required
+            ), f"NVCC {path} version {version} is lower than {least_version_required}"
             return path, version
     raise RuntimeError("Cannot find any available NVCC compiler")
 
@@ -115,13 +117,18 @@ def put(path, data, is_binary=False):
 
 def build(name: str, arg_defs: tuple, code: str) -> Runtime:
     # Compiler flags
-    cpp_standard = int(os.getenv("DG_NVCC_OVERRIDE_CPP_STANDARD", 20))
-    nvcc_flags = [f"-std=c++{cpp_standard}", "-shared", "-O3", "--expt-relaxed-constexpr", "--expt-extended-lambda",
-                  "-gencode=arch=compute_90a,code=sm_90a",
-                  "--ptxas-options=--register-usage-level=10" + (",--verbose" if "DG_PTXAS_VERBOSE" in os.environ else ""),
-                  # Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases
-                  "--diag-suppress=39,174,177,940"]
-    cxx_flags = ["-fPIC", "-O3", "-Wno-deprecated-declarations", "-Wno-abi", "-fconcepts"]
+    nvcc_flags = [
+        "-std=c++17",
+        "-shared",
+        "-O3",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+        "-gencode=arch=compute_90a,code=sm_90a",
+        "--ptxas-options=--register-usage-level=10" + (",--verbose" if "DG_PTXAS_VERBOSE" in os.environ else ""),
+        # Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases
+        "--diag-suppress=177,174,940",
+    ]
+    cxx_flags = ["-fPIC", "-O3", "-Wno-deprecated-declarations", "-Wno-abi"]
     flags = [*nvcc_flags, f'--compiler-options={",".join(cxx_flags)}']
     include_dirs = [get_jit_include_dir()]
 
@@ -148,12 +155,8 @@ def build(name: str, arg_defs: tuple, code: str) -> Runtime:
     # Compile into a temporary SO file
     so_path = f"{path}/kernel.so"
     tmp_so_path = f"{make_tmp_dir()}/nvcc.tmp.{str(uuid.uuid4())}.{hash_to_hex(so_path)}.so"
-
     # Compile
-    command = [get_nvcc_compiler()[0],
-               src_path, "-o", tmp_so_path,
-               *flags,
-               *[f"-I{d}" for d in include_dirs]]
+    command = [get_nvcc_compiler()[0], src_path, "-o", tmp_so_path, *flags, *[f"-I{d}" for d in include_dirs]]
     if os.getenv("DG_JIT_DEBUG", None) or os.getenv("DG_JIT_PRINT_NVCC_COMMAND", False):
         print(f"Compiling JIT runtime {name} with command {command}")
     return_code = subprocess.check_call(command)

ops/csrc/fp8/deep_gemm/jit/interleave_ffma.py

@@ -93,7 +93,7 @@ def parse_registers(line):
 
 
 def modify_segment(m, name, ffma_lines):
-    num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2
+    num_lines = len(ffma_lines)
     assert num_lines % 2 == 0
 
     le_bytes, new_le_bytes = [], []

ops/csrc/fp8/deep_gemm/jit/runtime.py

@@ -23,6 +23,8 @@ from typing import Optional
 import paddle
 from paddle import Tensor
 
+from .template import map_ctype
+
 
 class Runtime:
     def __init__(self, path: str) -> None:
@@ -44,52 +46,24 @@ class Runtime:
 
     def __call__(self, *args) -> int:
         # Load SO file
-        if self.lib is None:
+        if self.lib is None or self.args is None:
             self.lib = ctypes.CDLL(os.path.join(self.path, "kernel.so"))
+            with open(os.path.join(self.path, "kernel.args"), "r") as f:
+                self.args = eval(f.read())
+
+        # Check args and launch
+        assert len(args) == len(self.args), f"Expected {len(self.args)} arguments, got {len(args)}"
+        cargs = []
+        for arg, (name, dtype) in zip(args, self.args):
+            if isinstance(arg, Tensor):
+                assert arg.dtype == dtype, f"Expected tensor dtype `{dtype}` for `{name}`, got `{arg.dtype}`"
+            else:
+                assert isinstance(arg, dtype), f"Expected built-in type `{dtype}` for `{name}`, got `{type(arg)}`"
+            cargs.append(map_ctype(arg))
 
-        if len(args) == 9:
-            cargs = [
-                ctypes.c_void_p(args[0].data_ptr()),
-                ctypes.c_void_p(args[1].data_ptr()),
-                ctypes.c_void_p(args[2].data_ptr()),
-                ctypes.c_void_p(args[3].data_ptr()),
-                ctypes.c_void_p(args[4].data_ptr()),
-                ctypes.c_int(args[5]),
-                ctypes.c_void_p(args[6].cuda_stream),
-                ctypes.c_int(args[7]),
-                ctypes.c_int(args[8]),
-            ]
-        elif len(args) == 10:
-            cargs = [
-                ctypes.c_void_p(args[0].data_ptr()),
-                ctypes.c_void_p(args[1].data_ptr()),
-                ctypes.c_void_p(args[2].data_ptr()),
-                ctypes.c_void_p(args[3].data_ptr()),
-                ctypes.c_void_p(args[4].data_ptr()),
-                ctypes.c_void_p(args[5].data_ptr()),
-                ctypes.c_int(args[6]),
-                ctypes.c_void_p(args[7].cuda_stream),
-                ctypes.c_int(args[8]),
-                ctypes.c_int(args[9]),
-            ]
-        elif len(args) == 11:
-            cargs = [
-                ctypes.c_void_p(args[0].data_ptr()),
-                ctypes.c_void_p(args[1].data_ptr()),
-                ctypes.c_void_p(args[2].data_ptr()),
-                ctypes.c_void_p(args[3].data_ptr()),
-                ctypes.c_void_p(args[4].data_ptr()),
-                ctypes.c_void_p(args[5].data_ptr()),
-                ctypes.c_int(args[6]),
-                ctypes.c_int(args[7]),
-                ctypes.c_void_p(args[8].cuda_stream),
-                ctypes.c_int(args[9]),
-                ctypes.c_int(args[10]),
-            ]
-        else:
-            raise ValueError("Invalid number of arguments")
         return_code = ctypes.c_int(0)
         self.lib.launch(*cargs, ctypes.byref(return_code))
+        return return_code.value
 
 
 class RuntimeCache:

ops/csrc/fp8/deep_gemm/jit_kernels/gemm.py

@@ -16,16 +16,22 @@
 # Copyright (c) 2025 DeepSeek
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
-import math
-from functools import lru_cache
+import functools
 from typing import Tuple
+
 import paddle
 from paddle import Tensor
+
 from .tuner import jit_tuner
-from .utils import get_num_sms, ceil_div, get_col_major_tma_aligned_tensor, get_m_alignment_for_contiguous_layout
+from .utils import (
+    ceil_div,
+    get_col_major_tma_aligned_tensor,
+    get_m_alignment_for_contiguous_layout,
+    get_num_sms,
+)
 
 # C++ code templates
-includes = ('"deep_gemm/fp8_gemm.cuh"', )
+includes = ('"deep_gemm/fp8_gemm.cuh"',)
 template = """
 using namespace deep_gemm;
 
@@ -33,58 +39,32 @@ using namespace deep_gemm;
 constexpr auto N = {N}, K = {K};
 constexpr auto BLOCK_M = {BLOCK_M};
 constexpr auto BLOCK_N = {BLOCK_N};
-constexpr auto BLOCK_K = 128;
-constexpr auto BLOCK_N_PADDING = {BLOCK_N_PADDING};
-constexpr auto kSwizzleDMode = {SWIZZLE_D_MODE};
-constexpr auto kNumGroups = 1;
 constexpr auto kNumStages = {NUM_STAGES};
 constexpr auto kNumTMAMulticast = {NUM_TMA_MULTICAST};
-constexpr auto kIsTMAMulticastOnA = {IS_TMA_MULTICAST_ON_A};
 
 // Make a templated GEMM
-using gemm_t = Gemm<N, K, BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_N_PADDING, kSwizzleDMode, kNumGroups, kNumStages, kNumTMAMulticast, kIsTMAMulticastOnA, GemmType::Normal>;
+using GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, 1, kNumStages, kNumTMAMulticast, GemmType::Normal>;
 
 // Launch kernel
-auto tma_a_desc = gemm_t::make_2d_tma_a_desc(lhs, m);
-auto tma_b_desc = gemm_t::make_2d_tma_b_desc(rhs);
-auto tma_scales_a_desc = gemm_t::make_2d_tma_scales_a_desc(lhs_scales, m);
-auto tma_d_desc = gemm_t::make_2d_tma_d_desc(out, m);
-gemm_t::run(out, rhs_scales, nullptr,
-            m,
-            tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
-            stream, num_sms, smem_size);
+auto tma_a_desc = GemmType::make_2d_tma_a_desc(lhs, m);
+auto tma_b_desc = GemmType::make_2d_tma_b_desc(rhs);
+auto tma_scales_a_desc = GemmType::make_2d_tma_scales_a_desc(lhs_scales, m);
+auto tma_d_desc = GemmType::make_2d_tma_d_desc(out, m);
+GemmType::run(out, rhs_scales, nullptr,
+              m,
+              tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
+              stream, num_sms, smem_size);
 """
 
 
-def is_tma_multicast_legal(shape_dim: int, block_dim: int, num_tma_multicast: int, num_sms: int) -> bool:
+def is_tma_multicast_legal(n: int, block_n: int, num_tma_multicast: int, num_sms: int) -> bool:
     if num_tma_multicast == 1:
         return True
-    return (shape_dim % (block_dim * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
+    return (n % (block_n * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
 
 
-def get_swizzle_mode(block_n: int) -> int:
-    # TODO: remove some candidates if slow
-    elem_size = 2
-    for mode_bytes in (128, 64, 32):
-        if (block_n * elem_size) % mode_bytes == 0:
-            return mode_bytes
-    return 0
-
-
-def get_block_n_padding_for_smem_d(block_n: int) -> int:
-    # NOTES: padding is for solving bank conflicts, but wastes shared memory space
-    elem_size, requirement = 2, (4, 8)
-    bank_stride = (block_n * elem_size) // 4
-    padding = (requirement[0] - bank_stride) % requirement[1]
-    return (((padding + requirement[1]) if padding < 0 else padding) * 4) // elem_size
-
-
-def get_smem_config(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128) -> Tuple[int, int, int]:
-    # Try swizzle first, as it does not waste shared memory
-    swizzle_mode = get_swizzle_mode(block_n)
-    block_n_padding = get_block_n_padding_for_smem_d(block_n) if swizzle_mode == 0 else 0
-
-    smem_d = block_m * (block_n + block_n_padding) * 2
+def get_smem_size(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128) -> int:
+    smem_d = block_m * block_n * 2
     smem_a_per_stage = block_m * block_k
     smem_scales_a_per_stage = block_m * 4
     smem_b_per_stage = block_n * block_k
@@ -98,22 +78,18 @@ def get_smem_config(num_stages: int, k: int, block_m: int, block_n: int, block_k
     smem_size += num_stages * smem_b_per_stage
     smem_size += ceil_div(smem_scales_b * (1 if block_k % block_n == 0 else 2), 8) * 8
     smem_size += smem_barrier
-
-    # Swizzle and padding are not compatible
-    assert int(swizzle_mode > 0) + int(block_n_padding > 0) <= 1
-
-    return smem_size, swizzle_mode, block_n_padding
+    return smem_size
 
 
-@lru_cache(maxsize=None)
-def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
-                     is_grouped_contiguous: bool = False, is_grouped_masked: bool = False) -> \
-        Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]]:
+def get_best_configs(
+    m: int, n: int, k: int, num_groups: int, num_sms: int, is_grouped_contiguous: bool = False
+) -> Tuple[int, int, int, int, int]:
     if not is_grouped_contiguous:
-        block_ms = (64, 128, 256)
+        # TODO: for some cases, smaller M block is better, add them into tuning space
+        block_ms = (64 if m <= 64 else 128,)
     else:
-        block_ms = (get_m_alignment_for_contiguous_layout(), )
-    block_ns = tuple(range(16, 129, 8)) + (144, 160, )
+        block_ms = (get_m_alignment_for_contiguous_layout(),)
+    block_ns = tuple(range(16, 129, 8))
 
     fix_wave_saturate = lambda x: num_sms if x == 0 else x
     get_num_waves = lambda bm, bn: (ceil_div(ceil_div(m, bm) * ceil_div(n, bn) * num_groups, num_sms) if bm else None)
@@ -122,9 +98,7 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
     # Decide block sizes by waves
     best_block_m, best_block_n = None, None
     for block_m in block_ms:
-        # NOTES: the block sizes can not be too large, so at least one dim less than 128
-        for block_n in filter(lambda bn: block_m <= 128 or bn <= 128, block_ns):
-
+        for block_n in block_ns:
             success = False
             num_waves, best_num_waves = get_num_waves(block_m, block_n), get_num_waves(best_block_m, best_block_n)
             if best_block_m is None or best_block_n is None:
@@ -135,76 +109,48 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
                 # Check last wave utilization
                 util = get_last_wave_util(block_m, block_n)
                 best_util = get_last_wave_util(best_block_m, best_block_n)
-                success = util > best_util
-                if util == best_util:
-                    # Case 1: same `block_m`, smaller `block_n` (wasted)
-                    success |= block_m == best_block_m and block_n < best_block_n
-                    # Case 2: same `block_n`, smaller `block_m` (wasted)
-                    success |= block_n == best_block_n and block_m < best_block_m
-                    # Case 3: different for both `block_m` and `block_n`, `block_n` larger is better
-                    success |= block_m != best_block_m and block_n > best_block_n
+                success = util > best_util or (
+                    util == best_util
+                    and (block_m > best_block_m or (block_m == best_block_m and block_n < best_block_n))
+                )
             best_block_m, best_block_n = (block_m, block_n) if success else (best_block_m, best_block_n)
     assert best_block_m is not None and best_block_n is not None
 
     # Always pick the longest one
     # NOTES: for double B scales, the best number of stages may be reduced
-    best_num_stages, best_smem_config, sm90_capacity = None, None, 232448
-    stage_candidates = (8, 7, 6, 5, 4, 3)
-    if 128 % best_block_n != 0 and 128 // math.gcd(128, best_block_n) <= 4:
-        # Unrolling both stages and `num_former_iters` will cause large code size
-        stage_candidates = (4, 3)
-    for num_stages in stage_candidates:
-        best_smem_config = get_smem_config(num_stages, k, best_block_m, best_block_n)
-        if best_smem_config[0] <= sm90_capacity:
+    best_num_stages, best_smem_size, sm90_capacity = None, None, 232448
+    for num_stages in (6, 5, 4) if 128 % best_block_n != 0 else (8, 7, 6, 5, 4):
+        best_smem_size = get_smem_size(num_stages, k, best_block_m, best_block_n)
+        if best_smem_size <= sm90_capacity:
             best_num_stages = num_stages
             break
-    assert best_smem_config is not None
     assert best_num_stages is not None
 
-    # Decide the number of TMA multicast and whether broadcast on A
-    best_tma_multicast_config = (1, 1)
+    # Decide the number of TMA multicast
+    best_num_tma_multicast = 1
+    if m >= 1024 and is_tma_multicast_legal(n, best_block_n, 2, num_sms) and num_groups == 1:
+        best_num_tma_multicast = 2
 
-    # Try to multicast on the larger block side first
-    is_dense_gemm = (not is_grouped_contiguous) and (not is_grouped_masked)
-    is_multicast_legal = {
-        'A': is_tma_multicast_legal(n, best_block_n, 2, num_sms),
-        'B': is_tma_multicast_legal(m, best_block_m, 2, num_sms) and is_dense_gemm,
-    }
-    for i in ('A', 'B') if best_block_m > best_block_n else ('B', 'A'):
-        if m >= 512 and is_multicast_legal[i]:
-            best_tma_multicast_config = (2, int(i == 'A'))
-            break
+    return best_block_m, best_block_n, best_num_stages, best_num_tma_multicast, best_smem_size
 
-    # Recompute the minimal number of SMs required
-    # NOTES: less L2 cache usage and less GPU frequency drop
-    num_waves = get_num_waves(best_block_m, best_block_n)
-    num_min_sms = ceil_div(ceil_div(m, best_block_m) * ceil_div(n, best_block_n) * num_groups, num_waves)
-    num_min_sms = ceil_div(num_min_sms, best_tma_multicast_config[0]) * best_tma_multicast_config[0]
-    assert num_min_sms <= num_sms
 
-    return num_min_sms, best_block_m, best_block_n, best_num_stages, best_tma_multicast_config, best_smem_config
-
-@lru_cache()
-def auto_tuning_with_compilation(m, n, k, num_sms):
+@functools.lru_cache()
+def auto_tuning_with_compilation(m, n, k):
     global includes, template
-    if num_sms is None:
-        num_sms = get_num_sms()
-    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(m, n, k, 1, num_sms)
+    num_sms = get_num_sms()
+    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(m, n, k, 1, num_sms)
     runtime = jit_tuner.compile_and_tune(
         m,
         n,
         k,
         name="gemm_fp8_fp8_bf16_nt",
         keys={
-            'SWIZZLE_D_MODE': smem_config[1],
-            'BLOCK_N_PADDING': smem_config[2],
             "BLOCK_M": block_m,
             "BLOCK_N": block_n,
             "K": k,
             "N": n,
             "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": tma_multicast_config[0],
-            'IS_TMA_MULTICAST_ON_A': tma_multicast_config[1],
+            "NUM_TMA_MULTICAST": num_tma_multicast,
         },
         space=(),
         includes=includes,
@@ -221,10 +167,10 @@ def auto_tuning_with_compilation(m, n, k, num_sms):
         ),
         template=template,
     )
-    return runtime, num_sms, smem_config
+    return runtime, num_sms, smem_size
 
 
-def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, num_sms=132) -> None:
+def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor) -> None:
     """
     Do a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
     LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
@@ -264,7 +210,7 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor],
     # Do nothing if `m` is zero
     if m == 0:
         return
-    runtime, num_sms, smem_config = auto_tuning_with_compilation(m, n, k, num_sms)
-    args = (lhs, lhs_scales, rhs, rhs_scales, out, m, paddle.device.current_stream().stream_base, num_sms, smem_config[0])
+    runtime, num_sms, smem_size = auto_tuning_with_compilation(m, n, k)
+    args = (lhs, lhs_scales, rhs, rhs_scales, out, m, paddle.device.cuda.current_stream(), num_sms, smem_size)
     # Run the kernel.
     runtime(*args)

ops/csrc/fp8/deep_gemm/jit_kernels/m_grouped_gemm.py

@@ -17,17 +17,17 @@
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
 import functools
-from functools import reduce
 from typing import Tuple
 
 import paddle
 from paddle import Tensor
-from .gemm import get_best_configs, get_block_n_padding_for_smem_d
+
+from .gemm import get_best_configs
 from .tuner import jit_tuner
 from .utils import get_col_major_tma_aligned_tensor, get_num_sms
 
 # C++ code templates
-includes = ('"deep_gemm/fp8_gemm.cuh"', )
+includes = ('"deep_gemm/fp8_gemm.cuh"',)
 template = """
 using namespace deep_gemm;
 
@@ -35,35 +35,30 @@ using namespace deep_gemm;
 constexpr auto N = {N}, K = {K};
 constexpr auto BLOCK_M = {BLOCK_M};
 constexpr auto BLOCK_N = {BLOCK_N};
-constexpr auto BLOCK_K = 128;
-constexpr auto BLOCK_N_PADDING = {BLOCK_N_PADDING};
-constexpr auto kSwizzleDMode = {SWIZZLE_D_MODE};
-constexpr auto kNumGroups = {NUM_GROUPS};
 constexpr auto kNumStages = {NUM_STAGES};
 constexpr auto kNumTMAMulticast = {NUM_TMA_MULTICAST};
-constexpr auto kIsTMAMulticastOnA = {IS_TMA_MULTICAST_ON_A};
 
 // Make a templated grouped GEMM
-using gemm_t = Gemm<N, K, BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_N_PADDING, kSwizzleDMode, kNumGroups, kNumStages, kNumTMAMulticast, kIsTMAMulticastOnA, GemmType::{GEMM_TYPE}>;
+using GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, {NUM_GROUPS}, kNumStages, kNumTMAMulticast, GemmType::{GEMM_TYPE}>;
 
 // Launch kernel
-auto tma_a_desc = gemm_t::make_2d_tma_a_desc(lhs, m);
-auto tma_b_desc = gemm_t::make_2d_tma_b_desc(rhs);
-auto tma_scales_a_desc = gemm_t::make_2d_tma_scales_a_desc(lhs_scales, m);
-auto tma_d_desc = gemm_t::make_2d_tma_d_desc(out, m);
-gemm_t::run(out, rhs_scales, grouped_layout,
-            m,
-            tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
-            stream, num_sms, smem_size);
+auto tma_a_desc = GemmType::make_2d_tma_a_desc(lhs, m);
+auto tma_b_desc = GemmType::make_2d_tma_b_desc(rhs);
+auto tma_scales_a_desc = GemmType::make_2d_tma_scales_a_desc(lhs_scales, m);
+auto tma_d_desc = GemmType::make_2d_tma_d_desc(out, m);
+GemmType::run(out, rhs_scales, grouped_layout,
+              m,
+              tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
+              stream, num_sms, smem_size);
 """
 
 
 @functools.lru_cache()
 def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms):
     global includes, template
-    if num_sms is None:
-        num_sms = get_num_sms()
-    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(m, n, k, 1, num_sms, is_grouped_contiguous=True)
+    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(
+        m, n, k, 1, num_sms, is_grouped_contiguous=True
+    )
     runtime = jit_tuner.compile_and_tune(
         m,
         n,
@@ -72,15 +67,12 @@ def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, nu
         keys={
             "BLOCK_M": block_m,
             "BLOCK_N": block_n,
-            "SWIZZLE_D_MODE": smem_config[1],
-            "BLOCK_N_PADDING": smem_config[2],
             "GEMM_TYPE": "GroupedContiguous",
             "K": k,
             "N": n,
             "NUM_GROUPS": num_groups,
             "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": tma_multicast_config[0],
-            "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+            "NUM_TMA_MULTICAST": num_tma_multicast,
         },
         space=(),
         includes=includes,
@@ -99,11 +91,11 @@ def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, nu
         ),
         template=template,
     )
-    return runtime, num_sms, smem_config
+    return runtime, num_sms, smem_size
 
 
 def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
-    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, m_indices: Tensor, num_sms=132
+    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, m_indices: Tensor
 ) -> None:
     """
     Do a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
@@ -131,8 +123,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
     m, k = lhs.shape
     num_groups, n, k_ = rhs.shape
     m_, n_ = out.shape
-    m_shape = m_indices.shape
-    m__ = reduce(lambda x, y: x * y, m_shape)
+    m__ = m_indices.numel()
     # Type and shape checks
     assert m == m_ == m__ and k == k_ and n == n_
     assert lhs_scales.shape == [m, (k + 127) // 128]
@@ -153,7 +144,8 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
         return
     # Auto-tuning with compilation
     global includes, template
-    runtime, num_sms, smem_config = auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms)
+    num_sms = get_num_sms()
+    runtime, num_sms, smem_size = auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms)
 
     args = (
         lhs,
@@ -164,68 +156,15 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
         m_indices,
         m,
         num_groups,
-        paddle.device.current_stream().stream_base,
+        paddle.device.cuda.current_stream(),
         num_sms,
-        smem_config[0],
+        smem_size,
     )
     runtime(*args)
 
 
-@functools.lru_cache()
-def auto_tuning_with_compilation_grouped_gemm_masked(m, expected_m, n, k, num_groups, num_sms):
-    # Auto-tuning with compilation
-    global includes, template
-    if num_sms is None:
-        num_sms = get_num_sms()
-    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(
-        expected_m, n, k, num_groups, num_sms, is_grouped_masked=True
-    )
-
-    # Extra checks for TMA store
-    if num_groups > 1 and m > block_m:
-        while m % block_m != 0 and block_m > 128:
-            block_m = block_m // 2
-        assert (
-            m % block_m == 0
-        ), f"For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})"
-
-    runtime = jit_tuner.compile_and_tune_group_gemm_masked(
-        name="m_grouped_gemm_fp8_fp8_bf16_nt",
-        keys={
-            "N": n,
-            "K": k,
-            "BLOCK_M": block_m,
-            "BLOCK_N": block_n,
-            'SWIZZLE_D_MODE': smem_config[1],
-            'BLOCK_N_PADDING': smem_config[2],
-            "NUM_GROUPS": num_groups,
-            "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": tma_multicast_config[0],
-            'IS_TMA_MULTICAST_ON_A': tma_multicast_config[1],
-            "GEMM_TYPE": "GroupedMasked",
-        },
-        space=(),
-        includes=includes,
-        arg_defs=(
-            ("lhs", paddle.float8_e4m3fn),
-            ("lhs_scales", paddle.float32),
-            ("rhs", paddle.float8_e4m3fn),
-            ("rhs_scales", paddle.float32),
-            ("out", paddle.bfloat16),
-            ("grouped_layout", paddle.int32),
-            ("m", int),
-            ("stream", paddle.device.cuda.Stream),
-            ("num_sms", int),
-            ("smem_size", int),
-        ),
-        template=template,
-    )
-
-    return runtime, num_sms, smem_config
-
-
 def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
-    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, masked_m: Tensor, expected_m: int, num_sms=132
+    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, masked_m: Tensor, expected_m: int
 ) -> None:
     """
     Do a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
@@ -252,8 +191,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
     num_groups, m, k = lhs.shape
     num_groups_, n, k_ = rhs.shape
     num_groups__, m_, n_ = out.shape
-    masked_m_shape = masked_m.shape
-    num_groups___ = reduce(lambda x, y: x * y, masked_m_shape)
+    num_groups___ = masked_m.numel()
 
     # Type and shape checks
     assert num_groups == num_groups_ == num_groups__ == num_groups___
@@ -271,10 +209,20 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
     # LHS scales must be transposed for TMA load, but not for RHS scales
     lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
     assert rhs_scales.is_contiguous()
-    runtime, num_sms, smem_config = auto_tuning_with_compilation_grouped_gemm_masked(
-        m, expected_m, n, k, num_groups, num_sms
+
+    # Auto-tuning with compilation
+    global includes, template
+    num_sms = get_num_sms()
+    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(
+        expected_m, n, k, num_groups, num_sms
     )
 
+    # Extra checks for TMA store
+    if num_groups > 1 and m > block_m:
+        assert (
+            m % block_m == 0
+        ), f"For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})"
+
     args = (
         lhs,
         lhs_scales,
@@ -283,9 +231,38 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
         out,
         masked_m,
         m,
-        paddle.device.current_stream().stream_base,
+        paddle.device.cuda.current_stream(),
         num_sms,
-        smem_config[0],
+        smem_size,
+    )
+    runtime = jit_tuner.compile_and_tune_group_gemm_masked(
+        name="m_grouped_gemm_fp8_fp8_bf16_nt",
+        keys={
+            "N": n,
+            "K": k,
+            "BLOCK_M": block_m,
+            "BLOCK_N": block_n,
+            "NUM_GROUPS": num_groups,
+            "NUM_STAGES": num_stages,
+            "NUM_TMA_MULTICAST": num_tma_multicast,
+            "GEMM_TYPE": "GroupedMasked",
+        },
+        space=(),
+        includes=includes,
+        arg_defs=(
+            ("lhs", paddle.float8_e4m3fn),
+            ("lhs_scales", paddle.float32),
+            ("rhs", paddle.float8_e4m3fn),
+            ("rhs_scales", paddle.float32),
+            ("out", paddle.bfloat16),
+            ("grouped_layout", paddle.int32),
+            ("m", int),
+            ("stream", paddle.device.cuda.Stream),
+            ("num_sms", int),
+            ("smem_size", int),
+        ),
+        template=template,
+        args=args,
     )
 
     # Run the kernel

ops/csrc/fp8/deep_gemm/jit_kernels/tuner.py

@@ -20,6 +20,8 @@ import copy
 import os
 from typing import Any, Dict
 
+import paddle
+
 from ..jit import Runtime, build, cpp_format, generate
 
 
@@ -35,6 +37,7 @@ class JITTuner:
         includes: tuple,
         arg_defs: tuple,
         template: str,
+        args: tuple,
     ) -> Runtime:
         # NOTES: we always assume the space and template will not change
         # We also assume the GPU device will not be changed
@@ -50,6 +53,7 @@ class JITTuner:
             print(f"Auto-tuning JIT kernel {name} with keys {keys}")
 
         assert signature not in self.tuned
+        assert args is not None
         space = (dict(),) if len(space) == 0 else space
 
         kernels = []
@@ -64,7 +68,32 @@ class JITTuner:
 
         best_runtime, best_time, best_keys = None, None, None
         for runtime, tuned_keys in kernels:
-            elapsed_time = 0
+            if len(space) > 1:
+                # Check kernel validity
+                return_code = runtime(*args)
+                if return_code != 0:
+                    # Pass illegal kernels, e.g. insufficient shared memory capacity
+                    if os.getenv("DG_JIT_DEBUG", None):
+                        print(
+                            f"Illegal JIT kernel {name} with keys {keys} and tuned keys {tuned_keys}: error code {return_code}"
+                        )
+                    continue
+
+                # Measure performance with L2 flush and a large GEMM kernel before to reduce overhead between kernels
+                start_event = paddle.device.cuda.Event(enable_timing=True)
+                end_event = paddle.device.cuda.Event(enable_timing=True)
+                paddle.empty(int(256e6 // 4), dtype=paddle.int32).zero_()
+                paddle.randn((8192, 8192), dtype=paddle.float32, device="cuda") @ paddle.randn(
+                    (8192, 8192), dtype=paddle.float32
+                )
+                start_event.record()
+                for i in range(20):
+                    assert runtime(*args) == 0
+                end_event.record()
+                end_event.synchronize()
+                elapsed_time = start_event.elapsed_time(end_event)
+            else:
+                elapsed_time = 0
 
             # Compare if better
             if best_time is None or elapsed_time < best_time:
@@ -90,6 +119,7 @@ class JITTuner:
         includes: tuple,
         arg_defs: tuple,
         template: str,
+        # args: tuple,
     ) -> Runtime:
         # NOTES: we always assume the space and template will not change
         # We also assume the GPU device will not be changed

ops/csrc/fp8/deep_gemm/jit_kernels/utils.py

@@ -109,13 +109,6 @@ def get_col_major_tma_aligned_tensor(x: Tensor) -> Tensor:
     assert x.dim() in (2, 3)
     remove_dim = False
     if x.dim() == 2:
-        m, n = x.shape
-
-        aligned_m = get_tma_aligned_size(m, x.element_size())
-
-        if aligned_m == m and x.strides[0] == 1 and x.strides[1] == aligned_m:
-            return x
-
         x, remove_dim = x.unsqueeze(0), True
 
     b, m, n = x.shape

ops/csrc/fp8/setup.py

@@ -92,7 +92,7 @@ if __name__ == "__main__":
         packages=["deep_gemm", "deep_gemm/jit", "deep_gemm/jit_kernels"],
         package_data={
             "deep_gemm": [
-                "include/deep_gemm/*",
+                "include/deep_gemm/**/*",
                 "include/cute/**/*",
                 "include/cutlass/**/*",
             ]

ops/csrc/fp8/tests/test_core.py

@@ -114,8 +114,10 @@ def construct_grouped(
 
 def test_gemm() -> None:
     print("Testing GEMM:")
-    for m in (64, 128, 4096):
-        for k, n in [(7168, 2112), (1536, 24576), (512, 32768), (16384, 7168), (7168, 4096), (2048, 7168)]:
+    for m in (64,):
+        for k, n in [
+            (7168, 2112),
+        ]:
             x_fp8, y_fp8, out, ref_out = construct(m, k, n)
             deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, y_fp8, out)
             diff = calc_diff(out, ref_out)
@@ -127,7 +129,7 @@ def test_gemm() -> None:
 def test_m_grouped_gemm_contiguous() -> None:
     print("Testing grouped contiguous GEMM:")
 
-    for num_groups, m, k, n in ((8, 4096, 7168, 4096), (8, 4096, 2048, 7168), (4, 8192, 2048, 7168), (4, 8192, 7168, 4096), ):
+    for num_groups, m, k, n in ((4, 8192, 7168, 4096),):
         # TODO: make a stronger test
         x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=False)
         m_indices = paddle.arange(0, num_groups, dtype=paddle.int32)
@@ -142,9 +144,9 @@ def test_m_grouped_gemm_contiguous() -> None:
 
 def test_m_grouped_gemm_masked() -> None:
     print("Testing grouped masked GEMM:")
-    
-    for num_groups, m in ((1, 1024), (2, 512), (4, 256)):
-        for k, n in ((7168, 4096), (2048, 7168), ):
+
+    for num_groups, m in ((1, 1024),):
+        for k, n in ((7168, 4096),):
             # Test correctness
             masked_m_candidates = list(filter(lambda candidate: candidate <= m, (64, 128, 192, 256, 320, 384)))
             for i in range(10):
@@ -156,7 +158,7 @@ def test_m_grouped_gemm_masked() -> None:
                 masked_m_float = paddle.cast(masked_m, "float32")
                 masked_m_mean = paddle.mean(masked_m_float)
                 masked_m_mean_int = paddle.cast(masked_m_mean, "int32")
-                expected_m = min(int(masked_m_mean_int + 1), m)
+                expected_m = min(masked_m_mean_int + 1, m)
                 deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m)
                 for j in range(num_groups):
                     diff = calc_diff(out[j, : masked_m[j].item()], ref_out[j, : masked_m[j].item()])

ops/csrc/setup.py

@@ -53,6 +53,20 @@ def run_single(func):
     p.join()
 
 
+def run_multi(func_list):
+    processes = []
+    for func in func_list:
+        processes.append(multiprocessing.Process(target=func))
+        processes.append(multiprocessing.Process(target=func))
+        processes.append(multiprocessing.Process(target=func))
+
+    for p in processes:
+        p.start()
+
+    for p in processes:
+        p.join()
+
+
 cc_flag = get_gencode_flags(compiled_all=False)
 cc = get_sm_version()
 
@@ -237,14 +251,17 @@ def setup_paddle_bwd_ops():
         ext_modules=CUDAExtension(
             include_dirs=paddle_includes,
             sources=sources,
-            extra_compile_args={}
         ),
     )
 
 
 if __name__ == "__main__":
-    setup_fast_ln()
-    setup_fused_ln()
-    setup_causal_conv1d()
-    setup_selective_scan()
-    setup_paddle_bwd_ops()
+    run_multi(
+        [
+            setup_fast_ln,
+            setup_fused_ln,
+            setup_causal_conv1d,
+            setup_selective_scan,
+            setup_paddle_bwd_ops,
+        ],
+    )

paddlenlp/data/dist_dataloader.py

@@ -94,22 +94,22 @@ class DistDataLoader(paddle.io.DataLoader):
 
         if self._need_data:
             self._dataloader = paddle.io.DataLoader(
-                dataset=dataset,
-                feed_list=feed_list,
-                places=places,
-                return_list=return_list,
-                batch_sampler=batch_sampler,
-                batch_size=batch_size,
-                shuffle=shuffle,
-                drop_last=drop_last,
-                collate_fn=collate_fn,
-                num_workers=num_workers,
-                use_buffer_reader=use_buffer_reader,
-                prefetch_factor=prefetch_factor,
-                use_shared_memory=use_shared_memory,
-                timeout=timeout,
-                worker_init_fn=worker_init_fn,
-                persistent_workers=persistent_workers,
+                dataset,
+                feed_list,
+                places,
+                return_list,
+                batch_sampler,
+                batch_size,
+                shuffle,
+                drop_last,
+                collate_fn,
+                num_workers,
+                use_buffer_reader,
+                prefetch_factor,
+                use_shared_memory,
+                timeout,
+                worker_init_fn,
+                persistent_workers,
             )
 
             self._lazy_dataloader_iter = None

paddlenlp/datasets/rlhf_datasets/protocol.py

@@ -1,5 +1,4 @@
 # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,19 +16,15 @@
 Implement base data transfer protocol between any two functions, modules.
 We can subclass Protocol to define more detailed batch info with specific keys
 """
-
 import copy
 from dataclasses import dataclass, field
-from typing import Dict, List, Sequence, Union
+from typing import Dict, List, Union
 
 import numpy as np
 import paddle
 import pandas as pd
 from paddle.io import DataLoader
 
-from ...utils import logger
-from ...utils.nested import flatten_list
-
 __all__ = [
     "DataProto",
     "union_tensor_dict",
@@ -54,36 +49,7 @@ class TensorDict:
         self._tensors[key] = tensor
 
     def __getitem__(self, key):
-        if isinstance(key, str):
-            return self._tensors[key]
-        elif isinstance(key, slice):
-            strides = [1] if key.step is None else [key.step]
-            tensor_dict_slice = {
-                k: paddle.strided_slice(v, axes=[0], starts=[key.start], ends=[key.stop], strides=strides)
-                for k, v in self._tensors.items()
-            }
-            batch_size = tensor_dict_slice[list(tensor_dict_slice.keys())[0]].shape[: self.num_batch_dims]
-            return TensorDict(tensor_dict_slice, batch_size=batch_size, num_batch_dims=self.num_batch_dims)
-        else:
-            raise KeyError(f"Unsupported key type: {type(key)}")
-
-    def __contains__(self, key):
-        return key in self._tensors
-
-    def update(self, other):
-        """Updates the tensor dictionary of the current object.
-
-        Args:
-            other (TensorDict or dict): A TensorDict or a dictionary containing tensors to update with.
-
-        Returns:
-            None
-        """
-        if isinstance(other, TensorDict):
-            other = other._tensors
-        for key, value in other.items():
-            # consistency was not checked
-            self._tensors[key] = value
+        return self._tensors[key]
 
     def keys(self):
         return self._tensors.keys()
@@ -96,218 +62,16 @@ class TensorDict:
             self._tensors[key] = self._tensors[key].to(device)
         return self
 
-    def pop(self, key: str, default=None):
-        """Removes the value associated with the specified key from the TensorDict and returns it.
-
-        Args:
-            key (str): The key to remove from the TensorDict.
-            default (Optional[Any]): The default value to return if the key is not found. Defaults to None.
-
-        Returns:
-            Any: The value corresponding to the key if it exists; otherwise, returns the default value if provided.
-
-        Raises:
-            KeyError: If the key does not exist and no default value is provided.
-        """
-        if key in self._tensors:
-            value = self._tensors.pop(key)
-            return value
-        elif default is not None:
-            return default
-        else:
-            raise KeyError(f"Key '{key}' not found in TensorDict and no default value provided.")
-
-    def select(self, *keys, default=None):
-        """Selects tensors by the given keys and returns a new TensorDict object.
-
-        Args:
-            *keys (tuple): One or more keys to select.
-            default (Any, optional): The default value to return if a key does not exist. Defaults to None.
-
-        Returns:
-            TensorDict: A new TensorDict containing the selected tensors.
-
-        Raises:
-            KeyError: If a key does not exist and no default value is provided.
-        """
-        selected_tensors = {}
-        for key in keys:
-            if key in self._tensors:
-                selected_tensors[key] = self._tensors[key]
-            elif default is not None:
-                selected_tensors[key] = default
-            else:
-                raise KeyError(f"Key '{key}' not found in TensorDict and no default value provided.")
-
-        batch_size = (
-            list(selected_tensors[list(selected_tensors.keys())[0]].shape[: self.num_batch_dims])
-            if selected_tensors
-            else None
-        )
-        return TensorDict(selected_tensors, batch_size=batch_size, num_batch_dims=self.num_batch_dims)
-
-    def rename_key_(self, old_keys, new_keys):
-        """Renames keys in the TensorDict for use with the `rename` method of DataProto.
-        The naming convention is directly adopted from the official PyTorch implementation.
-
-        Args:
-            old_keys (str or List[str]): The key or list of keys to be renamed.
-            new_keys (str or List[str]): The new key or list of new keys.
-
-        Returns:
-            self: The updated TensorDict instance with renamed keys.
-
-        Raises:
-            ValueError: If the lengths of `old_keys` and `new_keys` do not match.
-            KeyError: If any key in `old_keys` does not exist in the TensorDict,
-                        or if any key in `new_keys` already exists in the TensorDict.
-        """
-        if isinstance(old_keys, (str,)):
-            old_keys = [old_keys]
-        if isinstance(new_keys, (str,)):
-            new_keys = [new_keys]
-        if len(old_keys) != len(new_keys):
-            raise ValueError("old_keys and new_keys must have the same length.")
-
-        for old_key, new_key in zip(old_keys, new_keys):
-            if old_key not in self._tensors:
-                raise KeyError(f"Key '{old_key}' not found in TensorDict.")
-            if new_key in self._tensors:
-                raise KeyError(f"Key '{new_key}' already exists in TensorDict.")
-            self._tensors[new_key] = self._tensors.pop(old_key)
-        return self
-
-    @classmethod
-    def concat(cls, tensordict_list, axis=0):
-        """Concatenate multiple TensorDict objects along a specified axis.
-
-        Args:
-            cls (class): The TensorDict class.
-            tensordict_list (list of TensorDict): The list of TensorDict objects to concatenate.
-            axis (int, optional): The axis to concatenate along, default is 0.
-
-        Returns:
-            TensorDict: The concatenated TensorDict object.
-
-        Raises:
-            ValueError: If tensordict_list is empty, or the keys of TensorDict objects are different,
-                        or the dimensions other than the concatenation axis do not match, a ValueError will be raised.
-        """
-        if not tensordict_list:
-            raise ValueError("tensordict_list must not be empty")
-
-        first_tensordict = tensordict_list[0]
-        first_keys = first_tensordict.keys()
-        first_shapes = {key: tensor.shape for key, tensor in first_tensordict.items()}
-
-        for tensordict in tensordict_list:
-            if tensordict.keys() != first_keys:
-                raise ValueError("All TensorDict objects must have the same keys")
-            for key in first_keys:
-                if (
-                    tensordict[key].shape[:axis] + tensordict[key].shape[axis + 1 :]
-                    != first_shapes[key][:axis] + first_shapes[key][axis + 1 :]
-                ):
-                    raise ValueError(f"Shapes of tensor '{key}' do not match except on concatenation axis {axis}")
-
-        concatenated_tensors = {
-            key: paddle.concat([tensordict[key] for tensordict in tensordict_list], axis=axis) for key in first_keys
-        }
-
-        batch_size = concatenated_tensors[list(concatenated_tensors.keys())[0]].shape[
-            : tensordict_list[0].num_batch_dims
-        ]
-        return cls(concatenated_tensors, batch_size=batch_size, num_batch_dims=tensordict_list[0].num_batch_dims)
-
-    def get(self, key, default=None):
-        return self._tensors.get(key, default)
-
-
-def tensor_or_tensordict_concat(
-    x: Union[Sequence[paddle.Tensor], Sequence[TensorDict]],
-    axis: int | paddle.Tensor = 0,
-    name: str | None = None,
-):
-    """Concatenates multiple tensors or TensorDicts along the specified axis.
-
-    Args:
-        x (Union[Sequence[paddle.Tensor], Sequence[TensorDict]]): A sequence of tensors or TensorDicts to concatenate.
-        axis (int or paddle.Tensor, optional): The axis along which to concatenate. Defaults to 0.
-            If `x` is a sequence of tensors, `axis` should be an integer;
-            if `x` is a sequence of TensorDicts, `axis` should be a paddle.Tensor. Defaults to 0.
-        name (str, optional): The name of the operation (optional). Defaults to None.
-
-    Returns:
-        Union[paddle.Tensor, TensorDict]: The concatenated tensor or TensorDict.
-    """
-
-    def is_tensor_sequence():
-        if isinstance(x[0], paddle.Tensor):
-            return True
-        else:
-            return False
-
-    if not is_tensor_sequence() and paddle.in_dynamic_mode():
-        return TensorDict.concat(x)
-    else:
-        return paddle.concat(x, axis, name)
-
-
-def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:
-    """Union two tensordicts."""
-    assert (
-        tensor_dict1.batch_size == tensor_dict2.batch_size
-    ), f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}"
-    for key in tensor_dict2.keys():
-        if key not in tensor_dict1.keys():
-            tensor_dict1[key] = tensor_dict2[key]
-        else:
-            assert tensor_dict1[key].equal(
-                tensor_dict2[key]
-            ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object"
-
-    return tensor_dict1
-
-
-def union_numpy_dict(tensor_dict1: dict[np.ndarray], tensor_dict2: dict[np.ndarray]) -> dict[np.ndarray]:
-    """Merges two dictionaries containing NumPy arrays.
-
-    Args:
-        tensor_dict1 (dict[np.ndarray]): The first dictionary containing NumPy arrays.
-        tensor_dict2 (dict[np.ndarray]): The second dictionary containing NumPy arrays.
-
-    Returns:
-        dict[np.ndarray]: The merged dictionary.
-
-    Raises:
-        AssertionError: If the values corresponding to the same key in both dictionaries are not NumPy arrays,
-                        or if the arrays are not equal (considering NaNs and object types).
-    """
-    for key, val in tensor_dict2.items():
-        if key in tensor_dict1:
-            assert isinstance(tensor_dict2[key], np.ndarray)
-            assert isinstance(tensor_dict1[key], np.ndarray)
-            # to properly deal with nan and object type
-            assert pd.DataFrame(tensor_dict2[key]).equals(
-                pd.DataFrame(tensor_dict1[key])
-            ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object"
-        tensor_dict1[key] = val
-
-    return tensor_dict1
-
 
 def union_two_dict(dict1: Dict, dict2: Dict):
-    """Merges two dictionaries. Raises an error if the same key exists in both dictionaries but maps to different objects.
+    """Union two dict. Will throw an error if there is an item not the same object with the same key.
 
     Args:
-        dict1 (Dict): The first dictionary.
-        dict2 (Dict): The second dictionary.
+        dict1:
+        dict2:
 
     Returns:
-        Dict: The merged dictionary.
 
-    Raises:
-        AssertionError: If the same key exists in both dictionaries with different values.
     """
     for key, val in dict2.items():
         if key in dict1:
@@ -344,39 +108,42 @@ def pad_dataproto_to_divisor(data: "DataProto", size_divisor: int):
 
 
 def unpad_dataproto(data: "DataProto", pad_size):
-    """Removes the padded parts from the given DataProto object.
-
-    Args:
-        data (DataProto): The DataProto object containing the data.
-        pad_size (int): The size of the padding to remove.
-
-    Returns:
-        DataProto: The DataProto object with the padding removed.
-    """
     if pad_size != 0:
         data = data[:-pad_size]
     return data
 
 
+def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:
+    """Union two tensordicts."""
+    assert (
+        tensor_dict1.batch_size == tensor_dict2.batch_size
+    ), f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}"
+    for key in tensor_dict2.keys():
+        if key not in tensor_dict1.keys():
+            tensor_dict1[key] = tensor_dict2[key]
+        else:
+            assert tensor_dict1[key].equal(
+                tensor_dict2[key]
+            ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object"
+
+    return tensor_dict1
+
+
+def union_numpy_dict(tensor_dict1: dict[np.ndarray], tensor_dict2: dict[np.ndarray]) -> dict[np.ndarray]:
+    for key, val in tensor_dict2.items():
+        if key in tensor_dict1:
+            assert isinstance(tensor_dict2[key], np.ndarray)
+            assert isinstance(tensor_dict1[key], np.ndarray)
+            # to properly deal with nan and object type
+            assert pd.DataFrame(tensor_dict2[key]).equals(
+                pd.DataFrame(tensor_dict1[key])
+            ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object"
+        tensor_dict1[key] = val
+
+    return tensor_dict1
+
+
 def list_of_dict_to_dict_of_list(list_of_dict: list[dict]):
-    """Converts a list of dictionaries into a dictionary of lists.
-
-    Args:
-        list_of_dict (list[dict]): A list containing dictionaries.
-
-    Returns:
-        dict: A dictionary where each key maps to a list of values collected from the input dictionaries.
-
-    Raises:
-        AssertionError: Raised if any key in the dictionaries is not present in the final output dictionary during iteration.
-
-    Examples:
-        >>> list_of_dict = [{'name': 'Alice', 'age': 25}, {'name': 'Bob', 'age': 30}]
-        >>> result = list_of_dict_to_dict_of_list(list_of_dict)
-        >>> print(result)
-        {'name': ['Alice', 'Bob'], 'age': [25, 30]}
-    """
-
     if len(list_of_dict) == 0:
         return {}
     keys = list_of_dict[0].keys()
@@ -428,15 +195,6 @@ def unfold_batch_dim(data: "DataProto", batch_dims=2):
 
 
 def collate_fn(x: list["DataProtoItem"]):
-    """Combines a list of DataProtoItem objects into a single batch.
-
-    Args:
-        x (list["DataProtoItem"]): A list of DataProtoItem objects.
-
-    Returns:
-        DataProto: A DataProto object containing the combined batch data and non-tensor batch data.
-    """
-
     batch = []
     non_tensor_batch = []
     for data in x:
@@ -449,26 +207,6 @@ def collate_fn(x: list["DataProtoItem"]):
     return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)
 
 
-def make_eos_mask(response_id, eos_token_ids=0, dtype=paddle.int64):
-    """
-    end of sentence token can be int or list: 1 or [1, 2]
-    e.g. eos_token=1
-    response_id: [0, 0, 2, 42, 3, 5, 1, 0, 0]
-    eos_mask:     [1, 1, 1, 1,  1, 1, 1, 0, 0]
-    """
-    if isinstance(eos_token_ids, int):
-        eos_token_ids = [eos_token_ids]
-
-    eos_mask = paddle.zeros_like(response_id, dtype=paddle.bool)
-    for token_id in eos_token_ids:
-        eos_mask |= response_id == token_id
-
-    eos_mask = eos_mask.to("int64")
-    eos_mask = (paddle.cumsum(eos_mask, axis=1) - eos_mask).to("bool")
-    eos_mask = paddle.logical_not(eos_mask).to(dtype)
-    return eos_mask
-
-
 @dataclass
 class DataProtoItem:
     # TODO(zhangchi.usc1992) add consistency check
@@ -501,49 +239,11 @@ class DataProto:
             return 0
 
     def __getitem__(self, item):
-        """
-        Enhanced indexing for DataProto objects.
-
-        Args:
-            item: Can be one of:
-                - int: A single index
-                - slice: A slice object (start:stop:step)
-                - list: A list of indices
-                - numpy.ndarray: An array of indices
-                - torch.Tensor: A tensor of indices
-
-        Returns:
-            DataProto: For all indexing types except single integers
-            DataProtoItem: Only for single integer indices
-        """
-        # Case 1: Slice object - use the slice method
-        if isinstance(item, slice):
-            return self.slice(item.start, item.stop, item.step)
-
-        # Case 2: List, numpy array, or torch tensor - use sel_idxs
-        elif isinstance(item, (list, np.ndarray, paddle.Tensor)):
-            return self.select_idxs(item)
-
-        # Case 3: Single integer - return DataProtoItem for backward compatibility
-        elif isinstance(item, (int, np.integer)):
-            tensor_data = self.batch[item]
-            non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()}
-            return_type = DataProto if isinstance(item, slice) else DataProtoItem
-            return return_type(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
-
-        # Case 4: Unsupported type
-        else:
-            raise TypeError(f"Indexing with {type(item)} is not supported")
+        tensor_data = self.batch[item]
+        non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()}
+        return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
 
     def print_size(self, prefix=""):
-        """Prints the sizes of `tensordict` and `non_tensor_batch`.
-
-        Args:
-            prefix (str, optional): A prefix string to identify the log message. Defaults to an empty string.
-
-        Returns:
-            None
-        """
         size_of_tensordict = 0
         for key, tensor in self.batch.items():
             size_of_tensordict += tensor.element_size() * tensor.numel()
@@ -586,19 +286,6 @@ class DataProto:
 
     @classmethod
     def from_single_dict(cls, data: Dict[str, Union[paddle.Tensor, np.ndarray]], meta_info=None):
-        """Converts a dictionary containing Paddle tensors and NumPy arrays into a DataProto object.
-
-        Args:
-            cls (class): The class invoking this method.
-            data (Dict[str, Union[paddle.Tensor, np.ndarray]]): A dictionary containing Paddle tensors and NumPy arrays.
-            meta_info (Any, optional): Optional metadata. Defaults to None.
-
-        Returns:
-            DataProto: A DataProto object containing `tensors` and `non_tensors`.
-
-        Raises:
-            ValueError: If the input `data` contains unsupported data types.
-        """
         tensors = {}
         non_tensors = {}
 
@@ -649,117 +336,6 @@ class DataProto:
         tensor_dict = TensorDict(source=tensors, batch_size=batch_size, num_batch_dims=num_batch_dims)
         return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)
 
-    @staticmethod
-    def concatenate_tensors(
-        gathered_list: List[Union[paddle.Tensor, np.ndarray, List]],
-        data_parallel_group=None,
-        sharding_parallel_group=None,
-    ) -> Union[paddle.Tensor, np.ndarray]:
-        """Concatenates a list of tensors/arrays that have been gathered from
-        different distributed ranks. Handles both Paddle Tensors and NumPy arrays,
-        and accounts for nested lists from `all_gather_object` with multiple ranks.
-
-        Args:
-            gathered_list (List[Union[paddle.Tensor, np.ndarray, List]]):
-                A list containing tensors or arrays gathered from all ranks.
-                This list might be nested if `all_gather_object` was used with
-                multiple groups (e.g., sharding then data parallel).
-            data_parallel_group (Optional[paddle.distributed.collective.Group]):
-                The data parallel communication group.
-            sharding_parallel_group (Optional[paddle.distributed.collective.Group]):
-                The sharding parallel communication group.
-
-        Returns:
-            Union[paddle.Tensor, np.ndarray]: The concatenated tensor/array.
-                                              If the input was Paddle Tensor, returns Paddle Tensor.
-                                              If the input was NumPy array, returns NumPy array.
-        """
-        if not gathered_list:
-            raise ValueError("Input gathered_list cannot be empty.")
-        is_paddle_tensor = isinstance(gathered_list[0], paddle.Tensor)
-        if is_paddle_tensor:
-            return paddle.concat(gathered_list, axis=0)
-        else:
-            return np.concatenate(flatten_list(gathered_list), axis=0)
-
-    @staticmethod
-    def pad_tensor(tensor_list, pad_index, dtype, padding_side):
-        """Pads a list of tensors or numpy arrays to the same size.
-
-        Args:
-            tensor_list (list): A list of tensors or numpy arrays. All elements must have the same first dimension (shape[0]).
-            pad_index (int): The value used for padding.
-            dtype (paddle.DataType): The data type of the resulting padded tensor.
-            padding_side (str): The padding direction. Use 'right' to pad on the right, or 'left' to pad on the left.
-
-        Returns:
-            Tensor or ndarray: The padded tensor or numpy array with shape (data_num, max_size).
-
-        Raises:
-            ValueError: If `padding_side` is not 'right' or 'left'.
-        """
-        max_size = max([i.shape[-1] for i in tensor_list])
-        data_num = sum([i.shape[0] for i in tensor_list])
-        if isinstance(tensor_list[0], paddle.Tensor):
-            new_tensor = paddle.full((data_num, max_size), pad_index, dtype=dtype)
-        elif isinstance(tensor_list[0], np.ndarray):
-            new_tensor = np.full((data_num, max_size), pad_index, dtype=dtype)
-
-        offset = 0
-        for idx, i in enumerate(tensor_list):
-            # new_tensor[offset : offset + i.shape[0], : i.shape[-1]] = i
-            data_length = i.shape[-1]
-
-            if padding_side == "right":
-                new_tensor[offset : offset + i.shape[0], :data_length] = i
-            elif padding_side == "left":
-                new_tensor[offset : offset + i.shape[0], -data_length:] = i
-            else:
-                raise ValueError("padding_side must be 'right' or 'left'")
-            offset += i.shape[0]
-        return new_tensor
-
-    @staticmethod
-    def pad_or_concat_tensor_list(tensor_list, pad_index, key):
-        """Pads or concatenates a given list of tensors.
-
-        Args:
-            tensor_list (list): The list of tensors to be processed.
-            pad_index (int): The value used for padding.
-            key (str): The key identifying the tensor list.
-
-        Returns:
-            Union[np.ndarray, torch.Tensor]: The processed tensor or list of tensors.
-        """
-        left_padding_key = ("prompt", "label_ids")
-        pad = False if len(flatten_list(tensor_list)[0].shape) == 1 else True
-        padding_side = "left" if (key in left_padding_key) else "right"
-        dtype = flatten_list(tensor_list)[0].dtype
-        if not pad:
-            return DataProto.concatenate_tensors(tensor_list)
-        else:
-            return DataProto.pad_tensor(tensor_list, pad_index, dtype, padding_side)
-
-    @staticmethod
-    def pad_batch_data(tensor_list: List[paddle.Tensor] = None, pad_token_id=None) -> List[paddle.Tensor]:
-        """Pads a batch of tensors.
-
-        Args:
-            tensor_list (List[paddle.Tensor], optional): A list of tensors to be padded. Defaults to None.
-            pad_token_id (int, optional): The token ID used for padding. Defaults to None.
-
-        Returns:
-            List[paddle.Tensor]: The list of padded tensors.
-        """
-        tensor_list = [paddle.unsqueeze(v, axis=0) if v.ndim == 1 else v for v in tensor_list]
-        padded_tensors = DataProto.pad_tensor(
-            tensor_list,
-            pad_index=pad_token_id,
-            dtype=tensor_list[0].dtype,
-            padding_side="right",
-        )
-        return padded_tensors
-
     def to(self, device) -> "DataProto":
         """move the batch to device
 
@@ -809,87 +385,6 @@ class DataProto:
 
         return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info)
 
-    def select_idxs(self, idxs):
-        """
-        Select specific indices from the DataProto.
-
-        Args:
-            idxs (torch.Tensor or numpy.ndarray or list): Indices to select
-
-        Returns:
-            DataProto: A new DataProto containing only the selected indices
-        """
-        if isinstance(idxs, list):
-            idxs = paddle.tensor(idxs, dtype=paddle.int32)
-
-        if isinstance(idxs, np.ndarray):
-            idxs_np = idxs
-            idxs_paddle = paddle.from_numpy(idxs)
-        else:  # torch.Tensor
-            idxs_paddle = idxs
-            idxs_np = idxs.numpy()
-
-        if self.batch is not None:
-            # Use TensorDict's built-in indexing capabilities
-            selected_batch = TensorDict(
-                source={key: tensor[idxs_paddle] for key, tensor in self.batch.items()},
-                batch_size=(idxs_paddle.shape[0],),
-            )
-        else:
-            selected_batch = None
-
-        selected_non_tensor = {}
-        for key, val in self.non_tensor_batch.items():
-            selected_non_tensor[key] = val[idxs_np]
-
-        return DataProto(batch=selected_batch, non_tensor_batch=selected_non_tensor, meta_info=self.meta_info)
-
-    def slice(self, start=None, end=None, step=None):
-        """
-        Slice the DataProto and return a new DataProto object.
-        This is an improved version of direct slicing which returns a DataProtoItem.
-
-        Args:
-            start (int, optional): Start index. Defaults to None (start from beginning).
-            end (int, optional): End index (exclusive). Defaults to None (go to end).
-            step (int, optional): Step size. Defaults to None (step=1).
-
-        Returns:
-            DataProto: A new DataProto containing the sliced data
-
-        Examples:
-            # Using the slice method directly
-            sliced_data = data_proto.slice(10, 20)
-
-            # Using enhanced indexing (returns DataProto)
-            sliced_data = data_proto[10:20]
-            sliced_data = data_proto[::2]  # Every other element
-
-            # Using list indexing (returns DataProto)
-            indices = [1, 5, 10]
-            selected_data = data_proto[indices]
-
-            # Single index still returns DataProtoItem
-            single_item = data_proto[5]
-        """
-        # Create a slice object
-        slice_obj = slice(start, end, step)
-
-        # Handle the batch data
-        if self.batch is not None:
-            # Use TensorDict's built-in slicing capabilities
-            sliced_batch = self.batch[slice_obj]
-        else:
-            sliced_batch = None
-
-        # Handle the non-tensor batch data
-        sliced_non_tensor = {}
-        for key, val in self.non_tensor_batch.items():
-            sliced_non_tensor[key] = val[slice_obj]
-
-        # Return a new DataProto object
-        return DataProto(batch=sliced_batch, non_tensor_batch=sliced_non_tensor, meta_info=self.meta_info)
-
     def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> "DataProto":
         """Pop a subset of the DataProto via `batch_keys` and `meta_info_keys`
 
@@ -923,21 +418,8 @@ class DataProto:
         return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)
 
     def rename(self, old_keys=None, new_keys=None) -> "DataProto":
-        """Rename the keys in the batch.
-
-        Args:
-            old_keys (Optional[Union[str, List[str]]]): The old keys to be renamed.
-            new_keys (Optional[Union[str, List[str]]]): The new keys after renaming.
-
-        Returns:
-            DataProto: Returns the instance itself.
-
-        Raises:
-            TypeError: If old_keys or new_keys is not a string or a list.
-            ValueError: If the lengths of old_keys and new_keys are not the same.
-
-        Note:
-            This function only renames the keys in the batch.
+        """
+        Note that this function only rename the key in the batch
         """
 
         def validate_input(keys):
@@ -965,7 +447,6 @@ class DataProto:
     def union(self, other: "DataProto") -> "DataProto":
         """Union with another DataProto. Union batch and meta_info separately.
         Throw an error if
-
         - there are conflict keys in batch and they are not equal
         - the batch size of two data batch is not the same
         - there are conflict keys in meta_info and they are not the same.
@@ -1063,7 +544,7 @@ class DataProto:
         for batch in data:
             batch_lst.append(batch.batch)
         if batch_lst[0] is not None:
-            new_batch = tensor_or_tensordict_concat(batch_lst, axis=0)
+            new_batch = paddle.concat(batch_lst, axis=0)
         else:
             new_batch = None
 
@@ -1096,8 +577,7 @@ class DataProto:
             if interleave:
                 # Interleave the data
                 repeated_tensors = {
-                    key: paddle.repeat_interleave(tensor, repeats=repeat_times, axis=0)
-                    for key, tensor in self.batch.items()
+                    key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items()
                 }
             else:
                 # Stack the data
@@ -1108,7 +588,7 @@ class DataProto:
 
             repeated_batch = TensorDict(
                 source=repeated_tensors,
-                batch_size=[self.batch.batch_size[0] * repeat_times],
+                batch_size=(self.batch.batch_size[0] * repeat_times,),
             )
         else:
             repeated_batch = None
@@ -1125,36 +605,3 @@ class DataProto:
             non_tensor_batch=repeated_non_tensor_batch,
             meta_info=self.meta_info,
         )
-
-    def split_batch_into_micro_batches(self, batch_size, pad_token_id=0) -> List["DataProto"]:
-        """
-        Splits total_batch into micro-batches of size `batch_size`.
-
-        Args:
-            total_batch (dict): Dictionary containing full batched tensors.
-            batch_size (int): Micro batch size per device.
-
-        Returns:
-            list of dict: A list of micro-batches.
-        """
-        micro_batches = []
-        num_micro_batches = self.batch["input_ids"].shape[0] // batch_size
-        if self.batch["input_ids"].shape[0] % batch_size != 0:
-            num_micro_batches += 1
-        if num_micro_batches <= 0:
-            logger.warning(
-                "The total batch size is smaller than the batch size, please consider using a smaller batch size or a larger global_batch_size."
-            )
-            num_micro_batches = 1
-
-        for i in range(num_micro_batches):
-            micro_batch = {}
-            for key, data in self.batch.items():
-                micro_batch[key] = data[i * batch_size : (i + 1) * batch_size]
-
-            # if os.getenv("PROCESS_PROMPT_AND_RESPONSE", "1").lower() in ["1", "t", "true", "yes", "y"]:
-            #     micro_batch = process_prompt_and_response(micro_batch=micro_batch, pad_token_id=pad_token_id)
-
-            micro_batches.append(DataProto.from_single_dict(micro_batch))
-
-        return micro_batches

paddlenlp/experimental/transformers/deepseek_v2/modeling.py

@@ -47,6 +47,9 @@ from paddlenlp.transformers.deepseek_v2.modeling import (
     yarn_get_mscale,
     yarn_linear_ramp_mask,
 )
+from paddlenlp.transformers.model_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+)
 from paddlenlp.transformers.model_utils import (
     dy2st_nocheck_guard_context,
     register_base_model,
@@ -263,6 +266,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
         self.weight_block_size = config.weight_block_size
         self.moe_quant_type = config.moe_quant_type
         self.rope_theta = config.rope_theta
+        self.return_full_hidden_states = config.get("return_full_hidden_states", False)
 
         self.use_weight_only = False
         self.weightonly_group_size = -1
@@ -587,6 +591,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
         speculate_config = SpeculateConfig(
             speculate_method=config.get("speculate_method", None),
             speculate_max_draft_token_num=config.get("speculate_max_draft_token_num", 5),
+            return_full_hidden_states=config.get("return_full_hidden_states", False),
         )
 
         transformer_config = FusedMultiTransformerConfig(
@@ -617,9 +622,9 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
             rotary_emb=self.rotary_emb,
             norm_type="rmsnorm",
             rank_id=config.tensor_parallel_rank,
-            append_attn=config.append_attn,
             moe_config=moe_config,
             mla_config=mla_config,
+            append_attn=config.append_attn,
             speculate_config=speculate_config,
         )
 
@@ -1284,7 +1289,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
                 inputs_embeds = inputs_embeds.reshape([-1, inputs_embeds.shape[2]])
 
         with dy2st_nocheck_guard_context():
-            hidden_states, full_hidden_states = self.transformer_block(
+            hidden_states, _ = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1296,7 +1301,13 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return hidden_states, full_hidden_states
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+            cum_offsets=cum_offsets,
+        )
 
 
 @register_base_model
@@ -1956,7 +1967,7 @@ class MTPDeepseekV2BlockInferenceModel(DeepseekV2BlockInferenceModel):
         inputs_embeds = self.eh_proj(inputs_embeds)
 
         with dy2st_nocheck_guard_context():
-            hidden_states, full_hidden_states = self.transformer_block(
+            hidden_states, _ = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1969,7 +1980,12 @@ class MTPDeepseekV2BlockInferenceModel(DeepseekV2BlockInferenceModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return hidden_states, full_hidden_states
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
 
 
 class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, DeepseekV2PretrainedModel):
@@ -2196,7 +2212,7 @@ class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, De
         draft_tokens=None,
         output_padding_offset=None,
     ):
-        hidden_states, full_hidden_states = self.deepseek_v2(
+        outputs = self.deepseek_v2(
             input_ids,
             inputs_embeds=inputs_embeds,
             src_mask=src_mask,
@@ -2214,7 +2230,21 @@ class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, De
             draft_tokens=draft_tokens,
             output_padding_offset=output_padding_offset,
         )
+        if self.return_full_hidden_states:
+            from paddlenlp_ops import rebuild_padding_v2
 
+            full_hidden_states = outputs[0]
+            cum_offsets = outputs[1]
+            hidden_states = rebuild_padding_v2(
+                full_hidden_states,
+                cum_offsets,
+                seq_lens_decoder,
+                seq_lens_encoder,
+                output_padding_offset,
+                self.max_seq_len,
+            )
+        else:
+            hidden_states = outputs[0]
         logits = self.lm_head(
             hidden_states,
             tensor_parallel_output=False,
@@ -2224,6 +2254,8 @@ class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, De
         else:
             return logits
 
+        return logits
+
     @paddle.no_grad()
     def set_state_dict(self, state_dict):
         if "lm_head.weight" in state_dict:
@@ -2331,7 +2363,7 @@ class MTPDeepseekV2ForCausalLMBlockInferenceModel(DeepseekV2ForCausalLMBlockInfe
         output_padding_offset=None,
         pre_hidden_states=None,
     ):
-        hidden_states, _ = self.mtp(
+        outputs = self.mtp(
             input_ids,
             src_mask=src_mask,
             caches=caches,
@@ -2350,6 +2382,8 @@ class MTPDeepseekV2ForCausalLMBlockInferenceModel(DeepseekV2ForCausalLMBlockInfe
             pre_hidden_states=pre_hidden_states,
         )
 
+        hidden_states = outputs[0]
+
         logits = self.lm_head(
             hidden_states,
             tensor_parallel_output=False,

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -1481,8 +1481,7 @@ class FusedMultiTransformerBase(Layer):
         self.pre_process(**kwargs)
         kwargs["cum_offsets"] = cum_offsets
 
-        kv_cache_reuse = kwargs.get("kv_cache_reuse", None)
-        if caches is not None and kv_cache_reuse is not True:
+        if caches is not None:
             assert len(caches) == len(self.linear_weights) or len(caches) == 2 * len(self.linear_weights)
 
         assert self.num_layers == len(self.linear_weights)
@@ -1590,7 +1589,7 @@ class FusedMultiTransformerBase(Layer):
         kwargs["input_ids"] = input_ids
 
         out = self.post_process(**kwargs)
-        return out, kwargs["multi_block_output"]
+        return out, caches
 
 
 class FusedMultiTransformerPostLayernorm(FusedMultiTransformerBase):
@@ -3173,17 +3172,10 @@ class FusedBlockMultiTransformer(FusedMultiTransformerBase):
             k_dequant_scales = kwargs.get("k_dequant_scales", None)
             v_dequant_scales = kwargs.get("v_dequant_scales", None)
 
-            kv_cache_reuse = kwargs.get("kv_cache_reuse", None)
-            if kv_cache_reuse:
-                k_cache_index = 0
-                v_cache_index = 1
-            else:
-                k_cache_index = 2 * i
-                v_cache_index = 2 * i + 1
             fmha_out = paddle.incubate.nn.functional.block_multihead_attention(
                 qkv_out,
-                caches[k_cache_index],
-                caches[v_cache_index],
+                caches[2 * i],
+                caches[2 * i + 1],
                 kwargs.get("seq_lens_encoder", None),
                 kwargs.get("seq_lens_decoder", None),
                 kwargs.get("seq_lens_this_time", None),

paddlenlp/experimental/transformers/generation_utils.py

@@ -295,7 +295,7 @@ class GenerationInferenceModel(GenerationMixin):
             model_inputs = self.prepare_inputs_for_generation(input_ids, cache_kvs, **args)
             return self(**model_inputs)
 
-        def _post_process_(logits, top_p, temperature, step_idx_ori, model_kwargs):
+        def _post_process_(outputs, top_p, temperature, step_idx_ori, model_kwargs):
             cache = model_kwargs.get("cache", None)
             just_decoder = model_kwargs["seq_len_encoder"] == 0
             if cache is None:  # first decoder
@@ -314,6 +314,7 @@ class GenerationInferenceModel(GenerationMixin):
                 step_idx,
                 model_kwargs["stop_flags"],
             )
+            logits = outputs[0] if isinstance(outputs, tuple) else outputs
 
             logits = paddle.cast(logits, paddle.float32)
             logits = logits_processors(model_kwargs["all_input_ids"], logits, decoding_step=step_idx_ori)
@@ -372,7 +373,7 @@ class GenerationInferenceModel(GenerationMixin):
         outputs = _forward_(**model_kwargs)
         # first decoder
         next_tokens, model_kwargs = _post_process_(
-            outputs[0] if isinstance(outputs, tuple) else outputs,
+            outputs,
             top_p,
             temperature,
             step_idx_ori,
@@ -388,9 +389,8 @@ class GenerationInferenceModel(GenerationMixin):
             paddle.sum(paddle.cast(model_kwargs["stop_flags"], "int64")),
             model_kwargs["stop_nums"],
         ):
-            outputs = _forward_(**model_kwargs)
             next_tokens, model_kwargs = _post_process_(
-                outputs[0] if isinstance(outputs, tuple) else outputs,
+                _forward_(**model_kwargs),
                 top_p,
                 temperature,
                 step_idx_ori,
@@ -692,7 +692,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             return self(**model_inputs)
 
         def _post_process_(
-            logits,
+            outputs,
             top_k,
             top_p,
             penalty_score,
@@ -702,7 +702,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             model_kwargs,
         ):
             step_idx = model_kwargs["step_idx"]
-            logits = paddle.cast(logits, paddle.float32)
+            logits = paddle.cast(outputs, paddle.float32)
 
             from paddlenlp_ops import set_preids_token_penalty_multi_scores
 
@@ -777,7 +777,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
         outputs = _forward_(**model_kwargs)  # [bs, 1, dim_embed]
         # first decoder
         next_tokens = _post_process_(
-            outputs[0] if isinstance(outputs, tuple) else outputs,
+            outputs,
             top_k,
             top_p,
             penalty_score,
@@ -806,7 +806,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             return self(**model_inputs)
 
         def _post_process_(
-            logits,
+            outputs,
             top_k,
             top_p,
             penalty_score,
@@ -816,7 +816,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             model_kwargs,
         ):
             step_idx = model_kwargs["step_idx"]
-            logits = paddle.cast(logits, paddle.float32)
+            logits = paddle.cast(outputs, paddle.float32)
 
             from paddlenlp_ops import speculate_get_token_penalty_multi_scores
 
@@ -959,7 +959,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             return self(**model_inputs)
 
         def _post_process_(
-            logits,
+            outputs,
             top_k,
             top_p,
             penalty_score,
@@ -968,7 +968,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             temperature,
             model_kwargs,
         ):
-            logits = paddle.cast(logits, paddle.float32)
+            logits = paddle.cast(outputs, paddle.float32)
 
             probs = F.softmax(logits)
 
@@ -1191,7 +1191,7 @@ class GenerationAvxInferenceModel(GenerationMixin):
             model_inputs = self.prepare_inputs_for_generation(input_ids, **args)
             return self(**model_inputs)
 
-        def _post_process_(logits, top_p, temperature, step_idx_ori, model_kwargs):
+        def _post_process_(outputs, top_p, temperature, step_idx_ori, model_kwargs):
             cache = model_kwargs.get("cache", None)
             just_decoder = model_kwargs["seq_len_encoder"] == 0
             if cache is None:  # first decoder
@@ -1211,6 +1211,7 @@ class GenerationAvxInferenceModel(GenerationMixin):
                 step_idx,
                 model_kwargs["stop_flags"],
             )
+            logits = outputs[0] if isinstance(outputs, tuple) else outputs
             logits = paddle.cast(logits, paddle.float32)
             logits = logits_processors(model_kwargs["all_input_ids"], logits, decoding_step=step_idx_ori)
 

paddlenlp/experimental/transformers/llama/modeling.py

@@ -150,15 +150,13 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
         ln_scale_attrs = [paddle.ParamAttr(name="fusellama.{}.ln_scale".format(i)) for i in range(self.num_layers)]
         qkv_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.qkv_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.qkv_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
         out_proj_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.out_proj_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.out_proj_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
@@ -167,22 +165,19 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
         ]
         gate_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.gate_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.gate_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
         up_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.up_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.up_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
         ffn2_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.ffn2_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.ffn2_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
@@ -262,12 +257,19 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
             # merge batch and seq_len dimension.
             inputs_embeds = inputs_embeds.reshape([batch * seq_len, hidden_dim])
 
+        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 inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         hidden_states = inputs_embeds
 
         # decoder layers
+        all_hidden_states = () if output_hidden_states else None
         with dy2st_nocheck_guard_context():
             hidden_states = self.transformer_block(
                 input_ids=input_ids,
@@ -278,7 +280,11 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return hidden_states
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, None, all_hidden_states, None] if v is not None)
 
     @paddle.no_grad()
     # avx
@@ -333,8 +339,7 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
                 )  # reshape(3, self.num_attention_heself.hidden_sizeads // self.config.tensor_parallel_degree, head_size, )
             if "llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx) in state_dict.keys():
                 concated_ffn1_weight = np.concatenate(
-                    split_fn(state_dict["llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx)]),
-                    axis=-1,
+                    split_fn(state_dict["llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx)]), axis=-1
                 )
             else:
                 unfused_state_dict["mlp.gate_proj.weight"] = state_dict[
@@ -342,11 +347,7 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
                 ]
                 unfused_state_dict["mlp.up_proj.weight"] = state_dict["llama.layers.{}.mlp.up_proj.weight".format(idx)]
                 concated_ffn1_weight = np.concatenate(
-                    [
-                        unfused_state_dict["mlp.gate_proj.weight"],
-                        unfused_state_dict["mlp.up_proj.weight"],
-                    ],
-                    axis=-1,
+                    [unfused_state_dict["mlp.gate_proj.weight"], unfused_state_dict["mlp.up_proj.weight"]], axis=-1
                 )
             gate_up_list = split_fn(concated_ffn1_weight)
             gate_weight_tensor = paddle.to_tensor(gate_up_list[0])
@@ -400,6 +401,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
         self.epsilon = config.rms_norm_eps
         self.max_position_embeddings = config.max_position_embeddings
         self.quant_type = config.get("quant_type", "")
+        self.return_full_hidden_states = config.get("return_full_hidden_states", False)
 
         self.rope_theta = config.rope_theta
         self.use_neox = True
@@ -485,15 +487,13 @@ class LlamaInferenceModel(LlamaPretrainedModel):
         ln_scale_attrs = [paddle.ParamAttr(name="fusellama.{}.ln_scale".format(i)) for i in range(self.num_layers)]
         qkv_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.qkv_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.qkv_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
         out_proj_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.out_proj_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.out_proj_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
@@ -503,30 +503,26 @@ class LlamaInferenceModel(LlamaPretrainedModel):
         if "fp8" in self.quant_type:
             ffn1_0_weight_attrs = [
                 paddle.ParamAttr(
-                    name="fusellama.{}.ffn1_0_weight".format(i),
-                    initializer=paddle.nn.initializer.Constant(value=0),
+                    name="fusellama.{}.ffn1_0_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
                 )
                 for i in range(self.num_layers)
             ]
             ffn1_1_weight_attrs = [
                 paddle.ParamAttr(
-                    name="fusellama.{}.ffn1_1_weight".format(i),
-                    initializer=paddle.nn.initializer.Constant(value=0),
+                    name="fusellama.{}.ffn1_1_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
                 )
                 for i in range(self.num_layers)
             ]
         else:
             ffn1_weight_attrs = [
                 paddle.ParamAttr(
-                    name="fusellama.{}.ffn1_weight".format(i),
-                    initializer=paddle.nn.initializer.Constant(value=0),
+                    name="fusellama.{}.ffn1_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
                 )
                 for i in range(self.num_layers)
             ]
         ffn2_weight_attrs = [
             paddle.ParamAttr(
-                name="fusellama.{}.ffn2_weight".format(i),
-                initializer=paddle.nn.initializer.Constant(value=0),
+                name="fusellama.{}.ffn2_weight".format(i), initializer=paddle.nn.initializer.Constant(value=0)
             )
             for i in range(self.num_layers)
         ]
@@ -621,6 +617,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
         speculate_config = SpeculateConfig(
             speculate_method=config.get("speculate_method", None),
             speculate_max_draft_token_num=config.get("speculate_max_draft_token_num", 5),
+            return_full_hidden_states=config.get("return_full_hidden_states", False),
         )
 
         hpu_config = HpuConfig(
@@ -800,12 +797,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
         from paddlenlp_ops import fused_get_rotary_embedding
 
         new_rope = fused_get_rotary_embedding(
-            input_ids,
-            position_ids,
-            self.head_dim_shape_tensor,
-            position_offset,
-            self.rope_theta,
-            self.use_neox,
+            input_ids, position_ids, self.head_dim_shape_tensor, position_offset, self.rope_theta, self.use_neox
         )
 
         with dy2st_nocheck_guard_context():
@@ -821,7 +813,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                 seq_lens=seq_lens,
                 rotary_embs=new_rope,
                 rotary_emb_dims=1,
-                time_step=(paddle.increment(paddle.shape(attention_mask)[-1], -1) if is_decoder else None),
+                time_step=paddle.increment(paddle.shape(attention_mask)[-1], -1) if is_decoder else None,
             )
         hidden_states = self.norm(hidden_states)
 
@@ -852,18 +844,14 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                 weight_scale_json_path = resolve_file_path(self.quant_model_path, "weight_scales.json")
                 if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
                     act_scale_json_path = resolve_file_path(
-                        self.quant_model_path,
-                        f"act_scales_{self.config.tensor_parallel_rank}.json",
+                        self.quant_model_path, f"act_scales_{self.config.tensor_parallel_rank}.json"
                     )
                     weight_scale_json_path = resolve_file_path(
-                        self.quant_model_path,
-                        f"weight_scales_{self.config.tensor_parallel_rank}.json",
+                        self.quant_model_path, f"weight_scales_{self.config.tensor_parallel_rank}.json"
                     )
 
                 act_scales = ActScalesLoader(
-                    act_scale_json_path,
-                    act_scale_map_dict,
-                    num_of_layers=self.config.num_hidden_layers,
+                    act_scale_json_path, act_scale_map_dict, num_of_layers=self.config.num_hidden_layers
                 )
 
                 weight_scales = PerTensorWeightScalesLoader(
@@ -894,17 +882,13 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                     weight_scale_json_path = resolve_file_path(self.quant_model_path, "weight_scales.json")
                     if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
                         act_scale_json_path = resolve_file_path(
-                            self.quant_model_path,
-                            f"act_scales_{self.config.tensor_parallel_rank}.json",
+                            self.quant_model_path, f"act_scales_{self.config.tensor_parallel_rank}.json"
                         )
                         weight_scale_json_path = resolve_file_path(
-                            self.quant_model_path,
-                            f"weight_scales_{self.config.tensor_parallel_rank}.json",
+                            self.quant_model_path, f"weight_scales_{self.config.tensor_parallel_rank}.json"
                         )
                     act_scale_loader = ActScalesLoader(
-                        act_scale_json_path,
-                        act_scale_map_dict,
-                        num_of_layers=self.config.num_hidden_layers,
+                        act_scale_json_path, act_scale_map_dict, num_of_layers=self.config.num_hidden_layers
                     )
                     weight_scales_loader = WeightScalesLoader(
                         weight_scale_json_path,
@@ -985,8 +969,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                 cache_scale_json_path = resolve_file_path(self.quant_model_path, "cachekv_scales.json")
                 if self.config.tensor_parallel_degree > 1 and not self.config.single_card_ptq:
                     cache_scale_json_path = resolve_file_path(
-                        self.quant_model_path,
-                        f"cachekv_scales_{self.config.tensor_parallel_rank}.json",
+                        self.quant_model_path, f"cachekv_scales_{self.config.tensor_parallel_rank}.json"
                     )
                 cache_scales_loader = CacheScaleLoader(
                     cache_scale_json_path,
@@ -1125,9 +1108,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
             if self.use_weight_only:
                 qkv_weight_tensor = paddle.transpose(qkv_weight_tensor, perm=[1, 0])
                 qkv_quanted_weight_tensor, qkv_weight_scale_tensor = weight_quantize(
-                    qkv_weight_tensor.cpu(),
-                    algo=self.quant_algo,
-                    group_size=self.weightonly_group_size,
+                    qkv_weight_tensor.cpu(), algo=self.quant_algo, group_size=self.weightonly_group_size
                 )
                 self.transformer_block.qkv_weights[idx].set_value(qkv_quanted_weight_tensor.cuda())
                 self.transformer_block.qkv_weights_scale[idx].set_value(qkv_weight_scale_tensor.cuda())
@@ -1145,9 +1126,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
             ).cast(paddle.get_default_dtype())
             if self.use_weight_only:
                 linear_quanted_weight_tensor, linear_weight_scale_tensor = weight_quantize(
-                    linear_weight_tensor.cpu(),
-                    algo=self.quant_algo,
-                    group_size=self.weightonly_group_size,
+                    linear_weight_tensor.cpu(), algo=self.quant_algo, group_size=self.weightonly_group_size
                 )
                 self.transformer_block.linear_weights[idx].set_value(linear_quanted_weight_tensor.cuda())
                 self.transformer_block.linear_weights_scale[idx].set_value(linear_weight_scale_tensor.cuda())
@@ -1194,8 +1173,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
 
             if "llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx) in state_dict.keys():
                 concated_ffn1_weight = np.concatenate(
-                    split_fn(state_dict["llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx)]),
-                    axis=-1,
+                    split_fn(state_dict["llama.layers.{}.mlp.gate_up_fused_proj.weight".format(idx)]), axis=-1
                 )
             else:
                 unfused_state_dict["mlp.gate_proj.weight"] = state_dict[
@@ -1203,19 +1181,13 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                 ]
                 unfused_state_dict["mlp.up_proj.weight"] = state_dict["llama.layers.{}.mlp.up_proj.weight".format(idx)]
                 concated_ffn1_weight = np.concatenate(
-                    [
-                        unfused_state_dict["mlp.gate_proj.weight"],
-                        unfused_state_dict["mlp.up_proj.weight"],
-                    ],
-                    axis=-1,
+                    [unfused_state_dict["mlp.gate_proj.weight"], unfused_state_dict["mlp.up_proj.weight"]], axis=-1
                 )
             ffn1_weight_tensor = paddle.to_tensor(concated_ffn1_weight).cast(paddle.get_default_dtype())
 
             if self.use_weight_only:
                 ffn1_quanted_weight_tensor, ffn1_weight_scale_tensor = weight_quantize(
-                    ffn1_weight_tensor.cpu(),
-                    algo=self.quant_algo,
-                    group_size=self.weightonly_group_size,
+                    ffn1_weight_tensor.cpu(), algo=self.quant_algo, group_size=self.weightonly_group_size
                 )
                 self.transformer_block.ffn1_weights[idx].set_value(ffn1_quanted_weight_tensor.cuda())
                 self.transformer_block.ffn1_weights_scale[idx].set_value(ffn1_weight_scale_tensor.cuda())
@@ -1242,10 +1214,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                     )
                 else:
                     self.transformer_block.ffn1_weights[idx].set_value(
-                        paddle.cast(
-                            paddle.to_tensor(concated_ffn1_weight).transpose((1, 0)),
-                            w_dtype,
-                        )
+                        paddle.cast(paddle.to_tensor(concated_ffn1_weight).transpose((1, 0)), w_dtype)
                     )
             else:
                 self.transformer_block.ffn1_weights[idx].set_value(ffn1_weight_tensor)
@@ -1255,9 +1224,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
             )
             if self.use_weight_only:
                 ffn2_quanted_weight_tensor, ffn2_weight_scale_tensor = weight_quantize(
-                    ffn2_weight_tensor.cpu(),
-                    algo=self.quant_algo,
-                    group_size=self.weightonly_group_size,
+                    ffn2_weight_tensor.cpu(), algo=self.quant_algo, group_size=self.weightonly_group_size
                 )
                 self.transformer_block.ffn2_weights[idx].set_value(ffn2_quanted_weight_tensor.cuda())
                 self.transformer_block.ffn2_weights_scale[idx].set_value(ffn2_weight_scale_tensor.cuda())
@@ -1277,8 +1244,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                 if paddle.is_compiled_with_rocm():
                     self.transformer_block.ffn2_weights[idx].set_value(
                         paddle.cast(
-                            paddle.to_tensor(state_dict["llama.layers.{}.mlp.down_proj.weight".format(idx)]),
-                            w_dtype,
+                            paddle.to_tensor(state_dict["llama.layers.{}.mlp.down_proj.weight".format(idx)]), w_dtype
                         )
                     )
                 else:
@@ -1344,8 +1310,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                     if self.use_fake_parameter:
                         if "llama.layers.{}.input_layernorm.bias".format(idx) not in state_dict:
                             state_dict["llama.layers.{}.input_layernorm.bias".format(idx)] = paddle.zeros(
-                                shape=[self.hidden_size],
-                                dtype=paddle.get_default_dtype(),
+                                shape=[self.hidden_size], dtype=paddle.get_default_dtype()
                             )
                             state_dict["llama.layers.{}.post_attention_layernorm.bias".format(idx)] = paddle.zeros(
                                 [self.hidden_size], dtype=paddle.get_default_dtype()
@@ -1363,12 +1328,10 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                                 dtype=paddle.get_default_dtype(),
                             )
                             unfused_state_dict["mlp.gate_proj.bias"] = paddle.zeros(
-                                shape=[self.intermediate_size],
-                                dtype=paddle.get_default_dtype(),
+                                shape=[self.intermediate_size], dtype=paddle.get_default_dtype()
                             )
                             unfused_state_dict["mlp.up_proj.bias"] = paddle.zeros(
-                                shape=[self.intermediate_size],
-                                dtype=paddle.get_default_dtype(),
+                                shape=[self.intermediate_size], dtype=paddle.get_default_dtype()
                             )
                     else:
                         unfused_state_dict["self_attn.q_proj.bias"] = state_dict[
@@ -1406,11 +1369,7 @@ class LlamaInferenceModel(LlamaPretrainedModel):
                     self.transformer_block.qkv_biases[idx].set_value(paddle.to_tensor(concated_qkv_biases))
 
                     concated_ffn1_bias = np.concatenate(
-                        [
-                            unfused_state_dict["mlp.gate_proj.bias"],
-                            unfused_state_dict["mlp.up_proj.bias"],
-                        ],
-                        axis=-1,
+                        [unfused_state_dict["mlp.gate_proj.bias"], unfused_state_dict["mlp.up_proj.bias"]], axis=-1
                     )
 
                     self.transformer_block.ffn1_biases[idx].set_value(paddle.to_tensor(concated_ffn1_bias))
@@ -1458,23 +1417,10 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
         token_num = paddle.sum(seq_lens_this_time)
         from paddlenlp_ops import get_padding_offset_v2
 
-        cum_offsets_now = paddle.cast(cum_offsets_now, dtype="int32")
-
         ids_remove_padding, cum_offsets, padding_offset, cu_seqlens_q, cu_seqlens_k = get_padding_offset_v2(
-            input_ids,
-            cum_offsets_now,
-            token_num,
-            seq_lens_this_time,
-            draft_tokens,
-            seq_lens_encoder,
-        )
-        return (
-            ids_remove_padding,
-            padding_offset,
-            cum_offsets,
-            cu_seqlens_q,
-            cu_seqlens_k,
+            input_ids, cum_offsets_now, token_num, seq_lens_this_time, draft_tokens, seq_lens_encoder
         )
+        return ids_remove_padding, padding_offset, cum_offsets, cu_seqlens_q, cu_seqlens_k
 
     def forward(
         self,
@@ -1533,21 +1479,13 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
 
             # whether speculative decoding or not
             if draft_tokens is None:
-                (
-                    ids_remove_padding,
-                    padding_offset,
-                    cum_offsets,
-                    cu_seqlens_q,
-                    cu_seqlens_k,
-                ) = self.remove_padding(input_ids, seq_lens_this_time)
+                ids_remove_padding, padding_offset, cum_offsets, cu_seqlens_q, cu_seqlens_k = self.remove_padding(
+                    input_ids, seq_lens_this_time
+                )
             else:
-                (
-                    ids_remove_padding,
-                    padding_offset,
-                    cum_offsets,
-                    cu_seqlens_q,
-                    cu_seqlens_k,
-                ) = self.remove_padding(input_ids, seq_lens_this_time, draft_tokens, seq_lens_encoder)
+                ids_remove_padding, padding_offset, cum_offsets, cu_seqlens_q, cu_seqlens_k = self.remove_padding(
+                    input_ids, seq_lens_this_time, draft_tokens, seq_lens_encoder
+                )
 
             kwargs["cu_seqlens_q"] = cu_seqlens_q
             kwargs["cu_seqlens_k"] = cu_seqlens_k
@@ -1558,7 +1496,7 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
         inputs_embeds = self.embed_tokens(ids_remove_padding)
 
         with dy2st_nocheck_guard_context():
-            hidden_states, full_hidden_states = self.transformer_block(
+            hidden_states, _ = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1570,11 +1508,17 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return hidden_states, full_hidden_states
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+            cum_offsets=cum_offsets,
+        )
 
 
 @register_base_model
-class EagleForLlamaBlockInferenceModel(LlamaBlockInferenceModel):
+class EagleForLlamaInferenceModel(LlamaBlockInferenceModel):
     def __init__(self, config: LlamaConfig):
         self.append_attn = config.append_attn
         super().__init__(config)
@@ -1584,11 +1528,7 @@ class EagleForLlamaBlockInferenceModel(LlamaBlockInferenceModel):
 
         if config.tensor_parallel_degree > 1:
             self.fc = ColumnParallelLinear(
-                self.hidden_size * 2,
-                self.hidden_size,
-                has_bias=True,
-                gather_output=True,
-                fuse_matmul_bias=True,
+                self.hidden_size * 2, self.hidden_size, has_bias=True, gather_output=True, fuse_matmul_bias=True
             )
         else:
             self.fc = nn.Linear(self.hidden_size * 2, self.hidden_size, bias_attr=True)
@@ -1748,11 +1688,7 @@ class LlamaForCausalLMInferenceModel(GenerationInferenceModel, LlamaPretrainedMo
         super().__init__(config)
         self.llama = LlamaInferenceModel(config)
         if config.tie_word_embeddings:
-            self.lm_head = LlamaLMHead(
-                config,
-                embedding_weights=self.llama.embed_tokens.weight,
-                transpose_y=True,
-            )
+            self.lm_head = LlamaLMHead(config, embedding_weights=self.llama.embed_tokens.weight, transpose_y=True)
             self.tie_weights()
         else:
             self.lm_head = LlamaLMHead(config)
@@ -1950,11 +1886,7 @@ class LlamaForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, LlamaPr
 
         self.llama = LlamaBlockInferenceModel(config)
         if config.tie_word_embeddings:
-            self.lm_head = LlamaLMHead(
-                config,
-                embedding_weights=self.llama.embed_tokens.weight,
-                transpose_y=True,
-            )
+            self.lm_head = LlamaLMHead(config, embedding_weights=self.llama.embed_tokens.weight, transpose_y=True)
             self.tie_weights()
         else:
             self.lm_head = LlamaLMHead(config)
@@ -2146,9 +2078,8 @@ class LlamaForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, LlamaPr
         excess_blocks=None,
         draft_tokens=None,
         output_padding_offset=None,
-        **kwargs,
     ):
-        hidden_states, full_hidden_states = self.llama(
+        outputs = self.llama(
             input_ids,
             src_mask=src_mask,
             caches=caches,
@@ -2165,9 +2096,24 @@ class LlamaForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, LlamaPr
             excess_blocks=excess_blocks,
             draft_tokens=draft_tokens,
             output_padding_offset=output_padding_offset,
-            **kwargs,
         )
+        # hidden_states = outputs[0]
+        if self.return_full_hidden_states:
+            from paddlenlp_ops import rebuild_padding_v2
 
+            # full_hidden_states = outputs[1]
+            full_hidden_states = outputs[0]
+            cum_offsets = outputs[1]
+            hidden_states = rebuild_padding_v2(
+                full_hidden_states,
+                cum_offsets,
+                seq_lens_decoder,
+                seq_lens_encoder,
+                output_padding_offset,
+                self.max_seq_len,
+            )
+        else:
+            hidden_states = outputs[0]
         logits = self.lm_head(
             hidden_states,
             tensor_parallel_output=False,
@@ -2193,13 +2139,9 @@ class EagleLlamaForCausalLMBlockInferenceModel(LlamaForCausalLMBlockInferenceMod
         self.verify_window = config.get("speculate_verify_window", 2)
         self.max_seq_len = config.max_seq_len
 
-        self.eagle = EagleForLlamaBlockInferenceModel(config)
+        self.eagle = EagleForLlamaInferenceModel(config)
         if config.tie_word_embeddings:
-            self.lm_head = LlamaLMHead(
-                config,
-                embedding_weights=self.llama.embed_tokens.weight,
-                transpose_y=True,
-            )
+            self.lm_head = LlamaLMHead(config, embedding_weights=self.llama.embed_tokens.weight, transpose_y=True)
             self.tie_weights()
         else:
             self.lm_head = LlamaLMHead(config)

paddlenlp/experimental/transformers/mistral/__init__.py

@@ -1,15 +0,0 @@
-# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .modeling import *

paddlenlp/experimental/transformers/qwen2/modeling.py

@@ -54,6 +54,7 @@ from paddlenlp.transformers import Qwen2Config, Qwen2PretrainedModel
 from paddlenlp.transformers.conversion_utils import split_param_func
 from paddlenlp.transformers.model_outputs import (  # CausalLMOutputWithCrossAttentions,
     BaseModelOutputWithPast,
+    BaseModelOutputWithPastAndCrossAttentions,
     CausalLMOutputWithPast,
 )
 from paddlenlp.transformers.model_utils import (
@@ -1324,7 +1325,7 @@ class Qwen2BlockInferenceModel(Qwen2InferenceModel):
                 inputs_embeds = inputs_embeds.reshape([-1, inputs_embeds.shape[2]])
 
         with dy2st_nocheck_guard_context():
-            hidden_states, full_hidden_states = self.transformer_block(
+            hidden_states, _ = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1336,7 +1337,12 @@ class Qwen2BlockInferenceModel(Qwen2InferenceModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return hidden_states, full_hidden_states
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
 
 
 class Qwen2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, Qwen2PretrainedModel):
@@ -1353,7 +1359,6 @@ class Qwen2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, Qwen2Pr
         self.max_candidate_len = config.get("speculate_max_candidate_len", 5)
         self.verify_window = config.get("speculate_verify_window", 2)
         self.max_seq_len = config.max_seq_len
-        self.return_full_hidden_states = config.get("return_full_hidden_states", False)
 
         self.qwen2 = Qwen2BlockInferenceModel(config, base_model_prefix)
         if config.tie_word_embeddings:
@@ -1550,7 +1555,7 @@ class Qwen2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, Qwen2Pr
         draft_tokens=None,
         output_padding_offset=None,
     ):
-        hidden_states, full_hidden_states = self.qwen2(
+        outputs = self.qwen2(
             input_ids,
             inputs_embeds=inputs_embeds,
             src_mask=src_mask,
@@ -1570,15 +1575,13 @@ class Qwen2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, Qwen2Pr
             output_padding_offset=output_padding_offset,
         )
 
+        hidden_states = outputs[0]
         logits = self.lm_head(
             hidden_states,
             tensor_parallel_output=False,
         )
 
-        if self.return_full_hidden_states:
-            return logits, full_hidden_states
-        else:
-            return logits
+        return logits
 
     @paddle.no_grad()
     def set_state_dict(self, state_dict):

paddlenlp/mergekit/merge_config.py

@@ -69,7 +69,6 @@ class MergeConfig:
     rescale: bool = field(default=True, metadata={"help": "Rescale the weights after sparsifying."})
     reserve_p: float = field(default=0.7, metadata={"help": "Random reserve probability for the sparsify model."})
     epsilon: float = field(default=0.14, metadata={"help": "Epsilon value for magprune."})
-    remove_keys: List[str] = field(default=None, metadata={"help": "Keys to be removed from the model."})
 
     def __post_init__(self):
         self.config_check()
@@ -117,9 +116,6 @@ class MergeConfig:
                         f"Error: reserve_p +- epsilon/2 must be in the range (0, 1). reserve_p + epsilon/2 = {self.reserve_p + self.epsilon / 2 }, reserve_p - epsilon/2 = {self.reserve_p - self.epsilon / 2 }"
                     )
 
-        if self.remove_keys is None:
-            self.remove_keys = []
-
     @property
     def __dict__(self):
         return asdict(self)

paddlenlp/mergekit/merge_model.py

@@ -15,7 +15,7 @@ import json
 import math
 import os
 import shutil
-from multiprocessing import Process, Queue
+from multiprocessing import Process
 
 import numpy as np
 import paddle
@@ -54,11 +54,6 @@ SPARSIFY_MERGE_MAPPING = {
 }
 
 
-def run_merge_in_process(merger_obj, base_index, shard_file, lora_config, file_type_list, key_list, q):
-    remove_nbytes = merger_obj.shard_lora_merge(base_index, shard_file, lora_config, file_type_list, key_list)
-    q.put(remove_nbytes)
-
-
 class MergeModel:
     def __init__(self, merge_config):
         self.reset_merge_model(merge_config=merge_config)
@@ -87,7 +82,6 @@ class MergeModel:
         # init merge method
         sparsify_method = SparsifyMethod(self.merge_config)
         self.merge_method = MergeMethod(merge_config, sparsify_method)
-        self.remove_keys = getattr(merge_config, "remove_keys", [])
 
     def merge_model(self):
         if self.merge_config.lora_model_path is not None:
@@ -109,10 +103,7 @@ class MergeModel:
             else:
                 src_path = self.merge_config.model_path_list[0]
             for file in self.merge_config.copy_file_list:
-                if file == "config.json":
-                    src_file = os.path.join(self.merge_config.lora_model_path, file)
-                else:
-                    src_file = os.path.join(src_path, file)
+                src_file = os.path.join(src_path, file)
                 dst_file = os.path.join(self.merge_config.output_path, file)
                 if os.path.isfile(src_file):
                     shutil.copy2(src_file, dst_file)
@@ -244,7 +235,7 @@ class MergeModel:
             logger.info(f"Model index file saved in {save_index_file}.")
             self.merge_config.save_pretrained(self.merge_config.output_path)
 
-    def get_model_state_dict(self, model_path, file_type, key_list=None, file=None, need_remove=False):
+    def get_model_state_dict(self, model_path, file_type, key_list=None, file=None):
         if file_type == "safetensors":
             state_dict = {}
             with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
@@ -293,14 +284,6 @@ class MergeModel:
                         state_dict[k] = f.get_tensor(k)
         else:
             raise ValueError(f"Unsupported file_type: {file_type}")
-
-        if need_remove:
-            remove_nbytes = 0
-            for k in list(state_dict.keys()):
-                if k in self.remove_keys:
-                    tensor = state_dict.pop(k)
-                    remove_nbytes += int(np.prod(tensor.shape) * self.numpy_dtype_map[str(tensor.dtype)])
-            return state_dict, remove_nbytes
         return state_dict
 
     def get_safetensor_index(self, model_path, file_type):
@@ -574,12 +557,8 @@ class MergeModel:
         merge_state_dict = {}
         lora_state_dict = self.get_model_state_dict(self.merge_config.lora_model_path, file_type_list[0])
         logger.info("Load LoRA weight successfully.")
-        base_state_dict, remove_nbytes = self.get_model_state_dict(
-            self.merge_config.base_model_path,
-            file_type_list[1],
-            key_list=key_list,
-            file=file,
-            need_remove=True,
+        base_state_dict = self.get_model_state_dict(
+            self.merge_config.base_model_path, file_type_list[1], key_list=key_list, file=file
         )
         logger.info("Load model weight successfully.")
         if not lora_config.rslora:
@@ -621,7 +600,6 @@ class MergeModel:
             metadata={"format": "np"},
         )
         logger.info(f"Model weights saved in {save_file_name}.")
-        return remove_nbytes
 
     def merge_safetensor_lora_model(self, file_type_list):
         # Load index
@@ -635,66 +613,45 @@ class MergeModel:
         index["weight_map"] = {}
 
         # LoRA Merge
-        old_key_list = list(base_index["weight_map"].keys())
-        key_list = [k for k in old_key_list if k not in self.remove_keys]
-        total_remove_nbytes = 0
+        key_list = list(base_index["weight_map"].keys())
         if not self.is_cpu:
             rank = dist.get_rank()
             file_list = sorted(list(set(base_index["weight_map"].values())))
-
             if file_type_list[-1] == "safetensors" and len(file_list) >= dist.get_world_size():
                 positions = divide_positions(len(file_list), dist.get_world_size())
                 logger.info(f"Merging file list: {file_list[positions[rank] : positions[rank + 1]]}")
                 for shard_file in file_list[positions[rank] : positions[rank + 1]]:
                     logger.info(f"Start merging tensor in {shard_file}")
-                    remove_nbytes = self.shard_lora_merge(
-                        base_index, shard_file, lora_config, file_type_list, file=shard_file
-                    )
-                    total_remove_nbytes += remove_nbytes
+                    self.shard_lora_merge(base_index, shard_file, lora_config, file_type_list, file=shard_file)
                 index["weight_map"] = base_index["weight_map"]
-                for k in list(index["weight_map"].keys()):
-                    if k in self.remove_keys:
-                        del index["weight_map"][k]
             else:
                 divided_key_list = divide_lora_key_list(key_list, dist.get_world_size(), lora_config)
                 local_keys = divided_key_list[rank]
                 shard_file = (
                     f"{self.merge_config.merge_prefix}-{rank+1:05d}-of-{dist.get_world_size():05d}.safetensors"
                 )
-                remove_nbytes = self.shard_lora_merge(
-                    base_index, shard_file, lora_config, file_type_list, key_list=local_keys
-                )
-                total_remove_nbytes += remove_nbytes
+                self.shard_lora_merge(base_index, shard_file, lora_config, file_type_list, key_list=local_keys)
                 for i in range(len(divided_key_list)):
                     shard_file = (
                         f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{dist.get_world_size():05d}.safetensors"
                     )
                     for k in divided_key_list[i]:
                         index["weight_map"][k] = shard_file
-
-            # update total_remove_nbytes in main process
-            total_remove_nbytes_tensor = paddle.to_tensor(total_remove_nbytes)
-            if dist.get_world_size() > 1:
-                dist.reduce(total_remove_nbytes_tensor, dst=0)  # op=ReduceOp.SUM
-            total_remove_nbytes = total_remove_nbytes_tensor.item()
         else:
             divided_key_list = divide_lora_key_list(key_list, self.merge_config.n_process, lora_config)
             threads = []
-            queue = Queue()
             for i in range(len(divided_key_list)):
                 shard_file = (
                     f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{self.merge_config.n_process:05d}.safetensors"
                 )
                 t = Process(
-                    target=run_merge_in_process,
+                    target=self.shard_lora_merge,
                     args=(
-                        self,
                         base_index,  # base index
                         shard_file,  # shard_file name
                         lora_config,
                         file_type_list,
                         divided_key_list[i],  # key_list
-                        queue,
                     ),
                 )
                 threads.append(t)
@@ -706,12 +663,8 @@ class MergeModel:
             for t in threads:
                 t.join()
 
-            while not queue.empty():
-                total_remove_nbytes += queue.get()
-
         # Save safe index file
         if paddle.distributed.get_rank() == 0:
-            index["metadata"]["total_size"] -= total_remove_nbytes
             save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
             with open(save_index_file, "w", encoding="utf-8") as f:
                 f.write(json.dumps(index, indent=2) + "\n")
@@ -722,11 +675,7 @@ class MergeModel:
         # Load & check state dict
         lora_state_dict = self.get_model_state_dict(self.merge_config.lora_model_path, file_type_list[0])
         logger.info("Load LoRA weight successfully.")
-        base_state_dict, remove_nbytes = self.get_model_state_dict(
-            self.merge_config.base_model_path,
-            file_type_list[1],
-            need_remove=True,
-        )
+        base_state_dict = self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[1])
         logger.info("Load model weight successfully.")
         for key in lora_state_dict.keys():
             if "lora_A" in key:

paddlenlp/peft/__init__.py

@@ -12,10 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .dislora import DisLoRAConfig, DisLoRALinear, DisLoRAModel
 from .lokr import LoKrConfig, LoKrModel
 from .lora import LoRAAutoConfig, LoRAAutoModel, LoRAConfig, LoRAModel
 from .prefix import PrefixConfig, PrefixModelForCausalLM
 from .reft import ReFTModel
-from .tare import TAREModel
 from .vera import VeRAConfig, VeRAModel

paddlenlp/peft/dislora/__init__.py

@@ -1,19 +0,0 @@
-# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .dislora_config import DisLoRAConfig
-from .dislora_layer import DisLoRALinear
-from .dislora_model import DisLoRAModel
-
-__all__ = ["DisLoRAConfig", "DisLoRAModel", "DisLoRALinear"]

paddlenlp/peft/dislora/dislora_config.py

@@ -1,160 +0,0 @@
-# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-import os
-from dataclasses import asdict, dataclass, field
-from typing import List, Optional, Union
-
-from ...utils.env import DISLORA_CONFIG_NAME
-
-
-@dataclass
-class DisLoRAConfig:
-    """
-    This is the configuration class to store the configuration of a [`DisLoRAModel`].
-    Args:
-        target_modules (`Union[List[str],str]`): The names of the modules to apply DisLoRA to.
-        trainable_modules (`List[str]`): The names of the modules to train when applying DisLoRA.
-        dislora_alpha (`float`): The alpha parameter for DisLoRA scaling.
-        merge_weights (`bool`):
-            Whether to merge the weights of the DisLoRA layers with the base transfoisrmer model in `eval` mode.
-    """
-
-    base_model_name_or_path: Optional[str] = field(
-        default=None, metadata={"help": "The name of the base model to use."}
-    )
-    r: int = field(default=8, metadata={"help": "DisLoRA attention dimension"})
-    target_modules: Optional[Union[List[str], str]] = field(
-        default=None,
-        metadata={
-            "help": "List of module names or regex expression of the module names to replace with DisLoRA."
-            "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$' "
-        },
-    )
-    trainable_modules: Optional[List[str]] = field(
-        default=None,
-        metadata={
-            "help": "List of module names or regex expression of the module names to train when applying with DisLoRA."
-            "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$' "
-        },
-    )
-    dislora_alpha: int = field(default=12, metadata={"help": "DisLoRA alpha"})
-    dislora_dropout: float = field(default=0.0, metadata={"help": "DisLoRA dropout"})
-    merge_weights: bool = field(
-        default=False, metadata={"help": "Merge weights of the original model and the DisLoRA model"}
-    )
-    trainable_bias: Optional[str] = field(
-        default=None, metadata={"help": "Define trainable bias parameters for the DisLoRA model."}
-    )
-
-    tensor_parallel_degree: int = field(default=-1, metadata={"help": "1 for not use tensor parallel"})
-    dtype: Optional[str] = field(default=None, metadata={"help": "The data type of tensor"})
-
-    dash_flag: int = field(  # characteristic
-        default=50,
-        metadata={"help": "The number of preheating steps before introducing additional low-rank updates"},
-    )
-
-    s_tsd: int = field(  # characteristic
-        default=8,
-        metadata={"help": "The number of top-k singular vectors dynamically selected after preheating"},
-    )
-
-    ortho_lambda: float = field(  # characteristic
-        default=1,
-        metadata={"help": "The weight of orthogonal regularization loss"},
-    )
-    prefer_small_sigma: bool = field(
-        default=True,
-        metadata={"help": "Whether to prioritize the smallest singular value in the top-k selection process"},
-    )
-
-    def __post_init__(self):
-
-        if self.target_modules is None:
-            raise ValueError("The target_modules must be specified as a string or a list of strings.")
-        if self.r <= 0:
-            raise ValueError("The rank r of LoRA must be greater than 0.")
-        if self.dislora_alpha <= 0:
-            raise ValueError("dislora_alpha must be greater than 0")
-        if self.r < self.s_tsd:
-            raise ValueError("The rank r of LoRA must be larger than the number of top-k singular values.")
-
-    @property
-    def scaling(self):
-        return self.dislora_alpha / self.r
-
-    @property
-    def __dict__(self):
-        return asdict(self)
-
-    def to_dict(self):
-        return self.__dict__
-
-    def save_pretrained(self, save_directory):
-        r"""
-        This method saves the configuration of your adapter model in a directory.
-        Args:
-            save_directory (`str`):
-                The directory where the configuration will be saved.
-        """
-        if os.path.isfile(save_directory):
-            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
-
-        os.makedirs(save_directory, exist_ok=True)
-
-        output_dict = self.__dict__
-        output_dict["scaling"] = self.scaling
-        output_path = os.path.join(save_directory, DISLORA_CONFIG_NAME)
-
-        # save it
-        with open(output_path, "w") as writer:
-            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-        r"""
-        This method loads the configuration of your adapter model from a directory.
-        Args:
-            pretrained_model_name_or_path (`str`):
-                The directory or the hub-id where the configuration is saved.
-            **kwargs:
-                Additional keyword arguments passed along to the child class initialization.
-        """
-        if os.path.isfile(os.path.join(pretrained_model_name_or_path, DISLORA_CONFIG_NAME)):
-            config_file = os.path.join(pretrained_model_name_or_path, DISLORA_CONFIG_NAME)
-        else:
-            raise ValueError(f"Can't find dislora_config.json at '{pretrained_model_name_or_path}'")
-
-        loaded_attributes = cls.from_json_file(config_file)
-        loaded_attributes.pop("scaling", None)
-
-        merged_kwargs = {**loaded_attributes, **kwargs}
-        config = cls(**merged_kwargs)
-
-        return config
-
-    @classmethod
-    def from_json_file(cls, path_json_file):
-        r"""
-        Loads a configuration file from a json file.
-        Args:
-            path_json_file (`str`):
-                The path to the json file.
-        """
-        with open(path_json_file, "r") as file:
-            json_object = json.load(file)
-
-        return json_object

paddlenlp/peft/dislora/dislora_layer.py

@@ -1,312 +0,0 @@
-# Copyright (c) 2023 PaddlePaddle Authors. 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
-
-
-import warnings
-from typing import Union
-
-import paddle
-import paddle.nn as nn
-
-
-class DisLoRALinear(nn.Linear):
-    """
-    Paddle implementation of Direct Low-Rank Adaptation (DisLoRA) layer.
-    DisLoRA decomposes W into backbone (W_prin) and task-specific (W_res) subspaces via SVD,
-    further identifying task-specific directions (W_TSD) for fine tuning.
-    """
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        r: int = 8,
-        dislora_alpha: int = 8,
-        dislora_dropout: float = 0.0,
-        dash_flag: int = 50,
-        s_tsd: int = 8,
-        prefer_small_sigma: bool = True,
-        merge_weights: bool = False,
-        init_lora_weights: Union[bool, str] = True,
-        **kwargs
-    ):
-
-        if r <= 0:
-            raise ValueError(f"`r` must be a positive integer, got {r}")
-        if s_tsd <= 0:
-            raise ValueError(f"`s_tsd` must be a positive integer, got {s_tsd}")
-
-        nn.Linear.__init__(self, in_features, out_features, **kwargs)
-
-        original_weight = self.weight.clone()
-        original_bias = self.bias.clone() if self.bias is not None else None
-
-        self.base_dtype = original_weight.dtype
-
-        delattr(self, "weight")
-        if hasattr(self, "bias") and self.bias is not None:
-            delattr(self, "bias")
-
-        self.weight = self.create_parameter(
-            shape=[in_features, out_features],
-            default_initializer=nn.initializer.Assign(original_weight),
-            dtype=self.base_dtype,
-            attr=paddle.ParamAttr(trainable=False),
-        )
-
-        if original_bias is not None:
-            self.bias = self.create_parameter(
-                shape=[out_features],
-                default_initializer=nn.initializer.Assign(original_bias),
-                dtype=self.base_dtype,
-                attr=paddle.ParamAttr(trainable=True),
-            )
-        else:
-            self.bias = None
-
-        self.r = r
-        self.dislora_alpha = dislora_alpha
-        self.scaling = dislora_alpha / r
-        self.dislora_dropout = nn.Dropout(p=dislora_dropout) if dislora_dropout > 0.0 else nn.Identity()
-        self.dash_flag = dash_flag
-        self.s_tsd = s_tsd
-        self.prefer_small_sigma = prefer_small_sigma
-        self.merge_weights = merge_weights
-        self.init_lora_weights = init_lora_weights
-
-        self._disable_adapters = False
-        self.merged = False
-
-        self.register_buffer("step", paddle.to_tensor(0, dtype="int64"))
-
-        self.U = None
-        self.S = None
-        self.Vh = None
-
-        self.Direc_Ur = nn.Linear(r, out_features, bias_attr=False)
-        self.Direc_Sr = self.create_parameter(
-            shape=[r], default_initializer=nn.initializer.Constant(0.0), dtype=self.base_dtype
-        )
-        self.Direc_Vhr = nn.Linear(in_features, r, bias_attr=False)
-        self.Direc_Ur.weight.stop_gradient = False
-        self.Direc_Sr.stop_gradient = False
-        self.Direc_Vhr.weight.stop_gradient = False
-
-        self.Direc_Utsd = nn.Linear(s_tsd, out_features, bias_attr=False)
-        self.Direc_Stsd = self.create_parameter(
-            shape=[s_tsd], default_initializer=nn.initializer.Constant(0.0), dtype=self.base_dtype
-        )
-        self.Direc_Vhtsd = nn.Linear(in_features, s_tsd, bias_attr=False)
-
-        self.Direc_Utsd.weight.stop_gradient = True
-        self.Direc_Vhtsd.weight.stop_gradient = True
-
-        self._align_dtypes()
-
-        if init_lora_weights:
-            self._init_lora_weights()
-
-    def _align_dtypes(self):
-        """Ensure that the data types of all parameters are consistent with those of the base layer."""
-        target_dtype = self.base_dtype
-
-        if self.Direc_Ur.weight.dtype != target_dtype:
-            self.Direc_Ur.weight.set_value(self.Direc_Ur.weight.astype(target_dtype))
-        if self.Direc_Vhr.weight.dtype != target_dtype:
-            self.Direc_Vhr.weight.set_value(self.Direc_Vhr.weight.astype(target_dtype))
-        if self.Direc_Utsd.weight.dtype != target_dtype:
-            self.Direc_Utsd.weight.set_value(self.Direc_Utsd.weight.astype(target_dtype))
-        if self.Direc_Vhtsd.weight.dtype != target_dtype:
-            self.Direc_Vhtsd.weight.set_value(self.Direc_Vhtsd.weight.astype(target_dtype))
-        if self.Direc_Sr.dtype != target_dtype:
-            self.Direc_Sr.set_value(self.Direc_Sr.astype(target_dtype))
-        if self.Direc_Stsd.dtype != target_dtype:
-            self.Direc_Stsd.set_value(self.Direc_Stsd.astype(target_dtype))
-
-    def _init_lora_weights(self):
-        """
-        Initialize LoRA weights using SVD
-        Decompose the original weight W into W_prin (frozen backbone) + W_res (trainable residual)
-        Note: The shape of the Linear weight in PaddlePaddle is [in_features, out_features]
-        """
-        weight_float32 = self.weight.astype("float32")
-
-        weight_transposed = weight_float32.T
-
-        U, S, Vh = paddle.linalg.svd(weight_transposed, full_matrices=False)
-
-        self.U = U.astype(self.base_dtype)
-        self.S = S.astype(self.base_dtype)
-        self.Vh = Vh.astype(self.base_dtype)
-
-        if self.prefer_small_sigma:
-            _, indices = paddle.topk(S, self.r, largest=False)
-        else:
-            _, indices = paddle.topk(S, self.r, largest=True)
-
-        self.Direc_Ur.weight.set_value(U[:, indices].T.astype(self.base_dtype))
-        self.Direc_Sr.set_value(S[indices].astype(self.base_dtype))
-
-        self.Direc_Vhr.weight.set_value(Vh[indices, :].T.astype(self.base_dtype))
-        self.Direc_Ur.weight.stop_gradient = False
-        self.Direc_Sr.stop_gradient = False
-        self.Direc_Vhr.weight.stop_gradient = False
-        self.Direc_Stsd.stop_gradient = False
-
-        S_diag = paddle.diag(self.Direc_Sr)  # [r, r]
-        W_res_T = self.Direc_Ur.weight.T @ S_diag @ self.Direc_Vhr.weight.T  # [out_features, in_features]
-        W_res = W_res_T.T * self.scaling  # [in_features, out_features]
-
-        if W_res.shape != self.weight.shape:
-            raise ValueError(f"Expected W_res shape {self.weight.shape}, but got {W_res.shape}.")
-
-        self.weight.set_value(self.weight - W_res.astype(self.base_dtype))
-        self.weight.stop_gradient = True
-
-    def forward(self, x: paddle.Tensor) -> paddle.Tensor:
-        """
-        Forward propagation: W_prin @ x + W_res @ x + W_TSD @ x
-        - W_prin is calculated through the base_layer
-        - W_res is calculated through the trainable LoRA structure
-        - W_TSD is calculated through the frozen dynamic vector (after warmup)
-        """
-        if self._disable_adapters:
-            if self.merged:
-                self.unmerge()
-            return super().forward(x)
-
-        if self.merged:
-            return super().forward(x)
-
-        result = super().forward(x)
-
-        temp = self.dislora_dropout(x)
-        temp = self.Direc_Vhr(temp)
-        temp = temp * self.Direc_Sr
-        temp = self.Direc_Ur(temp)
-        result += temp * self.scaling
-
-        if self.step < self.dash_flag:
-            pass
-        elif self.step == self.dash_flag:
-            self._initialize_dynamic_vectors()
-        else:
-            temp = self.dislora_dropout(x)
-            temp = self.Direc_Vhtsd(temp)
-            temp = temp * self.Direc_Stsd
-            temp = self.Direc_Utsd(temp)
-            result += temp * self.scaling
-
-        if self.training:
-            with paddle.no_grad():
-                self.step += 1
-
-        return result
-
-    def _initialize_dynamic_vectors(self):
-        """
-        After the warm-up steps, initialize the dynamic singular vector W_TSD.
-        Based on the current change of W_res, select the most important s_tsd directions.
-        """
-        with paddle.no_grad():
-
-            S_diag = paddle.diag(self.Direc_Sr)  # [r, r]
-            deltaW_T = self.Direc_Ur.weight.T @ S_diag @ self.Direc_Vhr.weight.T  # [out_features, in_features]
-
-            delta_sigma = paddle.diag(self.U.T @ deltaW_T @ self.Vh.T)
-
-            top_indices = self.calculate_change_rate(
-                self.S, delta_sigma, self.s_tsd, largest=not self.prefer_small_sigma
-            )
-
-            self.Direc_Utsd.weight.set_value(self.U[:, top_indices].T.astype(self.base_dtype))
-            self.Direc_Stsd.set_value(self.S[top_indices].astype(self.base_dtype))
-            self.Direc_Vhtsd.weight.set_value(self.Vh[top_indices, :].T.astype(self.base_dtype))
-
-            self.Direc_Utsd.weight.stop_gradient = True
-            self.Direc_Vhtsd.weight.stop_gradient = True
-
-    def calculate_change_rate(self, a: paddle.Tensor, b: paddle.Tensor, s: int, largest: bool = True) -> paddle.Tensor:
-        """
-        Calculate the rate of change of singular values and
-        select the top-s index change_rate = |b| / (|a| + eps)
-        """
-        with paddle.no_grad():
-
-            change_rate = paddle.abs(b) / (paddle.abs(a) + 1e-8)
-
-            _, top_s_indices = paddle.topk(change_rate, s, largest=largest)
-        return top_s_indices
-
-    def merge(self):
-        """
-        Merge the trainable W_res into the base weights.
-        After merging: base_layer.weight = W_prin + W_res
-        Note: W_TSD remains frozen and does not participate in the merge.
-        """
-        if self.merged:
-            warnings.warn("Already merged. Nothing to do.")
-            return
-
-        if self.r > 0:
-
-            delta_weight = self.get_delta_weight()
-            orig_weights = self.weight.clone()
-            orig_weights += delta_weight
-            self.weight.set_value(orig_weights)
-
-        self.merged = True
-
-    def unmerge(self):
-        """
-        Remove the merging of W_res from the base weights.
-        After the merging is removed: base_layer.weight = W_prin
-        """
-        if not self.merged:
-            warnings.warn("Already unmerged. Nothing to do.")
-            return
-
-        if self.r > 0:
-            delta_weight = self.get_delta_weight()
-            self.weight.set_value(self.weight - delta_weight)
-
-        self.merged = False
-
-    def get_delta_weight(self) -> paddle.Tensor:
-        """
-        Calculate the trainable LoRA incremental weights
-        It consists of two parts:
-        1. W_res = Ur @ diag(Sr) @ Vhr * scaling (transposed)
-        2. W_tsd = Utsd @ diag(Stsd) @ Vhtsd * scaling (transposed)
-        Return the incremental weights with the shape of [in_features, out_features]
-        """
-
-        S_diag_r = paddle.diag(self.Direc_Sr)  # [r, r]
-        delta_weight_T = self.Direc_Ur.weight.T @ S_diag_r @ self.Direc_Vhr.weight.T  # [out_features, in_features]
-        delta_weight = delta_weight_T.T * self.scaling  # [in_features, out_features]
-
-        if not paddle.all(self.Direc_Stsd == 0.0):
-            S_diag_tsd = paddle.diag(self.Direc_Stsd)  # [s_tsd, s_tsd]
-            delta_weight_tsd_T = (
-                self.Direc_Utsd.weight.T @ S_diag_tsd @ self.Direc_Vhtsd.weight.T
-            )  # [out_features, in_features]
-            delta_weight += delta_weight_tsd_T.T * self.scaling  # [in_features, out_features]
-
-        return delta_weight.astype(self.base_dtype)
-
-    def enable_adapters(self):
-        """Enable the adapter"""
-        self._disable_adapters = False
-
-    def disable_adapters(self):
-        """Disable adapter"""
-        self._disable_adapters = True
-
-    def __repr__(self) -> str:
-        rep = super().__repr__()
-        return rep

paddlenlp/peft/dislora/dislora_model.py

@@ -1,446 +0,0 @@
-# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import copy
-import os
-import re
-from collections import OrderedDict
-from typing import Dict, Union
-
-import numpy as np
-import paddle
-import paddle.nn as nn
-from paddle.distributed.fleet.meta_parallel import PipelineLayer
-
-from paddlenlp.transformers import AutoConfig, PretrainedModel
-from paddlenlp.transformers.model_utils import _add_variant, dtype_guard
-from paddlenlp.utils.log import logger
-
-from ...utils.env import DISLORA_WEIGHTS_NAME
-from .dislora_config import DisLoRAConfig
-
-
-def get_dislora_layers():
-    from .dislora_layer import DisLoRALinear
-
-    return {
-        "DisLoRALinear": DisLoRALinear,
-    }
-
-
-dislora_layers = get_dislora_layers()
-DisLoRALinear = dislora_layers["DisLoRALinear"]
-AVAILABLE_LAYERS = [
-    DisLoRALinear,
-]
-
-
-class DisLoRAModel(nn.Layer):
-    restore_layer_map: Dict[nn.Layer, nn.Layer] = {
-        DisLoRALinear: nn.Linear,
-    }
-
-    def __init__(self, model, dislora_config: DisLoRAConfig) -> None:
-        super().__init__()
-        self.model_config = AutoConfig.from_pretrained(dislora_config.base_model_name_or_path)
-        self.quantized = False
-        self.dislora_config = dislora_config
-        self.dislora_split_mapping = {}
-        if self.dislora_config.dtype is None:
-            self.dislora_config.dtype = paddle.get_default_dtype()
-        with dtype_guard(self.dislora_config.dtype):
-            self.model = self.get_dislora_model(model, dislora_config)
-        self.is_pipelinemodel = False
-        if issubclass(type(self.model), PipelineLayer):
-            raise NotImplementedError("dislora don't support pipeline parallel now")
-        if dislora_config.tensor_parallel_degree > 1:
-            self.dislora_config.tensor_parallel_degree = -1
-            self.model.config.tensor_parallel_degree = -1
-            raise NotImplementedError("dislora don't support tensor parallel now")
-        # currently tensor_parallel_degree should all be set to -1.
-        self.forward = self.model.forward
-
-        logger.info("Mark only dislora and trainable_module as trainable.")
-        self.mark_only_dislora_as_trainable()
-
-    @classmethod
-    def from_pretrained(cls, model, dislora_path, **kwargs):
-        dislora_config = kwargs.pop("dislora_config", None)
-        # init dislora config & dislora model
-        if not isinstance(dislora_config, DisLoRAConfig):
-            dislora_config = DisLoRAConfig.from_pretrained(dislora_path)
-        # define a new variable to conserve original lora_config.tensor_parallel_degree value which will update while initializing lora model
-        dislora_config_tensor_parallel_degree = dislora_config.tensor_parallel_degree
-        dislora_model = cls(model, dislora_config)
-
-        # define dislora weight name
-        dislora_weight_name = DISLORA_WEIGHTS_NAME
-
-        # load and set dislora weight parameter
-        dislora_weight_path = os.path.join(dislora_path, dislora_weight_name)
-        if os.path.exists(dislora_weight_path):
-            # load dislora weight parameter
-            dislora_state_dict = paddle.load(dislora_weight_path, return_numpy=True)
-            logger.info(f"Loading the DisLoRA weights from {dislora_weight_path}")
-
-            if (
-                dislora_config_tensor_parallel_degree > 1
-                and dislora_config_tensor_parallel_degree != model.config.tensor_parallel_degree
-            ):
-                raise NotImplementedError(
-                    f"{dislora_config_tensor_parallel_degree} is not equal to {model.config.tensor_parallel_degree}. Please merge DisLoRA weights first."
-                )
-            # set dislora state dict
-            dislora_model.set_state_dict(dislora_state_dict)
-        else:
-            logger.error(f"DisLoRA weights not found under {dislora_path}, creating DisLoRA weights from scratch")
-
-        return dislora_model
-
-    def set_state_dict(self, state_dict):
-        import warnings
-
-        warnings.filterwarnings(
-            action="ignore", message=".*Skip loading for.*", category=Warning, lineno=0, append=False
-        )
-        self.model.set_state_dict(state_dict)
-        logger.info("Load dislora weight successfully")
-
-    def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = False, **kwargs):
-        logger.info("save dislora pretrained")
-        save_model_config = kwargs.get("save_model_config", True)
-
-        variant = kwargs.get("variant", None)
-        is_main_process = kwargs.get("is_main_process", paddle.distributed.get_rank() == 0)
-
-        assert not os.path.isfile(
-            save_directory
-        ), f"Saving directory ({save_directory}) should be a directory, not a file"
-        os.makedirs(save_directory, exist_ok=True)
-
-        dislora_config_to_save = DisLoRAConfig(**self.dislora_config.to_dict())
-        trainable_state_dict = self.get_trainable_state_dict()
-
-        # save dislora weight
-        dislora_weight_name = _add_variant(DISLORA_WEIGHTS_NAME, variant)
-        weight_filename = os.path.join(save_directory, dislora_weight_name)
-        paddle.save(trainable_state_dict, weight_filename)
-
-        # save dislora config
-        if is_main_process:
-            dislora_config_to_save.save_pretrained(save_directory)
-            if save_model_config:
-                model_config_to_save = copy.deepcopy(self.model.config)
-                if merge_tensor_parallel:
-                    model_config_to_save.tensor_parallel_degree = -1
-                model_config_to_save.save_pretrained(save_directory)
-
-    def _find_and_replace_module(self, model, module_name, dislora_config):
-
-        if any(dislora_keyword in module_name.lower() for dislora_keyword in ["dislora", "direc_"]):
-            logger.debug(f"Skipping {module_name} - appears to be a DisLoRA submodule")
-            return
-
-        try:
-            parent_module = model
-            attribute_chain = module_name.split(".")
-            for name in attribute_chain[:-1]:
-                parent_module = getattr(parent_module, name)
-            module = getattr(parent_module, attribute_chain[-1])
-        except AttributeError as e:
-            logger.error(f"Cannot access module {module_name}: {e}")
-            raise ValueError(f"Cannot access target module {module_name}: {e}")
-
-        if isinstance(module, nn.Linear):
-            logger.debug(f"Converting {module_name} from nn.Linear to DisLoRALinear")
-
-            try:
-                dislora_module = DisLoRALinear(
-                    in_features=module.weight.shape[0],
-                    out_features=module.weight.shape[1],
-                    r=dislora_config.r,
-                    dislora_alpha=dislora_config.dislora_alpha,
-                    dislora_dropout=dislora_config.dislora_dropout,
-                    dash_flag=dislora_config.dash_flag,
-                    s_tsd=dislora_config.s_tsd,
-                    prefer_small_sigma=dislora_config.prefer_small_sigma,
-                    merge_weights=dislora_config.merge_weights,
-                    bias_attr=False if module.bias is None else None,
-                    init_lora_weights=False,
-                )
-
-                dislora_module.weight.set_value(module.weight)
-                if module.bias is not None:
-                    dislora_module.bias.set_value(module.bias)
-
-                dislora_module._init_lora_weights()
-
-                setattr(parent_module, attribute_chain[-1], dislora_module)
-                logger.debug(f"Successfully replaced {module_name}")
-
-            except Exception as e:
-                logger.error(f"Failed to create DisLoRALinear for {module_name}: {e}")
-                raise ValueError(f"Failed to create DisLoRALinear for {module_name}: {e}")
-
-        elif isinstance(module, DisLoRALinear):
-            logger.debug(f"Module {module_name} is already a DisLoRALinear, skipping")
-
-        else:
-
-            module_type = type(module).__name__
-            if any(keyword in module_name.lower() for keyword in ["dislora_dropout", "direc_"]):
-                logger.debug(f"Skipping DisLoRA submodule {module_name} ({module_type})")
-                return
-            else:
-
-                error_msg = f"Target module {module_name} is {module_type}, not nn.Linear. DisLoRA can only replace nn.Linear modules."
-                logger.error(f"Cannot replace {module_name}: expected nn.Linear, got {module_type}")
-                raise ValueError(error_msg)
-
-    def _find_and_restore_module(self, module_name):
-        parent_module = self.model
-        attribute_chain = module_name.split(".")
-        for name in attribute_chain[:-1]:
-            parent_module = getattr(parent_module, name)
-        module = getattr(parent_module, attribute_chain[-1])
-        original_model_class = self.restore_layer_map[module.__class__]
-        original_module = original_model_class(in_features=module.weight.shape[0], out_features=module.weight.shape[1])
-        original_module.weight = module.weight
-
-        if isinstance(module, DisLoRALinear):
-            if not module.merged:
-                complete_weight = module.weight + module.get_delta_weight()
-                original_module.weight.set_value(complete_weight)
-            else:
-                original_module.weight.set_value(module.weight)
-        else:
-            original_module.weight.set_value(module.weight)
-
-        if module.bias is not None:
-            original_module.bias.set_value(module.bias)
-
-        setattr(parent_module, attribute_chain[-1], original_module)
-
-    def get_trainable_state_dict(self):
-        """
-        Obtain the required state dictionary to be saved, including:
-        1. Trainable parameters (stop_gradient = False)
-        2. Main weight W_prin (although frozen, must be saved)
-        3. TSD direction parameters (although frozen, must be saved)
-        4. QAT-related parameters
-        """
-        trainable_state_dict = OrderedDict()
-        for name, weight in self.model.state_dict().items():
-            # Save trainable parameters and QAT parameters
-            if not weight.stop_gradient or "activation_quanter" in name or "weight_quanter" in name:
-                trainable_state_dict[name] = weight
-            # Save the main branch weight W_prin (for critical fixes)
-            elif "weight" in name and any(layer_name in name for layer_name in [".weight"]) and "Direc_" not in name:
-                trainable_state_dict[name] = weight
-                logger.debug(f"Saving backbone weight: {name}")
-            # Save all TSD parameters (excluding Direc_Stsd)
-            elif any(tsd_param in name for tsd_param in ["Direc_Utsd", "Direc_Vhtsd"]):
-                trainable_state_dict[name] = weight
-                logger.debug(f"Saving TSD parameter: {name}")
-            # Save the bias parameters (if any)
-            elif "bias" in name and "Direc_" not in name:
-                trainable_state_dict[name] = weight
-                logger.debug(f"Saving bias parameter: {name}")
-
-        return trainable_state_dict
-
-    def print_trainable_parameters(self) -> None:
-        freeze_numel = 0
-        trainable_numel = 0
-        for _, weight in self.model.state_dict().items():
-            if weight.stop_gradient:
-                freeze_numel += np.prod(weight.shape)
-            else:
-                trainable_numel += np.prod(weight.shape)
-        logger.debug(
-            f"Frozen parameters: {freeze_numel:.2e} || Trainable parameters:{trainable_numel:.2e} || Total parameters:{freeze_numel+trainable_numel:.2e}|| Trainable:{trainable_numel / (freeze_numel+trainable_numel):.2%}"
-        )
-
-    def mark_only_dislora_as_trainable(self) -> None:
-        """
-        Mark only the parameters related to DisLoRA as trainable, while ensuring that the TSD parameters remain in a frozen state.
-        """
-
-        for full_param_name, weight in self.model.state_dict().items():
-
-            is_dislora_layer = any(
-                re.fullmatch(target_module, full_param_name.rsplit(".", 1)[0])
-                for target_module in self.dislora_config.target_modules
-            )
-
-            if is_dislora_layer:
-                param_name = full_param_name.split(".")[-1]
-
-                if param_name == "weight" and "Direc_" not in full_param_name:
-                    weight.stop_gradient = True
-                    logger.debug(f"Freezing backbone weight: {full_param_name}")
-
-                elif param_name == "bias" and "Direc_" not in full_param_name:
-                    if self.dislora_config.trainable_bias in ["dislora", "all"]:
-                        weight.stop_gradient = False
-                        logger.debug(f"Setting bias as trainable: {full_param_name}")
-                    else:
-                        weight.stop_gradient = True
-                        logger.debug(f"Freezing bias: {full_param_name}")
-
-                elif any(tsd_param in full_param_name for tsd_param in ["Direc_Utsd", "Direc_Vhtsd"]):
-                    weight.stop_gradient = True
-                    logger.debug(f"Keeping TSD parameter frozen: {full_param_name}")
-
-                elif any(
-                    trainable_param in full_param_name
-                    for trainable_param in ["Direc_Ur", "Direc_Sr", "Direc_Vhr", "Direc_Stsd"]
-                ):
-                    weight.stop_gradient = False
-                    logger.debug(f"Setting DisLoRA parameter as trainable: {full_param_name}")
-
-                else:
-                    weight.stop_gradient = True
-                    logger.debug(f"Freezing other parameter: {full_param_name}")
-
-            else:
-                param_name = full_param_name.split(".")[-1]
-                if self.dislora_config.trainable_bias == "all" and param_name == "bias":
-                    weight.stop_gradient = False
-                    logger.debug(f"Setting bias as trainable in non-DisLoRA layer: {full_param_name}")
-                else:
-                    weight.stop_gradient = True
-                    logger.debug(f"Freezing parameter in non-DisLoRA layer: {full_param_name}")
-
-        if self.dislora_config.trainable_modules is not None:
-            for full_param_name, weight in self.model.state_dict().items():
-                if any(
-                    re.fullmatch(trainable_module, full_param_name)
-                    for trainable_module in self.dislora_config.trainable_modules
-                ):
-
-                    if not any(tsd_param in full_param_name for tsd_param in ["Direc_Utsd", "Direc_Vhtsd"]):
-                        weight.stop_gradient = False
-                        logger.debug(f"Setting additional trainable module parameter: {full_param_name}")
-                    else:
-                        logger.warning(
-                            f"TSD parameter {full_param_name} matched trainable_modules pattern but kept frozen"
-                        )
-
-    def get_dislora_model(self, model: Union[PretrainedModel, nn.Layer], dislora_config: DisLoRAConfig):
-        """
-        Iterate all base model layers, change target modules to DisLoRALayer.
-        """
-        if dislora_config.target_modules is None:
-            return model
-        else:
-            target_modules = dislora_config.target_modules
-
-        target_module_names = []
-
-        existing_dislora_paths = set()
-        for module_name, module in model.named_sublayers():
-            if isinstance(module, DisLoRALinear):
-                existing_dislora_paths.add(module_name)
-
-        for target_module in target_modules:
-            for module_name, module in model.named_sublayers():
-
-                if re.fullmatch(target_module, module_name):
-
-                    if not isinstance(module, DisLoRALinear):
-
-                        is_submodule = any(
-                            module_name.startswith(dislora_path + ".") for dislora_path in existing_dislora_paths
-                        )
-
-                        if not is_submodule:
-                            target_module_names.append(module_name)
-                        else:
-                            logger.debug(f"Skipping {module_name} - it's a submodule of existing DisLoRA module")
-                    else:
-                        logger.debug(f"Skipping {module_name} - already a DisLoRA module")
-
-        for module_name in target_module_names:
-            try:
-                self._find_and_replace_module(model, module_name, dislora_config)
-                logger.debug(f"Replaced {module_name} with DisLoRALinear")
-            except ValueError as e:
-                raise e
-            except Exception as e:
-
-                logger.warning(f"Failed to replace {module_name}: {e}")
-
-        return model
-
-    def restore_original_model(self):
-        # make sure W and dislora weights are not merged before we restore the original model
-        for layer_name, layer in self.model.named_sublayers():
-            if isinstance(layer, DisLoRALinear):
-                self._find_and_restore_module(layer_name)
-        return self.model
-
-    def __getattr__(self, name: str):
-        """
-        Forward missing attributes to the wrapped module.
-        """
-        try:
-            return super().__getattr__(name)  # defer to nn.Layer's logic
-        except AttributeError:
-            return getattr(self.model, name)
-
-    def train(self):
-        self.training = True
-        self.model.training = True
-        for layer in self.model.sublayers():
-            layer.training = True
-            layer.train()
-
-    def eval(self):
-        self.training = False
-        self.model.training = False
-        for layer in self.model.sublayers():
-            layer.training = False
-            layer.eval()
-
-    def disable_dislora(self):
-        """
-        Disable the DisLoRA adapter
-        """
-        for _, layer in self.model.named_sublayers():
-            if isinstance(layer, DisLoRALinear):
-                layer.disable_adapters()
-
-    def enable_dislora(self):
-        """
-        Enable the DisLoRA adapter
-        """
-        for _, layer in self.model.named_sublayers():
-            if isinstance(layer, DisLoRALinear):
-                layer.enable_adapters()
-
-    def merge(self):
-        for _, layer in self.model.named_sublayers():
-            if any(isinstance(layer, dislora_layer) for dislora_layer in AVAILABLE_LAYERS):
-                layer.merge()
-
-    def unmerge(self):
-        for _, layer in self.model.named_sublayers():
-            if any(isinstance(layer, dislora_layer) for dislora_layer in AVAILABLE_LAYERS):
-                layer.unmerge()
-
-    def get_model_config(
-        self,
-    ):
-        return self.model_config.to_dict()

paddlenlp/peft/lora/lora_config.py

@@ -77,8 +77,6 @@ class LoRAConfig:
     rslora: bool = field(default=False, metadata={"help": "Whether to use RsLoRA"})
     pissa: bool = field(default=False, metadata={"help": "Whether to use Pissa: https://arxiv.org/pdf/2404.02948.pdf"})
     loraga: bool = field(default=False, metadata={"help": "Whether to LoRA-GA"})
-    nola: bool = field(default=False, metadata={"help": "Whether to use Nola: https://arxiv.org/pdf/2310.02556"})
-    nola_basis_num: int = field(default=1, metadata={"help": "When use nola, the number of basis"})
     use_mora: bool = field(
         default=False, metadata={"help": "Whether to use MoRA: https://arxiv.org/pdf/2405.12130.pdf"}
     )
@@ -96,12 +94,6 @@ class LoRAConfig:
         default=False,
         metadata={"help": "Whether to use mos lora."},
     )
-    mixer_num: int = field(
-        default=1,
-        metadata={
-            "help": "Num of mixer matrices. Mixer matrices will be added between the LoRA_A and LoRA_B matrices, as referenced in the paper https://arxiv.org/abs/2411.00039."
-        },
-    )
     lorapro: bool = field(default=False, metadata={"help": "Whether to use LoRA-PRO"})
 
     def __post_init__(self):

paddlenlp/peft/lora/lora_layers.py

@@ -63,10 +63,7 @@ class LoRALinear(nn.Linear):
         rslora: bool = False,
         lora_plus_scale: float = 1.0,
         pissa: bool = False,
-        nola: bool = False,
-        nola_basis_num: int = 1,
         lora_use_mixer: bool = False,
-        mixer_num: int = 1,
         use_mora: bool = False,
         lorapro: bool = False,
         mp_moe: bool = False,
@@ -87,10 +84,7 @@ class LoRALinear(nn.Linear):
         # Mark the weight as unmerged
         self.merged = False
         self.pissa = pissa
-        self.nola = nola
-        self.nola_basis_num = nola_basis_num
         self.lora_use_mixer = lora_use_mixer
-        self.mixer_num = mixer_num
         self.lorapro = lorapro
 
         # Actual trainable parameters
@@ -124,20 +118,14 @@ class LoRALinear(nn.Linear):
                 ),
             )
             if self.lora_use_mixer:
-                for i in range(self.mixer_num):
-                    key = "lora_mixer_" + str(i)
-                    setattr(
-                        self,
-                        key,
-                        self.create_parameter(
-                            shape=[r, r],
-                            dtype=self._dtype,
-                            is_bias=False,
-                            default_initializer=nn.initializer.KaimingUniform(
-                                negative_slope=math.sqrt(5), nonlinearity="leaky_relu"
-                            ),
-                        ),
-                    )
+                self.lora_AB = self.create_parameter(
+                    shape=[r, r],
+                    dtype=self._dtype,
+                    is_bias=False,
+                    default_initializer=nn.initializer.KaimingUniform(
+                        negative_slope=math.sqrt(5), nonlinearity="leaky_relu"
+                    ),
+                )
             self.lora_B = self.create_parameter(
                 shape=[r, out_features],
                 dtype=self._dtype,
@@ -148,32 +136,6 @@ class LoRALinear(nn.Linear):
                 ),
             )
         self.apply_pissa = False
-        if nola:
-            # Initialize placeholders for NOLA parameters
-            self.nola_basis_A = self.create_parameter(
-                shape=[nola_basis_num, in_features, r],
-                dtype=self._dtype,
-                is_bias=False,
-            )
-            self.nola_basis_A.stop_gradient = True
-            self.nola_basis_B = self.create_parameter(
-                shape=[nola_basis_num, r, out_features],
-                dtype=self._dtype,
-                is_bias=False,
-            )
-            self.nola_basis_B.stop_gradient = True
-            self.nola_alpha = self.create_parameter(
-                shape=[nola_basis_num],
-                dtype=self._dtype,
-                is_bias=False,
-                default_initializer=nn.initializer.Constant(value=0.0),
-            )
-            self.nola_beta = self.create_parameter(
-                shape=[nola_basis_num],
-                dtype=self._dtype,
-                is_bias=False,
-                default_initializer=nn.initializer.Constant(value=0.0),
-            )
         if use_mora or pissa:
             self.scaling = 1.0
         elif not rslora:
@@ -209,16 +171,6 @@ class LoRALinear(nn.Linear):
         weight = res.astype(dtype)
         self.weight.set_value(weight)
 
-    def get_nola_lora_matrices(self):
-        """Compute LoRA matrices A and B from NOLA basis and coefficients."""
-        if not self.nola:
-            return self.lora_A, self.lora_B
-        # Compute A = sum(alpha_i * A_i)
-        lora_A = paddle.einsum("k,kir->ir", self.nola_alpha, self.nola_basis_A)  # [in_features, r]
-        # Compute B = sum(beta_j * B_j)
-        lora_B = paddle.einsum("k,kro->ro", self.nola_beta, self.nola_basis_B)  # [r, out_features]
-        return lora_A, lora_B
-
     def rope_init(self):
         if self.cos is None or self.sin is None:
             inv_freq = 1.0 / (10000 ** (paddle.arange(0, self.r, 2, dtype=paddle.float32) / self.r))
@@ -269,7 +221,7 @@ class LoRALinear(nn.Linear):
         if self.lora_use_mixer:
             lora_A = lora_A if lora_A is not None else self.lora_A
             lora_B = lora_B if lora_B is not None else self.lora_B
-            lora_AB = lora_AB if lora_AB is not None else self.get_mixer_params(0)
+            lora_AB = lora_AB if lora_AB is not None else self.lora_AB
             delta_weight = lora_A @ lora_AB @ lora_B * self.scaling
         elif self.use_mora:
             lora_A = lora_A if lora_A is not None else self.lora_A
@@ -297,9 +249,6 @@ class LoRALinear(nn.Linear):
                 w = w[: self.out_features]
             final_weight = w
             delta_weight = final_weight.T
-        elif self.nola:
-            lora_A, lora_B = self.get_nola_lora_matrices()
-            delta_weight = lora_A @ lora_B * self.scaling
         else:
             lora_A = lora_A if lora_A is not None else self.lora_A
             lora_B = lora_B if lora_B is not None else self.lora_B
@@ -307,25 +256,18 @@ class LoRALinear(nn.Linear):
 
         return delta_weight
 
-    def get_mixer_params(self, index):
-        key = "lora_mixer_" + str(index)
-        if index == self.mixer_num - 1:
-            return getattr(self, key)
-        else:
-            return getattr(self, key) @ self.get_mixer_params(index + 1)
-
     def merge(self):
         if not self.merged:
             delta_weight = self.get_delta_weight()
             new_weight = self.weight + delta_weight
-            self.weight.set_value(new_weight.astype(self.weight.dtype))
+            self.weight.set_value(new_weight)
             self.merged = True
 
     def unmerge(self):
         if self.merged:
             delta_weight = self.get_delta_weight()
             new_weight = self.weight - delta_weight
-            self.weight.set_value(new_weight.astype(self.weight.dtype))
+            self.weight.set_value(new_weight)
             self.merged = False
 
     def forward(self, input: paddle.Tensor, *args, **kwargs):
@@ -342,17 +284,10 @@ class LoRALinear(nn.Linear):
             input = self.lora_dropout(input)
             mora_out = self._apply_mora(input)
             result += mora_out
-        elif self.nola:
-            result = F.linear(x=input, weight=self.weight, bias=self.bias, name=self.name)
-            input = self.lora_dropout(input)
-            lora_A, lora_B = self.get_nola_lora_matrices()
-            result += (self.lora_dropout(input) @ lora_A @ lora_B) * self.scaling
         else:
             result = F.linear(x=input, weight=self.weight, bias=self.bias, name=self.name)
             if self.lora_use_mixer:
-                result += (
-                    self.lora_dropout(input) @ self.lora_A @ self.get_mixer_params(0) @ self.lora_B
-                ) * self.scaling
+                result += (self.lora_dropout(input) @ self.lora_A @ self.lora_AB @ self.lora_B) * self.scaling
             else:
                 result += (self.lora_dropout(input) @ self.lora_A @ self.lora_B) * self.scaling
         return result
@@ -375,16 +310,14 @@ class RowParallelLoRALinear(RowParallelLinear):
         use_quick_lora: bool = False,
         pissa: bool = False,
         use_mora: bool = False,
-        nola: bool = False,
-        nola_basis_num: int = 1,
         **kwargs
     ):
         RowParallelLinear.__init__(self, in_features, out_features, **kwargs)
         if not isinstance(r, int) or r <= 0:
             raise ValueError("Lora rank r should be a positive integer")
 
-        if pissa or use_mora or nola:
-            raise ValueError("Pissa, Mora or NoLA is not supported in model parallel by now")
+        if pissa or use_mora:
+            raise ValueError("Pissa or Mora is not supported in model parallel by now")
 
         self.r = r
         self.lora_alpha = lora_alpha
@@ -643,16 +576,14 @@ class ColumnParallelLoRALinear(ColumnParallelLinear):
         use_quick_lora: bool = False,
         pissa: bool = False,
         use_mora: bool = False,
-        nola: bool = False,
-        nola_basis_num: int = 1,
         **kwargs
     ):
         ColumnParallelLinear.__init__(self, in_features, out_features, **kwargs)
         if not isinstance(r, int) or r <= 0:
             raise ValueError("Lora rank r should be a positive integer")
 
-        if pissa or use_mora or nola:
-            raise ValueError("Pissa, Mora or NoLA is not supported in model parallel by now")
+        if pissa or use_mora:
+            raise ValueError("Pissa or Mora is not supported in model parallel by now")
 
         self.r = r
         self.lora_alpha = lora_alpha

paddlenlp/peft/lora/lora_model.py

@@ -482,12 +482,9 @@ class LoRAModel(nn.Layer):
                 rslora=lora_config.rslora,
                 lora_plus_scale=lora_config.lora_plus_scale,
                 pissa=lora_config.pissa,
-                nola=lora_config.nola,
-                nola_basis_num=lora_config.nola_basis_num,
                 bias_attr=False if module.bias is None else None,
                 use_quick_lora=lora_config.use_quick_lora,
                 lora_use_mixer=lora_config.lora_use_mixer,
-                mixer_num=lora_config.mixer_num,
                 use_mora=lora_config.use_mora,
                 mp_moe=getattr(module.weight, "mp_moe", False),
                 is_distributed=getattr(module.weight, "is_distributed", False),
@@ -523,8 +520,6 @@ class LoRAModel(nn.Layer):
                 rslora=lora_config.rslora,
                 lora_plus_scale=lora_config.lora_plus_scale,
                 pissa=lora_config.pissa,
-                nola=lora_config.nola,
-                nola_basis_num=lora_config.nola_basis_num,
                 lora_A_weight_attr=paddle.ParamAttr(
                     initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
                 ),
@@ -551,8 +546,6 @@ class LoRAModel(nn.Layer):
                 rslora=lora_config.rslora,
                 lora_plus_scale=lora_config.lora_plus_scale,
                 pissa=lora_config.pissa,
-                nola=lora_config.nola,
-                nola_basis_num=lora_config.nola_basis_num,
                 use_quick_lora=lora_config.use_quick_lora,
             )
             # Lora column parallel will spilt lora A matrix

paddlenlp/peft/tare/__init__.py

@@ -1,17 +0,0 @@
-# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .tare_model import TAREModel
-
-__all__ = ["TAREModel"]

paddlenlp/peft/tare/tare_model.py

@@ -1,336 +0,0 @@
-# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import re
-import time
-
-import paddle
-import paddle.nn.functional as F
-
-
-def cast_all_params_(layer: paddle.nn.Layer, dtype):
-    """把整模里所有 Parameter 的 dtype 原地转成 dtype。"""
-    for _, p in layer.named_parameters():
-        if p.dtype != dtype:
-            p.set_value(p.astype(dtype))
-
-
-def device2str(type=None, index=None, *, device=None):
-    type = device if device else type
-    if isinstance(type, int):
-        type = f"gpu:{type}"
-    elif isinstance(type, str):
-        if "cuda" in type:
-            type = type.replace("cuda", "gpu")
-        if "cpu" in type:
-            type = "cpu"
-        elif index is not None:
-            type = f"{type}:{index}"
-    elif isinstance(type, paddle.CPUPlace) or (type is None):
-        type = "cpu"
-    elif isinstance(type, paddle.CUDAPlace):
-        type = f"gpu:{type.get_device_id()}"
-
-    return type
-
-
-def _get_token(obj, tok: str):
-    """
-    根据 tok 从 obj 里取下一层：
-    - name[idx] 形式：先 getattr(obj, name) 再 [idx]
-    - 纯数字：认为 obj 是 list-like，直接 [int(tok)]
-    - 普通标识符：getattr(obj, tok)
-    """
-    m = re.fullmatch(r"([A-Za-z_]\w*)\[(\d+)\]", tok)
-    if m:
-        name, idx = m.group(1), int(m.group(2))
-        container = getattr(obj, name)
-        return container[idx]
-    if tok.isdigit():
-        return obj[int(tok)]
-    return getattr(obj, tok)
-
-
-def _resolve_parent_and_name(root, path: str):
-    """
-    沿着点路径解析，返回 (parent, last_name_or_index_str, obj)
-    例如 "llama.layers.0.mlp.down_proj" ->
-      parent=<mlp Layer>, last="down_proj", obj=<Linear>
-    """
-    parts = path.split(".") if path else []
-    if not parts:
-        raise AttributeError(f"Empty path for resolve: {path}")
-    cur = root
-    parent = None
-    last = None
-    for tok in parts:
-        parent = cur
-        last = tok
-        cur = _get_token(cur, tok)
-    return parent, last, cur
-
-
-def _set_child(parent, name_or_idx: str, new_obj):
-    """
-    把 parent 的子模块替换为 new_obj。
-    - 如果 name_or_idx 是数字，认为 parent 是 list-like，做 parent[int] = new_obj
-    - 否则 setattr(parent, name_or_idx, new_obj)
-    """
-    print("parent:", parent)
-    print("name_or_idx:", name_or_idx)
-    print("new_obj", new_obj)
-    if name_or_idx.isdigit():
-        print("Yes!!!!!!")
-        parent[int(name_or_idx)] = new_obj
-    else:
-        setattr(parent, name_or_idx, new_obj)
-
-
-device = device2str("cuda" if paddle.device.cuda.device_count() >= 1 else "cpu")
-target_dict = {}
-total_parameter = 0
-paddle.set_device(device)
-
-
-class _DeltaUnit(paddle.nn.Layer):
-    def __init__(self, hidden_size, dtype):
-        super().__init__()
-        self.activation_scaling = self.create_parameter(
-            shape=[1, hidden_size],
-            dtype=dtype,
-            default_initializer=paddle.nn.initializer.Constant(1.0),
-        )
-        self.activation_bias = self.create_parameter(
-            shape=[1, hidden_size],
-            dtype=dtype,
-            default_initializer=paddle.nn.initializer.Constant(0.0),
-        )
-
-
-class ActivationLayer(paddle.nn.Layer):
-    def __init__(self, hidden_size, update_layer, layer_type="all", op_position="ffn", is_llama=False, n=8, k=7):
-        super().__init__()
-        self.update_layer = update_layer
-        self.layer_type = layer_type
-        self.op_position = op_position
-        self.n = n
-        self.k = k
-        self.fc = paddle.nn.Linear(in_features=5120, out_features=self.n)
-        if is_llama:
-            self.weight_type = paddle.bfloat16
-        else:
-            self.weight_type = paddle.float16
-        if self.layer_type == "all":
-            self.delta_vector = paddle.nn.LayerList(
-                _DeltaUnit(hidden_size, dtype=self.weight_type) for _ in range(self.n)
-            )
-        elif self.layer_type == "scaling":
-            self.delta_vector = paddle.nn.ParameterDict(
-                parameters={"activation_scaling": paddle.nn.parameter.Parameter(paddle.ones(1, hidden_size))}
-            )
-        elif self.layer_type == "bias":
-            self.delta_vector = paddle.nn.ParameterDict(
-                parameters={"activation_bias": paddle.nn.parameter.Parameter(paddle.zeros(1, hidden_size))}
-            )
-        elif self.layer_type == "ln":
-            self.delta_vector = paddle.nn.ParameterDict(
-                parameters={
-                    "activation_ln": paddle.nn.LayerNorm(normalized_shape=hidden_size),
-                    "activation_scaling": paddle.nn.parameter.Parameter(paddle.ones(1, hidden_size)),
-                    "activation_bias": paddle.nn.parameter.Parameter(paddle.zeros(1, hidden_size)),
-                }
-            )
-        self.weight = paddle.rand(shape=[1], dtype=self.weight_type)
-
-    def forward(self, x, input_tensor=None):
-        # self.delta_vector.to(self.weight_type).to(device)
-        if self.op_position == "res" or self.op_position == "res_with_attn" or self.op_position == "res_with_res":
-            hidden_states = self.update_layer(x, input_tensor)
-        else:
-            hidden_states = self.update_layer(x)
-        if self.layer_type == "all":
-            hidden_states_new = self.fc(hidden_states)
-            # hidden_states_updated = hidden_states.clone()
-            # print('hidden_states_new:',hidden_states_new.shape)
-            topk_values, topk_indices = paddle.topk(hidden_states_new, self.k, axis=-1)
-            # topk_weights = paddle.softmax(topk_values, axis=-1)
-            topk_weights = F.softmax(topk_values, axis=-1)
-
-            scaling_tensor = paddle.stack([v.activation_scaling for v in self.delta_vector], axis=0).squeeze(
-                1
-            )  # [8, hidden]
-            bias_tensor = paddle.stack([v.activation_bias for v in self.delta_vector], axis=0).squeeze(
-                1
-            )  # [8, hidden]
-
-            hidden_states_updated = paddle.zeros_like(hidden_states)  # [..., hidden]
-            for i in range(self.k):
-                idx = topk_indices[..., i]  # [...,]
-                s_i = scaling_tensor[idx]  # [..., hidden]
-                b_i = bias_tensor[idx]  # [..., hidden]
-                hd_i = hidden_states * s_i + b_i  # [..., hidden]
-                w_i = topk_weights[..., i : i + 1]  # [..., 1]
-                hidden_states_updated = hidden_states_updated + hd_i * w_i  # 广播到 [..., hidden]
-
-            hidden_states = hidden_states_updated.squeeze(-1)
-        elif self.layer_type == "scaling":
-            hidden_states = hidden_states * self.delta_vector["activation_scaling"]
-        elif self.layer_type == "bias":
-            hidden_states = hidden_states + self.delta_vector["activation_bias"]
-        elif self.layer_type == "ln":
-            hidden_states = hidden_states * self.delta_vector["activation_scaling"]
-            hidden_states = hidden_states + self.delta_vector["activation_bias"]
-            hidden_states = self.delta_vector["activation_ln"](hidden_states)
-        if self.op_position == "res_with_res":
-            hidden_states = hidden_states + x
-        return hidden_states
-
-
-class TAREModel(paddle.nn.Layer):
-    _no_split_modules = ["LlamaDecoderLayer"]
-
-    def __init__(self, base_model, op_position="ffn", layer_type="all", exclude_layers=[], only_lora=False, n=8, k=7):
-        print("初始化模型！！！！！！")
-        time.sleep(5)
-        super().__init__()
-        self.base_model = base_model
-        self.model_type = "llama-7b"
-        self.layer_type = layer_type
-        self.op_position = op_position
-        self.exclude_layers = exclude_layers
-        self.n = n
-        self.k = k
-        if exclude_layers:
-            pattern_str = "|".join(map(str, exclude_layers))
-            pattern = re.compile("\\b(?:" + pattern_str + ")\\b")
-        self.frozen_model()
-        if only_lora is False:
-            key_list = [key for key, _ in base_model.named_sublayers(include_self=True)]
-            for key in key_list:
-                if exclude_layers:
-                    match = pattern.search(key)
-                    if match:
-                        continue
-                if self.check_update(key):
-                    self.replace_layer(key)
-        print(self.print_trainable_parameters())
-
-    def check_update(self, key):
-        if self.op_position == "ffn":
-            return self.match_substring(key)
-
-    def generate(self, **args):
-        return self.base_model.generate(**args)
-
-    def replace_layer(self, key):
-        # 打印一下，便于定位真实前缀（有的模型是 self.llama，有的是 self.model）
-        print("base_model:", self.base_model)
-        print("key:", key)
-
-        # 逐段解析，拿到 旧模块 与 其父模块
-        parent, last_name, old_mod = _resolve_parent_and_name(self.base_model, key)
-        print("replaced_module(old):", old_mod)
-
-        # 构造你的替换模块；确保 forward 签名与 old_mod 兼容
-        new_module = ActivationLayer(
-            hidden_size=self.base_model.config.hidden_size,
-            update_layer=old_mod,
-            layer_type=self.layer_type,
-            op_position=self.op_position,
-            is_llama=True,
-            n=self.n,
-            k=self.k,
-        )
-
-        # 覆盖到父模块上
-        _set_child(parent, last_name, new_module)
-
-    def print_trainable_parameters(self):
-        total_parameters = 0
-        trainable_parameters = 0
-
-        for name, param in self.base_model.named_parameters():
-            # 统计参数量
-            try:
-                n = int(paddle.numel(param))
-            except Exception:
-                # 兜底
-                n = 1
-                for d in param.shape:
-                    n *= int(d)
-            total_parameters += n
-
-            # 在 Paddle 中，用 stop_gradient 判断是否可训练
-            is_trainable = not getattr(param, "stop_gradient", True)
-            if is_trainable:
-                trainable_parameters += n
-
-        pct = (100.0 * trainable_parameters / total_parameters) if total_parameters else 0.0
-        return {
-            "total_para:": total_parameters,
-            "trainable_para: ": trainable_parameters,
-            "trainable%:": f"{pct:.4f}",
-        }
-
-    def frozen_model(self):
-        for name, param in self.base_model.named_parameters():
-            if "lora" in name:
-                param.stop_gradient = not True
-            else:
-                param.stop_gradient = not False
-
-    def match_substring(self, input_string):
-        pattern = "down_proj"
-        pattern2 = "c_proj"
-        match = re.search(pattern, input_string) or re.search(pattern2, input_string)
-        if match:
-            return True
-        else:
-            return False
-
-    def forward(self, input_ids, attention_mask=None, labels=None, position_ids=None):
-        output = self.base_model(
-            input_ids=input_ids, attention_mask=attention_mask, labels=labels, position_ids=position_ids
-        )
-        return output
-
-    def load_model(self, save_path):
-        save_sd = paddle.load(str(save_path))
-        model_sd = self.base_model.state_dict()
-        loaded = 0
-        for k, v in save_sd.items():
-            # print(k)
-            if k in model_sd:
-                tgt = model_sd[k]
-                if v.dtype != tgt.dtype:
-                    v = v.astype(tgt.dtype)
-                assert list(v.shape) == list(tgt.shape), f"{k} shape {v.shape} vs {tgt.shape}"
-                tgt.set_value(v)
-                tgt.stop_gradient = False  # 若这些参数要训练
-                loaded += 1
-        print(f"loaded {loaded} tensors")
-        time.sleep(5)
-
-    def get_save_dict(self):
-        state_dict = self.base_model.state_dict()
-        for k in state_dict:
-            if "activation_" in k or "lora" in k or "fc" in k:
-                print(k)
-
-        save_dict = {k: state_dict[k] for k in state_dict if "activation_" in k or "lora" in k or "fc" in k}
-        return save_dict
-
-    def save_model(self, save_path):
-        save_dict = self.get_save_dict()
-        paddle.save(obj=save_dict, path=save_path)

paddlenlp/quantization/hadamard_utils.py

@@ -14,8 +14,6 @@
 
 import paddle
 
-from paddlenlp.utils import infohub
-
 
 def matmul_hadU(X):
 
@@ -33,37 +31,22 @@ def matmul_hadU(X):
     return input.reshape(X.shape)
 
 
-def create_hadamard_matrix(block_size, dtype):
-    Q = paddle.diag(paddle.ones((block_size), dtype=dtype))
-    block = matmul_hadU(Q)
-    return block
-
-
-def hadamard_matmul(input, side, hadamard_matrix, block_size):
-    # left -> H.T@input right -> input@H
-    origin_shape = input.shape
-    input = input.reshape([-1, origin_shape[-1]])
-    if side == "left":
-        # H.T@input -> (input.T@H).T
-        input = input.transpose([1, 0])
-    block_num = input.shape[-1] // block_size
-    output = input.reshape([-1, block_num, block_size]) @ hadamard_matrix
-    output = output.reshape([-1, block_num * block_size])
-    if side == "left":
-        output = output.transpose([1, 0])
-    output = output.reshape(origin_shape)
-
-    return output
-
-
-def apply_hadamard_matmul(x, side, block_size):
-    if getattr(infohub, "hadamard") is None:
-        setattr(infohub, "hadamard", {})
-
-    if block_size in infohub.hadamard:
-        hadamard_matrix = infohub.hadamard[block_size]
+def random_hadamard_matrix(size, dtype, is_block=False):
+    if not is_block:
+        A = paddle.randint(low=0, high=2, shape=[size, size]).astype("float32") * 2 - 1
+        Q, _ = paddle.linalg.qr(A)
+        return Q.astype(dtype), 1
     else:
-        hadamard_matrix = create_hadamard_matrix(block_size, x.dtype)
-        infohub.hadamard[block_size] = hadamard_matrix
-    target_x = hadamard_matmul(x, side, hadamard_matrix, block_size)
-    return target_x
+        num_blocks = size
+        while not (num_blocks % 2):
+            num_blocks = num_blocks // 2
+        block_size = size // num_blocks
+        Q = paddle.diag(paddle.ones((block_size,), dtype="float32"))
+        block = matmul_hadU(Q)
+        large_matrix = paddle.zeros([size, size])
+
+        for i in range(num_blocks):
+            start_row = i * block_size
+            start_col = i * block_size
+            large_matrix[start_row : start_row + block_size, start_col : start_col + block_size] = block
+    return large_matrix.cast(dtype), block_size

paddlenlp/quantization/qat_utils.py

@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
 
 import paddle
 from paddle.autograd import PyLayer
 
-from .hadamard_utils import apply_hadamard_matmul
+from paddlenlp.utils import infohub
+
+from .hadamard_utils import random_hadamard_matrix
 
 try:
     from transformer_engine import transformer_engine_paddle as tex
@@ -28,156 +31,151 @@ try:
         paddle.float8_e4m3fn: tex.DType.kFloat8E4M3,
         paddle.float8_e5m2: tex.DType.kFloat8E5M2,
     }
-    SUPPORT_TE = True
+    USE_FP8_GEMM = True
 except ImportError:
-    SUPPORT_TE = False
-
-from paddle.linalg import fp8_fp8_half_gemm_fused
-
-QMIN_QMAX_MAPPING = {
-    "a8w8linear_activation": (-128, 127),
-    "a8w4linear_activation": (-128, 127),
-    "a8w8linear_weight": (-128, 127),
-    "a8w4linear_weight": (-8, 7),
-    "float8_e4m3fn": (-488, 488),
-    "float8_e5m2": (-57344, 57344),
-}
+    USE_FP8_GEMM = False
 
 
-def quantize(
-    x,
-    weight_quantize_algo,
-    tensor_type,
-    quantization_config,
-    side="right",
-    apply_hadamard=False,
-    act_scale=None,
-    state=0,
-    training=False,
-    group=None,
-):
-    if apply_hadamard:
-        target_x = apply_hadamard_matmul(x, side, quantization_config.hadamard_block_size)
-        hadamard_scale = quantization_config.hadamard_block_size
+def quantize_tensorwise(x, quantization_config=None, bit_length=8, state=0, training=False, act_scale=None):
+    qmax = (1 << (bit_length - 1)) - 1
+    qmin = -1 * qmax - 1
+    if quantization_config.apply_hadamard:
+        target_x = x @ infohub.hadamard[x.shape[-1]][0]
     else:
-        target_x, hadamard_scale = x, 1.0
-    if weight_quantize_algo in ["fp8linear"]:
-        qmin, qmax = QMIN_QMAX_MAPPING[quantization_config.fp8_format[tensor_type]]
-    else:
-        qmin, qmax = QMIN_QMAX_MAPPING[weight_quantize_algo + "_" + tensor_type]
-    if tensor_type == "activation":
-        if act_scale is not None:
-            if training:
-                scale = (paddle.max(paddle.abs(target_x)) / qmax + quantization_config.scale_epsilon).reshape([1])
-                if group is not None:
-                    paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX, group=group, sync_op=True)
-                if state < quantization_config.apply_online_actscale_step:
-                    act_scale[:] = (state * act_scale + scale) / (state + 1)
-                else:
-                    scale = (
-                        1 - quantization_config.actscale_moving_rate
-                    ) * act_scale + quantization_config.actscale_moving_rate * scale
-                    act_scale[:] = scale
+        target_x = x
+
+    if act_scale is not None:
+        if training:
+            scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+            if state < quantization_config.skip_first_act_scale_step:
+                act_scale.set_value((state * act_scale + scale) / (state + 1))
             else:
+                act_scale.set_value(
+                    (1 - quantization_config.moving_rate) * act_scale + quantization_config.moving_rate * scale
+                )
                 scale = act_scale
         else:
-            scale = (paddle.max(paddle.abs(target_x)) / qmax + quantization_config.scale_epsilon).reshape([1])
-        if weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
-            quant_x = paddle.clip((target_x / scale).round(), qmin, qmax).astype("int8")
-        elif weight_quantize_algo in ["fp8linear"]:
-            quant_x = (target_x / scale).astype(quantization_config.fp8_format[tensor_type])
-        else:
-            raise NotImplementedError(f"Unknown {weight_quantize_algo}.")
-    elif tensor_type == "weight":
-        if weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
-            # channelwise
-            scale = paddle.max(paddle.abs(target_x), axis=0, keepdim=True) / qmax + quantization_config.scale_epsilon
-            if group is not None:
-                paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX, group=group, sync_op=True)
-            quant_x = paddle.clip((target_x / scale).round(), qmin, qmax).astype("int8").T
-            scale = scale.squeeze(0) / hadamard_scale
-        elif weight_quantize_algo in ["fp8linear"]:
-            scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.scale_epsilon
-            if group is not None:
-                paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX, group=group, sync_op=True)
-            quant_x = (target_x / scale).astype(quantization_config.fp8_format[tensor_type]).view("int8").T
-            scale = (scale / hadamard_scale).reshape([1])
-        else:
-            raise NotImplementedError(f"Unknown {weight_quantize_algo}.")
-    elif tensor_type == "grad_output":
-        if weight_quantize_algo in ["fp8linear"]:
-            scale = (paddle.max(paddle.abs(target_x)) / qmax + quantization_config.scale_epsilon).reshape([1])
-            quant_x = (target_x / scale).astype(quantization_config.fp8_format[tensor_type])
-            scale = scale / hadamard_scale
-        else:
-            raise NotImplementedError(f"Unknown {weight_quantize_algo}.")
+            scale = act_scale
     else:
-        raise NotImplementedError(f"Unknown {tensor_type}.")
-    scale.stop_gradient = True
-    return quant_x, scale
+        scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+
+    x_int8 = paddle.clip((target_x / scale).round(), qmin, qmax).astype("int8")
+    return x_int8, scale
 
 
-def dequantize(
-    quant_x, scale, tensor_type, weight_quantize_algo, quantization_config, apply_hadamard=False, side="left"
-):
-    if tensor_type == "weight":
-        if weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
-            x = quant_x.T.astype(scale.dtype)
-        elif weight_quantize_algo in ["fp8linear"]:
-            x = quant_x.view(quantization_config.fp8_format[tensor_type]).T.astype(scale.dtype)
-        else:
-            raise NotImplementedError(f"Unknown weight_quantize_algo: {weight_quantize_algo}")
-        if apply_hadamard:
-            x = apply_hadamard_matmul(x, side, quantization_config.hadamard_block_size)
-        x *= scale
-    else:
-        raise NotImplementedError(f"Unknown {tensor_type}.")
+def dequantize_tensorwise(x_int8, scale, apply_hadamard=False):
+    x = x_int8.astype(scale.dtype) * scale
+    if apply_hadamard:
+        x = x @ infohub.hadamard[x.shape[-1]][0].T
     return x
 
 
-def int8_forward(
-    x,
-    quant_w,
-    scale_w,
-    weight_quantize_algo,
-    bias=None,
-    quantization_config=None,
-    state=0,
-    training=False,
-    act_scale=None,
-    group=None,
-):
-    quant_x, scale_x = quantize(
-        x=x,
-        weight_quantize_algo=weight_quantize_algo,
-        tensor_type="activation",
-        quantization_config=quantization_config,
-        side="right",
-        apply_hadamard=quantization_config.apply_hadamard,
-        act_scale=act_scale,
-        state=state,
-        training=training,
-        group=group,
-    )
+def fp8_quantize_tensorwise(x, tensor_type, quantization_config=None, state=0, training=False, act_scale=None):
+    assert tensor_type in ["weight", "activation", "grad_output"], "Only support weight, activation and grad_output"
+    fp8_format = quantization_config.fp8_format[tensor_type]
+    qmin, qmax = (-448, 448) if fp8_format == "float8_e4m3fn" else (-57344, 57344)
+    tensor_type_to_shape_index = {"weight": 0, "activation": -1, "grad_output": -2}
 
-    out = paddle.matmul(quant_x, quant_w.T).astype(scale_w.dtype) * (scale_x * scale_w)
+    if quantization_config is not None and quantization_config.apply_hadamard:
+        if getattr(infohub, "hadamard") is None:
+            setattr(infohub, "hadamard", {})
+
+        hadamard_matrix_shape = x.shape[tensor_type_to_shape_index[tensor_type]]
+        if hadamard_matrix_shape in infohub.hadamard:
+            hadamard_matrix, block_size = infohub.hadamard[hadamard_matrix_shape]
+        else:
+            hadamard_matrix, block_size = random_hadamard_matrix(hadamard_matrix_shape, x.dtype, is_block=True)
+            infohub.hadamard[hadamard_matrix_shape] = (hadamard_matrix, block_size)
+        target_x = hadamard_matrix.T @ x if tensor_type in ["weight", "grad_output"] else x @ hadamard_matrix
+    else:
+        target_x = x
+        block_size = 1
+
+    if act_scale is not None:
+        if training:
+            scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+            if state < quantization_config.skip_first_act_scale_step:
+                act_scale.set_value((state * act_scale + scale) / (state + 1))
+            else:
+                act_scale.set_value(
+                    (1 - quantization_config.moving_rate) * act_scale + quantization_config.moving_rate * scale
+                )
+                # scale = act_scale
+        else:
+            scale = act_scale
+    else:
+        scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+
+    x_fp8 = target_x / scale
+    x_fp8 = x_fp8.astype(fp8_format).view("int8")
+    x_fp8 = x_fp8.T if tensor_type == "weight" else x_fp8
+    scale.stop_gradient = True
+    scale = scale / block_size if tensor_type in ["weight", "grad_output"] else scale
+    return x_fp8, scale
+
+
+def fp8_dequantize_tensorwise(x_fp8, scale, tensor_type, quantization_config=None):
+    x_fp8 = x_fp8.view(quantization_config.fp8_format[tensor_type])
+    x_fp8 = x_fp8.T if tensor_type == "weight" else x_fp8
+    x = x_fp8.astype(scale.dtype) * scale
+    if quantization_config.apply_hadamard:
+        hadamard_matrix_shape = x.shape[0] if tensor_type == "weight" else x.shape[-1]
+        hadamard_matrix, _ = infohub.hadamard[hadamard_matrix_shape]
+        x = hadamard_matrix @ x if tensor_type == "weight" else x @ hadamard_matrix.T
+    return x
+
+
+def quantize_channelwise(w, apply_hadamard=False, bit_length=8):
+    qmax = (1 << (bit_length - 1)) - 1
+    qmin = -1 * qmax - 1
+    if apply_hadamard:
+        if getattr(infohub, "hadamard") is None:
+            setattr(infohub, "hadamard", {})
+        if w.shape[0] in infohub.hadamard:
+            hadamard_matrix, block_size = infohub.hadamard[w.shape[0]]
+        else:
+            hadamard_matrix, block_size = random_hadamard_matrix(w.shape[0], w.dtype, is_block=True)
+            infohub.hadamard[w.shape[0]] = (hadamard_matrix, block_size)
+        w = hadamard_matrix.T @ w
+    else:
+        block_size = 1
+    scale = paddle.max(paddle.abs(w), axis=0, keepdim=True) / qmax
+    w_int8 = paddle.clip((w / scale).round(), qmin, qmax).astype("int8")
+    scale.stop_gradient = True
+    return w_int8.T, scale.squeeze(0) / block_size
+
+
+def dequantize_channelwise(w_int8, scale, apply_hadamard=False):
+    w = w_int8.T.astype(scale.dtype) * scale
+    if apply_hadamard:
+        w = infohub.hadamard[w_int8.shape[1]][0] @ w
+    return w
+
+
+def a8w8_forward(
+    x, w_int8, w_scale=None, bias=None, dtype=None, quantization_config=None, state=0, training=False, act_scale=None
+):
+    x_int8, x_scale = quantize_tensorwise(
+        x, quantization_config, bit_length=8, state=state, training=training, act_scale=act_scale
+    )
+    out = paddle.matmul(x_int8, w_int8.T).astype(dtype) * (x_scale * w_scale.unsqueeze(0))
     if bias is not None:
         out += bias
-    return out, quant_x, scale_x
+    return out, x_int8, x_scale
 
 
-def int8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale):
+def a8w8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale):
     if not ctx.x_stop_gradient:
-        qdq_weight = dequantize(
-            quant_weight,
-            quant_scale,
-            "weight",
-            ctx.weight_quantize_algo,
-            ctx.quantization_config,
-            ctx.quantization_config.apply_hadamard,
-            "left",
-        )
-        input_grad = paddle.matmul(grad_output, qdq_weight.T)
+        if ctx.quantization_config.quant_input_grad:
+            grad_output_int8, grad_output_scale = quantize_tensorwise(grad_output * quant_scale)
+            input_grad = paddle.matmul(grad_output_int8, quant_weight).astype(ctx.dtype) * grad_output_scale
+            if ctx.quantization_config.apply_hadamard:
+                input_grad = input_grad @ infohub.hadamard[quant_weight.shape[-1]][0].T
+        else:
+            qdq_weight = dequantize_channelwise(
+                quant_weight, quant_scale, apply_hadamard=ctx.quantization_config.apply_hadamard
+            )
+            input_grad = paddle.matmul(grad_output, qdq_weight.T)
     else:
         input_grad = None
 
@@ -195,59 +193,64 @@ def int8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_sca
 
 
 def fp8_forward(
-    x,
-    w_fp8,
-    w_scale,
-    weight_quantize_algo,
-    bias=None,
-    dtype=None,
-    quantization_config=None,
-    state=0,
-    training=False,
-    act_scale=None,
-    group=None,
+    x, w_fp8, w_scale=None, bias=None, dtype=None, quantization_config=None, state=0, training=False, act_scale=None
 ):
-    x_fp8, x_scale = quantize(
+    x_fp8, x_scale = fp8_quantize_tensorwise(
         x,
-        weight_quantize_algo,
-        "activation",
-        quantization_config,
-        side="right",
-        apply_hadamard=quantization_config.apply_hadamard,
-        act_scale=act_scale,
+        tensor_type="activation",
+        quantization_config=quantization_config,
         state=state,
         training=training,
-        group=group,
+        act_scale=act_scale,
     )
+    x_fp8 = x_fp8.view(quantization_config.fp8_format["activation"])
     w_fp8 = w_fp8.view(quantization_config.fp8_format["weight"])
-    out = fp8_fp8_half_gemm_fused(
-        x_fp8,
-        w_fp8,
-        transpose_x=False,
-        transpose_y=True,
-        bias=bias,
-        scale=x_scale * w_scale,
-        output_dtype=dtype,
-    )
+
+    if USE_FP8_GEMM:
+        x_shape = x_fp8.shape
+        x_fp8 = x_fp8.view((-1, x_fp8.shape[-1]))
+        fwd_scales = paddle.stack([x_scale.astype("float32"), w_scale.astype("float32")])
+        out, _ = fp8_gemm(
+            A=w_fp8,
+            A_scale_inv=fwd_scales,
+            A_fp8_tensor=FP8FwdTensors.GEMM1_WEIGHT,
+            A_dtype=TE_DType[w_fp8.dtype],
+            B=x_fp8,
+            B_scale_inv=fwd_scales,
+            B_fp8_tensor=FP8FwdTensors.GEMM1_INPUT,
+            B_dtype=TE_DType[x_fp8.dtype],
+            out_dtype=dtype,
+            workspace=get_workspace(),
+            bias=bias,
+            use_bias=True if bias is not None else False,
+            use_split_accumulator=True,
+        )
+        x_fp8 = x_fp8.view(x_shape)
+        out = out.view((*x_shape[:-1], -1))
+    else:
+        x = x_fp8.astype(dtype) * x_scale
+        w = w_fp8.astype(dtype) * w_scale
+        out = paddle.matmul(x, w.T).astype(dtype)
+        if bias is not None:
+            out += bias
+
     return out, x_fp8, x_scale
 
 
 def fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale):
     if not ctx.x_stop_gradient:
         if ctx.quantization_config.quant_input_grad:
-            grad_output_fp8, grad_output_scale = quantize(
+            grad_output_fp8, grad_output_scale = fp8_quantize_tensorwise(
                 grad_output,
-                ctx.weight_quantize_algo,
-                "grad_output",
-                ctx.quantization_config,
-                side="left",
-                apply_hadamard=False,
+                tensor_type="grad_output",
+                quantization_config=ctx.quantization_config,
             )
+            grad_output_fp8 = grad_output_fp8.view(ctx.quantization_config.fp8_format["grad_output"])
             quant_weight = quant_weight.view(ctx.quantization_config.fp8_format["weight"])
-            if SUPPORT_TE:
+            if USE_FP8_GEMM:
                 grad_output_shape = grad_output_fp8.shape
                 grad_output_fp8 = grad_output_fp8.view((-1, grad_output_fp8.shape[-1]))
-                fwd_scales = paddle.concat([x_scale.astype("float32"), quant_scale.astype("float32")])
+                fwd_scales = paddle.stack([x_scale.astype("float32"), quant_scale.astype("float32")])
                 bwd_scales = grad_output_scale[None].astype("float32")
                 input_grad, _ = fp8_gemm(
                     A=quant_weight.T,
@@ -268,16 +271,11 @@ def fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scal
                 weight_ = quant_weight.astype(ctx.dtype) * quant_scale
                 input_grad = paddle.matmul(grad_output_, weight_).astype(ctx.dtype)
             if ctx.quantization_config.apply_hadamard:
-                input_grad = apply_hadamard_matmul(input_grad, "right", ctx.quantization_config.hadamard_block_size)
+                input_grad = infohub.hadamard[grad_output.shape[-2]][0] @ input_grad
+                input_grad = input_grad @ infohub.hadamard[quant_weight.shape[-1]][0].T
         else:
-            qdq_weight = dequantize(
-                quant_weight,
-                quant_scale,
-                "weight",
-                ctx.weight_quantize_algo,
-                ctx.quantization_config,
-                apply_hadamard=ctx.quantization_config.apply_hadamard,
-                side="left",
+            qdq_weight = fp8_dequantize_tensorwise(
+                quant_weight, quant_scale, tensor_type="weight", quantization_config=ctx.quantization_config
             )
             input_grad = paddle.matmul(grad_output, qdq_weight.T)
     else:
@@ -285,36 +283,20 @@ def fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scal
 
     if not ctx.w_stop_gradient:
         if ctx.quantization_config.quant_weight_grad:
-            grad_output_fp8, grad_output_scale = quantize(
-                x=grad_output,
-                weight_quantize_algo=ctx.weight_quantize_algo,
+            quantization_config_ = deepcopy(ctx.quantization_config)
+            quantization_config_.apply_hadamard = False
+            grad_output_fp8, grad_output_scale = fp8_quantize_tensorwise(
+                grad_output,
                 tensor_type="grad_output",
-                quantization_config=ctx.quantization_config,
-                apply_hadamard=False,
+                quantization_config=quantization_config_,
             )
-            if SUPPORT_TE:
+            grad_output_fp8 = grad_output_fp8.view(ctx.quantization_config.fp8_format["grad_output"])
+            quant_x = quant_x.view(ctx.quantization_config.fp8_format["activation"])
+            if USE_FP8_GEMM:
                 quant_x = quant_x.view((-1, quant_x.shape[-1]))
                 grad_output_fp8 = grad_output_fp8.view((-1, grad_output_fp8.shape[-1]))
-                fwd_scales = paddle.concat([x_scale.astype("float32"), quant_scale.astype("float32")])
+                fwd_scales = paddle.stack([x_scale.astype("float32"), quant_scale.astype("float32")])
                 bwd_scales = grad_output_scale[None].astype("float32")
-                # FP8 gemm need k % 16 = 0
-                ALIGNMENT_SIZE = 16
-
-                def pad_tensor_to_multiple(tensor, dtype):
-                    current_size = tensor.shape[0]
-                    padding_size = ALIGNMENT_SIZE - current_size % ALIGNMENT_SIZE
-                    # Create padding zeros with matching shape and dtype
-                    padding_shape = [padding_size, tensor.shape[1]]
-                    padding = paddle.zeros(padding_shape, dtype=dtype)
-                    padded_tensor = paddle.concat([tensor, padding], axis=0)
-                    return padded_tensor
-
-                if quant_x.shape[0] % ALIGNMENT_SIZE != 0:
-                    quant_x = pad_tensor_to_multiple(quant_x, ctx.quantization_config.fp8_format["activation"])
-                    grad_output_fp8 = pad_tensor_to_multiple(
-                        grad_output_fp8, ctx.quantization_config.fp8_format["grad_output"]
-                    )
-
                 weight_grad, _ = fp8_gemm(
                     A=grad_output_fp8.T,
                     A_scale_inv=bwd_scales,
@@ -339,8 +321,9 @@ def fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scal
                         grad_output_.reshape([-1, grad_output_.shape[-1]]),
                     ).astype(ctx.dtype)
             if ctx.quantization_config.apply_hadamard:
-                weight_grad = weight_grad / ctx.quantization_config.hadamard_block_size
-                weight_grad = apply_hadamard_matmul(weight_grad, "left", ctx.quantization_config.hadamard_block_size)
+                hadamard_matrix, block_size = infohub.hadamard[quant_x.shape[-1]]
+                weight_grad = weight_grad / block_size
+                weight_grad = hadamard_matrix @ weight_grad
         else:
             if len(x.shape) == 2:
                 weight_grad = paddle.matmul(x.transpose([1, 0]), grad_output)
@@ -367,39 +350,33 @@ class QATFunc(PyLayer):
         state,
         training,
         act_scale,
-        weight_quantize_algo,
-        group,
     ):
         quant_x, x_scale = None, None
-        if weight_quantize_algo in ["fp8linear"]:
+        if quantization_config.weight_quantize_algo in ["fp8linear"]:
             output, quant_x, x_scale = fp8_forward(
                 x,
                 quant_weight,
                 w_scale=quant_scale,
-                weight_quantize_algo=weight_quantize_algo,
                 bias=bias,
                 dtype=dtype,
                 quantization_config=quantization_config,
                 state=state,
                 training=training,
                 act_scale=act_scale,
-                group=group,
             )
         else:
-            output, quant_x, x_scale = int8_forward(
+            output, quant_x, x_scale = a8w8_forward(
                 x,
-                quant_w=quant_weight,
-                scale_w=quant_scale,
-                weight_quantize_algo=weight_quantize_algo,
+                quant_weight,
+                w_scale=quant_scale,
                 bias=bias,
+                dtype=dtype,
                 quantization_config=quantization_config,
                 state=state,
                 training=training,
                 act_scale=act_scale,
-                group=group,
             )
         ctx.quantization_config = quantization_config
-        ctx.weight_quantize_algo = weight_quantize_algo
         ctx.dtype = dtype
         ctx.x_stop_gradient = x.stop_gradient
         ctx.w_stop_gradient = quant_weight.stop_gradient
@@ -422,7 +399,7 @@ class QATFunc(PyLayer):
         if ctx.quantization_config.weight_quantize_algo in ["fp8linear"]:
             input_grad, weight_grad = fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale)
         else:
-            input_grad, weight_grad = int8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale)
+            input_grad, weight_grad = a8w8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale)
 
         if not ctx.b_stop_gradient:
             bias_grad = grad_output.sum(axis=[0, 1])

paddlenlp/quantization/quantization_config.py

@@ -20,6 +20,7 @@ quant_inference_mapping = {"avg": "abs_max", "abs_max_channel_wise": "abs_max_ch
 fp8_format_mapping = {
     "hybrid": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e5m2"},
     "e4m3": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e4m3fn"},
+    "e5m2": {"weight": "float8_e5m2", "activation": "float8_e5m2", "grad_output": "float8_e5m2"},
 }
 
 
@@ -40,7 +41,6 @@ class QuantizationConfig:
         qlora_weight_double_quant_block_size: Block size for quant_scale of weight quant_scale.
         weight_quant_method: The method for weight quantization.
         act_quant_method: The method for activation quantization.
-        apply_online_actscale_step: Use online (per-step) activation scales for the first N steps. During these steps, activation scales are also collected to compute their mean for later use.
     """
 
     def __init__(
@@ -64,14 +64,13 @@ class QuantizationConfig:
         dtype=None,
         ignore_modules=None,
         group_size=-1,
-        apply_hadamard=False,
-        hadamard_block_size=32,
+        apply_hadamard=True,
         quant_input_grad=False,
         quant_weight_grad=False,
-        apply_online_actscale_step=200,
-        actscale_moving_rate=0.01,
+        skip_first_act_scale_step=20,
+        moving_rate=0.01,
+        epsilon=1e-8,
         fp8_format_type="hybrid",
-        scale_epsilon=1e-8,
         **kwargs,
     ):
         if weight_quantize_algo is not None:
@@ -140,17 +139,12 @@ class QuantizationConfig:
         self.ignore_modules = ignore_modules
         self.group_size = group_size
         self.apply_hadamard = apply_hadamard
-        self.hadamard_block_size = hadamard_block_size
         self.quant_input_grad = quant_input_grad
         self.quant_weight_grad = quant_weight_grad
-        self.apply_online_actscale_step = apply_online_actscale_step
-        self.actscale_moving_rate = actscale_moving_rate
-        self.fp8_format_type = fp8_format_type
-        self.scale_epsilon = scale_epsilon
-
-    @property
-    def fp8_format(self):
-        return fp8_format_mapping[self.fp8_format_type]
+        self.skip_first_act_scale_step = skip_first_act_scale_step
+        self.moving_rate = moving_rate
+        self.epsilon = epsilon
+        self.fp8_format = fp8_format_mapping[fp8_format_type]
 
     def is_weight_quantize(self):
         if isinstance(self.weight_quantize_algo, dict):
@@ -220,7 +214,6 @@ class QuantizationConfig:
         config_dict = self.to_dict()
 
         # get the default config dict
-
         default_config_dict = QuantizationConfig().to_dict()
 
         serializable_config_dict = {}

paddlenlp/quantization/quantization_linear.py

@@ -15,7 +15,6 @@
 import paddle
 import paddle.nn as nn
 from paddle.autograd import PyLayer
-from paddle.distributed import fleet
 from paddle.distributed.fleet.base import topology as tp
 from paddle.distributed.fleet.layers.mpu import mp_ops
 from paddle.distributed.fleet.utils.sequence_parallel_utils import (
@@ -24,8 +23,6 @@ from paddle.distributed.fleet.utils.sequence_parallel_utils import (
 )
 from paddle.nn.quant import llm_int8_linear, weight_dequantize, weight_only_linear
 
-from paddlenlp.utils import infohub
-
 from .qat_utils import QATFunc
 
 try:
@@ -212,20 +209,10 @@ def quant_weight_linear(
 ):
     if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
 
-        state, training, act_scale, group = act_state
+        state, training, act_scale = act_state
 
         return QATFunc.apply(
-            x,
-            quant_weight,
-            bias,
-            quant_scale,
-            quantization_config,
-            dtype,
-            state,
-            training,
-            act_scale,
-            weight_quantize_algo,
-            group,
+            x, quant_weight, bias, quant_scale, quantization_config, dtype, state, training, act_scale
         )
     else:
         return QuantizationLinearFunc.apply(
@@ -241,22 +228,6 @@ def quant_weight_linear(
         )
 
 
-def get_act_scale_group(is_row=False):
-    if paddle.distributed.is_initialized():
-        if getattr(infohub, "scale_group") is None:
-            hcg = fleet.get_hybrid_communicate_group()
-            rank = hcg._dp_degree * hcg._sharding_degree
-            group_no_row = hcg.create_fuse_group(["data", "sharding"])[1] if rank > 1 else None
-            rank *= hcg._mp_degree
-            group_row = hcg.create_fuse_group(["data", "sharding", "model"])[1] if rank > 1 else None
-
-            setattr(infohub, "scale_group", [group_no_row, group_row])
-        group = infohub.scale_group[1] if is_row else infohub.scale_group[0]
-    else:
-        group = None
-    return group
-
-
 class QuantizationLinear(nn.Layer):
     """Quantization Linear layer."""
 
@@ -307,10 +278,9 @@ class QuantizationLinear(nn.Layer):
                 raise NotImplementedError("Not yet support grouwise weightonly quantization.")
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                 self.act_scale = self.create_parameter(
-                    shape=[1], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
+                    shape=[], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
                 self.act_scale.stop_gradient = True
-                self.group = get_act_scale_group()
 
         elif self.weight_quantize_algo in ["fp4", "nf4"]:
             if qlora_weight_linear is None:
@@ -370,7 +340,6 @@ class QuantizationLinear(nn.Layer):
             for p in self.parameters():
                 p.is_distributed = is_distributed
                 p.mp_moe = mp_moe
-        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
 
     def forward(self, x):
         output = quant_weight_linear(
@@ -385,7 +354,7 @@ class QuantizationLinear(nn.Layer):
             if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
             else None,
             bias=self.bias,
-            act_state=(self.state, self.training, self.act_scale, self.group)
+            act_state=(self.state, self.training, self.act_scale)
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
@@ -465,10 +434,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
                     is_bias=False,
                 )
                 self.quant_scale.stop_gradient = True
-                if self.weight_quantize_algo in ["fp8linear", "a8w4linear", "a8w8linear"]:
-                    self.quant_scale.is_distributed = False
-                else:
-                    self.quant_scale.is_distributed = True if self.is_mp else False
+                self.quant_scale.is_distributed = True if self.is_mp else False
                 if self.quant_scale.is_distributed:
                     self.quant_scale.split_axis = 0
             else:
@@ -476,11 +442,10 @@ class ColumnParallelQuantizationLinear(nn.Layer):
                 raise NotImplementedError("Not yet support grouwise weightonly quantization.")
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                 self.act_scale = self.create_parameter(
-                    shape=[1], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
+                    shape=[], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
-                self.act_scale.is_distributed = False
+                self.act_scale.is_distributed = True if self.is_mp else False
                 self.act_scale.stop_gradient = True
-                self.group = get_act_scale_group()
         else:
             raise NotImplementedError(f"Not yet support weight_quantize_algo: {self.weight_quantize_algo}")
         if bias_attr is False:
@@ -495,7 +460,6 @@ class ColumnParallelQuantizationLinear(nn.Layer):
             self.bias.is_distributed = True if self.is_mp else False
             if self.bias.is_distributed:
                 self.bias.split_axis = 0
-        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
 
     def forward(self, x):
         if self.is_mp:
@@ -522,7 +486,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
             if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
             else None,
             bias=self.bias,
-            act_state=(self.state, self.training, self.act_scale, self.group)
+            act_state=(self.state, self.training, self.act_scale)
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
@@ -609,10 +573,7 @@ class RowParallelQuantizationLinear(nn.Layer):
                     is_bias=False,
                 )
                 self.quant_scale.stop_gradient = True
-                if self.weight_quantize_algo in ["fp8linear", "a8w4linear", "a8w8linear"]:
-                    self.quant_scale.is_distributed = False
-                else:
-                    self.quant_scale.is_distributed = True if self.is_mp else False
+                self.quant_scale.is_distributed = True if self.is_mp else False
                 if self.quant_scale.is_distributed:
                     self.quant_scale.split_axis = 0
             else:
@@ -622,9 +583,8 @@ class RowParallelQuantizationLinear(nn.Layer):
                 self.act_scale = self.create_parameter(
                     shape=[1], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
-                self.act_scale.is_distributed = False
+                self.act_scale.is_distributed = True if self.is_mp else False
                 self.act_scale.stop_gradient = True
-                self.group = get_act_scale_group(is_row=True)
         else:
             raise NotImplementedError(f"Not yet support weight_quantize_algo: {self.weight_quantize_algo}")
 
@@ -638,8 +598,6 @@ class RowParallelQuantizationLinear(nn.Layer):
                 is_bias=True,
             )
 
-        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
-
     def forward(self, x):
         if self.input_is_parallel or (not self.is_mp):
             input_parallel = x
@@ -661,7 +619,7 @@ class RowParallelQuantizationLinear(nn.Layer):
                 if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
                 else None,
                 bias=None,
-                act_state=(self.state, self.training, self.act_scale, self.group)
+                act_state=(self.state, self.training, self.act_scale)
                 if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )
@@ -689,7 +647,7 @@ class RowParallelQuantizationLinear(nn.Layer):
                 if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
                 else None,
                 bias=self.bias,
-                act_state=(self.state, self.training, self.act_scale, self.group)
+                act_state=(self.state, self.training, self.act_scale)
                 if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )

paddlenlp/quantization/quantization_utils.py

@@ -34,7 +34,7 @@ except:
     qlora_weight_quantize = None
 
 from ..utils.log import logger
-from .qat_utils import quantize
+from .qat_utils import fp8_quantize_tensorwise, quantize_channelwise
 from .quantization_linear import (
     ColumnParallelQuantizationLinear,
     QuantizationLinear,
@@ -155,17 +155,25 @@ def convert_to_weight_quantize_state_dict(state_dict, name, quantization_config,
     if weight_name in state_dict:
         # gpu weight_quantize will fix in future
         target_weight = state_dict.pop(weight_name).cast(dtype).cuda()
-
-        if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
-            quant_weight, quant_scale = quantize(
-                target_weight,
-                weight_quantize_algo,
-                "weight",
-                quantization_config,
-                side="left",
-                apply_hadamard=quantization_config.apply_hadamard,
+        if weight_quantize_algo in ["a8w8linear"]:
+            quant_weight, quant_scale = quantize_channelwise(
+                target_weight, quantization_config.apply_hadamard, bit_length=8
             )
-            act_scale = paddle.ones([1], dtype=dtype).cuda()
+            act_scale = paddle.zeros([], dtype="bfloat16").cuda()
+            act_scale.stop_gradient = True
+            state_dict[act_scale_name] = act_scale
+        elif weight_quantize_algo in ["a8w4linear"]:
+            quant_weight, quant_scale = quantize_channelwise(
+                target_weight, quantization_config.apply_hadamard, bit_length=4
+            )
+            act_scale = paddle.zeros([], dtype="bfloat16").cuda()
+            act_scale.stop_gradient = True
+            state_dict[act_scale_name] = act_scale
+        elif weight_quantize_algo in ["fp8linear"]:
+            quant_weight, quant_scale = fp8_quantize_tensorwise(
+                target_weight, tensor_type="weight", quantization_config=quantization_config
+            )
+            act_scale = paddle.zeros([], dtype="bfloat16").cuda()
             act_scale.stop_gradient = True
             state_dict[act_scale_name] = act_scale
         else:

paddlenlp/rl/algos/advantage.py

@@ -20,32 +20,57 @@ import paddle
 from ..utils.comm_utils import masked_whiten
 
 
-@paddle.no_grad()
 def compute_gae_advantage_return(
     token_level_rewards: paddle.Tensor,
     values: paddle.Tensor,
     sequence_mask: paddle.Tensor,
+    start: int,
     gamma: paddle.Tensor,
     lam: paddle.Tensor,
+    use_tgt_len_return: bool = True,
 ) -> Tuple[paddle.Tensor, paddle.Tensor]:
     """Compute advantages and returns using Generalized Advantage Estimation (GAE)."""
     # Modified from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py
-
     lastgaelam = 0.0
     advantages_reversed = []
     gen_len = token_level_rewards.shape[-1]
 
     values = values * sequence_mask
     token_level_rewards = token_level_rewards * sequence_mask
-
-    for t in reversed(range(0, gen_len)):
+    if use_tgt_len_return and start > 0:
+        # consistent with Beaver
+        # values length is src+tgt-1, start is src-1, return length is tgt
+        pass
+    elif use_tgt_len_return:
+        # values length is tgt, start is 0, return length is tgt
+        assert start == 0
+    else:
+        # values length is src+tgt-1, start is src-1, return length is src+tgt-1
+        pass
+    for t in reversed(range(start, gen_len)):  # pylint: disable=invalid-name
         next_values = values[:, t + 1] if t < gen_len - 1 else 0.0
         delta = token_level_rewards[:, t] + gamma * next_values - values[:, t]
         lastgaelam = delta + gamma * lam * lastgaelam
         advantages_reversed.append(lastgaelam)
     advantages = paddle.stack(advantages_reversed[::-1], axis=1)
 
-    returns = advantages + values.contiguous()
+    returns = advantages + values[:, start:].contiguous()
+
+    if not use_tgt_len_return:
+        advantages = paddle.concat(
+            [
+                paddle.zeros([advantages.shape[0], start], dtype=advantages.dtype),
+                advantages,
+            ],
+            axis=-1,
+        )
+        returns = paddle.concat(
+            [
+                paddle.zeros([returns.shape[0], start], dtype=returns.dtype),
+                returns,
+            ],
+            axis=-1,
+        )
 
     return advantages.detach(), returns
 
@@ -133,7 +158,6 @@ def compute_reinforce_plus_plus_advantages_and_returns(
     return advantages, returns
 
 
-@paddle.no_grad()
 def add_kl_divergence_regularization(
     prompt: paddle.Tensor,  # size = (B, S) # pylint: disable=unused-argument
     log_probs: paddle.Tensor,  # size = (B, L)

paddlenlp/rl/models/ppo_model_utils.py

@@ -30,12 +30,6 @@ from paddle.distributed import fleet
 from paddle.distributed.fleet.layers.mpu import mp_ops
 from paddle.distributed.fleet.meta_parallel import ParallelCrossEntropy
 
-try:
-    from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
-except:
-    pass
-
-
 from ...transformers.llama.modeling import (
     LlamaPretrainingCriterion as PretrainingCriterion,
 )
@@ -452,16 +446,6 @@ class RLHFPPOMixedLoss(nn.Layer):
 
         else:
             hidden_states, weight, bias, transpose_y = logits
-            if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel:
-                hidden_states = GatherOp.apply(hidden_states)
-                hidden_states = hidden_states.reshape(
-                    [
-                        input_ids.shape[0],
-                        -1,
-                        hidden_states.shape[-1],
-                    ]
-                )
-
             if use_remove_padding:
                 input_ids = raw_input_ids
                 if pad_size > 0:
@@ -615,7 +599,7 @@ class RLHFPPOMixedLoss(nn.Layer):
 
 @merge_fwd_labels
 class RLHFValueLoss(nn.Layer):
-    def __init__(self, config, clip_range_value=5.0, use_fp32_compute=False):
+    def __init__(self, config, clip_range_value=5.0):
         """
         Initializes the `ClipRewardRange` object.
 
@@ -633,7 +617,6 @@ class RLHFValueLoss(nn.Layer):
         super().__init__()
         self.clip_range_value = clip_range_value
         self.config = config
-        self.use_fp32_compute = use_fp32_compute
 
     def critic_loss_fn(
         self,
@@ -653,48 +636,54 @@ class RLHFValueLoss(nn.Layer):
         vf_loss2 = paddle.square(values_clipped - returns)
         return 0.5 * paddle.sum(paddle.maximum(vf_loss1, vf_loss2) * mask) / mask.sum()
 
-    def forward(
-        self,
-        reward_values,
-        old_reward_values,
-        reward_returns,
-        sequence_mask,
-        response_start=0,
-        # for varlen flaskmask
-        pad_size=0,
-        raw_input_ids=None,
-        indices=None,
-        raw_input_shape=None,
-        input_ids_rmpad_rolled=None,
-    ):
-        """ """
-        reward_values = reward_values[0].squeeze(0)
-        if self.config.sequence_parallel:
-            from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
+    def forward(self, reward_values, old_reward_values, reward_returns, sequence_mask):
+        """
+        计算奖励值的损失函数。
+        如果输入的奖励值和旧奖励值的长度相同，则使用给定的序列掩码来确定有效长度。
+        如果输入的奖励值的长度比旧奖励值少一个，则将最后一个元素视为与输入IDs一致的填充，并删除它。
+        否则，奖励值只有tgt长度。
 
-            reward_values = GatherOp.apply(reward_values)
+        Args:
+            reward_values (paddle.Tensor, list of paddle.Tensor or None, optional): 奖励值，可以是单个张量或列表中的多个张量。默认为None。
+            old_reward_values (paddle.Tensor, optional): 旧奖励值。
+            reward_returns (paddle.Tensor, optional): 奖励返回值。
+            sequence_mask (paddle.Tensor, optional): 序列掩码。
 
-        use_remove_padding = indices is not None
-        if use_remove_padding:
-            if pad_size > 0:
-                reward_values = reward_values[:-pad_size, :]
+        Returns:
+            paddle.Tensor, float32: 奖励值的损失函数。
 
-            from ..utils.bert_padding import pad_input
-
-            reward_values = pad_input(
-                reward_values.squeeze(0), indices, batch=raw_input_shape[0], seqlen=raw_input_shape[1]
-            ).contiguous()
-
-        if self.use_fp32_compute and reward_values.dtype != paddle.float32:
-            reward_values = reward_values.cast(paddle.float32)
-
-        reward_values = reward_values.squeeze(axis=-1)[:, response_start:-1]
+        Raises:
+            ValueError: 当奖励值和旧奖励值的长度不匹配时引发。
+        """
+        reward_values = reward_values if isinstance(reward_values, paddle.Tensor) else reward_values[0]
+        reward_values = reward_values.squeeze(axis=-1)[:, :-1]
+        if reward_values.shape[1] == old_reward_values.shape[1]:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # src+tgt-1 length, valid length is determined by sequence_mask
+            pass
+        elif reward_values.shape[1] < old_reward_values.shape[1]:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # src+tgt length and the last one is a padding to be consistent
+            # with input_ids
+            assert reward_values.shape[1] == old_reward_values.shape[1] - 1
+            reward_values = paddle.concat(
+                [
+                    reward_values,
+                    paddle.zeros([reward_values.shape[0], 1], dtype=reward_values.dtype),
+                ],
+                -1,
+            )
+        else:
+            # labels (old_reward_values, reward_returns, sequence_mask) has
+            # tgt length
+            reward_values = reward_values[:, -old_reward_values.shape[1] :]
         reward_critic_loss = self.critic_loss_fn(
             reward_values,
             old_reward_values,
             reward_returns,
             sequence_mask,
         )
+
         return reward_critic_loss
 
 

paddlenlp/rl/models/reward_utils.py

@@ -1,41 +0,0 @@
-# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import re
-
-
-def extract_solution(solution_str, method="strict"):
-    assert method in ["strict", "flexible"]
-
-    if method == "strict":
-        # this also tests the formatting of the model
-        solutions = re.findall("#### (\\-?[0-9\\.\\,]+)", solution_str)
-        if len(solutions) == 0:
-            final_answer = None
-        else:
-            # take the last solution
-            final_answer = solutions[-1].replace(",", "").replace("$", "")
-    elif method == "flexible":
-        answer = re.findall("(\\-?[0-9\\.\\,]+)", solution_str)
-        final_answer = None
-        if len(answer) == 0:
-            # no reward is there is no answer
-            pass
-        else:
-            invalid_str = ["", "."]
-            # find the last number that is not '.'
-            for final_answer in reversed(answer):
-                if final_answer not in invalid_str:
-                    break
-    return final_answer

paddlenlp/rl/models/score_model.py

@@ -18,7 +18,6 @@ from typing import Any, Optional
 
 import paddle
 from paddle import nn
-from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
 
 from ...transformers import PretrainedConfig, PretrainedModel
 from ...transformers.auto import AutoModel
@@ -41,35 +40,10 @@ class AutoModelForScore(ScoreModelMixin, PretrainedModel):
             TypeError: If the config is not an instance of `PretrainedConfig`.
         """
         super().__init__(config)
+        self.model = AutoModel.from_config(config)
 
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, *model_args, **kwargs):
-        config = kwargs.pop("config", None)
-        model = cls(config, **kwargs)
-        model.config = config
-        model.model = AutoModel.from_pretrained(pretrained_model_name_or_path, config=config)
-        model.init_score_head(config, hidden_size=config.hidden_size, **kwargs)
-        model.head_init_weights()
-        return model
-
-    def head_init_weights(self):
-        self.score_head.weight.set_value(
-            paddle.tensor.normal(
-                mean=0.0,
-                std=self.config.initializer_range,
-                shape=self.score_head.weight.shape,
-            )
-        )
-        if hasattr(self.score_head, "bias") and isinstance(self.score_head.bias, paddle.Tensor):
-            self.score_head.bias.set_value(paddle.zeros_like(self.score_head.bias))
-
-    @classmethod
-    def from_config(cls, config, **kwargs):
-        model = cls(config, **kwargs)
-        model.model = AutoModel.from_config(config)
-        model.config = config
-        model.init_score_head(config, hidden_size=config.hidden_size, **kwargs)
-        return model
+        # config.architectures = [self.__class__.__name__]
+        self.init_score_head(config, hidden_size=config.hidden_size, **kwargs)
 
     def get_input_embeddings(self) -> Optional[nn.Embedding]:
         """
@@ -117,6 +91,7 @@ class AutoModelForScore(ScoreModelMixin, PretrainedModel):
     def forward(
         self,
         input_ids: paddle.Tensor,
+        attention_mask: paddle.Tensor,
         position_ids: paddle.Tensor | None = None,
         past_key_values: list[paddle.Tensor] | None = None,
         inputs_embeds: paddle.Tensor | None = None,
@@ -124,8 +99,6 @@ class AutoModelForScore(ScoreModelMixin, PretrainedModel):
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
-        attn_mask_startend_row_indices: paddle.Tensor | None = None,
-        **kwargs
     ) -> tuple[paddle.Tensor, paddle.Tensor] | ScoreModelOutput:
         """
         Forward pass of the sentence.
@@ -157,7 +130,7 @@ class AutoModelForScore(ScoreModelMixin, PretrainedModel):
             AssertionError:
                 Raised when `attention_mask` is not None.
         """
-
+        assert attention_mask is not None
         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
@@ -166,26 +139,19 @@ class AutoModelForScore(ScoreModelMixin, PretrainedModel):
 
         outputs = self.model(
             input_ids=input_ids,
+            attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            attn_mask_startend_row_indices=attn_mask_startend_row_indices,
             return_dict=return_dict,
         )
-        critic_hidden_states = outputs[0]  # size = (B, L, E)
-
-        if self.config.sequence_parallel:
-            gather_hidden_states = GatherOp.apply(critic_hidden_states)
-
-        hidden_states = paddle.reshape_(
-            gather_hidden_states, [-1, position_ids.shape[1], gather_hidden_states.shape[-1]]
-        )
-
+        hidden_states = outputs[0]  # size = (B, L, E)
         return self.get_score(
             hidden_states,
+            attention_mask=attention_mask,
             position_ids=position_ids,
             return_dict=return_dict,
         )

paddlenlp/rl/models/score_model_utils.py

@@ -53,7 +53,6 @@ class ScoreModelMixin:
     normalize_function: NormalizeFunction = "affine"
     _initialized: bool = False
 
-    @classmethod
     def init_score_head(self, config: PretrainedConfig, hidden_size: int, **kwargs: Any) -> None:
         """Initialize the score head."""
         if self._initialized:
@@ -114,13 +113,11 @@ class ScoreModelMixin:
     def get_score(
         self,
         hidden_state: paddle.Tensor,  # size = (B, L, E)
-        position_ids: paddle.Tensor | None = None,  # size = (B, L)
-        attn_mask_startend_row_indices: paddle.Tensor | None = None,  # size = (B, 1), (B, 2), (B, 3) or (B, 4)
-        return_dict: bool | None = None,
         attention_mask: paddle.Tensor | None = None,  # size = (B, L)
+        position_ids: paddle.Tensor | None = None,  # size = (B, L)
+        return_dict: bool | None = None,
     ) -> ScoreModelOutput:
         """Forward pass of the score model."""
-        hidden_state = hidden_state.cast(paddle.float32)
         scores = self.score_head(hidden_state)  # size = (B, L, D)
         if scores.dtype != hidden_state.dtype:  # EB rm cast to float32
             scores = scores.cast(hidden_state.dtype)
@@ -173,6 +170,7 @@ class ScoreModelMixin:
 
         if self.do_normalize:
             scores = self.normalizer.normalize(scores)
+            end_score = self.normalizer.normalize(end_score)
 
         if not return_dict:
             return scores, end_score
@@ -182,8 +180,6 @@ class ScoreModelMixin:
             end_scores=end_score,  # size = (B, D)
         )
 
-        return scores
-
     def set_normalize(self, mode: bool = True) -> None:
         """
         Set whether to normalize the input data, default is True.

paddlenlp/rl/trainer/actor_trainer.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import uuid
-from typing import List
+from typing import Any, Dict, List
 
 import numpy as np
 import paddle
@@ -21,7 +21,6 @@ from paddle.distributed import fleet
 from paddle.distributed.fleet.layers.mpu import mp_ops
 from paddle.distributed.fleet.meta_parallel import ParallelCrossEntropy
 
-from ...datasets.rlhf_datasets.protocol import DataProto
 from ..models.ppo_model_utils import (
     RLHFPPOMixedLoss,
     create_startend_row_indices,
@@ -37,7 +36,7 @@ class ActorReferenceTrainerBase(RLTrainer):
     loss_identifier = lambda self, inputs: "actor_loss"
 
     @paddle.no_grad()
-    def compute_logprob(self, batch: DataProto, key) -> DataProto:
+    def compute_logprob(self, input_ids: paddle.Tensor, position_ids: paddle.Tensor = None, **kwargs) -> paddle.Tensor:
         """
         Computes the log probability of each token during the rollout process.
 
@@ -61,15 +60,11 @@ class ActorReferenceTrainerBase(RLTrainer):
         Raises:
             None.
         """
-        input_ids = batch.batch["input_ids"]
-        position_ids = batch.batch["position_ids"]
-        prompt = batch.batch.get("prompt", None)
         if self.args.use_fused_head_and_loss_fn:
             return self.compute_fused_logprob(
                 input_ids=input_ids,
-                key=key,
                 position_ids=position_ids,
-                prompt=prompt,
+                **kwargs,
             )
 
         log_probs_list = []
@@ -80,7 +75,7 @@ class ActorReferenceTrainerBase(RLTrainer):
         # Pipe model outputs a logits tensor with LMHead, while non-pipe model
         # outputs a tuple with logits tensor as the only one element.
         startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
-        response_start = prompt.shape[-1] - 1 if prompt is not None else 0
+        response_start = kwargs["prompt"].shape[-1] - 1 if "prompt" in kwargs else 0
 
         for i in range(num_batches):
             # Calculate the start and end indices for the current batch
@@ -160,14 +155,11 @@ class ActorReferenceTrainerBase(RLTrainer):
             logits = None
             paddle.device.cuda.empty_cache()
 
-        return DataProto.from_single_dict(
-            {key: paddle.concat(log_probs_list, axis=0)}, meta_info={"temperature": self.args.temperature}
-        )
+        return paddle.concat(log_probs_list, axis=0)
 
-    @paddle.no_grad()
     def compute_fused_logprob(
-        self, input_ids: paddle.Tensor, key, position_ids: paddle.Tensor = None, prompt=None, loop_chunk_size=1024
-    ) -> DataProto:
+        self, input_ids: paddle.Tensor, position_ids: paddle.Tensor = None, loop_chunk_size=1024, **kwargs
+    ):
         log_probs_list = []
         batch_size, sequence_length = input_ids.shape
         per_device_logprob_batch_size = self.args.per_device_logprob_batch_size
@@ -176,7 +168,7 @@ class ActorReferenceTrainerBase(RLTrainer):
         # Pipe model outputs a logits tensor with LMHead, while non-pipe model
         # outputs a tuple with logits tensor as the only one element.
         startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
-        response_start = (prompt.shape[-1] - 1) if prompt is not None else 0
+        response_start = kwargs["prompt"].shape[-1] - 1 if "prompt" in kwargs else 0
 
         num_embeddings = self.model.config.vocab_size
         tensor_parallel_degree = self.model.config.tensor_parallel_degree
@@ -219,7 +211,6 @@ class ActorReferenceTrainerBase(RLTrainer):
                 position_ids=current_position_ids,
                 attn_mask_startend_row_indices=current_startend_row_indices,
             )
-
             self.model.training = False
 
             if self.args.use_remove_padding:
@@ -311,22 +302,20 @@ class ActorReferenceTrainerBase(RLTrainer):
             log_prob_chunks = None
             paddle.device.cuda.empty_cache()
 
-        return DataProto.from_single_dict(
-            {key: paddle.concat(log_probs_list, axis=0)}, meta_info={"temperature": self.args.temperature}
-        )
+        return paddle.concat(log_probs_list, axis=0)
 
-    def update_actor(self, rl_batch: DataProto) -> DataProto:
+    def update_actor(self, rl_batch: Dict[str, paddle.Tensor]) -> Dict[str, Any]:
         # inputs shared by policy and value trainer
-        input_ids = rl_batch.batch["input_ids"].contiguous()  # length: src+tgt
-        position_ids = rl_batch.batch["position_ids"]  # length: src+tgt
-        sequence_mask = rl_batch.batch["eos_mask"]  # length: tgt(-1)
+        input_ids = rl_batch["input_ids"].contiguous()  # length: src+tgt
+        position_ids = rl_batch["position_ids"]  # length: src+tgt
+        sequence_mask = rl_batch["eos_mask"]  # length: tgt(-1)
         if self.args.use_fp32_compute and sequence_mask.dtype != paddle.float32:
             sequence_mask = sequence_mask.cast(paddle.float32)
         # inputs used by policy trainer
-        old_log_probs = rl_batch.batch["log_probs"]  # length: tgt(-1)
-        reward_advantages = rl_batch.batch["reward_advantages"]  # length: tgt(-1)
+        old_log_probs = rl_batch["log_probs"]  # length: tgt(-1)
+        reward_advantages = rl_batch["reward_advantages"]  # length: tgt(-1)
 
-        response_start = rl_batch.batch["prompt"].shape[-1] - 1
+        response_start = rl_batch["prompt"].shape[-1] - 1
 
         attn_mask_startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
         policy_trainer_inputs = {
@@ -340,7 +329,7 @@ class ActorReferenceTrainerBase(RLTrainer):
         }
 
         if self.args.rl_algorithm == "grpo":
-            policy_trainer_inputs.update({"ref_log_probs": rl_batch.batch["ref_log_probs"]})
+            policy_trainer_inputs.update({"ref_log_probs": rl_batch["ref_log_probs"]})
         else:
             policy_trainer_inputs.update({"ref_log_probs": None})
 
@@ -348,65 +337,61 @@ class ActorReferenceTrainerBase(RLTrainer):
 
         # metric
         with paddle.no_grad():
-            rewards = rl_batch.batch["rewards"].mean()
-            ori_rewards = rl_batch.batch["ori_rewards"].mean()
+            rewards = rl_batch["rewards"].mean()
+            ori_rewards = rl_batch["ori_rewards"].mean()
             mask_cast = sequence_mask.cast(paddle.float32)
             if self.args.rl_algorithm in ["ppo", "reinforce_plus_plus"]:
-                kl_rewards = (rl_batch.batch["kl_rewards"] * mask_cast).sum() / mask_cast.sum()
-                rewards_with_kl = (rl_batch.batch["rewards_with_kl"] * mask_cast).sum() / mask_cast.sum()
+                kl_rewards = (rl_batch["kl_rewards"] * mask_cast).sum() / mask_cast.sum()
+                rewards_with_kl = (rl_batch["rewards_with_kl"] * mask_cast).sum() / mask_cast.sum()
                 if self.args.rl_algorithm == "ppo":
-                    values = (rl_batch.batch["reward_values"] * mask_cast).sum() / mask_cast.sum()
-                returns = (rl_batch.batch["reward_returns"] * mask_cast).sum() / mask_cast.sum()
-            ref_log_probs = rl_batch.batch["ref_log_probs"]
+                    values = (rl_batch["reward_values"] * mask_cast).sum() / mask_cast.sum()
+                returns = (rl_batch["reward_returns"] * mask_cast).sum() / mask_cast.sum()
+            ref_log_probs = rl_batch["ref_log_probs"]
             kl_divergence = ((old_log_probs - ref_log_probs) * mask_cast).sum() / mask_cast.sum()
             mean_generated_length = mask_cast.sum(axis=-1).mean()
             max_generated_length = mask_cast.sum(axis=-1).max()
             min_generated_length = mask_cast.sum(axis=-1).min()
 
-        return DataProto(
-            meta_info={
-                "metrics": {
-                    # when using PipelienParallel, the loss returned is 0 when not reach
-                    # accumulated step and the loss returned at accumulated step is a
-                    # mixed loss.
-                    "train_policy_loss": actor_loss,
-                    **(
-                        {
-                            "train_pure_policy_loss": self.info_buffer.get("pure_policy_loss"),
-                            "train_kl_loss": self.info_buffer.get("kl_loss"),
-                            "train_entropy_loss": self.info_buffer.get("entropy_loss"),
-                        }
-                        if self.args.rl_algorithm == "grpo"
-                        else {}
-                    ),
-                    "train_reward": ori_rewards,  # use original reward to log
-                    **(
-                        {
-                            "train_norm_reward": rewards,
-                            "train_kl_reward": kl_rewards,
-                            "train_norm_reward_with_kl": rewards_with_kl,
-                            "train_pure_policy_loss": self.info_buffer.get("pure_policy_loss"),
-                            "train_entropy_loss": self.info_buffer.get("entropy_loss"),
-                            **({"train_values": values} if self.args.rl_algorithm == "ppo" else {}),
-                            "train_returns": returns,
-                        }
-                        if self.args.rl_algorithm in ["ppo", "reinforce_plus_plus"]
-                        else {}
-                    ),
-                    "train_kl_divergence": kl_divergence,
-                    "train_mean_generated_length": mean_generated_length,
-                    "train_max_generated_length": max_generated_length,
-                    "train_min_generated_length": min_generated_length,
+        return {
+            # when using PipelienParallel, the loss returned is 0 when not reach
+            # accumulated step and the loss returned at accumulated step is a
+            # mixed loss.
+            "train_policy_loss": actor_loss,
+            **(
+                {
+                    "train_pure_policy_loss": self.info_buffer.get("pure_policy_loss"),
+                    "train_kl_loss": self.info_buffer.get("kl_loss"),
+                    "train_entropy_loss": self.info_buffer.get("entropy_loss"),
                 }
-            }
-        )
+                if self.args.rl_algorithm == "grpo"
+                else {}
+            ),
+            "train_reward": ori_rewards,  # use original reward to log
+            **(
+                {
+                    "train_norm_reward": rewards,
+                    "train_kl_reward": kl_rewards,
+                    "train_norm_reward_with_kl": rewards_with_kl,
+                    "train_pure_policy_loss": self.info_buffer.get("pure_policy_loss"),
+                    "train_entropy_loss": self.info_buffer.get("entropy_loss"),
+                    **({"train_values": values} if self.args.rl_algorithm == "ppo" else {}),
+                    "train_returns": returns,
+                }
+                if self.args.rl_algorithm in ["ppo", "reinforce_plus_plus"]
+                else {}
+            ),
+            "train_kl_divergence": kl_divergence,
+            "train_mean_generated_length": mean_generated_length,
+            "train_max_generated_length": max_generated_length,
+            "train_min_generated_length": min_generated_length,
+        }
 
 
 class ActorReferenceTrainer(ActorReferenceTrainerBase):
     @paddle.no_grad()
-    def generate_sequences(self, prompt_only_batch: DataProto, do_eval=False) -> List[DataProto]:
+    def generate_sequences(self, prompt_only_batch: Dict, do_eval=False) -> List[Dict[str, Any]]:
         """Rollout a batch of experiences."""
-        input_ids = prompt_only_batch.batch["input_ids"]
+        input_ids = prompt_only_batch["input_ids"]
 
         repeat_num = 1 if do_eval else self.args.rollout_n
 
@@ -422,8 +407,8 @@ class ActorReferenceTrainer(ActorReferenceTrainerBase):
         if repeat_num > 1:
             input_ids = input_ids.repeat_interleave(repeat_num, axis=0)
 
-        if self.args.use_rm_server or self.args.use_rule_reward:
-            label_ids = prompt_only_batch.batch["label_ids"]
+        if self.args.use_rm_server:
+            label_ids = prompt_only_batch["label_ids"]
             if repeat_num > 1:
                 label_ids = label_ids.repeat_interleave(repeat_num, axis=0)
 
@@ -432,17 +417,11 @@ class ActorReferenceTrainer(ActorReferenceTrainerBase):
             sequences = sequences.transpose([1, 0, 2])
         # prompt, sequence, attention_mask
         return [
-            DataProto.from_single_dict(
-                {
-                    "prompt": input_ids[idx * len(seq) : (idx + 1) * len(seq)],  # src prompt
-                    "input_ids": seq,
-                    **(
-                        {"label_ids": label_ids[idx * len(seq) : (idx + 1) * len(seq)]}
-                        if self.args.use_rm_server or self.args.use_rule_reward
-                        else {}
-                    ),  # tgt response
-                    "index": np.array([str(uuid.uuid4())] * len(seq), dtype=object),
-                }
-            )
+            {
+                "prompt": input_ids,
+                "input_ids": seq,
+                **({"label_ids": label_ids[idx * len(seq) : (idx + 1) * len(seq)]} if self.args.use_rm_server else {}),
+                "index": np.array([str(uuid.uuid4())] * len(seq), dtype=object),
+            }
             for idx, seq in enumerate(sequences)
         ]

paddlenlp/rl/trainer/critic_trainer.py

@@ -15,15 +15,15 @@
 
 from __future__ import annotations
 
+from typing import Any, Dict
+
 import paddle
 
-from ...datasets.rlhf_datasets.protocol import DataProto
 from ...transformers import PretrainedTokenizer
 from ..models.ppo_model_utils import RLHFValueLoss, create_startend_row_indices
-
-# from ..utils.comm_utils import CriticStages
-# from ..utils.offload_utils import reload_and_offload_scope
-# from ..utils.timer_utils import TimerScope
+from ..utils.comm_utils import CriticStages
+from ..utils.offload_utils import reload_and_offload_scope
+from ..utils.timer_utils import TimerScope
 from .rl_trainer import RLTrainer
 
 
@@ -33,79 +33,28 @@ class CriticTrainer(RLTrainer):
     # define loss name for logging
     loss_identifier = lambda self, inputs: "reward_critic_loss"
 
-    @paddle.no_grad()
     def compute_value(
         self,
-        batch: DataProto,
+        input_ids: paddle.Tensor,
+        position_ids: paddle.Tensor = None,
         input_ids_tokenizer: PretrainedTokenizer = None,
-    ) -> DataProto:
-        self.model.eval()
+        **kwargs,
+    ) -> Dict[str, paddle.Tensor]:
+        # TODO: confirm actor_tokenizer or reward_tokenizer or critic_tokenizer
+        # need retokenize?
+        attn_mask_startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
+        reward_value = self.model(
+            input_ids,
+            attention_mask=None,
+            position_ids=position_ids,
+            attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+        )[0]
+        reward_value = reward_value.squeeze(axis=-1)
+        reward_value = reward_value[:, :-1]
 
-        input_ids = batch.batch["input_ids"]
-        position_ids = batch.batch["position_ids"]
+        return reward_value
 
-        values_list = []
-        batch_size, sequence_length = input_ids.shape
-        per_device_value_batch_size = self.args.per_device_value_batch_size
-        num_batches = (batch_size + per_device_value_batch_size - 1) // per_device_value_batch_size
-        startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
-        response_start = batch.batch["prompt"].shape[-1] - 1 if "prompt" in batch.batch else 0
-        for i in range(num_batches):
-            start_index = i * per_device_value_batch_size
-            end_index = min(start_index + per_device_value_batch_size, batch_size)
-
-            # Extract the current batch
-            current_input_ids = input_ids[start_index:end_index]
-            current_startend_row_indices = (
-                startend_row_indices[start_index:end_index] if startend_row_indices is not None else None
-            )
-            current_position_ids = position_ids[start_index:end_index] if position_ids is not None else None
-            if self.args.use_remove_padding:
-                from ..utils.bert_padding import prepare_flashmask_inputs
-
-                update_inputs = prepare_flashmask_inputs(
-                    current_input_ids,
-                    current_position_ids,
-                    self.tokenizer.pad_token_id,
-                    self.model.config.sequence_parallel,
-                    self.model.config.tensor_parallel_degree,
-                )
-                current_input_ids = update_inputs["input_ids"]
-                current_position_ids = update_inputs["position_ids"]
-                current_startend_row_indices = update_inputs["attn_mask_startend_row_indices"]
-                indices = update_inputs["indices"]
-                raw_input_shape = update_inputs["raw_input_shape"]
-                pad_size = update_inputs["pad_size"]
-            reward_value = self.model(
-                current_input_ids,
-                position_ids=current_position_ids,
-                attn_mask_startend_row_indices=current_startend_row_indices,
-                use_cache=False,
-            )[0]
-            reward_value = reward_value.squeeze(0)
-            if self.model.config.sequence_parallel:
-                from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-                    GatherOp,
-                )
-
-                reward_value = GatherOp.apply(reward_value)
-
-            if self.args.use_remove_padding:
-                from ..utils.bert_padding import pad_input
-
-                if pad_size > 0:
-                    reward_value = reward_value[:-pad_size, :]
-                reward_value = pad_input(
-                    reward_value.squeeze(0).unsqueeze(-1), indices, batch=raw_input_shape[0], seqlen=raw_input_shape[1]
-                ).squeeze(-1)
-                reward_value = reward_value[:, response_start:-1].contiguous()
-            values_list.append(reward_value.squeeze(-1))
-            reward_value = None
-            paddle.device.cuda.empty_cache()
-
-        return DataProto.from_single_dict({"reward_values": paddle.concat(values_list, axis=0)})
-
-    def update_critic(self, rl_batch: DataProto) -> DataProto:
+    def update_critic(self, rl_batch: Dict[str, paddle.Tensor]) -> Dict[str, Any]:
         """
         Update the parameters of the critic (reward function).
 
@@ -120,30 +69,30 @@ class CriticTrainer(RLTrainer):
         Returns (Dict[str, Any]):
             - train_value_loss (float): Training loss of the critic (reward function).
         """
-        self.model.train()
         # Inputs shared by policy and value trainer
-        input_ids = rl_batch.batch["input_ids"].contiguous()  # length: src+tgt
-        position_ids = rl_batch.batch["position_ids"]  # length: src+tgt
-        sequence_mask = rl_batch.batch["eos_mask"]  # length: src+tgt(-1)
-        if self.args.use_fp32_compute and sequence_mask.dtype != paddle.float32:
-            sequence_mask = sequence_mask.cast(paddle.float32)
+        input_ids = rl_batch["input_ids"].contiguous()  # length: src+tgt
+        attention_mask = rl_batch["attention_mask"]  # length: src+tgt
+        position_ids = rl_batch["position_ids"]  # length: src+tgt
+        sequence_mask = rl_batch["sequence_mask"]  # length: src+tgt(-1)
 
         # Inputs used by value trainer
-        old_reward_values = rl_batch.batch["reward_values"]  # length: src+tgt(-1)
-        reward_returns = rl_batch.batch["reward_returns"]  # length: src+tgt(-1)
+        old_reward_values = rl_batch["reward_values"]  # length: src+tgt(-1)
+        reward_returns = rl_batch["reward_returns"]  # length: src+tgt(-1)
 
-        attn_mask_startend_row_indices = create_startend_row_indices(input_ids, self.tokenizer.pad_token_id)
         value_trainer_inputs = {
             "input_ids": input_ids,
+            "attention_mask": attention_mask,
             "position_ids": position_ids,
             "old_reward_values": old_reward_values,
             "reward_returns": reward_returns,
             "sequence_mask": sequence_mask,
-            "response_start": rl_batch.batch["prompt"].shape[-1] - 1,
-            "attn_mask_startend_row_indices": attn_mask_startend_row_indices,
         }
 
-        reward_critic_loss = self.full_training_step(**value_trainer_inputs)
+        with TimerScope(
+            self.timers, CriticStages.MODEL_ENABLE_DISABLE, minus_names=[CriticStages.CRITIC_TRAINING_STEP]
+        ):
+            with reload_and_offload_scope(self, self.model, self.optimizer):
+                with TimerScope(self.timers, CriticStages.CRITIC_TRAINING_STEP):
+                    reward_critic_loss = self.full_training_step(**value_trainer_inputs)
 
-        # return DataProto(meta_info={"metrics": {"train_value_loss": reward_critic_loss}})
-        return {"train_value_loss": reward_critic_loss}
+        return reward_critic_loss

paddlenlp/rl/trainer/reward_trainer.py

@@ -25,7 +25,6 @@ from paddle.distributed import fleet
 from paddle.io import Dataset
 
 from ...data import DataCollator
-from ...datasets.rlhf_datasets.protocol import DataProto
 from ...trainer.trainer import (
     EvalPrediction,
     TrainerCallback,
@@ -34,7 +33,6 @@ from ...trainer.trainer import (
 )
 from ...transformers import PretrainedModel, PretrainedTokenizer
 from ..models.ppo_model_utils import create_startend_row_indices
-from ..models.reward_utils import extract_solution
 from .rl_trainer import RLTrainer
 from .trainer_utils import batch_retokenize
 
@@ -57,10 +55,7 @@ class RewardTrainer(RLTrainer):
         preprocess_logits_for_metrics: Optional[Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor]] = None,
         reward_server: str = None,
     ):
-        if args.use_rule_reward:
-            self.args = args
-            self.tokenizer = tokenizer
-        elif args.use_rm_server:
+        if args.use_rm_server:
             assert isinstance(model, str), "reward trainer need a str (http://xxx:port) for request"
             self.args = args
             self.tokenizer = tokenizer
@@ -84,27 +79,13 @@ class RewardTrainer(RLTrainer):
     @paddle.no_grad()
     def compute_reward(
         self,
-        batch: DataProto,
+        input_ids: paddle.Tensor,
+        position_ids: paddle.Tensor = None,
         input_ids_tokenizer: PretrainedTokenizer = None,
+        label_ids: paddle.Tensor = None,
+        **kwargs,
     ) -> Dict[str, paddle.Tensor]:
-        input_ids = batch.batch["input_ids"]
-        position_ids = batch.batch["position_ids"]
-        label_ids = batch.batch["label_ids"]
-        prompt = batch.batch["prompt"]
-
-        if self.args.use_rule_reward:
-            prompt_len = prompt.shape[-1]
-            if label_ids is None:
-                raise ValueError("Rule-based reward needs labels.")
-            tgt = input_ids_tokenizer.batch_decode(label_ids, skip_special_tokens=False)
-            response = input_ids_tokenizer.batch_decode(input_ids[:, prompt_len:], skip_special_tokens=False)
-            ground_truth = [i.replace(self.tokenizer.pad_token, "") for i in tgt]
-            response = [i.replace(self.tokenizer.pad_token, "") for i in response]
-            reward_score = self.compute_score(
-                solution=response,
-                ground_truth=ground_truth,
-            )
-        elif not self.args.use_rm_server:
+        if not self.args.use_rm_server:
             if self.tokenizer is not input_ids_tokenizer:
                 # right padding
                 reward_tokenize_output = batch_retokenize(
@@ -126,7 +107,7 @@ class RewardTrainer(RLTrainer):
                 position_ids=reward_position_ids,
             )[1]
         else:
-            prompt_len = prompt.shape[-1]
+            prompt_len = kwargs["prompt"].shape[-1]
             if label_ids is None:
                 raise ValueError("Rule-based reward needs labels.")
             src = input_ids_tokenizer.batch_decode(input_ids[:, :prompt_len], skip_special_tokens=False)
@@ -141,6 +122,8 @@ class RewardTrainer(RLTrainer):
         reward_score = reward_score.squeeze(axis=-1)
 
         return reward_score
+        # if self.args.rl_algorithm in ["grpo", "reinforce_plus_plus"]:
+        #     return {"rewards": reward_score}
 
     def request_reward_server(self, src, tgt, response):
         data = {"src": src, "tgt": tgt, "response": response}
@@ -182,30 +165,4 @@ class RewardTrainer(RLTrainer):
             paddle.distributed.barrier(tp_group)
             paddle.distributed.broadcast(reward_score, src=tp_group.ranks[0], group=tp_group)
 
-        reward_score = reward_score.unsqueeze(-1)
-        return reward_score
-
-    def compute_score(self, solution, ground_truth, method="strict", format_score=0.0, score=1.0):
-        """The scoring function for GSM8k.
-
-        Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual
-        Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.
-
-        Args:
-            solution: the solution text
-            ground_truth: the ground truth
-            method: the method to extract the solution, choices are 'strict' and 'flexible'
-            format_score: the score for the format
-            score: the score for the correct answer
-        """
-        reward_tensor = paddle.zeros((len(solution), 1), dtype=paddle.float32)
-        for i in range(len(solution)):
-            answer = extract_solution(solution_str=solution[i], method=method)
-            if answer is None:
-                reward_tensor[i] = 0
-            else:
-                if answer == ground_truth[i]:
-                    reward_tensor[i] = score
-                else:
-                    reward_tensor[i] = format_score
-        return reward_tensor
+        return reward_score.unsqueeze(-1)

paddlenlp/rl/utils/comm_utils.py

@@ -15,7 +15,6 @@
 import sys
 from collections import defaultdict
 from enum import Enum, auto
-from functools import wraps
 
 import numpy as np
 import paddle
@@ -23,7 +22,6 @@ import paddle.distributed as dist
 from paddle import nn
 from paddle.distributed import fleet
 
-from ...datasets.rlhf_datasets.protocol import DataProto, make_eos_mask
 from ...trainer.trainer import Trainer, logger
 from ...utils.nested import flatten_list, nested_broadcast_tensor_with_empty
 from ..models.ppo_model_utils import make_position_ids_from_input_ids
@@ -751,6 +749,29 @@ def masked_whiten(values, mask, shift_mean=True):
     return whitened
 
 
+def pad_tensor(tensor_list, pad_index=0.0, dtype="bfloat16", padding_side="right"):
+    max_size = max([i.shape[-1] for i in tensor_list])
+    data_num = sum([i.shape[0] for i in tensor_list])
+    if isinstance(tensor_list[0], paddle.Tensor):
+        new_tensor = paddle.full((data_num, max_size), pad_index, dtype=dtype)
+    elif isinstance(tensor_list[0], np.ndarray):
+        new_tensor = np.full((data_num, max_size), pad_index, dtype=dtype)
+
+    offset = 0
+    for idx, i in enumerate(tensor_list):
+        # new_tensor[offset : offset + i.shape[0], : i.shape[-1]] = i
+        data_length = i.shape[-1]
+
+        if padding_side == "right":
+            new_tensor[offset : offset + i.shape[0], :data_length] = i
+        elif padding_side == "left":
+            new_tensor[offset : offset + i.shape[0], -data_length:] = i
+        else:
+            raise ValueError("padding_side must be 'right' or 'left'")
+        offset += i.shape[0]
+    return new_tensor
+
+
 def gather_and_pad(tensor, dp_group=None, sd_group=None, pad_index=0.0, pad=True, padding_side="right"):
     """Gather tensor from all devices."""
 
@@ -773,7 +794,7 @@ def gather_and_pad(tensor, dp_group=None, sd_group=None, pad_index=0.0, pad=True
             else:
                 return np.concatenate(tensor, axis=0)
         else:
-            return DataProto.pad_tensor(tensor, pad_index=pad_index, dtype=dtype, padding_side=padding_side)
+            return pad_tensor(tensor, pad_index=pad_index, dtype=dtype, padding_side=padding_side)
 
     def map_func(weight):
         if isinstance(weight, paddle.Tensor):
@@ -806,10 +827,49 @@ def gather_and_pad(tensor, dp_group=None, sd_group=None, pad_index=0.0, pad=True
         else:
             return np.concatenate(flatten_list(gathered_tensor), axis=0)
     else:
-        return DataProto.pad_tensor(gathered_tensor, pad_index=pad_index, dtype=dtype, padding_side=padding_side)
+        return pad_tensor(gathered_tensor, pad_index=pad_index, dtype=dtype, padding_side=padding_side)
 
 
-def filter_valid_reward_groups(combined_batch: DataProto, total_batch, rollout_n, variance_threshold=1e-6):
+def combine_micro_batches_into_batch(micro_batches, pad_token_id=0):
+    """combine micro batches to get a complete batch"""
+    if not isinstance(micro_batches, list):
+        return micro_batches
+    combined_batch = {}
+
+    for micro_batch in micro_batches:
+        for key, value in micro_batch.items():
+            if isinstance(value, list):
+                if isinstance(value[0], paddle.Tensor):
+                    if key == "label_ids":
+                        value = [paddle.unsqueeze(v, axis=0) if v.ndim == 1 else v for v in value]
+                        concat_value = pad_tensor(
+                            value,
+                            pad_index=pad_token_id,
+                            dtype=value[0].dtype,
+                            padding_side="left",
+                        )
+                    else:
+                        concat_value = paddle.concat(value, axis=0)
+                elif isinstance(value[0], np.ndarray):
+                    concat_value = np.concatenate(value, axis=0)
+                combined_batch.setdefault(key, []).append(concat_value)
+            else:
+                combined_batch.setdefault(key, []).append(value)
+
+    for key, values in combined_batch.items():
+        if len(combined_batch[key][0].shape) > 1:
+            pad_index = pad_token_id
+            padding_side = "left" if (key == "prompt" or key == "label_ids") else "right"
+            combined_batch[key] = gather_and_pad(values, pad_index=pad_index, padding_side=padding_side)
+        elif isinstance(values[0], paddle.Tensor):
+            combined_batch[key] = paddle.concat(values, axis=0)
+        elif isinstance(values[0], np.ndarray):
+            combined_batch[key] = np.concatenate(values, axis=0)
+
+    return combined_batch
+
+
+def filter_valid_reward_groups(combined_batch, total_batch, rollout_n, variance_threshold=1e-6):
     """
     Filters out invalid prompt groups based on reward variance, and appends the valid samples to total_batch.
 
@@ -827,14 +887,10 @@ def filter_valid_reward_groups(combined_batch: DataProto, total_batch, rollout_n
     """
 
     # Choose the reward key to filter by
-    select_key = (
-        "rewards_before_length_penalty"
-        if "rewards_before_length_penalty" in combined_batch.batch.keys()
-        else "rewards"
-    )
+    select_key = "rewards_before_length_penalty" if "rewards_before_length_penalty" in combined_batch else "rewards"
 
-    rewards = combined_batch.batch[select_key].flatten()  # paddle.Tensor
-    indices = combined_batch.non_tensor_batch["index"].flatten()  # numpy.ndarray
+    rewards = combined_batch[select_key].flatten()  # paddle.Tensor
+    indices = combined_batch["index"].flatten()  # numpy.ndarray
 
     # Group by prompt index
     group_map = defaultdict(list)
@@ -856,11 +912,8 @@ def filter_valid_reward_groups(combined_batch: DataProto, total_batch, rollout_n
 
     # Select only valid samples for each key and append to total_batch
     valid_indices = np.array(valid_indices, dtype=int)
-    for key in combined_batch.batch.keys():
-        filtered = combined_batch.batch[key][valid_indices]
-        total_batch[key].append(filtered)
-    for key in combined_batch.non_tensor_batch.keys():
-        filtered = combined_batch.non_tensor_batch[key][valid_indices]
+    for key in combined_batch:
+        filtered = combined_batch[key][valid_indices]
         total_batch[key].append(filtered)
 
     return total_batch, num_valid_prompts
@@ -983,93 +1036,62 @@ def process_prompt_and_response(micro_batch, pad_token_id=0):
     return micro_batch
 
 
-def gather_tensor(tensor, dp_group=None, sd_group=None):
-    """Gather tensor from all devices."""
+def split_batch_into_micro_batches(total_batch, batch_size, pad_token_id=0):
+    """
+    Splits total_batch into micro-batches of size `batch_size`.
 
-    if not isinstance(tensor, list):
-        tensor = [tensor]
+    Args:
+        total_batch (dict): Dictionary containing full batched tensors.
+        batch_size (int): Micro batch size per device.
 
-    if isinstance(tensor[0], paddle.Tensor):
-        type = "tensor"
-    elif isinstance(tensor[0], np.ndarray):
-        type = "numpy"
-    else:
-        raise TypeError(f"{type(tensor[0])} is not supported for gather and pad")
+    Returns:
+        list of dict: A list of micro-batches.
+    """
+    micro_batches = []
+    num_micro_batches = total_batch["input_ids"].shape[0] // batch_size
+    if total_batch["input_ids"].shape[0] % batch_size != 0:
+        num_micro_batches += 1
+    if num_micro_batches <= 0:
+        logger.warning(
+            "The total batch size is smaller than the batch size, please consider using a smaller batch size or a larger global_batch_size."
+        )
+        num_micro_batches = 1
 
-    dtype = tensor[0].dtype
+    for i in range(num_micro_batches):
+        micro_batch = {}
+        for key, data in total_batch.items():
+            if isinstance(data, paddle.Tensor):
+                micro_batch[key] = data[i * batch_size : (i + 1) * batch_size]
+            elif isinstance(data, np.ndarray):
+                micro_batch[key] = data[i * batch_size : (i + 1) * batch_size]
+            elif isinstance(data, list):
+                micro_batch[key] = data[i * batch_size : (i + 1) * batch_size]
+            else:
+                raise TypeError(f"Unsupported data type for key {key}: {type(data)}")
 
-    if (dp_group is None and sd_group is None) or (dp_group.nranks == 1 and sd_group.nranks == 1):
-        return tensor
+        # if os.getenv("PROCESS_PROMPT_AND_RESPONSE", "1").lower() in ["1", "t", "true", "yes", "y"]:
+        #     micro_batch = process_prompt_and_response(micro_batch=micro_batch, pad_token_id=pad_token_id)
 
-    def map_func(weight):
-        if isinstance(weight, paddle.Tensor):
-            weight = weight.numpy()
-        return weight
+        micro_batches.append(micro_batch)
 
-    tensor = [map_func(i) for i in tensor]
-
-    sd_gathered_tensor = []
-    if sd_group.nranks > 1:
-        dist.all_gather_object(sd_gathered_tensor, tensor, group=sd_group)
-
-    dp_gathered_tensor = []
-    if dp_group.nranks > 1:
-        if len(sd_gathered_tensor) > 0:
-            tensor = sd_gathered_tensor
-        dist.all_gather_object(dp_gathered_tensor, tensor, group=dp_group)
-
-    if len(dp_gathered_tensor) > 0:
-        gathered_tensor = dp_gathered_tensor
-    else:
-        gathered_tensor = sd_gathered_tensor
-
-    if type == "tensor":
-        gathered_tensor = [paddle.to_tensor(i, dtype=dtype) for i in flatten_list(gathered_tensor)]
-
-    return gathered_tensor
+    return micro_batches
 
 
-def gather_tensor_list(dp_group=None, sd_group=None):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(tensors_list, *args, **kwargs):
-            gathered = gather_tensor(tensors_list, dp_group, sd_group)
-            return func(gathered, *args, **kwargs)
+def make_eos_mask(response_id, eos_token_ids=0, dtype=paddle.int64):
+    """
+    end of sentence token can be int or list: 1 or [1, 2]
+    e.g. eos_token=1
+    response_id: [0, 0, 2, 42, 3, 5, 1, 0, 0]
+    eos_mask:     [1, 1, 1, 1,  1, 1, 1, 0, 0]
+    """
+    if isinstance(eos_token_ids, int):
+        eos_token_ids = [eos_token_ids]
 
-        return wrapper
+    eos_mask = paddle.zeros_like(response_id, dtype=paddle.bool)
+    for token_id in eos_token_ids:
+        eos_mask |= response_id == token_id
 
-    return decorator
-
-
-def gather_and_pad_dataproto(batch, dp_group, sd_group, eos_token_ids, pad_token_id, select_keys) -> "DataProto":
-    new_batch = {}
-    gather_then_pad_or_concat = gather_tensor_list(dp_group, sd_group)(DataProto.pad_or_concat_tensor_list)
-
-    if "eos_mask" in select_keys:
-        eos_mask = make_eos_mask(
-            batch.batch["input_ids"][:, batch.batch["prompt"].shape[-1] :],
-            eos_token_ids=eos_token_ids,
-        ).to(batch.batch["log_probs"].dtype)
-
-        new_batch["eos_mask"] = gather_then_pad_or_concat(eos_mask, pad_token_id, key="eos_mask")
-
-    for key in select_keys:
-        if key == "eos_mask":
-            continue
-        if key in batch.batch:
-            value = batch.batch[key]
-        elif key in batch.non_tensor_batch:
-            value = batch.non_tensor_batch[key]
-        else:
-            raise KeyError(f"{key} not found in batch or non_tensor_batch")
-
-        if not isinstance(value, list):
-            tensor_list = [value]
-        else:
-            tensor_list = value
-
-        gathered = gather_then_pad_or_concat(tensor_list, pad_token_id, key=key)
-
-        new_batch[key] = gathered
-
-    return DataProto.from_single_dict(new_batch)
+    eos_mask = eos_mask.to("int64")
+    eos_mask = (paddle.cumsum(eos_mask, axis=1) - eos_mask).to("bool")
+    eos_mask = paddle.logical_not(eos_mask).to(dtype)
+    return eos_mask

paddlenlp/rl/utils/config_utils.py

@@ -300,7 +300,6 @@ class TrainingArguments(TrainingArguments):
         metadata={"help": "Whether to use tgt for KL."},
     )
     use_rm_server: bool = field(default=False, metadata={"help": "Use reward server instead of reward model."})
-    use_rule_reward: bool = field(default=False, metadata={"help": "Use rule-based reward only for gsm8k, to date."})
     use_fp32_compute: bool = field(
         default=False, metadata={"help": "Use fp32 to compute xx_log_prob,rewards, advantages and loss."}
     )
@@ -383,7 +382,6 @@ class TrainingArguments(TrainingArguments):
             // self.per_device_train_batch_size
             // self.dataset_world_size
         )
-
         if self.gradient_accumulation_steps <= 0:
             logger.warning(
                 f"gradient_accumulation_steps: {self.gradient_accumulation_steps} must be greater than zero!"
@@ -445,14 +443,14 @@ class TrainingArguments(TrainingArguments):
             self.normalize_advantage = False
 
         max_per_device_eval_batch_size = (
-            self.global_batch_size * self.rollout_n * self.update_iters // self.dataset_world_size
+            self.global_mini_batch_size * self.rollout_n * self.update_iters // self.dataset_world_size
         )
         if self.per_device_eval_batch_size > max_per_device_eval_batch_size:
             logger.warning(
                 f"per_device_eval_batch_size: {self.per_device_eval_batch_size} is larger than "
-                f"global_batch_size: {self.global_batch_size} * rollout_n: "
+                f"global_mini_batch_size: {self.global_mini_batch_size} * rollout_n: "
                 f"{self.rollout_n} * update_iters: {self.update_iters}, which may cause infer error. "
-                f"We will set it to global_batch_size * rollout_n * update_iters // dataset_world_size!"
+                f"We will set it to global_mini_batch_size * rollout_n * update_iters // dataset_world_size!"
             )
             self.per_device_eval_batch_size = max_per_device_eval_batch_size
 
@@ -532,7 +530,7 @@ class ModelArgument:
     )
     actor_tokenizer_alpha: float = field(default=None, metadata={"help": "Tokenizer will tokenize randomly"})
     reward_tokenizer_alpha: float = field(default=None, metadata={"help": "Tokenizer will tokenize randomly"})
-    critic_tokenizer_alpha: float = field(default=None, metadata={"help": "Tokenizer will tokenize randomly"})
+    reward_critic_tokenizer_alpha: float = field(default=None, metadata={"help": "Tokenizer will tokenize randomly"})
     stage: str = field(default="PPO", metadata={"help": "The type of training."})
     critic_recompute_granularity: str = field(
         default="full",

paddlenlp/rl/utils/offload_utils.py

@@ -183,18 +183,14 @@ def reload_and_offload_scope(trainer, *args):
     offload_map = {
         trainer.actor_model: "train_model",
         trainer.reference_model: "freeze_model",
-        **(
-            {trainer.reward_model: "freeze_model"}
-            if not trainer.args.use_rm_server and not trainer.args.use_rule_reward
-            else {}
-        ),
+        **({trainer.reward_model: "freeze_model"} if not trainer.args.use_rm_server else {}),
         trainer.actor_trainer.optimizer: "optimizer",
     }
 
     if trainer.args.rl_algorithm == "ppo":
         offload_map.update(
             {
-                trainer.critic_model: "train_model",
+                trainer.reward_critic_model: "train_model",
                 trainer.critic_trainer.optimizer: "optimizer",
             }
         )
@@ -205,8 +201,8 @@ def reload_and_offload_scope(trainer, *args):
             # NOTE(gongenlei): for export_evaluate_model
             objs.append((trainer.actor_model, offload_map.get(trainer.actor_model, "")))
     if trainer.args.rl_algorithm == "ppo":
-        if trainer.critic_model not in [i for i, _ in objs]:
+        if trainer.reward_critic_model not in [i for i, _ in objs]:
             if getattr(trainer.critic_trainer, "_inner_eval_model", None) is not None:
                 # NOTE(gongenlei): for export_evaluate_model
-                objs.append((trainer.critic_model, offload_map.get(trainer.critic_model, "")))
+                objs.append((trainer.reward_critic_model, offload_map.get(trainer.reward_critic_model, "")))
     return OffloadController(objs)

paddlenlp/trainer/argparser.py

@@ -157,8 +157,7 @@ class PdArgumentParser(ArgumentParser):
             elif field.default is dataclasses.MISSING:
                 kwargs["required"] = True
         else:
-            kwargs["type"] = json.loads if origin_type is dict else field.type
-
+            kwargs["type"] = json.loads if field.type is dict else field.type
             if field.default is not dataclasses.MISSING:
                 kwargs["default"] = field.default
             elif field.default_factory is not dataclasses.MISSING:

paddlenlp/trainer/auto_trainer.py

@@ -16,7 +16,6 @@ import copy
 import os
 import random
 import time
-import types
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
@@ -25,15 +24,12 @@ import paddle.distributed as dist
 import paddle.distributed.auto_parallel.intermediate.parallelize as parallelize
 import paddle.nn as nn
 from paddle.distributed import fleet
-from paddle.distributed.auto_parallel._utils import _patch_grads_for_step
 from paddle.profiler.utils import switch_job_schedule_profiler
 from tqdm.auto import tqdm
 
 from paddlenlp.trainer import Trainer
 
-from ..transformers.context_parallel_utils import auto_split_sequence_dim_load_balance
-from ..transformers.model_utils import clean_model_class_name, unwrap_model
-from ..transformers.segment_parallel_utils import auto_split_inputs_sequence_dim
+from ..transformers.model_utils import unwrap_model
 from ..utils.batch_sampler import DistributedBatchSampler as NlpDistributedBatchSampler
 from ..utils.env import (
     PREFIX_CHECKPOINT_DIR,
@@ -50,9 +46,7 @@ from .trainer_utils import (  # set_hyrbid_parallel_seed,
     ShardingOption,
     TrainOutput,
     _exec_mode_guard,
-    check_auto_parallel_pipeline_support,
     get_last_checkpoint,
-    get_pp_schedule,
     has_length,
     speed_metrics,
 )
@@ -83,7 +77,6 @@ class AutoTrainer(Trainer):
                 kwargs.update({"criterion": loss_func})
         self.auto_dist_config = kwargs.pop("auto_dist_config", None)
         model = kwargs.get("model", None)
-        self.model_type = kwargs.pop("model_type", None)
         assert model is not None
         if kwargs.get("args", None) is not None and kwargs["args"].use_intermediate_api:
             if not parallelize.has_parallelized_model:
@@ -100,21 +93,12 @@ class AutoTrainer(Trainer):
             if not param._is_initialized() and param._init_func is not None:
                 param.initialize()
         kwargs["model"] = model
+
         super().__init__(*args, **kwargs)
         assert self.args.enable_auto_parallel
 
         self.global_mesh = fleet.auto.get_mesh()
         self.comm_group_in_pp = fleet.get_hybrid_communicate_group().get_pipe_parallel_group()
-        if self.args.pipeline_parallel_degree > 1 and check_auto_parallel_pipeline_support(self.model_type):
-            self.pp_schedule = get_pp_schedule(
-                model,
-                self.model_type,
-                self.args.n_microbatches,
-                self.criterion,
-                self.args.pipeline_schedule_mode,
-                self.args.pipeline_parallel_degree,
-                self.comm_group_in_pp,
-            )
         self._in_pir_mode = paddle.base.framework.get_flags("FLAGS_enable_pir_api")["FLAGS_enable_pir_api"]
 
     @classmethod
@@ -145,7 +129,6 @@ class AutoTrainer(Trainer):
             "data_sharding_parallel": training_args.dataset_world_size > 1,
             "sharding": training_args.sharding,
             "sharding_mesh_dim": training_args.sharding_parallel_mesh_dimension,
-            "context_parallel": training_args.context_parallel_degree > 1 or training_args.sep_parallel_degree > 1,
         }
         auto_dist_config = model._generate_auto_dist_config(auto_dist_degree)
         model = parallelize.parallelize_model(
@@ -229,16 +212,6 @@ class AutoTrainer(Trainer):
                 )
             else:
                 self.optimizer = dist.shard_optimizer(self.optimizer, None, self.args.gradient_accumulation_steps)
-            if (
-                hasattr(self.optimizer, "_enable_tensor_fusion")
-                and "enable_tensor_fusion" in self.args.sharding_parallel_config
-            ):
-                self.optimizer._enable_tensor_fusion()
-            if (
-                hasattr(self.optimizer, "_enable_sharding_overlap")
-                and "enable_overlap" in self.args.sharding_parallel_config
-            ):
-                self.optimizer._enable_sharding_overlap(model)
 
         if self.args.to_static:
             unified_strategy = dist.Strategy()
@@ -520,18 +493,6 @@ class AutoTrainer(Trainer):
             npu_accelerate_plugin(self.optimizer)
 
         model, dist_loader = self._wrap_for_auto(model, train_dataloader)
-
-        if (
-            dist.in_auto_parallel_align_mode()
-        ):  # When in auto parallel align mode, patching the optimizer step function
-
-            orig_step = (
-                self.optimizer.step.__func__ if hasattr(self.optimizer.step, "__func__") else self.optimizer.step
-            )
-            decorator = _patch_grads_for_step(amp_master_grad=self.args.amp_master_grad)
-            new_step = decorator(orig_step)
-            self.optimizer.__dict__["step"] = types.MethodType(new_step, self.optimizer)
-
         train_dataloader = dist_loader()
         if resume_from_checkpoint is not None:
             self._load_from_checkpoint(resume_from_checkpoint)
@@ -584,10 +545,7 @@ class AutoTrainer(Trainer):
                     if step_control % args.gradient_accumulation_steps == 0:
                         self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
                         self.timers and self.timers("forward-backward").start()
-                    if self.args.sep_parallel_degree > 1 and self.args.split_inputs_sequence_dim:
-                        inputs = auto_split_inputs_sequence_dim(inputs)
-                    if self.args.context_parallel_degree > 1 and self.args.split_inputs_sequence_dim:
-                        inputs = auto_split_sequence_dim_load_balance(inputs)
+
                     tr_loss_step = self.training_step(model, inputs)
 
                     with _exec_mode_guard("dynamic"):
@@ -745,56 +703,7 @@ class AutoTrainer(Trainer):
 
         return (loss, outputs) if return_outputs else loss
 
-    def compute_pipeline_loss(self, model, inputs, return_outputs=False):
-        """
-        How the loss is computed by Trainer. By default, all models return the loss in the first element.
-        Subclass and override for custom behavior.
-        """
-        if self.criterion is not None:
-            if "labels" in inputs:
-                labels = inputs.pop("labels")
-            elif "start_positions" in inputs and "end_positions" in inputs:
-                labels = (inputs.pop("start_positions"), inputs.pop("end_positions"))
-            elif self.args.label_names is not None:
-                labels = []
-                for label in self.label_names:
-                    labels.append(inputs.pop(label))
-                labels = tuple(labels)
-            elif "generator_labels" in inputs:
-                labels = inputs["generator_labels"]
-        else:
-            labels = None
-
-        pp_rank = self.comm_group_in_pp.rank
-        losses = []
-        if pp_rank == 0:  # 第一个pp_stage，参数传入数据流
-            self.pp_schedule.step(**inputs)  # 最后的pp_stage，参数传入label, 并输出loss
-        elif pp_rank == self.args.pipeline_parallel_degree - 1:
-            self.pp_schedule.step(target=labels, losses=losses)
-        else:
-            self.pp_schedule.step()
-
-        final_loss = None
-        if len(losses) != 0:
-            final_loss = paddle.stack(losses).mean()
-
-        return final_loss
-
-    def dynamic_auto_parallel_pipeline_training(
-        self, model: nn.Layer, inputs: Dict[str, Union[paddle.Tensor, Any]]
-    ) -> paddle.Tensor:
-        assert self.args.pipeline_parallel_degree > 1, "pipeline_parallel_degree must be greater than 1."
-        assert check_auto_parallel_pipeline_support(
-            self.model_type
-        ), "dynamic auto_parallel pipeline only supports special models"
-        with self.autocast_smart_context_manager():
-            loss = self.compute_pipeline_loss(model, inputs)
-
-        return loss
-
     def dynamic_training(self, model: nn.Layer, inputs: Dict[str, Union[paddle.Tensor, Any]]) -> paddle.Tensor:
-        if self.args.pipeline_parallel_degree > 1 and check_auto_parallel_pipeline_support(self.model_type):
-            return self.dynamic_auto_parallel_pipeline_training(model, inputs)
         with self.autocast_smart_context_manager():
             loss = self.compute_loss(model, inputs)
 
@@ -1001,19 +910,15 @@ class AutoTrainer(Trainer):
             config_to_save = copy.deepcopy(model_to_save.config)
             config_to_save.mp_degree = getattr(config_to_save, "config_to_save", 1)
             # Attach architecture to the config
-            config_to_save.architectures = [clean_model_class_name(model_to_save.__class__.__name__)]
+            config_to_save.architectures = [model_to_save.__class__.__name__]
 
             config_to_save.save_pretrained(output_dir)
             if self.model.can_generate():
                 model_to_save.generation_config.save_pretrained(output_dir)
 
         if self.args.should_save_model_state:
-            if state_dict is None:
-                self._save_ckpt_func(self.model.state_dict(), output_dir)
-                logger.info(f"Model weights saved in {output_dir}")
-            else:
-                self._save_ckpt_func(state_dict, output_dir)
-                logger.info(f"Model weights and optimizer states saved in {output_dir}")
+            self._save_ckpt_func(self.model.state_dict(), output_dir)
+            logger.info(f"Model weights and optimizer states saved in {output_dir}")
 
     def _load_from_checkpoint(self, resume_from_checkpoint=None):
 

paddlenlp/trainer/integrations.py

@@ -403,132 +403,81 @@ class WandbCallback(TrainerCallback):
 
 class SwanLabCallback(TrainerCallback):
     """
-    A [`TrainerCallback`] that logs metrics, media, model checkpoints to [SwanLab](https://swanlab.cn/).
+    A [`TrainerCallback`] that logs metrics, media to [Swanlab](https://swanlab.cn/).
     """
 
     def __init__(self):
-        if not is_swanlab_available():
-            raise RuntimeError("SwanLabCallback requires swanlab to be installed. Run `pip install swanlab`.")
-        import swanlab
+        has_swanlab = is_swanlab_available()
+        if not has_swanlab:
+            raise RuntimeError("SwanlabCallback requires swanlab to be installed. Run `pip install swanlab`.")
+        if has_swanlab:
+            import swanlab
+
+            self._swanlab = swanlab
 
-        self._swanlab = swanlab
         self._initialized = False
-        self._log_model = os.getenv("SWANLAB_LOG_MODEL", None)
 
     def setup(self, args, state, model, **kwargs):
         """
-        Setup the optional SwanLab (*swanlab*) integration.
-        One can subclass and override this method to customize the setup if needed. Find more information
-        [here](https://docs.swanlab.cn/guide_cloud/integration/integration-huggingface-transformers.html).
-        You can also override the following environment variables. Find more information about environment
-        variables [here](https://docs.swanlab.cn/en/api/environment-variable.html#environment-variables)
+        Setup the optional Swanlab integration.
+
+        One can subclass and override this method to customize the setup if needed.
+        variables:
         Environment:
-        - **SWANLAB_API_KEY** (`str`, *optional*, defaults to `None`):
-            Cloud API Key. During login, this environment variable is checked first. If it doesn't exist, the system
-            checks if the user is already logged in. If not, the login process is initiated.
-                - If a string is passed to the login interface, this environment variable is ignored.
-                - If the user is already logged in, this environment variable takes precedence over locally stored
-                login information.
-        - **SWANLAB_PROJECT** (`str`, *optional*, defaults to `None`):
-            Set this to a custom string to store results in a different project. If not specified, the name of the current
-            running directory is used.
-        - **SWANLAB_LOG_DIR** (`str`, *optional*, defaults to `swanlog`):
-            This environment variable specifies the storage path for log files when running in local mode.
-            By default, logs are saved in a folder named swanlog under the working directory.
-        - **SWANLAB_MODE** (`Literal["local", "cloud", "disabled"]`, *optional*, defaults to `cloud`):
-            SwanLab's parsing mode, which involves callbacks registered by the operator. Currently, there are three modes:
-            local, cloud, and disabled. Note: Case-sensitive. Find more information
-            [here](https://docs.swanlab.cn/en/api/py-init.html#swanlab-init)
-        - **SWANLAB_LOG_MODEL** (`str`, *optional*, defaults to `None`):
-            SwanLab does not currently support the save mode functionality.This feature will be available in a future
-            release
-        - **SWANLAB_WEB_HOST** (`str`, *optional*, defaults to `None`):
-            Web address for the SwanLab cloud environment for private version (its free)
-        - **SWANLAB_API_HOST** (`str`, *optional*, defaults to `None`):
-            API address for the SwanLab cloud environment for private version (its free)
+        - **SWANLAB_MODE** (`str`, *optional*, defaults to `"cloud"`):
+            Whether to use swanlab cloud, local or disabled. Set `SWANLAB_MODE="local"` to use local. Set `SWANLAB_MODE="disabled"` to disable.
+        - **SWANLAB_PROJECT** (`str`, *optional*, defaults to `"PaddleNLP"`):
+            Set this to a custom string to store results in a different project.
         """
+
+        if self._swanlab is None:
+            return
+
         self._initialized = True
 
         if state.is_world_process_zero:
-            logger.info('Automatic SwanLab logging enabled, to disable set os.environ["SWANLAB_MODE"] = "disabled"')
+            logger.info('Automatic Swanlab logging enabled, to disable set os.environ["SWANLAB_MODE"] = "disabled"')
+
             combined_dict = {**args.to_dict()}
 
             if hasattr(model, "config") and model.config is not None:
-                model_config = model.config if isinstance(model.config, dict) else model.config.to_dict()
+                model_config = model.config.to_dict()
                 combined_dict = {**model_config, **combined_dict}
-            if hasattr(model, "lora_config") and model.lora_config is not None:
-                lora_config = model.lora_config if isinstance(model.lora_config, dict) else model.lora_config.to_dict()
-                combined_dict = {**{"lora_config": lora_config}, **combined_dict}
+
             trial_name = state.trial_name
             init_args = {}
-            if trial_name is not None and args.run_name is not None:
-                init_args["experiment_name"] = f"{args.run_name}-{trial_name}"
-            elif args.run_name is not None:
-                init_args["experiment_name"] = args.run_name
-            elif trial_name is not None:
-                init_args["experiment_name"] = trial_name
-
-            # new add this for experiment_name
-            experiment_name = os.getenv("SWANLAB_EXP_NAME", None)
-            if experiment_name is not None:
-                init_args["experiment_name"] = experiment_name
-
-            init_args["project"] = os.getenv("SWANLAB_PROJECT", None)
-            if args.logging_dir is not None:
-                init_args["logdir"] = os.getenv("SWANLAB_LOG_DIR", args.logging_dir)
-
+            if trial_name is not None:
+                init_args["name"] = trial_name
+                init_args["group"] = args.run_name
+            else:
+                if not (args.run_name is None or args.run_name == args.output_dir):
+                    init_args["name"] = args.run_name
+            init_args["dir"] = args.logging_dir
             if self._swanlab.get_run() is None:
                 self._swanlab.init(
+                    project=os.getenv("SWANLAB_PROJECT", "PaddleNLP"),
                     **init_args,
                 )
-            # show paddlenlp logo!
-            self._swanlab.config["FRAMEWORK"] = "paddlenlp"
-            # add config parameters (run may have been created manually)
-            self._swanlab.config.update(combined_dict)
+            self._swanlab.config.update(combined_dict, allow_val_change=True)
 
     def on_train_begin(self, args, state, control, model=None, **kwargs):
+        if self._swanlab is None:
+            return
         if not self._initialized:
             self.setup(args, state, model, **kwargs)
 
-    def on_train_end(self, args, state, control, model=None, processing_class=None, **kwargs):
-        if self._log_model is not None and self._initialized and state.is_world_process_zero:
-            logger.warning(
-                "SwanLab does not currently support the save mode functionality. "
-                "This feature will be available in a future release."
-            )
+    def on_train_end(self, args, state, control, model=None, tokenizer=None, **kwargs):
+        if self._swanlab is None:
+            return
 
     def on_log(self, args, state, control, model=None, logs=None, **kwargs):
-        single_value_scalars = [
-            "train_runtime",
-            "train_samples_per_second",
-            "train_steps_per_second",
-            "train_loss",
-            "total_flos",
-        ]
-
+        if self._swanlab is None:
+            return
         if not self._initialized:
             self.setup(args, state, model)
         if state.is_world_process_zero:
-            for k, v in logs.items():
-                if k in single_value_scalars:
-                    self._swanlab.log({f"single_value/{k}": v}, step=state.global_step)
-            non_scalar_logs = {k: v for k, v in logs.items() if k not in single_value_scalars}
-            non_scalar_logs = rewrite_logs(non_scalar_logs)
-            self._swanlab.log({**non_scalar_logs, "train/global_step": state.global_step}, step=state.global_step)
-
-    def on_save(self, args, state, control, **kwargs):
-        if self._log_model is not None and self._initialized and state.is_world_process_zero:
-            logger.warning(
-                "SwanLab does not currently support the save mode functionality. "
-                "This feature will be available in a future release."
-            )
-
-    def on_predict(self, args, state, control, metrics, **kwargs):
-        if not self._initialized:
-            self.setup(args, state, **kwargs)
-        if state.is_world_process_zero:
-            metrics = rewrite_logs(metrics)
-            self._swanlab.log(metrics)
+            logs = rewrite_logs(logs)
+            self._swanlab.log({**logs, "train/global_step": state.global_step}, step=state.global_step)
 
 
 class AutoNLPCallback(TrainerCallback):

paddlenlp/trainer/trainer_callback.py

@@ -20,18 +20,10 @@ Callbacks to use with the Trainer class and customize the training loop.
 """
 import dataclasses
 import json
-import time
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 import numpy as np
-from paddle.distributed.fleet import fleet
-from paddle.distributed.fleet.utils.hybrid_parallel_util import (
-    fused_allreduce_gradients_with_group,
-)
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    is_sequence_parallel_parameter,
-)
 from tqdm.auto import tqdm
 
 from paddlenlp.utils.log import logger
@@ -157,7 +149,6 @@ class TrainerControl:
     should_save: bool = False
     should_evaluate: bool = False
     should_log: bool = False
-    should_save_hf: bool = False
 
     def _new_training(self):
         """Internal method that resets the variable for a new training."""
@@ -172,7 +163,6 @@ class TrainerControl:
         self.should_save = False
         self.should_evaluate = False
         self.should_log = False
-        self.should_save_hf = False
 
 
 class TrainerCallback:
@@ -308,12 +298,6 @@ class TrainerCallback:
         """
         pass
 
-    def on_save_hf(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        """
-        Event called after a huggingface checkpoint save.
-        """
-        pass
-
 
 class CallbackHandler(TrainerCallback):
     """Internal class that just calls the list of callbacks in order."""
@@ -394,7 +378,6 @@ class CallbackHandler(TrainerCallback):
         control.should_log = False
         control.should_evaluate = False
         control.should_save = False
-        control.should_save_hf = False
         return self.call_event("on_step_begin", args, state, control)
 
     def on_load_data_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, inputs: Dict):
@@ -427,10 +410,6 @@ class CallbackHandler(TrainerCallback):
     def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
         return self.call_event("on_prediction_step", args, state, control)
 
-    def on_save_hf(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
-        control.should_save_hf = False
-        return self.call_event("on_save_hf", args, state, control)
-
     def call_event(self, event, args, state, control, **kwargs):
         for callback in self.callbacks:
             result = getattr(callback, event)(
@@ -487,14 +466,6 @@ class DefaultFlowCallback(TrainerCallback):
         if state.global_step >= state.max_steps:
             control.should_training_stop = True
 
-        # Save hf
-        if (
-            args.save_strategy == IntervalStrategy.STEPS
-            and args.save_hf_steps > 0
-            and state.global_step % args.save_hf_steps == 0
-        ):
-            control.should_save_hf = True
-
         return control
 
     def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
@@ -638,33 +609,3 @@ class EarlyStoppingCallback(TrainerCallback):
         self.check_metric_value(args, state, control, metric_value)
         if self.early_stopping_patience_counter >= self.early_stopping_patience:
             control.should_training_stop = True
-
-
-class SPGradSyncCallback(TrainerCallback):
-    """
-    SPGradSyncCallback
-    只能在非 sharding stage2 的情况下使用。
-    开启sharding stage2 时，在 `on_optimizer_begin` 的时候 grad 已经被清空了
-    """
-
-    def __init__(self, model):
-        assert hasattr(fleet, "_hcg"), "must use MP when calling this Callback"
-        logger.info("using sp callback")
-        params = []
-        self.model = model
-        for n, p in model.named_parameters():
-            if is_sequence_parallel_parameter(p):
-                logger.info(f"register bw hook for:{n}")
-                params.append(p)
-
-        logger.info(f"#-sp-sync param:{len(params)}")
-        self._sp_params = params
-
-    def on_optimizer_begin(self, args, state, control, **kwargs):
-        """on_optimizer_begin"""
-        if self._sp_params:
-            now = time.time()
-            mp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
-            fused_allreduce_gradients_with_group(self._sp_params, group=mp_group, scale=1.0)  # sum not mean
-            another_time = time.time()
-            logger.info(f"sync gradients takes {another_time - now} time")

paddlenlp/trainer/trainer_utils.py

@@ -30,17 +30,12 @@ import threading
 import time
 from contextlib import contextmanager
 from enum import Enum
-from pathlib import Path
 from typing import Dict, List, NamedTuple, Optional, Tuple, Union
 
 import numpy as np
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
-from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
-    DygraphShardingOptimizer,
-    DygraphShardingOptimizerV2,
-)
 from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.io import IterableDataset
 from paddle.optimizer.lr import LambdaDecay
@@ -48,7 +43,6 @@ from paddle.optimizer.lr import LambdaDecay
 from paddlenlp.ops import Topology
 
 from ..trainer.argparser import strtobool
-from ..transformers import get_gpt_pp_schedule, get_llama_pp_schedule
 from ..transformers.tokenizer_utils_base import BatchEncoding
 from ..utils.env import PREFIX_CHECKPOINT_DIR, _re_checkpoint  # noqa for compatibility
 from ..utils.fault_tolerance import PDC_DOWNLOAD_ERROR
@@ -324,6 +318,8 @@ class OptimizerNames(ExplicitEnum):
     ADAFACTOR = "adafactor"
     ADAMW_MINI = "adamw_mini"
     ADAMW_CUSTOM = "adamw_custom"
+    ADAMW_16BIT_MOMENT = "adamw_16bit_moment"
+    AdamW_Qweight = "adamw_qweight"
 
 
 class ShardingOption(ExplicitEnum):
@@ -1259,189 +1255,3 @@ def download_recovery_ckpt_from_pdc(recovery_checkpoint_path, timeout):
         raise RuntimeError(
             f"{PDC_DOWNLOAD_ERROR}; Error occurred when trying to download checkpoint from PDC, recovery_checkpoint_path: {recovery_checkpoint_path}, timeout: {timeout}; error details: {PDCErrorMessageMap[result]}"
         )
-
-
-def check_auto_parallel_pipeline_support(model_type=None):
-    support_types = ["llama_pp", "gpt_pp"]
-    return model_type in support_types
-
-
-def get_pp_schedule(model, model_type, n_microbatches, loss_fn, mode, pp_degree, group):
-    assert check_auto_parallel_pipeline_support(model_type)
-    if model_type == "llama_pp":
-        return get_llama_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group)
-    elif model_type == "gpt_pp":
-        return get_gpt_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group)
-
-
-def parse_nccl_config_file(config_dir):
-    json_file = Path(config_dir)
-    if json_file.exists():
-        with open(json_file, "r") as file:
-            data = json.load(file)
-
-        def get_full_config_from_dict(comm_config):
-            assert type(comm_config) is dict
-            min_val = {
-                "ll_buffsize": 2**15,  # 32KB
-                "ll128_buffsize": 2**17,  # 128KB
-                "simple_buffsize": 2**17,  # 128KB
-            }
-            final_config = {}
-
-            # if user does not set group name, use the default name set by Paddle
-            if comm_config.get("name", None) is not None:
-                final_config["commName"] = comm_config["name"]
-            final_config["buffsize_align"] = comm_config.get("buffsize_align", 1024)
-            final_config["algoStr"] = comm_config.get("algo", "")
-            final_config["protoStr"] = comm_config.get("proto", "")
-            final_config["nchannels"] = comm_config.get("n_channels", -1)
-
-            # ll part
-            # -1 means using the default value
-            final_config["ll_buffsize"] = comm_config.get("ll_buffsize", -1)
-            # keep the buffsize > the min value
-            if final_config["ll_buffsize"] != -1:
-                final_config["ll_buffsize"] = max(final_config["ll_buffsize"], min_val["ll_buffsize"])
-
-            # ll128 part
-            final_config["ll128_buffsize"] = comm_config.get("ll128_buffsize", -1)
-            if final_config["ll128_buffsize"] != -1:
-                final_config["ll128_buffsize"] = max(final_config["ll128_buffsize"], min_val["ll128_buffsize"])
-
-            # simple part
-            final_config["simple_buffsize"] = comm_config.get("simple_buffsize", -1)
-            if final_config["simple_buffsize"] != -1:
-                final_config["simple_buffsize"] = max(final_config["simple_buffsize"], min_val["simple_buffsize"])
-
-            # set the buffer size of unused protocols to the minimum value
-            if final_config["protoStr"] != "":
-                protos = split_parallel_config(final_config["protoStr"].lower())
-                for proto in ["ll", "ll128", "simple"]:
-                    if proto not in protos:
-                        final_config[(proto + "_buffsize")] = min_val[(proto + "_buffsize")]
-
-            return final_config
-
-        for key in data.keys():
-            data[key] = get_full_config_from_dict(data[key])
-
-        return data
-    else:
-        raise FileNotFoundError(f"The argument file {json_file} does not exist.")
-
-
-def init_nccl_config(nccl_comm_group_config, strategy):
-    nccl_config = parse_nccl_config_file(nccl_comm_group_config)
-
-    def set_comm_config(configs, attr, dict_obj):
-        if strategy.hybrid_configs.get(configs, None) is None or dict_obj is None:
-            return
-        if not hasattr(strategy.hybrid_configs[configs], attr):
-            return
-        attr_obj = getattr(strategy.hybrid_configs[configs], attr)
-        for key, value in dict_obj.items():
-            if hasattr(attr_obj, key):
-                setattr(attr_obj, key, value)
-
-    set_comm_config("pp_configs", "coll_nccl_config", nccl_config.get("pp", None))
-    set_comm_config("pp_configs", "p2p_nccl_config", nccl_config.get("pp_p2p", None))
-    set_comm_config("pp_configs", "shared_nccl_config", nccl_config.get("pp_shared", None))
-    set_comm_config("mp_configs", "nccl_config", nccl_config.get("tp", None))
-    set_comm_config("sharding_configs", "nccl_config", nccl_config.get("sharding", None))
-    set_comm_config("sharding_configs", "check_nccl_config", nccl_config.get("sharding_check", None))
-    set_comm_config("dp_configs", "nccl_config", nccl_config.get("dp", None))
-    set_comm_config("dp_configs", "check_nccl_config", nccl_config.get("dp_check", None))
-    set_comm_config("sep_configs", "nccl_config", nccl_config.get("sep", None))
-    set_comm_config("dp_sep_configs", "nccl_config", nccl_config.get("dp_sep", None))
-    set_comm_config("pp_tp_configs", "nccl_config", nccl_config.get("pp_tp", None))
-    set_comm_config("ep_configs", "nccl_config", nccl_config.get("ep", None))
-    set_comm_config("ep_configs", "grad_nccl_config", nccl_config.get("ep_grad", None))
-    set_comm_config("moe_sharding_configs", "nccl_config", nccl_config.get("moe_sharding", None))
-    set_comm_config("moe_sharding_configs", "check_nccl_config", nccl_config.get("moe_sharding_check", None))
-    set_comm_config("default_comm_group_configs", "nccl_config", nccl_config.get("default", None))
-    return strategy
-
-
-def init_optimizer(optimizer, model_sharded_state_dict, state_dict_metadata):
-    """
-    Initialize the optimizer's states according to its type.
-
-    For DygraphShardingOptimizer (V1), initializes accumulators for local parameters.
-    For DygraphShardingOptimizerV2, manually initializes master weights and state dict for sharded parameters.
-    For other cases, initializes accumulators for all parameters.
-
-    Args:
-        optimizer: The optimizer instance to be initialized.
-    """
-    optimizer_state_names = [".moment1_0", ".moment2_0", ".beta1_pow_acc_0", ".beta2_pow_acc_0", ".w_0"]
-    inner_opt = getattr(optimizer, "_inner_opt", None)
-    static_to_struct_mapping = {}
-    model_sharded_state_dict = dict(sorted(model_sharded_state_dict.items()))
-    for k, v in model_sharded_state_dict.items():
-        if v.local_tensor.name not in static_to_struct_mapping:
-            static_to_struct_mapping[v.local_tensor.name] = k
-
-    if isinstance(inner_opt, DygraphShardingOptimizer):
-        local_params = optimizer._rank2params[optimizer._sharding_rank]
-        param_list = []
-        for param in local_params:
-            param_name = param.name
-            struct_name = static_to_struct_mapping[param_name]
-            if not any(struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names):
-                continue
-            param_list.append(param)
-        optimizer._create_accumulators(paddle.base.framework.default_main_program().global_block(), param_list)
-        return
-
-    elif isinstance(inner_opt, DygraphShardingOptimizerV2):
-
-        def init_param_optimizer_states(param_iter):
-            master_weights = {}
-            state_dict = {}
-            moments = ("moment1_0", "moment2_0")
-            betas = ("beta1_pow_acc_0", "beta2_pow_acc_0")
-            for static_name, shape, no_need_master_weights in param_iter:
-                if not no_need_master_weights:
-                    master_weights[static_name] = paddle.zeros(shape, dtype="float32")
-                    prefix = f"{static_name}_fp32_master_0_"
-                else:
-                    prefix = f"{static_name}_"
-
-                for moment in moments:
-                    key = f"{prefix}{moment}"
-                    state_dict[key] = paddle.zeros(shape, dtype="float32")
-                for beta in betas:
-                    key = f"{prefix}{beta}"
-                    state_dict[key] = paddle.zeros((1,), dtype="float32")
-            return master_weights, state_dict
-
-        def buffer_params():
-            for buffer in optimizer._comm_buffer_list:
-                for param_name, grad_view in buffer._sharding_param_grad_view.items():
-                    struct_name = static_to_struct_mapping[param_name]
-                    if not any(
-                        struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names
-                    ):
-                        continue
-                    param_begin = grad_view._param_begin
-                    param_end = grad_view._param_end
-                    shape = (param_end - param_begin,)
-                    no_need_master_weights = grad_view._param.dtype == paddle.float32
-
-                    if shape[0] > 0:
-                        yield param_name, shape, no_need_master_weights
-
-        master_weights, state_dict = init_param_optimizer_states(buffer_params())
-        state_dict["master_weights"] = master_weights
-        state_dict["LR_Scheduler"] = {"last_epoch": 1, "last_lr": 5e-06}
-        optimizer.set_state_dict(state_dict)
-        return
-    param_list = []
-    for param in optimizer._parameter_list:
-        param_name = param.name
-        struct_name = static_to_struct_mapping[param_name]
-        if not any(struct_name + state_name in state_dict_metadata for state_name in optimizer_state_names):
-            continue
-        param_list.append(param)
-    optimizer._create_accumulators(paddle.base.framework.default_main_program().global_block(), param_list)

paddlenlp/trainer/training_args.py

@@ -31,11 +31,6 @@ import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
 
-try:
-    from paddle.distributed.fleet.base.topology import message2nccl_config
-except ImportError:
-    message2nccl_config = None
-
 from ..utils.env import PREFIX_CHECKPOINT_DIR
 from ..utils.log import logger
 from ..utils.pdc_sdk import FLASH_DEVICE
@@ -44,7 +39,6 @@ from .trainer_utils import (
     OptimizerNames,
     SchedulerType,
     ShardingOption,
-    init_nccl_config,
     split_parallel_config,
 )
 
@@ -407,29 +401,6 @@ class TrainingArguments:
             Whether to release gradients during training. Default is `False`.
         ckpt_quant_stage (`str`, *optional*):
             Whether activate checkpoint quantization. O0: deactivate, O1: Int8 compression, O2: Int4 compression. (default: O0).
-        save_checkpoint_format (`str`, *optional*):
-            Specifies the format for saving checkpoints. Options are: None, 'sharding_io', 'unified_checkpoint', 'flex_checkpoint'. (default: None). This setting is ignored if the corresponding switch is configured.
-        load_checkpoint_format (`str`, *optional*):
-            Specifies the format for loading checkpoints. Options are: None, 'sharding_io', 'unified_checkpoint', 'flex_checkpoint'. (default: None). This setting is ignored if the corresponding switch is configured.
-        aoa_config (`Optional[dict[str, list[str]]]`, *optional*):
-            The AoA configuration of FlexCheckpoint, used to describe the mapping between model weights and the checkpoint content. Default is None.
-        save_hf_steps (`int`, *optional*, defaults to 500):
-            Number of updates steps before two huggingface checkpoint saves if `save_strategy="steps"`.
-        load_via_cpu (bool, optional):
-            Whether to load checkpoint data into CPU memory first before transferring to GPU.
-            This helps mitigate GPU memory shortage by staging data on the CPU and only moving required parts to the GPU on demand during communication.
-            Defaults to False.
-        load_from_hf (bool, optional):
-            Whether to load a checkpoint in the HuggingFace format.
-            Defaults to False.
-        flex_ckpt_comm_method (str, optional):
-            Communication method used for checkpoint resharding.
-            Choices are "send_recv", "broadcast", "multi_group_broadcast", and "grouped_send_recv".
-            Defaults to "broadcast".
-        replicate_saved_into_local (bool, optional):
-            Whether to save checkpoint replicas into local files in a distributed save/load system.
-            If set to True, replicas will be stored locally on each node/machine.
-            Defaults to False.
     """
 
     output_dir: str = field(
@@ -637,16 +608,6 @@ class TrainingArguments:
             )
         },
     )
-    sharding_offload_opt_buffersize_GB: int = field(
-        default=-1,
-        metadata={
-            "help": (
-                "Set the size of the optimizer offload buffer when need_hack_offload_optimizer() is True. This option only takes effect when "
-                "use DygraphShardingOptimizerV2. The default value is -1, which means that all of the optimizer states will be offloaded. Only "
-                "works when export HACK_OFFLOAD_OPTIMIZER=1. "
-            )
-        },
-    )
 
     save_sharded_model: bool = field(
         default=False,
@@ -666,16 +627,6 @@ class TrainingArguments:
         },
     )
 
-    load_sharded_model_remap_parameter_name: bool = field(
-        default=False,
-        metadata={"help": "Whether to remap parameter name when load_sharded_model = true."},
-    )
-
-    sharded_model_from_ema: bool = field(
-        default=False,
-        metadata={"help": "Whether to load sharded model from EMA."},
-    )
-
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={
@@ -825,7 +776,7 @@ class TrainingArguments:
                 "Following options are supported:\n"
                 "- pp_first. the topo order is dp, pp, sharding, mp \n"
                 "- sharding_first. the topo order is dp, sharding, pp, mp \n"
-                "Default is None, for sharding_first"
+                "Default is None, for pp_first"
             )
         },
     )
@@ -917,10 +868,6 @@ class TrainingArguments:
         default="adamw",
         metadata={"help": "The optimizer to use."},
     )
-    use_lowprecision_moment: bool = field(
-        default=False,
-        metadata={"help": "AdamW use 16bit moment as model parameter."},
-    )
     report_to: Optional[List[str]] = field(
         default=None, metadata={"help": "The list of integrations to report the results and logs to."}
     )
@@ -964,29 +911,6 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Whether to use async_save instead of paddle.save."},
     )
-    save_checkpoint_format: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Specifies the format used to save checkpoints. "
-                "Available options: 'sharding_io', 'unified_checkpoint', "
-                "'flex_checkpoint'."
-                "This setting is ignored if the corresponding switch is configured."
-            )
-        },
-    )
-
-    load_checkpoint_format: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Specifies the format used to load checkpoints. "
-                "Available options: 'sharding_io', 'unified_checkpoint', "
-                "'flex_checkpoint'."
-                "This setting is ignored if the corresponding switch is configured."
-            )
-        },
-    )
     ordered_save_group_size: int = field(
         default=0,
         metadata={
@@ -1072,16 +996,6 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Offload optimizer after optimizer.step()"},
     )
-    tensorwise_offload_optimizer: Optional[bool] = field(
-        default=False,
-        metadata={
-            "help": (
-                "Offload all optimizer states to CPU memory. "
-                "The corresponding parameters will only be loaded to GPU during optimizer step, "
-                "which reduces GPU memory usage but may increase step time."
-            )
-        },
-    )
     save_sharding_stage1_model_include_freeze_params: Optional[bool] = field(
         default=False, metadata={"help": "Save Sharding Stage1 Model Exclude Freeze Params"}
     )
@@ -1111,10 +1025,6 @@ class TrainingArguments:
         default=1,
         metadata={"help": "Interval between updating EMA parameters."},
     )
-    zcc_ema_loss_threshold: Optional[float] = field(
-        default=None,
-        metadata={"help": "If set not None, only do EMA when the training loss is smaller than the threshold value"},
-    )
     save_tokenizer: Optional[bool] = field(
         default=True,
         metadata={"help": "Save tokenizer to output_dir."},
@@ -1148,56 +1058,10 @@ class TrainingArguments:
         metadata={"help": "是否开启单路sharding时global norm通信拆分全局通信组为pp通信和mp通信分别做"},
     )
 
-    nccl_comm_group_config: Optional[str] = field(
-        default=None, metadata={"help": "NCCL中通信组的细粒度控制的配置文件路径, 默认值为None, 代表不启用此项配置"}
-    )
-
-    aoa_config: Optional[dict[str, list[str]]] = field(
-        default=None,
-        metadata={
-            "help": "The AoA configuration of FlexCheckpoint, used to describe the mapping between model weights and the checkpoint content. Default is None."
-        },
-    )
-
-    save_hf_steps: int = field(default=-1, metadata={"help": "Save huggingface checkpoint every X updates steps."})
-
-    load_via_cpu: Optional[bool] = field(
-        default=False,
-        metadata={
-            "help": "If True, loads checkpoint data to CPU first, then transfers required parts to GPU on demand to reduce GPU memory usage. Defaults to False."
-        },
-    )
-
-    load_from_hf: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to load a checkpoint in the HuggingFace format."},
-    )
-
-    flex_ckpt_comm_method: Optional[str] = field(
-        default="broadcast",
-        metadata={
-            "help": (
-                "Communication method used by FlexCheckpoint for checkpoint resharding. "
-                'Choices are "send_recv", "broadcast", "multi_group_broadcast", and "grouped_send_recv". '
-                'Default is "broadcast".'
-            )
-        },
-    )
-
-    replicate_saved_into_local: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to save replicas cross files in distributed save load system."},
-    )
-
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()
         if in_auto_parallel_align_mode():
-            # self.max_grad_norm = 0.0
-            # The current auto_hybrid_pp has aligned the handling of ClipGradByGlobalNorm with the original dygraph semi-auto parallel and dynamic manual-parallel modes and can correctly handle grad_clip, so it is no longer necessary to set max_grad_norm=0.0.
-            if self.max_grad_norm != 0.0:
-                warnings.warn(
-                    "max_grad_norm is not 0.0,We will execute ClipGradByGlobalNorm,if you want to disable it,please set max_grad_norm=0.0"
-                )
+            self.max_grad_norm = 0.0
             os.environ["FLAGS_max_inplace_grad_add"] = "65536"
             os.environ["FLAGS_embedding_deterministic"] = "1"
             os.environ["FLAGS_cudnn_deterministic"] = "1"
@@ -1293,8 +1157,7 @@ class TrainingArguments:
             raise ValueError("AdamW Mini currently doesn't support tensor parallelism.")
 
         self._post_init_parallel_degree()
-        self._post_init_save_checkpoint_format()
-        self._post_init_load_checkpoint_format()
+
         if self.to_static:
             assert world_size == 1 or self.enable_auto_parallel, (
                 "It's not supported for training in static mode except the following cases : "
@@ -1313,6 +1176,7 @@ class TrainingArguments:
 
         # use_hybrid_parallel
         if self.use_hybrid_parallel:
+
             if ShardingOption.OFFLOAD in self.sharding:
                 warnings.warn("`offload` is not supported NOW!")
 
@@ -1347,8 +1211,6 @@ class TrainingArguments:
                                 "use_dualpipev",
                                 "forward_backward_overlap_scheduler",
                                 "enable_dynamic_shape",
-                                "sync_moment",
-                                "sync_param",
                             ]:
                                 raise ValueError(
                                     f"Found unknown pipeline mode config {x}, accept config is disable_p2p_cache_shape, disable_partial_send_recv."
@@ -1401,18 +1263,6 @@ class TrainingArguments:
                         in pipeline_parallel_config,
                         "enable_dynamic_shape": "enable_dynamic_shape" in pipeline_parallel_config,
                     }
-
-                    pp_sync_param = "sync_param" in pipeline_parallel_config
-                    pp_sync_moment = "sync_moment" in pipeline_parallel_config
-
-                    if pp_sync_param:
-                        logger.info("setting pp sync_param")
-                        strategy.hybrid_configs["pp_configs"].sync_param = True
-
-                    if pp_sync_moment:
-                        logger.info("setting pp sync_moment")
-                        strategy.hybrid_configs["pp_configs"].sync_moment = True
-
                     if dygraph_pp_configs["dp_comm_overlap"]:
                         raise ValueError("overlap has accuracy issue")  # TODO: fix `overalap` + `delay_scale` issue
 
@@ -1509,17 +1359,6 @@ class TrainingArguments:
                         logger.warning("segment parallel is not supported!!!, Ignore it.")
                     return support_sep
 
-                def is_context_parallel_supported():
-                    import inspect
-
-                    members = [
-                        name for (name, date) in inspect.getmembers(fleet.base.topology.EPHybridCommunicateGroup)
-                    ]
-                    support_cp = "get_context_parallel_world_size" in members
-                    if not support_cp:
-                        logger.warning("context parallel is not supported!!! Ignore it.")
-                    return support_cp
-
                 if self.hybrid_parallel_topo_order == "pp_first":
                     if is_segment_parallel_supported():
                         order = ["dp", "pp", "sharding", "sep", "mp"]
@@ -1531,25 +1370,9 @@ class TrainingArguments:
                     else:
                         order = ["dp", "sharding", "pp", "mp"]
                 if self.use_expert_parallel:
-                    if self.moe_sharding_parallel_degree >= 1 and self.expert_parallel_degree > 1:
-                        if is_context_parallel_supported():
-                            order = ["sharding", "moe_sharding", "pp", "sep", "cp", "dp", "ep", "mp"]
-                        else:
-                            order = ["sharding", "moe_sharding", "pp", "sep", "dp", "ep", "mp"]
-                    else:
-                        order = ["sharding", "pp", "sep", "dp", "mp"]
+                    order = order[1:-1] + ["dp", "mp"]
 
-                if is_context_parallel_supported():
-                    hybrid_configs = {
-                        "dp_degree": self.data_parallel_degree,
-                        "mp_degree": self.tensor_parallel_degree,
-                        "pp_degree": self.pipeline_parallel_degree,
-                        "sharding_degree": self.sharding_parallel_degree,
-                        "sep_degree": self.sep_parallel_degree,
-                        "cp_degree": self.context_parallel_degree,
-                        "order": order,
-                    }
-                elif is_segment_parallel_supported():
+                if is_segment_parallel_supported():
                     hybrid_configs = {
                         "dp_degree": self.data_parallel_degree,
                         "mp_degree": self.tensor_parallel_degree,
@@ -1569,13 +1392,6 @@ class TrainingArguments:
                         "order": order,
                     }
 
-                if self.expert_parallel_degree > 1:
-                    assert (
-                        self.use_expert_parallel is True and self.moe_sharding_parallel_degree >= 0
-                    ), f"invalid expert_parallel_degree {self.expert_parallel_degree} and use_expert_paralle:{self.use_expert_parallel}."
-                    hybrid_configs["ep_degree"] = self.expert_parallel_degree
-                    hybrid_configs["moe_sharding_degree"] = self.moe_sharding_parallel_degree
-
                 try:
                     if self.split_norm_comm:
                         hybrid_configs["split_norm_comm"] = True
@@ -1629,11 +1445,6 @@ class TrainingArguments:
                                 self.sharding_comm_buffer_size_MB
                             )
 
-                        if hasattr(strategy.hybrid_configs["sharding_configs"], "offload_opt_buffer_size"):
-                            strategy.hybrid_configs["sharding_configs"].offload_opt_buffer_size = int(
-                                self.sharding_offload_opt_buffersize_GB
-                            )
-
                         if "split_param" in sharding_parallel_config:
                             strategy.hybrid_configs["sharding_configs"].split_param = True
                             assert self.amp_master_grad, "Currently sharding stage1 v2 only support amp_master_grad"
@@ -1695,6 +1506,12 @@ class TrainingArguments:
                         assert (
                             "split_param" in sharding_parallel_config
                         ), "split_param should be set when enable_stage1_allgather_overlap."
+                        use_casual_mask = os.getenv("USE_CASUAL_MASK", "False")
+                        assert use_casual_mask, "enable_stage1_allgather_overlap requires USE_CASUAL_MASK=True."
+                        assert self.logging_steps > 1, (
+                            "The logging_steps should be greater than 1 for enable_stage1_allgather_overlap, "
+                            f"but got logging_steps={self.logging_steps}."
+                        )
 
                     if "split_param" in sharding_parallel_config:
                         if ShardingOption.SHARD_OP not in self.sharding:
@@ -1703,12 +1520,12 @@ class TrainingArguments:
                             self.amp_master_grad
                         ), "If `split_param` in sharding_parallel_config, `amp_master_grad` must be True."
 
-                if self.nccl_comm_group_config is not None:
-                    strategy = init_nccl_config(self.nccl_comm_group_config, strategy)
-
                 fleet.init(is_collective=True, strategy=strategy)
                 logger.info(strategy)
 
+                if self.expert_parallel_degree > 1:
+                    self.add_moe_comm_group()
+
         elif self.enable_auto_parallel:
             self.tensor_parallel_degree = max(self.tensor_parallel_degree, 1)
             self.sep_parallel_degree = max(self.sep_parallel_degree, 1)
@@ -1722,7 +1539,10 @@ class TrainingArguments:
             if self.sharding_parallel_degree == -1:
                 if len(self.sharding) > 0:
                     self.sharding_parallel_degree = world_size // (
-                        self.tensor_parallel_degree * self.sep_parallel_degree * self.pipeline_parallel_degree
+                        self.tensor_parallel_degree
+                        * self.sep_parallel_degree
+                        * self.context_parallel_degree
+                        * self.pipeline_parallel_degree
                     )
 
             self.sharding_parallel_degree = max(self.sharding_parallel_degree, 1)
@@ -1902,25 +1722,12 @@ class TrainingArguments:
                 amp.custom_white_list = self.amp_custom_white_list if self.amp_custom_white_list is not None else []
 
             self.strategy = strategy
-            sep_degree = self.sep_parallel_degree if self.sep_parallel_degree > 1 else self.context_parallel_degree
             if self.hybrid_parallel_topo_order == "pp_first":
                 order = ["pp", "dp", "mp"]
-
-                degree = [
-                    self.pipeline_parallel_degree,
-                    self.dataset_world_size,
-                    self.tensor_parallel_degree,
-                ]
+                degree = [self.pipeline_parallel_degree, self.dataset_world_size, self.tensor_parallel_degree]
             elif self.hybrid_parallel_topo_order == "sharding_first":
                 order = ["dp", "pp", "mp"]
-                degree = [
-                    self.dataset_world_size,
-                    self.pipeline_parallel_degree,
-                    self.tensor_parallel_degree,
-                ]
-            if sep_degree > 1:
-                order.insert(-1, "sep")
-                degree.insert(-1, sep_degree)
+                degree = [self.dataset_world_size, self.pipeline_parallel_degree, self.tensor_parallel_degree]
             mesh_dims = list(zip(order, degree))
             fleet.auto.create_mesh(mesh_dims)
 
@@ -1939,7 +1746,6 @@ class TrainingArguments:
                 "dp_degree": self.dataset_world_size,
                 "mp_degree": self.tensor_parallel_degree,
                 "pp_degree": self.pipeline_parallel_degree,
-                "sep_degree": sep_degree,
                 "order": order,
             }
             fleet.init(is_collective=True, strategy=strategy)
@@ -1947,10 +1753,7 @@ class TrainingArguments:
         else:
             if world_size > 1:
                 if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized():
-                    if self.save_checkpoint_format in [
-                        "unified_checkpoint",
-                        "flex_checkpoint",
-                    ] or self.load_checkpoint_format in ["unified_checkpoint", "flex_checkpoint"]:
+                    if self.unified_checkpoint:
                         # DP use hybrid group
                         strategy = fleet.DistributedStrategy()
                         fleet.init(is_collective=True, strategy=strategy)
@@ -1958,20 +1761,16 @@ class TrainingArguments:
                         paddle.distributed.init_parallel_env()
 
         if (
-            (
-                self.save_checkpoint_format == "unified_checkpoint"
-                or self.load_checkpoint_format == "unified_checkpoint"
-            )
+            self.unified_checkpoint
             and self.sharding_parallel_degree > 0
             and ShardingOption.FULL_SHARD in self.sharding
         ):
             logger.warning(
-                "Unified checkpoint currently do not support sharding stage3, disabling unified_checkpoint format."
+                "Unified checkpoint currently do not support sharding stage3, set `unified_checkpoint` to False."
             )
-            self.save_checkpoint_format = None
-            self.load_checkpoint_format = None
+            self.unified_checkpoint = False
 
-        if self.save_checkpoint_format == "unified_checkpoint" or self.load_checkpoint_format == "unified_checkpoint":
+        if self.unified_checkpoint:
             unified_checkpoint_config = set(self.unified_checkpoint_config.split(" "))
             if sys.platform.startswith("win") and "async_save" in self.unified_checkpoint_config:
                 raise ValueError("Currently do not support asynchronous saving for Windows system!")
@@ -2095,7 +1894,7 @@ class TrainingArguments:
         assert (
             self.save_steps % self.zcc_ema_interval == 0
         ), f"save_steps[{self.save_steps}] must be divisible by zcc_ema_interval[{self.zcc_ema_interval}]"
-        if self.enable_zero_cost_checkpoint and self.zcc_save_ema_coef is not None:
+        if self.zcc_save_ema_coef is not None:
             assert (
                 self.zcc_workers_num == 1
             ), "EMA function in zero cost checkpoint mode does not support zcc_workers_num > 1 for now."
@@ -2154,7 +1953,10 @@ class TrainingArguments:
             if self.sharding_parallel_degree == -1:
                 if len(self.sharding) > 0:
                     self.sharding_parallel_degree = world_size // (
-                        tensor_parallel_degree * sep_parallel_degree * pipeline_parallel_degree
+                        tensor_parallel_degree
+                        * sep_parallel_degree
+                        * context_parallel_degree
+                        * pipeline_parallel_degree
                     )
 
             sharding_parallel_degree = max(self.sharding_parallel_degree, 1)
@@ -2162,28 +1964,19 @@ class TrainingArguments:
                 logger.warning("sharding_parallel_degree=1 means no sharding, please set sharding to empty!")
                 self.sharding = []
 
+            if sharding_parallel_degree > 1:
+                assert (
+                    sharding_parallel_degree % expert_parallel_degree == 0
+                ), f"sharding_parallel_degree should be divided by expert_parallel_degree, current sharding_parallel_degree: {sharding_parallel_degree}, expert_parallel_degree: {expert_parallel_degree}."
+
             self.data_parallel_degree = world_size // (
-                sharding_parallel_degree * tensor_parallel_degree * sep_parallel_degree * pipeline_parallel_degree
+                sharding_parallel_degree
+                * tensor_parallel_degree
+                * sep_parallel_degree
+                * context_parallel_degree
+                * pipeline_parallel_degree
             )
 
-            if expert_parallel_degree > 1:
-                assert (
-                    self.expert_tensor_parallel_degree <= 1
-                ), "expert_tensor_parallel_degree > 1 is not supported when expert_parallel_degree > 1"
-                moe_sharding_parallel_degree = world_size // (pipeline_parallel_degree * expert_parallel_degree)
-            else:
-                moe_sharding_parallel_degree = 1
-            moe_sharding_parallel_degree = max(moe_sharding_parallel_degree, 1)
-            if moe_sharding_parallel_degree > 1 and self.data_parallel_degree > 1:
-                raise NotImplementedError(
-                    f"Currently only support use expert_data_parallel strategy together with sharding_parallel strategy, but not with data_parallel strategy. But got data_parallel_degree: {self.data_parallel_degree}, expert_parallel_degree: {expert_parallel_degree}, moe_sharding_parallel_degree: {moe_sharding_parallel_degree}."
-                )
-
-            if sharding_parallel_degree > 1 and moe_sharding_parallel_degree > 1:
-                assert (
-                    sharding_parallel_degree % moe_sharding_parallel_degree == 0
-                ), f"sharding_parallel_degree should be divided by moe_sharding_parallel_degree, current sharding_parallel_degree: {sharding_parallel_degree}, moe_sharding_parallel_degree: {moe_sharding_parallel_degree}."
-
             assert not (
                 self.data_parallel_degree > 1 and expert_parallel_degree > 1
             ), f"Currently only support use expert_data_parallel strategy together with sharding_parallel strategy, but not with data_parallel strategy. Currently data_parallel_degree is {self.data_parallel_degree}."
@@ -2205,7 +1998,6 @@ class TrainingArguments:
                 self.context_parallel_degree = context_parallel_degree
                 self.expert_parallel_degree = expert_parallel_degree
                 self.expert_tensor_parallel_degree = expert_tensor_parallel_degree
-                self.moe_sharding_parallel_degree = moe_sharding_parallel_degree
 
             if not self.use_hybrid_parallel:
                 self.sharding = []
@@ -2218,38 +2010,13 @@ class TrainingArguments:
                 self.expert_tensor_parallel_degree = -1
 
         if self.hybrid_parallel_topo_order is None:
-            self.hybrid_parallel_topo_order = "sharding_first"
+            self.hybrid_parallel_topo_order = "pp_first"
         assert self.hybrid_parallel_topo_order in ["pp_first", "sharding_first"]
 
         if self.use_hybrid_parallel and self.enable_auto_parallel:
             self.use_hybrid_parallel = False
 
-    def _post_init_save_checkpoint_format(self):
-        if self.save_checkpoint_format:
-            valid_modes = ["unified_checkpoint", "sharding_io", "flex_checkpoint"]
-            assert (
-                self.save_checkpoint_format in valid_modes
-            ), f"Invalid save_checkpoint_format: {self.save_checkpoint_format}, Only these formats are allowed: {valid_modes}."
-        else:
-            if self.unified_checkpoint:
-                self.save_checkpoint_format = "unified_checkpoint"
-            elif self.save_sharded_model:
-                self.save_checkpoint_format = "sharding_io"
-
-    def _post_init_load_checkpoint_format(self):
-        if self.load_checkpoint_format:
-            valid_modes = ["unified_checkpoint", "sharding_io", "flex_checkpoint"]
-            assert (
-                self.load_checkpoint_format in valid_modes
-            ), f"Invalid load_checkpoint_format: {self.load_checkpoint_format}, Only these formats are allowed: {valid_modes}."
-        else:
-            if self.unified_checkpoint:
-                self.load_checkpoint_format = "unified_checkpoint"
-            elif self.load_sharded_model:
-                self.load_checkpoint_format = "sharding_io"
-
     def add_moe_comm_group(self):
-        hybrid_configs = fleet.fleet._user_defined_strategy.hybrid_configs
         hcg = fleet.get_hybrid_communicate_group()
         topo = hcg._topo
         sharding_parallel_groups = topo.get_comm_list("sharding")
@@ -2261,12 +2028,7 @@ class TrainingArguments:
             for i in range(experts_replicas):
                 rank_indices = list(range(i * self.expert_parallel_degree, (i + 1) * self.expert_parallel_degree))
                 ranks = [ranks_in_current_sharding_group[i] for i in rank_indices]
-                if message2nccl_config is not None and hybrid_configs.get("ep_configs", None) is not None:
-                    group = dist.new_group(
-                        ranks=ranks, nccl_config=message2nccl_config(hybrid_configs["ep_configs"].nccl_config, "ep")
-                    )
-                else:
-                    group = dist.new_group(ranks=ranks)
+                group = dist.new_group(ranks=ranks)
                 if dist.get_rank() in ranks:
                     assert not hasattr(hcg, "expert_parallel_group"), "expert_parallel_group can not be set repeate"
                     setattr(hcg, "expert_parallel_group", group)
@@ -2275,13 +2037,7 @@ class TrainingArguments:
             for i in range(self.expert_parallel_degree):
                 rank_indices = list(range(i, self.sharding_parallel_degree, self.expert_parallel_degree))
                 ranks = [ranks_in_current_sharding_group[i] for i in rank_indices]
-                if message2nccl_config is not None and hybrid_configs.get("ep_configs", None) is not None:
-                    group = dist.new_group(
-                        ranks=ranks,
-                        nccl_config=message2nccl_config(hybrid_configs["ep_configs"].grad_nccl_config, "ep_grad"),
-                    )
-                else:
-                    group = dist.new_group(ranks=ranks)
+                group = dist.new_group(ranks=ranks)
                 if dist.get_rank() in ranks:
                     assert not hasattr(hcg, "expert_grad_comm_group"), "expert_grad_comm_group can not be set repeate"
                     setattr(hcg, "expert_grad_comm_group", group)
@@ -2373,17 +2129,6 @@ class TrainingArguments:
         else:
             return 0
 
-    @property
-    def moe_sharding_parallel_rank(self):
-        if self.use_hybrid_parallel:
-            hcg = fleet.get_hybrid_communicate_group()
-            if hasattr(hcg, "get_moe_sharding_parallel_group"):
-                return max(hcg.get_moe_sharding_parallel_group().rank, 0)
-            else:
-                return 0
-        else:
-            return 0
-
     @property
     def tensor_parallel_rank(self):
         if self.use_hybrid_parallel:
@@ -2408,28 +2153,6 @@ class TrainingArguments:
         else:
             return 0
 
-    @property
-    def expert_parallel_rank(self):
-        if self.use_hybrid_parallel:
-            hcg = fleet.get_hybrid_communicate_group()
-            if hasattr(hcg, "get_expert_parallel_rank"):
-                return max(hcg.get_expert_parallel_rank(), 0)
-            else:
-                return 0
-        else:
-            return 0
-
-    @property
-    def context_parallel_rank(self):
-        if self.use_hybrid_parallel:
-            hcg = fleet.get_hybrid_communicate_group()
-            if hasattr(hcg, "get_context_parallel_rank"):
-                return max(hcg.get_context_parallel_rank(), 0)
-            else:
-                return 0
-        else:
-            return 0
-
     def _format_name(self, prefix, rank, degree):
         size = 2
         return f"{prefix}{rank:0>{size}d}"
@@ -2444,7 +2167,7 @@ class TrainingArguments:
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
             if self.sharding_parallel_degree > 1:
                 name.append(self._format_name("shard", self.sharding_parallel_rank, self.sharding_parallel_degree))
-            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
+            if self.use_expert_parallel:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
             return "_".join(name)
         else:
@@ -2460,10 +2183,8 @@ class TrainingArguments:
                 name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
-            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
+            if self.use_expert_parallel:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
-            if self.use_expert_parallel and self.expert_parallel_degree > 1:
-                name.append(self._format_name("moe_sharding", self.expert_parallel_rank, self.expert_parallel_degree))
             return "_".join(name)
 
         else:
@@ -2471,9 +2192,7 @@ class TrainingArguments:
                 return self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree)
             return None
 
-    def sharded_name_suffix(self, shard_id=None, pp_id=None, moe_id=None, sharding_parallel_degree=None):
-        if sharding_parallel_degree is None:
-            sharding_parallel_degree = self.sharding_parallel_degree
+    def sharded_name_suffix(self, shard_id=None, pp_id=None, moe_id=None):
         if self.use_hybrid_parallel:
             name = []
             if self.tensor_parallel_degree > 1:
@@ -2483,12 +2202,12 @@ class TrainingArguments:
                     pp_id = self.pipeline_parallel_rank
                 assert isinstance(pp_id, int)
                 name.append(self._format_name("pp", pp_id, self.pipeline_parallel_degree))
-            if sharding_parallel_degree > 1:
+            if self.sharding_parallel_degree > 1:
                 if shard_id is None:
                     shard_id = self.sharding_parallel_rank
                 assert isinstance(shard_id, int)
-                name.append(self._format_name("shard", shard_id, sharding_parallel_degree))
-            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
+                name.append(self._format_name("shard", shard_id, self.sharding_parallel_degree))
+            if self.use_expert_parallel:
                 if moe_id is None:
                     moe_id = self.data_parallel_rank
                 assert isinstance(moe_id, int)
@@ -2589,13 +2308,9 @@ class TrainingArguments:
                 return True
             elif self.enable_auto_parallel:
                 return True
-            elif self.save_checkpoint_format == "flex_checkpoint":
-                return False
             elif self.use_hybrid_parallel:
                 # save on dataset rank 0
-                return (
-                    self.sharding_parallel_rank == 0 and (self.data_parallel_rank == 0 or self.use_expert_parallel)
-                ) or (self.expert_parallel_degree > 1 and self.moe_sharding_parallel_rank == 0)
+                return self.sharding_parallel_rank == 0 and (self.data_parallel_rank == 0 or self.use_expert_parallel)
             else:
                 return self.process_index == 0 or self.use_expert_parallel
 
@@ -2611,16 +2326,16 @@ class TrainingArguments:
         if self.enable_auto_parallel:
             return False
         return (
-            ShardingOption.SHARD_OP in self.sharding
-            and self.sharding_parallel_degree > 1
-            and self.save_checkpoint_format == "sharding_io"
+            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.save_sharded_model
         )
 
     @property
     def should_load_sharding_stage1_model(self):
         if self.enable_auto_parallel:
             return False
-        return self.load_checkpoint_format == "sharding_io"
+        return (
+            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.load_sharded_model
+        )
 
     @property
     def should_load_dataset(self):

paddlenlp/trainer/unified_checkpoint/load_local.py

@@ -233,7 +233,7 @@ def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoin
                         )
                     else:
                         tp_actions = model.get_tensor_parallel_convert_actions(
-                            model.config, model_keys, ignore_error=True, is_optim=True
+                            model.config, model_keys, ignore_error=True
                         )
                     if not is_master_weights:
                         tp_actions = mapping_optimizer_tp_actions(tp_actions, expected_keys)

paddlenlp/trainer/unified_checkpoint/shared_memory_utils.py

@@ -37,7 +37,6 @@ dtype_mapping = {
     paddle.int32: np.int32,
     paddle.int64: np.int64,
     paddle.uint8: np.uint8,
-    paddle.int8: np.int8,
     paddle.bool: np.bool_,
     paddle.float16: np.float16,
     paddle.bfloat16: np.uint16,

paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py

@@ -648,7 +648,7 @@ def unified_optimizer_into_shards(
             tp_actions = model._get_tensor_parallel_convert_actions(model_keys, is_split=False, ignore_error=True)
         else:
             tp_actions = model.get_tensor_parallel_convert_actions(
-                model.config, model_keys, is_split=False, ignore_error=True, is_optim=True
+                model.config, model_keys, is_split=False, ignore_error=True
             )
         logger.info("Unified optimizer tensor parallel in shards")
         optim_state_dict = merge_tensor_parallel_for_optimizer(

paddlenlp/trainer/unified_checkpoint/utils.py

@@ -26,7 +26,6 @@ from paddlenlp.trainer.trainer_utils import ExplicitEnum, ShardingOption
 from paddlenlp.trainer.utils.helper import distributed_isfile
 from paddlenlp.transformers.model_utils import (
     PretrainedModel,
-    clean_model_class_name,
     get_parameter_dtype,
     unwrap_model,
 )
@@ -766,9 +765,9 @@ def save_model_config(model_to_save, save_directory):
     config_to_save = save_config(model_to_save)
     # Attach architecture to the config
     if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
-        config_to_save.architectures = [clean_model_class_name(model_to_save.model.__class__.__name__)]
+        config_to_save.architectures = [model_to_save.model.__class__.__name__]
     else:
-        config_to_save.architectures = [clean_model_class_name(model_to_save.__class__.__name__)]
+        config_to_save.architectures = [model_to_save.__class__.__name__]
 
     config_to_save.save_pretrained(save_directory)
     # save generation config