Delete .github/workflows/distribute-v100.yml

* change_save_meta_in_zcc_crossponding_to_paddle

* delete

* fix

* change_save_meta_in_zcc_crossponding_to_paddle and fix args

* fix deadcode

---------

Co-authored-by: liufengwei02 <liufengwei02@baidu.com>

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

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

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

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

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

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

@@ -0,0 +1,52 @@
+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,17 +1,80 @@
-name: Lint
+name: Codestyle Check
 
 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
-    runs-on: ubuntu-24.04
+    needs: check-bypass
+    if: ${{ needs.check-bypass.outputs.can-skip != 'true' }}
+    runs-on: [self-hosted, ernie-cpu]
     steps:
-      - uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
-      - name: checkout develop
+      - 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 \
+            -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
@@ -19,14 +82,32 @@ jobs:
           else
             echo "local develop branch exist, skipping"
           fi
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.10'
-          cache: 'pip' # caching pip dependencies
-      - name: Install dependencies
+          '
+          
+      - name: Setup Environment
         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: run the command
-        run: make lint
-    
+          '
+
+      - 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

.github/workflows/llm.yml

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

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

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

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

@@ -0,0 +1,204 @@
+# 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,3 +62,12 @@ 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/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h

@@ -69,7 +69,8 @@ void generic_moe_gemm_kernelLauncher(const T* A,
                                      cudaStream_t stream,
                                      int* kernel_occupancy = nullptr) {
   if (gemm_config.split_k_style != SplitKStyle::NO_SPLIT_K) {
-    PADDLE_FATAL("[MoeGemm] Grouped gemm does not support split-k");
+    throw std::runtime_error(
+        "[MoeGemm] Grouped gemm does not support split-k");
   }
 
 #ifdef PADDLE_CUDA_BF16
@@ -169,7 +170,7 @@ void generic_moe_gemm_kernelLauncher(const T* A,
   int occupancy = std::min(2, GemmGrouped::maximum_active_blocks());
 
   if (occupancy == 0) {
-    PADDLE_FATAL(
+    throw std::runtime_error(
         "[MoE Runner] GPU lacks the shared memory resources to run "
         "GroupedGEMM kernel");
   }
@@ -197,7 +198,7 @@ void generic_moe_gemm_kernelLauncher(const T* A,
   if (can_implement != cutlass::Status::kSuccess) {
     std::string err_msg = "MoEFC kernel will fail for params. Error: " +
                           std::string(cutlassGetStatusString(can_implement));
-    PADDLE_FATAL("[MoE Runner] " + err_msg);
+    throw std::runtime_error("[MoE Runner] " + err_msg);
   }
 
   auto init_status = gemm.initialize(args);
@@ -243,7 +244,7 @@ struct dispatch_stages {
     std::string err_msg = "Cutlass fpA_intB gemm. Not instantiates for arch " +
                           std::to_string(arch::kMinComputeCapability) +
                           " with stages set to " + std::to_string(Stages);
-    PADDLE_FATAL("[dispatch_stages::dispatch] " + err_msg);
+    throw std::runtime_error("[dispatch_stages::dispatch] " + err_msg);
   }
 };
 
@@ -394,7 +395,8 @@ void dispatch_gemm_config(const T* A,
     default:
       std::string err_msg = "dispatch_gemm_config does not support stages " +
                             std::to_string(gemm_config.stages);
-      PADDLE_FATAL("[MoE][dispatch_gemm_config] " + err_msg);
+      throw std::runtime_error(
+          "[MoE][dispatch_gemm_config] " + err_msg);
       break;
   }
 }
@@ -452,15 +454,16 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
     dispatch_gemm_config_macro(64, 128, 64, 32, 64, 64);
     dispatch_gemm_config_macro(128, 128, 64, 64, 32, 64);
     case CutlassTileConfig::Undefined:
-      PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+      throw std::runtime_error(
+          "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
       break;
     case CutlassTileConfig::ChooseWithHeuristic:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass] gemm config should have "
           "already been set by heuristic.");
       break;
     default:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass] Config is invalid for same "
           "type MoE tensorop GEMM.");
       break;
@@ -497,15 +500,16 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
       dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
       case CutlassTileConfig::Undefined:
-        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
         break;
       case CutlassTileConfig::ChooseWithHeuristic:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] gemm config should have "
             "already been set by heuristic.");
         break;
       default:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
             "mixed type tensorop GEMM.");
         break;
@@ -513,22 +517,27 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
   } else {
     switch (gemm_config.tile_config) {
       dispatch_gemm_config_macro(16, 128, 64, 16, 32, 64);
+      dispatch_gemm_config_macro(16, 256, 64, 16, 64, 64);
+      dispatch_gemm_config_macro(64, 64, 64, 32, 32, 64);
       dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+      dispatch_gemm_config_macro(128, 64, 64, 64, 32, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
       dispatch_gemm_config_macro(128, 128, 64, 64, 64, 64);
       dispatch_gemm_config_macro(128, 128, 64, 128, 32, 64);
       dispatch_gemm_config_macro(128, 256, 64, 64, 64, 64);
       dispatch_gemm_config_macro(64, 128, 64, 64, 32, 64);
+      dispatch_gemm_config_macro(256, 128, 64, 64, 64, 64);
       case CutlassTileConfig::Undefined:
-        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        throw std::runtime_error(
+            "[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
         break;
       case CutlassTileConfig::ChooseWithHeuristic:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] gemm config should have "
             "already been set by heuristic.");
         break;
       default:
-        PADDLE_FATAL(
+        throw std::runtime_error(
             "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
             "mixed type tensorop GEMM.");
         break;
@@ -561,17 +570,17 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
   switch (gemm_config.tile_config) {
     dispatch_gemm_config_macro(128, 128, 8, 64, 64, 8);
     case CutlassTileConfig::Undefined:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config "
           "undefined.");
       break;
     case CutlassTileConfig::ChooseWithHeuristic:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config should "
           "have already been set by heuristic.");
       break;
     default:
-      PADDLE_FATAL(
+      throw std::runtime_error(
           "[dispatch_moe_gemm_to_cutlass][SIMT] Unsupported config "
           "for float MoE gemm.");
       break;

csrc/gpu/moe/preprocess_for_moe.cu

@@ -140,12 +140,14 @@ 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("errors");
+      printf("【ERROR】, please add the current number of experts in the case");
     }
 
     const int block_threads = 256;

csrc/gpu/unittest/test_get_padding_offset_v2.py

@@ -64,6 +64,5 @@ 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()

csrc/gpu/update_inputs_v2.cu

@@ -42,8 +42,10 @@ __global__ void update_inputs_kernel_v2(
     const int bsz,
     const int max_bsz,
     const int input_ids_stride,
-    const int end_length) {
+    const int end_length,
+    const int Flag_truncated_return_eos) {
   int thread_idx = threadIdx.x;
+  bool output_len_truncated = false;
   // update step_idx and stop_flags
   if (thread_idx < max_bsz) {
     bool stop_flag = stop_flags[thread_idx];
@@ -52,6 +54,7 @@ __global__ void update_inputs_kernel_v2(
     }
     if (step_idx[thread_idx] >= max_dec_len[thread_idx]) {
       stop_flags[thread_idx] = true;
+      output_len_truncated = true;
     }
   }
   __syncthreads();
@@ -61,8 +64,13 @@ __global__ void update_inputs_kernel_v2(
       if (seq_lens_this_time[thread_idx] == 0) {
         next_tokens[thread_idx] = -1;
       } else {
-        next_tokens[thread_idx] = end_ids[0];
-        kwargs_next_tokens[thread_idx] = end_ids[0];
+        if (!Flag_truncated_return_eos && output_len_truncated) {
+          // output len truncated will not return eos for rl.
+          kwargs_next_tokens[thread_idx] = next_tokens[thread_idx];
+        }else{
+          next_tokens[thread_idx] = end_ids[0];
+          kwargs_next_tokens[thread_idx] = end_ids[0];
+        }
       }
     } else {
       kwargs_next_tokens[thread_idx] = next_tokens[thread_idx];
@@ -127,6 +135,15 @@ void UpdateInputesV2(const paddle::Tensor& stop_flags,
   const int end_length = end_ids.shape()[0];
 
   auto not_need_stop_gpu = not_need_stop.copy_to(stop_flags.place(), false);
+  int Flag_truncated_return_eos = 1;
+  if (const char* inference_truncated_return_eos_env_p =
+          std::getenv("INFERENCE_TRUNCATED_RETURN_EOS")) {
+      std::string inference_truncated_return_eos_env_str(
+          inference_truncated_return_eos_env_p);
+      int inference_truncated_return_eos_from_env =
+          std::stoi(inference_truncated_return_eos_env_str);
+      Flag_truncated_return_eos = inference_truncated_return_eos_from_env;
+  }
 
   update_inputs_kernel_v2<1024><<<1, 1024, 0, input_ids.stream()>>>(
     const_cast<bool*>(not_need_stop_gpu.data<bool>()),
@@ -145,7 +162,8 @@ void UpdateInputesV2(const paddle::Tensor& stop_flags,
     now_bsz,
     max_bsz,
     input_ids_stride,
-    end_length
+    end_length,
+    Flag_truncated_return_eos
   );
 
   auto not_need_stop_cpu = not_need_stop_gpu.copy_to(not_need_stop.place(), false);

docs/zh/index.rst

@@ -128,6 +128,10 @@
    :caption: 实践教程
 
    AI Studio Notebook <tutorials/overview>
+   大模型预训练新手指南 <llm/docs/pretrain_tutorial.md>
+   大模型精调新手指南 <llm/docs/finetune_tutorial.md>
+   大模型对齐新手指南 <llm/docs/alignment_tutorial.md>
+   大模型量化新手指南 <llm/docs/quantization_tutorial.md>
 
 .. toctree::
    :maxdepth: 1

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

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

docs/zh/llm/devices/intel_hpu/tests/README.md

@@ -0,0 +1 @@
+../../../../../../llm/devices/intel_hpu/tests/README.md

docs/zh/llm/docs/alignment_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/alignment_tutorial.md

docs/zh/llm/docs/finetune_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/finetune_tutorial.md

docs/zh/llm/docs/pretrain.md

@@ -0,0 +1 @@
+../../../../llm/docs/pretrain.md

docs/zh/llm/docs/pretrain_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/pretrain_tutorial.md

docs/zh/llm/docs/quantization_tutorial.md

@@ -0,0 +1 @@
+../../../../llm/docs/quantization_tutorial.md

docs/zh/locale/en/LC_MESSAGES/source/paddlenlp.taskflow.utils.po

@@ -276,7 +276,7 @@ msgid "word of current node."
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.BurkhardKellerTree:1
-msgid "Implementataion of BK-Tree"
+msgid "Implementation of BK-Tree"
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.BurkhardKellerTree.add:1
@@ -300,7 +300,7 @@ msgid "similar words."
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.TriedTree:1
-msgid "Implementataion of TriedTree"
+msgid "Implementation of TriedTree"
 msgstr ""
 
 #: of paddlenlp.taskflow.utils.TriedTree.add_word:1

docs/zh/locale/en/LC_MESSAGES/source/paddlenlp.transformers.squeezebert.tokenizer.po

@@ -50,7 +50,7 @@ msgid ""
 msgstr ""
 
 #: of paddlenlp.transformers.squeezebert.tokenizer.SqueezeBertTokenizer:11
-msgid "The special token for unkown words. Default: \"[UNK]\"."
+msgid "The special token for unknown words. Default: \"[UNK]\"."
 msgstr ""
 
 #: of paddlenlp.transformers.squeezebert.tokenizer.SqueezeBertTokenizer:13

llm/README.md

@@ -13,27 +13,9 @@
 - **大模型无损量化**。大模型套件预先集成了 PaddleSlim LLM.PTQ 量化算法，以及业界广泛采用的 GPTQ 和 AWQ 的 W4量化方法，成功实现了对主流大模型的无损量化处理，显著加速了模型的推理速度。
 - **高性能推理**。大模型套件的高性能推理模块内置了动态插入和全环节算子融合的高级策略，这极大地提升了并行推理的速度。同时，该模块隐藏了底层技术细节，为用户提供了开箱即用的高性能并行推理功能。
 
-
 ## 🛠️ 支持模型列表 🛠️
 
-| Model                                  | Pretrain | SFT | LoRA | Prefix Tuning | DPO/SimPO/ORPO/KTO | RLHF | Mergekit | Quantization | Torch convert |
-|----------------------------------------|----------|-----|------|---------------|----------------|------|-------|--------------|---------------|
-| [LLaMA](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)                | ✅        | ✅   | ✅    | ✅             | ✅             | ✅    | ✅    | ✅            | ✅             |
-| [Qwen](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen)                  | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | 🚧           | ✅             |
-| [Mixtral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mixtral)            | ✅        | ✅   | ✅    | ❌             | ✅             | 🚧   | ✅    | 🚧           | 🚧            |
-| [Mistral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mistral)            | ✅         | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | 🚧           | ✅             |
-| [Baichuan/Baichuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)   | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | ✅            | ✅             |
-| [ChatGLM-6B](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm)         | ✅        | ✅   | ✅    | ✅             | 🚧            | 🚧   | ✅    | ✅            | ❌             |
-| [ChatGLM2/ChatGLM3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2) | ✅        | ✅   | ✅    | ✅             | ✅             | 🚧   | ✅    | ✅            | ✅             |
-| [Bloom](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/bloom)                | ✅        | ✅   | ✅    | ✅             | 🚧            | 🚧   | ✅    | ✅            | ✅             |
-| [GPT-3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gpt-3)                | ✅        | ✅   | 🚧   | 🚧            | 🚧            | 🚧   | ✅    | 🚧           | ✅             |
-| [OPT](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/opt)                    | ✅       | ✅   | ✅    | 🚧            | 🚧            | 🚧   | ✅    | 🚧           | ✅             |
-| [Gemma](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gemma)                | ✅       | ✅   | ✅    | 🚧            | ✅            | 🚧   | ✅    | 🚧           | 🚧             |
-| [Yuan](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan)                  | ✅       | ✅   | ✅    | 🚧            | ✅            | 🚧   | ✅    | 🚧           | 🚧             |
-
-- ✅: Supported
-- 🚧: In Progress
-- ❌: Not Supported
+参考 [PaddleNLP 支持模型列表](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/README.md#模型支持)
 
 ## 🚀 快速开始 🚀
 

llm/alignment/rl/README.md

@@ -1,6 +1,6 @@
 # GRPO && REINFORCE++
 
-GRPO（Group Relative Policy Optimization，组相对策略优化）是 PPO（Proximal Policy Optimization，近端策略优化）算法的一种变体。与 PPO 不同，GRPO 省略了价值函数估计器。在 GRPO 中，对于每个状态 \(s\)，算法会从当前策略 \(\pi_{\theta_{t}}\) 中采样多个动作 \(a_{1}, \dots, a_{G}\)。然后，GRPO 计算这些动作相对于组内其他动作的“组相对优势”（group-relative advantage），以此作为优化策略的依据。
+GRPO（Group Relative Policy Optimization，组相对策略优化）是 PPO（Proximal Policy Optimization，近端策略优化）算法的一种变体。与 PPO 不同，GRPO 省略了价值函数估计器。在 GRPO 中，对于每个状态 $s$，算法会从当前策略 $\pi_{\theta_{t}}$ 中采样多个动作 $a_{1}, \dots, a_{G}$。然后，GRPO 计算这些动作相对于组内其他动作的“组相对优势”（group-relative advantage），以此作为优化策略的依据。
 REINFORCE++ 是经典 REINFORCE 算法的改进版本，通过融合 PPO 的关键优化技术并移除 Critic Model，实现了更加简洁高效的策略优化。相比于传统的 REINFORCE，REINFORCE++ 在 Token-Level KL 惩罚、PPO-Clip、优势标准化、奖励裁剪和奖励标准化等关键技术上进行了改进，从而提高了训练过程的效率和稳定性。
 
 以下是详细的使用文档和示例：
@@ -203,6 +203,7 @@ export FLAGS_cascade_attention_max_partition_size=2048
 
 python -u -m paddle.distributed.launch --devices "0,1,2,3" run_rl.py ../../config/qwen/reinforce_plus_plus_argument.yaml
 ```
+我们提供根据上述脚本可复现的[wandb 日志](https://api.wandb.ai/links/ainlp66-netflix/injcw3ra)。
 
 ### 在线监控
 在`grpo_argument.yaml`和`reinforce_plus_plus_argument.yaml`中设置的输出目录为`"logging_dir": "vdl_log"`, 可以通过以下命令查看训练过程

llm/alignment/rl/gsm8k_processor.py

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

@@ -32,7 +32,7 @@ from paddlenlp.rl.utils.config_utils import (
     TrainingArguments,
 )
 from paddlenlp.rl.utils.offload_utils import offload_tensor_to_cpu
-from paddlenlp.rl.utils.reshard_utils import init_rollout_env
+from paddlenlp.rl.utils.reshard_utils import ReshardController
 from paddlenlp.rl.utils.timer_utils import timers_scope_runtimer
 from paddlenlp.trainer import (
     EarlyStoppingCallback,
@@ -42,6 +42,7 @@ from paddlenlp.trainer import (
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForTokenClassification,
     AutoTokenizer,
     PretrainedConfig,
 )
@@ -81,6 +82,7 @@ def create_actor_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
+    reshard_controller: ReshardController = None,
 ):
     with timers_scope_runtimer("Actor model loading time"):
         # actor model
@@ -103,7 +105,7 @@ def create_actor_models(
         actor_model_config.use_sparse_head_and_loss_fn = False
         actor_model_config.seq_length = data_args.max_length
         actor_model_config.max_sequence_length = data_args.max_length
-        print(f"Loading Actor model with config:\n\t{actor_model_config}\n")
+        logger.info(f"Loading Actor model with config:\n\t{actor_model_config}\n")
 
         if not training_args.autotuner_benchmark:
             actor_model = AutoModelForCausalLM.from_pretrained(
@@ -113,18 +115,16 @@ def create_actor_models(
             actor_model = AutoModelForCausalLM.from_config(actor_model_config)
 
     with timers_scope_runtimer("Actor eval model loading time"):
-        if (
-            training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
-            or training_args.pipeline_parallel_degree > 1
-        ):
+        if reshard_controller is not None:
+            reshard_controller.set_rollout_env("[create actor eval model]")
             actor_eval_model_config = copy.deepcopy(actor_model_config)
             actor_eval_model_config.use_fused_head_and_loss_fn = False
-            with init_rollout_env(training_args.rollout_tensor_parallel_degree):
-                hcg = fleet.get_hybrid_communicate_group()
-                actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
-                actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
-                # TODO(gongenlei): lazy load lazy guard
-                actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            hcg = fleet.get_hybrid_communicate_group()
+            actor_eval_model_config.tensor_parallel_degree = hcg.get_model_parallel_world_size()
+            actor_eval_model_config.tensor_parallel_rank = hcg.get_model_parallel_rank()
+            # TODO(gongenlei): lazy load lazy guard
+            actor_eval_model = AutoModelForCausalLM.from_config(actor_eval_model_config)
+            reshard_controller.set_train_env("[after create actor eval model]")
         else:
             actor_eval_model = None
 
@@ -135,7 +135,6 @@ 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,
@@ -171,7 +170,7 @@ def create_reward_models(
         LlmMetaConfig.set_llm_config(reward_model_config, training_args)
         reward_model_config.max_position_embeddings = data_args.max_length
         reward_model_config.use_sparse_head_and_loss_fn = False
-        print(f"Loading Reward model with config:\n\t{reward_model_config}\n")
+        logger.info(f"Loading Reward model with config:\n\t{reward_model_config}\n")
 
         config = copy.deepcopy(reward_model_config)
         if training_args.eval_mode is not None:
@@ -211,46 +210,43 @@ def create_critic_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
-    reward_model,
 ):
     with timers_scope_runtimer("Critic model loading time"):
-        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,
+        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,
             )
-            if not training_args.autotuner_benchmark:
-                critic_model.set_state_dict(reward_model.state_dict())
         else:
-            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_model = AutoModelForTokenClassification.from_config(
+                critic_model_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.reward_critic_tokenizer_alpha,
+        tokenizer_alpha=model_args.critic_tokenizer_alpha,
         use_fast=True,
     )
     if critic_tokenizer.pad_token_id is None:
@@ -262,8 +258,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("Reward critic eval model loading time"):
-            critic_eval_model = AutoModelForScore.from_config(config)
+        with timers_scope_runtimer("Critic eval model loading time"):
+            critic_eval_model = AutoModelForTokenClassification.from_config(config)
     else:
         critic_eval_model = None
 
@@ -271,7 +267,7 @@ def create_critic_models(
 
 
 def create_rl_dataset(data_args, training_args, tokenizer):
-    requires_label = True if training_args.use_rm_server else False
+    requires_label = True if training_args.use_rm_server or training_args.use_rule_reward else False
     train_ds = RLHFDataset(
         dataset_name_or_path=data_args.train_datasets,
         tokenizer=tokenizer,
@@ -323,18 +319,27 @@ def main():
         max_sequence_length=data_args.max_length,
     )
 
-    actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
-        model_args, data_args, training_args, common_config
-    )
+    if (
+        training_args.rollout_tensor_parallel_degree != training_args.tensor_parallel_degree
+        or training_args.pipeline_parallel_degree > 1
+    ):
+        reshard_controller = ReshardController(tensor_parallel_degree=training_args.rollout_tensor_parallel_degree)
+    else:
+        reshard_controller = None
 
-    if not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
+    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:
         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, reward_model
+            model_args, data_args, training_args, common_config
         )
     else:
         critic_model, critic_eval_model, critic_tokenizer = None, None, None
@@ -348,7 +353,8 @@ def main():
         offload_tensor_to_cpu((reference_model, "freeze_model"))
 
         if training_args.rl_algorithm == "ppo":
-            offload_tensor_to_cpu((reward_model, "freeze_model"))
+            if not training_args.use_rm_server and not training_args.use_rule_reward:
+                offload_tensor_to_cpu((reward_model, "freeze_model"))
             if critic_eval_model is not None:
                 offload_tensor_to_cpu((critic_eval_model, "freeze_model"))
 
@@ -356,7 +362,14 @@ def main():
         paddle.device.cuda.empty_cache()
 
     def compute_metrics(eval_preds):
-        accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
+        '''
+        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()
         return {"accuracy": accuracy}
 
     try:
@@ -382,11 +395,12 @@ def main():
         data_collator=partial(
             collate_fn,
             pad_token_id=actor_tokenizer.pad_token_id,
-            requires_label=True if training_args.use_rm_server else False,
+            requires_label=True if training_args.use_rm_server or training_args.use_rule_reward 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)
         generation_config=generation_config,
+        reshard_controller=reshard_controller,
     )
 
     # TODO(gongenlei) resume_from_checkpoint is not ready

llm/auto_parallel/deepseek-v3/run_pretrain_auto.py

@@ -460,10 +460,25 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    if training_args.enable_linear_fused_grad_add:
-        from fused_layers import mock_layers
+    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()
+        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

llm/auto_parallel/gpt-3/run_pretrain_auto.py

@@ -38,6 +38,7 @@ from paddlenlp.transformers import (
     CosineAnnealingWithWarmupDecay,
     GPTConfig,
     GPTForCausalLMAuto,
+    GPTForCausalLMAutoPP,
     GPTForCausalLMNet,
     GPTPretrainingCriterionAuto,
     GPTPretrainingCriterionNet,
@@ -48,6 +49,7 @@ from paddlenlp.utils.tools import get_env_device
 
 MODEL_CLASSES = {
     "gpt": (GPTConfig, GPTForCausalLMAuto, GPTPretrainingCriterionAuto),
+    "gpt_pp": (GPTConfig, GPTForCausalLMAutoPP, GPTPretrainingCriterionAuto),
     "gpt_network": (GPTConfig, GPTForCausalLMNet, GPTPretrainingCriterionNet),
 }
 
@@ -72,6 +74,12 @@ 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."},
@@ -85,13 +93,18 @@ 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={
@@ -431,6 +444,26 @@ 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
 
@@ -525,6 +558,7 @@ 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
@@ -534,6 +568,14 @@ 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"
@@ -589,6 +631,7 @@ 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,6 +41,7 @@ from paddlenlp.transformers import (
     LinearAnnealingWithWarmupDecay,
     LlamaConfig,
     LlamaForCausalLM3DAuto,
+    LlamaForCausalLM3DAutoPP,
     LlamaForCausalLMNet,
     LlamaPretrainingCriterion3DAuto,
     LlamaPretrainingCriterionNet,
@@ -49,6 +50,7 @@ from paddlenlp.utils.log import logger
 
 MODEL_CLASSES = {
     "llama": (LlamaConfig, LlamaForCausalLM3DAuto, LlamaPretrainingCriterion3DAuto),
+    "llama_pp": (LlamaConfig, LlamaForCausalLM3DAutoPP, LlamaPretrainingCriterion3DAuto),
     "llama_network": (LlamaConfig, LlamaForCausalLMNet, LlamaPretrainingCriterionNet),
 }
 
@@ -86,13 +88,18 @@ 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."},
+    )
 
     def __post_init__(self):
         super().__post_init__()
@@ -403,12 +410,20 @@ 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,
             )
@@ -449,10 +464,25 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    if training_args.enable_linear_fused_grad_add:
-        from fused_layers import mock_layers
+    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()
+        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
@@ -540,6 +570,8 @@ 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
@@ -557,7 +589,11 @@ def main():
 
     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:
+    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()
@@ -619,6 +655,7 @@ 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

@@ -58,6 +58,8 @@ from paddlenlp.data.causal_dataset import (
 )
 from paddlenlp.trainer.utils.doc import add_start_docstrings
 
+os.environ["USE_CASUAL_MASK"] = "True"
+
 
 @dataclass
 @add_start_docstrings(AutoTrainingArguments.__doc__)
@@ -91,7 +93,8 @@ 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(
@@ -430,10 +433,25 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    if training_args.enable_linear_fused_grad_add:
-        from fused_layers import mock_layers
+    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()
+        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
@@ -530,11 +548,15 @@ def main():
 
     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:
+    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"
     if training_args.fp16_opt_level == "O2":

llm/benchmark/rl/README.md

@@ -0,0 +1,188 @@
+# Large Language Model Throughput Testing Framework
+
+This framework is designed to test the **throughput performance** of Large Language Models (LLMs) across different deployment methods: **offline inference** (using PaddlePaddle and PyTorch) and **online inference** (via API). It specifically focuses on evaluating the model's ability to handle **batch queries**, measuring throughput in tokens per second under various configurations and batch sizes.
+
+## Features
+
+  * **Diverse Deployment Method Support:** Tests LLMs deployed via online API, and offline inference with PaddlePaddle and PyTorch.
+  * **Batch Query Throughput Calculation:** Accurately measures throughput (tokens/s) for concurrent queries, providing insights into the model's performance under load.
+  * **Detailed Time Logging:** Records the total time for each batch processing operation.
+
+## How It Works
+
+This framework operates by sending batched requests to specified endpoints (API or local inference scripts) and collecting performance data on how the model generates responses for multiple queries simultaneously.
+
+1.  **Input Queries:** You provide a set of questions as test input, typically from a `.parquet` or text file.
+2.  **Batch Processing:** The framework groups these questions into batches of a specified `rollout_input_batch_size`.
+3.  **Generate Responses:** For each query within a batch, the framework requests the model to generate `rollout_n` responses.
+4.  **Time Measurement:** The total time from sending a batch of questions to receiving all corresponding responses is recorded.
+5.  **Token Statistics:** The total number of tokens generated across all responses in a batch is summed up.
+6.  **Throughput Calculation:** Throughput (tokens/s) is calculated by dividing the total tokens generated by the total time taken for the batch to complete.
+
+## Usage
+
+This section details how to run throughput tests for each deployment method using the provided shell scripts.
+
+### Data Preparation
+
+To run the tests, you'll first need to download and extract the `rl_data.tar.gz` archive, which contains the GSM8K dataset in a suitable format for testing.
+
+```
+cd llm/benchmark/rl
+wget https://paddle-qa.bj.bcebos.com/paddlenlp/rl_data.tar.gz     
+tar -zxvf rl_data.tar.gz 
+```
+
+Extract the contents of the archive. This will create a data folder containing the GSM8K dataset.
+
+### Online API Inference
+
+This script tests the throughput of a remote LLM API.
+
+**Configuration (`api_serve.sh`):**
+
+```bash
+output_dir="api_serve_results"
+
+python api_serve.py \
+    --openai_urls "your_url1" "your_url2"\
+    --api_keys "key1" "key2" \
+    --model "Qwen2.5-7B-Instruct-1M" \
+    --tokenizer "Qwen/Qwen2.5-7B-Instruct-1M" \
+    --input_file ./data/gsm8k/instruct/train.parquet \
+    --output_dir ${output_dir} \
+    --use_fastdeploy true \
+    --rollout_input_batch_size 8 \
+    --rollout_n 8 \
+    --top_p 1.0 \
+    --temperature 0.7 \
+    --max_dec_len 8192 \
+    --limit_rows 512
+```
+
+  * **`--openai_urls`**: URLs of the API endpoints to test.
+  * **`--api_keys`**: API keys for authentication (if required).
+  * **`--model`**: Name of the model being tested.
+  * **`--tokenizer`**: Path or name of the tokenizer.
+  * **`--input_file`**: Path to the input dataset file.
+  * **`--output_dir`**: Directory to save output results.
+  * **`--use_fastdeploy`**: Use FastDeploy if true, otherwise use vLLM (default: true).
+  * **`--rollout_input_batch_size`**: The batch size for API requests.
+  * **`--rollout_n`**: Number of responses to generate for each input query.
+  * **`--max_dec_len`**: Maximum decoding length for responses.
+  * **`--limit_rows`**: Limit the number of input rows processed.
+
+**Run command:**
+
+```bash
+bash scripts/api_serve.sh
+```
+
+### Offline PaddlePaddle Inference
+
+This script tests the throughput of an LLM using offline PaddlePaddle inference, potentially with distributed processing.
+
+**Configuration (`paddle_infer.sh`):**
+
+```bash
+unset PADDLE_TRAINERS_NUM
+unset PADDLE_ELASTIC_JOB_ID
+unset PADDLE_TRAINER_ENDPOINTS
+unset DISTRIBUTED_TRAINER_ENDPOINTS
+unset FLAGS_START_PORT
+unset PADDLE_ELASTIC_TIMEOUT
+
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP":$PYTHONPATH
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP/llm":$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_custom_allreduce=0
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+output_dir="pdpd_bf16_offline"
+
+python -u -m paddle.distributed.launch --log_dir ${output_dir}/logs --gpus ${CUDA_VISIBLE_DEVICES} paddle_infer.py \
+  --actor_model_name_or_path your_model_name \
+  --max_src_len 2048 \
+  --min_dec_len 32 \
+  --max_dec_len 30720 \
+  --top_p 1.0 \
+  --temperature 1.0 \
+  --rollout_input_batch_size 4 \
+  --rollout_n 8 \
+  --rollout_max_num_seqs 24 \
+  --rollout_quant_type "" \
+  --tensor_parallel_degree 4 \
+  --limit_rows 640 \
+  --input_file file.parquet \
+  --output_dir ${output_dir} > ./paddleinfer.log 2>&1
+```
+
+  * **`CUDA_VISIBLE_DEVICES`**: Specifies the GPUs to be used.
+  * **`paddle.distributed.launch`**: Initiates a distributed PaddlePaddle training/inference job.
+  * **`--actor_model_name_or_path`**: Path to the pre-trained model.
+  * **`--max_src_len`**: Maximum source sequence length.
+  * **`--rollout_input_batch_size`**: The batch size for inference.
+  * **`--rollout_n`**: Number of responses to generate for each input query.
+  * **`--tensor_parallel_degree`**: Degree of tensor parallelism for distributed inference.
+  * **`--input_file`**: Path to the input dataset file.
+  * **`--output_dir`**: Directory to save output results and logs.
+
+### Offline PyTorch Inference
+
+This script tests the throughput of an LLM using offline PyTorch inference.
+
+**Configuration (`torch_infer.sh`):**
+
+```bash
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+
+output_dir="vllm_bf16_offline_flashattn"
+
+python torch_infer.py \
+    --actor_model_name_or_path Qwen/Qwen2.5-7B-Instruct-1M \
+    --max_src_len 2048 \
+    --min_dec_len 32 \
+    --max_dec_len 30720 \
+    --top_p 1.0 \
+    --temperature 1.0 \
+    --rollout_input_batch_size 4 \
+    --rollout_n 8 \
+    --tensor_parallel_degree 4 \
+    --limit_rows 640 \
+    --input_file ./data/gsm8k/instruct/train.parquet \
+    --output_dir ${output_dir} \
+    --gpu_memory_utilization 0.8 > ./torchinferflashattn.log 2>&1
+```
+  * **`--gpu_memory_utilization`**: Fraction of GPU memory to be reserved for the model.
+
+-----
+
+## Output Results
+
+The `output_dir` contains the following files:  
+
+**1. Statistics Files**  
+• `dispersed_stats.csv`  
+
+  Per-batch request length and throughput statistics. Fields:  
+  `batch_index, rollout_lengths, min_length, max_length, avg_length, completion_time, throughput_tokens_per_sec`  
+
+• `global_stats.csv`  
+
+  Aggregated global metrics. Fields:  
+  `batch_index, min_response_tokens, max_response_tokens, avg_response_tokens, total_response_tokens, completion_time, throughput_tokens_per_sec`  
+
+**2. Detailed Records**  
+• `rollout_details.jsonl`  
+
+  Raw per-request outputs (JSON Lines format), including input/output text.
+

llm/benchmark/rl/api_serve.py

@@ -0,0 +1,391 @@
+# 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 argparse
+import asyncio
+import csv
+import json
+import math
+import time
+from dataclasses import dataclass, field
+from itertools import cycle
+from pathlib import Path
+from typing import List, Tuple
+
+import pandas as pd
+from openai import AsyncOpenAI
+from tqdm import tqdm
+from utils import RangeSet
+
+from paddlenlp.transformers import AutoTokenizer
+from paddlenlp.utils.log import logger
+
+
+@dataclass
+class RequestPayload:
+    prompt: str = ""
+    num_responses: int = 8
+    idx: int = 0
+
+
+@dataclass
+class ResponsePayload:
+    idx: int = 0
+    question: str = ""
+    question_token_length: int = 0
+    responses: List[str] = field(default_factory=list)
+    elapsed_times: List[float] = field(default_factory=list)
+    token_lengths: List[int] = field(default_factory=list)
+    total_length: int = 0
+
+
+class ApiTask:
+    def __init__(self, args):
+        self.args = args
+        self.model = args.model
+        self.tokenzizer = AutoTokenizer.from_pretrained(self.args.tokenizer, use_fast=True)
+        self.clients = cycle(
+            AsyncOpenAI(base_url=url, api_key=api) for url, api in zip(args.openai_urls, args.api_keys)
+        )
+        self._max_concurrency = args.max_concurrency
+        self.semaphore = asyncio.Semaphore(args.max_concurrency)
+        self.output_dir = Path(self.args.output_dir)
+
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._load_status()
+
+    def tokenize(self, response: ResponsePayload) -> ResponsePayload:
+        question = response.question
+        responses = response.responses
+        response.question_token_length = len(self.tokenzizer(question).input_ids)
+
+        for i, resp in enumerate(responses):
+            tokens = self.tokenzizer(resp).input_ids
+            length = len(tokens)
+            response.token_lengths.append(length)
+            response.total_length += length
+
+        return response
+
+    def get_active_tasks_count(self) -> int:
+        return self._max_concurrency - self.semaphore._value
+
+    def get_client(self) -> AsyncOpenAI:
+        # Returns an AsyncOpenAI client instance
+        return next(self.clients)
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        batch_prompts = []
+        for idx, prompt in enumerate(dataframe[self.args.prompt_key]):
+            batch_prompts.append(
+                RequestPayload(prompt=prompt[0]["content"], idx=idx, num_responses=self.args.rollout_n)
+            )
+            if len(batch_prompts) == self.args.rollout_input_batch_size:
+                yield batch_prompts
+                batch_prompts = []
+
+    async def fastdeploy_call(self, request: RequestPayload) -> Tuple[str, float]:
+        client = self.get_client()
+        try:
+            async with self.semaphore:
+                start_time = time.time()
+                response = await client.chat.completions.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": request.prompt}],
+                    temperature=self.args.temperature,
+                    top_p=self.args.top_p,
+                    max_tokens=self.args.max_response_length,
+                    n=1,
+                    stream=True,
+                    timeout=60*60,
+                    metadata={
+                        "training": True,
+                        "raw_request": False,
+                    }
+                ) 
+                # Streaming text is stored in a list of chunks
+                chunks = []
+                # Streaming responses
+                async for chunk in response:
+                    delta = chunk.choices[0].delta
+                    if delta and delta.content:
+                        chunks.append(delta.content)
+                text = "".join(chunks)
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                logger.debug("Streaming response took %.2f seconds", elapsed_time)
+                return text, round(elapsed_time, 2)
+
+        except Exception as e:
+            logger.error("Error while streaming: %s", e)
+            raise ValueError(e)
+
+    async def vllm_call(self, request: RequestPayload) -> Tuple[str, float]:
+        client = self.get_client()
+        try:
+            async with self.semaphore:
+                start_time = time.time()
+                response = await client.completions.create(
+                    model=self.model,
+                    prompt=request.prompt,
+                    temperature=self.args.temperature,
+                    top_p=self.args.top_p,
+                    max_tokens=self.args.max_dec_len,
+                    n=1,
+                    stream=True,
+                )
+                # Streaming text is stored in a list of chunks
+                chunks = []
+                # Streaming responses
+                async for chunk in response:
+                    if chunk.choices and chunk.choices[0].text:
+                        chunks.append(chunk.choices[0].text)
+                text = "".join(chunks)
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                logger.debug("Streaming response took %.2f seconds", elapsed_time)
+                return text, round(elapsed_time, 2)
+
+        except Exception as e:
+            logger.error("Error while streaming: %s", e)
+            raise ValueError(e)
+
+    async def group_call(self, request: RequestPayload) -> ResponsePayload:
+        """Performs n complete token generation rollouts for the given query."""
+        if self.args.use_fastdeploy == "true":
+            call = self.fastdeploy_call
+        else:
+            call = self.vllm_call
+
+        tasks = [call(request) for _ in range(request.num_responses)]
+
+        result = ResponsePayload()
+        result.idx = request.idx
+        result.question = request.prompt
+        for task, elapsed_time in await asyncio.gather(*tasks):
+            result.responses.append(task)
+            result.elapsed_times.append(elapsed_time)
+        return result
+
+    async def batch_call(self, requests: List[RequestPayload]) -> Tuple[List[ResponsePayload], int]:
+        """Batch execution requests"""
+        start_time = time.time()
+        batch_results = await asyncio.gather(*[self.group_call(request) for request in requests])
+        end_time = time.time()
+        batch_elapsed_time = end_time - start_time
+        logger.debug("total batch took %.2f seconds", batch_elapsed_time)
+        return batch_results, batch_elapsed_time
+
+    def dispersed_stats(self, responses: List[ResponsePayload], batch_elapsed_time: float, batch_index):
+        batch_group_pd = pd.DataFrame(responses)
+
+        dispersed_stats_dict = {
+            "batch_index": batch_index,
+            "rollout_lengths": batch_group_pd["token_lengths"].to_list(),
+            "min_length": batch_group_pd["token_lengths"].apply(lambda x: min(x)).tolist(),
+            "max_length": batch_group_pd["token_lengths"].apply(lambda x: max(x)).tolist(),
+            "avg_length": batch_group_pd["token_lengths"].apply(lambda x: sum(x) / len(x)).tolist(),
+            "completion_time": batch_elapsed_time,
+            "throughput_tokens_per_sec": batch_group_pd["token_lengths"].apply((lambda x: sum(x))).sum()
+            / batch_elapsed_time,
+        }
+
+        return dispersed_stats_dict
+
+    def global_stats(self, responses: List[ResponsePayload], batch_elapsed_time: float, batch_index):
+        dispersed_stats_dict = self.dispersed_stats(responses, batch_elapsed_time, batch_index)
+
+        total_response_tokens = 0
+        for lengths in dispersed_stats_dict["rollout_lengths"]:
+            total_response_tokens += sum(lengths)
+
+        global_stats_dict = {}
+        global_stats_dict["batch_index"] = dispersed_stats_dict["batch_index"]
+        global_stats_dict["min_response_tokens"] = min(dispersed_stats_dict["min_length"])
+        global_stats_dict["max_response_tokens"] = max(dispersed_stats_dict["max_length"])
+        global_stats_dict["avg_response_tokens"] = total_response_tokens / (
+            self.args.rollout_n * self.args.rollout_input_batch_size
+        )
+        global_stats_dict["total_response_tokens"] = total_response_tokens
+        global_stats_dict["group_max_response_tokens"] = dispersed_stats_dict["max_length"]
+        global_stats_dict["completion_time"] = dispersed_stats_dict["completion_time"]
+        global_stats_dict["throughput_tokens_per_sec"] = dispersed_stats_dict["throughput_tokens_per_sec"]
+
+        return global_stats_dict, dispersed_stats_dict
+
+    def execute(self):
+        dataframe = self.process_data(self.args.input_file)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+            global_writer = csv.writer(global_f)
+            dispersed_writer = csv.writer(dispersed_f)
+
+            if self.processed_set.processed_count <= 0:
+                global_writer.writerow(
+                    [
+                        "batch_index",
+                        "min_response_tokens",
+                        "max_response_tokens",
+                        "avg_response_tokens",
+                        "total_response_tokens",
+                        "group_max_response_tokens",
+                        "completion_time",
+                        "throughput_tokens_per_sec",
+                    ]
+                )
+                dispersed_writer.writerow(
+                    [
+                        "batch_index",
+                        "rollout_lengths",
+                        "min_length",
+                        "max_length",
+                        "avg_length",
+                        "completion_time",
+                        "throughput_tokens_per_sec",
+                    ]
+                )
+
+            for batch_index, input_ids in tqdm(
+                enumerate(self.batch_process(dataframe)),
+                total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                mininterval=0.1,
+            ):
+                if self.processed_set.contains(batch_index):
+                    continue
+
+                batch_results, batch_elapsed_time = asyncio.run(self.batch_call(input_ids))
+
+                for i in range(len(batch_results)):
+                    batch_results[i] = self.tokenize(batch_results[i])
+
+                global_stats_dict, dispersed_stats_dict = self.global_stats(
+                    batch_results, batch_elapsed_time, batch_index
+                )
+
+                global_writer.writerow(
+                    [
+                        batch_index,
+                        global_stats_dict["min_response_tokens"],
+                        global_stats_dict["max_response_tokens"],
+                        round(global_stats_dict["avg_response_tokens"], 2),
+                        global_stats_dict["total_response_tokens"],
+                        global_stats_dict["group_max_response_tokens"],
+                        round(global_stats_dict["completion_time"], 2),
+                        round(global_stats_dict["throughput_tokens_per_sec"], 2),
+                    ]
+                )
+
+                dispersed_writer.writerow(
+                    [
+                        batch_index,
+                        dispersed_stats_dict["rollout_lengths"],
+                        dispersed_stats_dict["min_length"],
+                        dispersed_stats_dict["max_length"],
+                        dispersed_stats_dict["avg_length"],
+                        round(dispersed_stats_dict["completion_time"], 2),
+                        round(dispersed_stats_dict["throughput_tokens_per_sec"], 2),
+                    ]
+                )
+
+                record = [
+                    {
+                        "batch_index": batch_index,
+                        "prompt_text": result.question,
+                        "rollouts": [
+                            {"response": res, "token_length": token_length}
+                            for res, token_length in zip(result.responses, result.token_lengths)
+                        ],
+                        "total_time": dispersed_stats_dict["completion_time"],
+                        "throughput_tokens_per_sec": dispersed_stats_dict["throughput_tokens_per_sec"],
+                    }
+                    for result in batch_results
+                ]
+
+                jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                global_f.flush()
+                dispersed_f.flush()
+                jsonl_f.flush()
+                self._save_status(batch_index)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process prompts with OpenAI clients.")
+    parser.add_argument("--openai_urls", type=str, nargs="+", required=True, help="List of OpenAI service URLs")
+    parser.add_argument(
+        "--api_keys", type=str, nargs="+", default=None, help="List of API keys (default: 'NONE' for each service)"
+    )
+    parser.add_argument("--model", type=str, required=True, help="Model name (e.g., Qwen2.5-7B-Instruct-1M)")
+    parser.add_argument(
+        "--tokenizer", type=str, required=True, help="Tokenizer name (e.g., Qwen/Qwen2.5-7B-Instruct-1M)"
+    )
+    parser.add_argument("--rollout_input_batch_size", type=int, default=4, help="Batch size for requests")
+    parser.add_argument("--rollout_n", type=int, default=8, help="Number of responses per request")
+    parser.add_argument(
+        "--prompt_key", type=str, default="prompt", help="Key in the DataFrame for prompts (default: 'prompt')"
+    )
+    parser.add_argument("--input_file", type=str, required=True, help="Path to the input Parquet file")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./api_infer_results",
+        help="Directory for output CSV files (default: './api_infer_results')",
+    )
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p sampling parameter for text generation")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature parameter for text generation")
+    parser.add_argument("--max_dec_len", type=int, default=1024 * 2, help="Maximum response length (in tokens)")
+    parser.add_argument("--max_concurrency", type=int, default=1000, help="Maximum concurrent connections")
+    parser.add_argument(
+        "--limit_rows", type=int, default=-1, help="Maximum number of rows to read from the dataset (-1 means all)"
+    )
+    parser.add_argument("--use_fastdeploy", type=str.lower, choices=["true", "false"], default="true", help="Engine selection (true=FastDeploy, false=vLLM, default: true)")
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    task = ApiTask(args)
+    task.execute()

llm/benchmark/rl/paddle_infer.py

@@ -0,0 +1,351 @@
+# 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 csv
+import json
+import math
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import paddle
+import pandas as pd
+from tqdm import tqdm
+from utils import RangeSet
+
+from paddlenlp.rl.trainer import process_row
+from paddlenlp.rl.utils import TrainingArguments
+from paddlenlp.rl.utils.infer_utils import get_policy_predictor, infer_guard
+from paddlenlp.trainer import PdArgumentParser, set_seed
+from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer
+from paddlenlp.trl.llm_utils import init_dist_env
+from paddlenlp.utils.log import logger
+
+
+@contextmanager
+def switch_level_context(level="ERROR"):
+    original_level = logger.logLevel
+    logger.set_level(level)
+
+    try:
+        yield
+    finally:
+        logger.set_level(original_level)
+
+
+def chunk(all_input_ids, size):
+    if size <= 0:
+        raise ValueError("Size must be greater than 0")
+    return [all_input_ids[i : i + size] for i in range(0, len(all_input_ids), size)]
+
+
+@dataclass
+class DumpyTrainingArguments(TrainingArguments):
+    actor_model_name_or_path: str = field(
+        default="Qwen/Qwen2.5-7B-Instruct-1M", metadata={"help": "pretrained model name or path"}
+    )
+    input_file: str = field(
+        default="kk/instruct/3-7ppl/combined.parquet", metadata={"help": "input the Parquet file path"}
+    )
+    limit_rows: int = field(
+        default=-1, metadata={"help": "Maximum number of rows to read from the dataset (-1 means all)"}
+    )
+    output_dir: str = field(default="./pt_infer_results", metadata={"help": "output directory path"})
+    rollout_input_batch_size: int = field(default=32, metadata={"help": "batch size for inference engine inputs"})
+    rollout_n: int = field(default=8, metadata={"help": "number of rollouts (for performance testing)"})
+    log_interval: int = field(default=1, metadata={"help": "logging interval (in batches)"})
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.dtype = "bfloat16"
+        if self.rollout_quant_type.lower() in ["bfloat16", "bf16"]:
+            self.rollout_quant_type = ""
+        elif self.rollout_quant_type.lower() in ["wint8", "weight_only_int8"]:
+            self.rollout_quant_type = "weight_only_int8"
+
+
+class DumpyInferenceTask:
+    def __init__(self, args):
+        self.args = args
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        # ​​Initialize output file path​
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._init_environment()
+        self._load_status()
+        self._load_model()
+        self._prepare_tokenizer()
+
+    def _init_environment(self):
+        self.tensor_parallel_rank, self.tensor_parallel_degree = init_dist_env()
+        self.use_fusemt = True
+        self.amp_dtype = self.args.dtype
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_model(self):
+        logger.info(f"Loading model from {self.args.actor_model_name_or_path}...")
+        start_time = time.time()
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.args.actor_model_name_or_path,
+            dtype=self.args.dtype,
+            low_cpu_mem_usage=True,
+            tensor_parallel_degree=self.tensor_parallel_degree,
+            tensor_parallel_rank=self.tensor_parallel_rank,
+        )
+        self.model.eval()
+        self.model.to(device=paddle.CUDAPinnedPlace())
+        logger.info(f"Model loaded in {time.time() - start_time:.2f}s")
+
+    def _prepare_tokenizer(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.actor_model_name_or_path, use_fast=True)
+        self.tokenizer.padding_side = "left"
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    @switch_level_context(level="ERROR")
+    def run_inference(self, input_ids, batch_index=0):
+        set_seed(42)
+        start_time = time.time()
+        output_ids = get_policy_predictor().predict(
+            input_ids=input_ids,
+            repeat_num=self.args.rollout_n,
+            temperature=self.args.temperature,
+            top_p=self.args.top_p,
+        )
+        if self.args.world_size > 1:
+            paddle.distributed.barrier()
+        end_time = time.time()
+        if self.args.should_log:
+            statistics = self.postprocess_data(input_ids, output_ids, batch_index=batch_index)
+            statistics["total_time"] = end_time - start_time
+            return statistics
+        return None
+
+    def postprocess_data(self, input_ids, output_ids, batch_index=0):
+        # Process prompts
+        global_prompt_tokens = []
+        global_prompt_tokens_len = []
+        group_prompt_texts = []
+        for row in input_ids:
+            row_ids = process_row(row, remove_value=self.tokenizer.pad_token_id, remove_side="both").tolist()
+            global_prompt_tokens.append(row_ids)
+            global_prompt_tokens_len.append(len(row_ids))
+            group_prompt_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Process responses
+        response_tokens = []
+        response_tokens_len = []
+        response_texts = []
+        for row in output_ids:
+            row_ids = process_row(row, remove_value=self.tokenizer.pad_token_id, remove_side="both").tolist()
+            response_tokens.append(row_ids)
+            response_tokens_len.append(len(row_ids))
+            response_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Group processing
+        group_response_tokens = chunk(response_tokens, self.args.rollout_n)
+        group_response_tokens_len = chunk(response_tokens_len, self.args.rollout_n)
+        group_response_texts = chunk(response_texts, self.args.rollout_n)
+
+        # Calculate statistics
+        global_response_tokens_total = sum(response_tokens_len)
+        global_response_tokens_len_min = min(response_tokens_len)
+        global_response_tokens_len_max = max(response_tokens_len)
+        global_response_tokens_len_mean = round(global_response_tokens_total / len(response_tokens_len), 2)
+
+        group_response_tokens_len_mean = [round(sum(x) / len(x), 2) for x in group_response_tokens_len]
+        group_response_tokens_len_max = [max(x) for x in group_response_tokens_len]
+        group_response_tokens_len_min = [min(x) for x in group_response_tokens_len]
+
+        return {
+            # Global prompt stats
+            "batch_index": batch_index,
+            "global_prompt_tokens_len": global_prompt_tokens_len,
+            "group_prompt_texts": group_prompt_texts,
+            # Group response stats
+            "group_response_tokens": group_response_tokens,
+            "group_response_tokens_len": group_response_tokens_len,
+            "group_response_texts": group_response_texts,
+            "group_response_tokens_len_min": group_response_tokens_len_min,
+            "group_response_tokens_len_max": group_response_tokens_len_max,
+            "group_response_tokens_len_mean": group_response_tokens_len_mean,
+            # Global response stats
+            "global_response_tokens_total": global_response_tokens_total,
+            "global_response_tokens_len_min": global_response_tokens_len_min,
+            "global_response_tokens_len_max": global_response_tokens_len_max,
+            "global_response_tokens_len_mean": global_response_tokens_len_mean,
+        }
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        all_input_ids = []
+
+        def _pad_batch(input_ids_list):
+            return self.tokenizer.pad({"input_ids": input_ids_list}, padding=True, return_tensors="pd")["input_ids"]
+
+        for idx, row in dataframe.iterrows():
+            prompt = row["prompt"].tolist()[0]["content"]
+            input_ids = self.tokenizer(prompt, add_special_tokens=False, return_attention_mask=False)["input_ids"]
+            all_input_ids.append(input_ids)
+
+            if len(all_input_ids) >= self.args.rollout_input_batch_size:
+                yield _pad_batch(all_input_ids)
+                all_input_ids = []
+
+        if all_input_ids:
+            yield _pad_batch(all_input_ids)
+
+    def execute(self):
+        """Main Execution Pipeline"""
+        # Data Preparation​
+        dataframe = self.process_data(self.args.input_file)
+
+        # Create Output Directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+
+            if self.args.should_log:
+                # Initialize the CSV writer
+                global_writer = csv.writer(global_f)
+                dispersed_writer = csv.writer(dispersed_f)
+                if self.processed_set.processed_count <= 0:
+                    global_writer.writerow(
+                        [
+                            "batch_index",
+                            "min_response_tokens",
+                            "max_response_tokens",
+                            "avg_response_tokens",
+                            "total_response_tokens",
+                            "group_max_response_tokens",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+                    dispersed_writer.writerow(
+                        [
+                            "batch_index",
+                            "rollout_lengths",
+                            "min_length",
+                            "max_length",
+                            "avg_length",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+
+            with infer_guard(self):
+                for batch_index, input_ids in tqdm(
+                    enumerate(self.batch_process(dataframe)),
+                    total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                    disable=not self.args.should_log,
+                ):
+                    if self.processed_set.contains(batch_index):
+                        continue
+
+                    statistics = self.run_inference(input_ids, batch_index=batch_index)
+
+                    if self.args.should_log:
+                        # Write global statistics
+                        total_time = round(statistics["total_time"], 2)
+                        total_tokens = statistics["global_response_tokens_total"]
+                        throughput = round(total_tokens / total_time if total_time > 0 else 0, 2)
+
+                        global_writer.writerow(
+                            [
+                                batch_index,
+                                statistics["global_response_tokens_len_min"],
+                                statistics["global_response_tokens_len_max"],
+                                statistics["global_response_tokens_len_mean"],
+                                total_tokens,
+                                statistics["group_response_tokens_len_max"],
+                                total_time,
+                                throughput,
+                            ]
+                        )
+
+                        dispersed_writer.writerow(
+                            [
+                                batch_index,
+                                statistics["group_response_tokens_len"],
+                                statistics["group_response_tokens_len_min"],
+                                statistics["group_response_tokens_len_max"],
+                                statistics["group_response_tokens_len_mean"],
+                                total_time,
+                                throughput,
+                            ]
+                        )
+
+                        # Write detailed records (one row per query)
+                        prompt_text = statistics["group_prompt_texts"]
+                        response_tokens = statistics["group_response_tokens"]
+                        response_texts = statistics["group_response_texts"]
+
+                        record = {
+                            "batch_index": batch_index,
+                            "prompt_text": prompt_text,
+                            "rollouts": [
+                                {"token_ids": tokens, "text": text, "length": len(tokens)}
+                                for tokens, text in zip(response_tokens, response_texts)
+                            ],
+                            "total_time": total_time,
+                            "throughput_tokens_per_sec": throughput,
+                        }
+                        jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                        global_f.flush()
+                        dispersed_f.flush()
+                        jsonl_f.flush()
+                        self._save_status(batch_index)
+
+
+def main():
+    parser = PdArgumentParser((DumpyTrainingArguments,))
+    args = parser.parse_args_into_dataclasses()[0]
+    task = DumpyInferenceTask(args)
+    task.execute()
+
+
+if __name__ == "__main__":
+    main()

llm/benchmark/rl/scripts/api_serve.sh

@@ -0,0 +1,16 @@
+output_dir="your_output_dir"
+
+python api_serve.py \
+    --openai_urls "your_url1" "your_url2"\
+    --api_keys "key1" "key2" \
+    --model "Qwen2.5-7B-Instruct-1M" \
+    --tokenizer "Qwen/Qwen2.5-7B-Instruct-1M" \
+    --use_fastdeploy true \
+    --input_file your_file \
+    --output_dir ${output_dir} \
+    --rollout_input_batch_size 8 \
+    --rollout_n 8 \
+    --top_p 1.0 \
+    --temperature 0.7 \
+    --max_dec_len 8192 \
+    --limit_rows 512

llm/benchmark/rl/scripts/paddle_infer.sh

@@ -0,0 +1,39 @@
+unset PADDLE_TRAINERS_NUM
+unset PADDLE_ELASTIC_JOB_ID
+unset PADDLE_TRAINER_ENDPOINTS
+unset DISTRIBUTED_TRAINER_ENDPOINTS
+unset FLAGS_START_PORT
+unset PADDLE_ELASTIC_TIMEOUT
+
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP":$PYTHONPATH
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP/llm":$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_custom_allreduce=0
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+output_dir="pdpd_bf16_offline"
+
+python -u -m paddle.distributed.launch --log_dir ${output_dir}/logs --gpus ${CUDA_VISIBLE_DEVICES} paddle_infer.py \
+  --actor_model_name_or_path your_model_name \
+  --max_src_len 2048 \
+  --min_dec_len 32 \
+  --max_dec_len 30720 \
+  --top_p 1.0 \
+  --temperature 1.0 \
+  --rollout_input_batch_size 4 \
+  --rollout_n 8 \
+  --rollout_max_num_seqs 24 \
+  --rollout_quant_type "" \
+  --tensor_parallel_degree 4 \
+  --limit_rows 640 \
+  --input_file file.parquet \
+  --output_dir ${output_dir} > ./paddleinfer.log 2>&1 

llm/benchmark/rl/scripts/torch_infer.sh

@@ -0,0 +1,18 @@
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+
+output_dir="vllm_bf16_offline_flashattn"
+
+nohup python torch_infer.py \
+    --actor_model_name_or_path your_model_name \
+    --max_src_len 2048 \
+    --min_dec_len 32 \
+    --max_dec_len 30720 \
+    --top_p 1.0 \
+    --temperature 1.0 \
+    --rollout_input_batch_size 4 \
+    --rollout_n 8 \
+    --tensor_parallel_degree 4 \
+    --limit_rows 640 \
+    --input_file file.parquet \
+    --output_dir ${output_dir} \
+    --gpu_memory_utilization 0.8 > ./torchinferflashattn.log 2>&1 &

llm/benchmark/rl/torch_infer.py

@@ -0,0 +1,339 @@
+# 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 argparse
+import csv
+import json
+import math
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import pandas as pd
+import tqdm
+from transformers import AutoTokenizer, logging
+from utils import RangeSet
+from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--actor_model_name_or_path",
+        type=str,
+        default="Qwen/Qwen2.5-7B-Instruct-1M",
+        help="Name or path of the pretrained model",
+    )
+    parser.add_argument("--input_file", type=str, default="./combined.parquet", help="Path to input parquet file")
+    parser.add_argument("--output_dir", type=str, default="./pt_infer_results", help="Path to output directory")
+    parser.add_argument(
+        "--rollout_input_batch_size", type=int, default=2, help="Batch size for each inference engine input"
+    )
+    parser.add_argument(
+        "--rollout_n", type=int, default=2, help="Number of repeated inference runs (for performance testing)"
+    )
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type for computation")
+    parser.add_argument("--tensor_parallel_degree", type=int, default=1, help="Degree of model parallelism")
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p sampling parameter for text generation")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature parameter for text generation")
+    parser.add_argument("--should_log", type=bool, default=True, help="Whether to print logging information")
+    parser.add_argument("--max_src_len", type=int, default=1024 * 2, help="Maximum prompt length (in tokens)")
+    parser.add_argument("--max_dec_len", type=int, default=1024 * 2, help="Maximum response length (in tokens)")
+    parser.add_argument("--min_dec_len", type=int, default=0, help="Minimum response length (in tokens)")
+    parser.add_argument(
+        "--limit_rows", type=int, default=-1, help="Maximum number of rows to read from the dataset (-1 means all)"
+    )
+    parser.add_argument("--gpu_memory_utilization", type=float, default=0.9, help="Percentage of GPU memory usage")
+    args = parser.parse_args()
+
+    return args
+
+
+def chunk(all_input_ids, size):
+    if size <= 0:
+        raise ValueError("Size must be greater than 0")
+    return [all_input_ids[i : i + size] for i in range(0, len(all_input_ids), size)]
+
+
+class DumpyInferenceTask:
+    def __init__(self, args):
+        self.args = args
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._load_status()
+        self._load_model()
+        self._prepare_tokenizer()
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_model(self):
+        logger.info(f"Loading model from {self.args.actor_model_name_or_path}...")
+        start_time = time.time()
+        self.model = LLM(
+            model=self.args.actor_model_name_or_path,
+            dtype=self.args.dtype,
+            enable_sleep_mode=False,
+            tensor_parallel_size=self.args.tensor_parallel_degree,
+            # distributed_executor_backend="external_launcher",
+            enforce_eager=False,
+            gpu_memory_utilization=self.args.gpu_memory_utilization,
+            disable_custom_all_reduce=True,
+            disable_mm_preprocessor_cache=True,
+            skip_tokenizer_init=False,
+            max_model_len=self.args.max_src_len + self.args.max_dec_len,
+            disable_log_stats=True,
+            max_num_batched_tokens=self.args.max_src_len + self.args.max_dec_len,
+            enable_chunked_prefill=True,
+            enable_prefix_caching=True,
+        )
+        self.sampling_params = SamplingParams(
+            top_p=self.args.top_p,
+            temperature=self.args.temperature,
+            min_tokens=self.args.min_dec_len,
+            max_tokens=self.args.max_dec_len,
+            n=self.args.rollout_n,
+            seed=42,
+        )
+        logger.info(f"Model loaded in {time.time() - start_time:.2f}s")
+
+    def _prepare_tokenizer(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.actor_model_name_or_path, use_fast=True)
+        self.tokenizer.padding_side = "left"
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    def run_inference(self, prompts, batch_index=0):
+        input_ids = [self.tokenizer(prompt, add_special_tokens=False)["input_ids"] for prompt in prompts]
+        token_prompt_ids = [TokensPrompt(prompt_token_ids=prompt) for prompt in input_ids]
+
+        start_time = time.time()
+        request_outputs = self.model.generate(
+            prompts=token_prompt_ids, sampling_params=self.sampling_params, use_tqdm=True
+        )
+        end_time = time.time()
+
+        batch_token_ids = []
+        for output in request_outputs:
+            for each_output in output.outputs:
+                batch_token_ids.append(each_output.token_ids)
+
+        if self.args.should_log:
+            statistics = self.postprocess_data(input_ids, batch_token_ids, batch_index=batch_index)
+            statistics["total_time"] = end_time - start_time
+            return statistics
+        return None
+
+    def postprocess_data(self, input_ids, output_ids, batch_index=0):
+        global_prompt_tokens = []
+        global_prompt_tokens_len = []
+        group_prompt_texts = []
+        for row in input_ids:
+            row_ids = row
+            global_prompt_tokens.append(row_ids)
+            global_prompt_tokens_len.append(len(row_ids))
+            group_prompt_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Process responses
+        response_tokens = []
+        response_tokens_len = []
+        response_texts = []
+        for row in output_ids:
+            row_ids = row
+            response_tokens.append(row_ids)
+            response_tokens_len.append(len(row_ids))
+            response_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Group processing
+        group_response_tokens = chunk(response_tokens, self.args.rollout_n)
+        group_response_tokens_len = chunk(response_tokens_len, self.args.rollout_n)
+        group_response_texts = chunk(response_texts, self.args.rollout_n)
+
+        # Calculate statistics
+        global_response_tokens_total = sum(response_tokens_len)
+        global_response_tokens_len_min = min(response_tokens_len)
+        global_response_tokens_len_max = max(response_tokens_len)
+        global_response_tokens_len_mean = round(global_response_tokens_total / len(response_tokens_len), 2)
+
+        group_response_tokens_len_mean = [round(sum(x) / len(x), 2) for x in group_response_tokens_len]
+        group_response_tokens_len_max = [max(x) for x in group_response_tokens_len]
+        group_response_tokens_len_min = [min(x) for x in group_response_tokens_len]
+
+        return {
+            # Global prompt stats
+            "batch_index": batch_index,
+            "global_prompt_tokens_len": global_prompt_tokens_len,
+            "group_prompt_texts": group_prompt_texts,
+            # Group response stats
+            "group_response_tokens": group_response_tokens,
+            "group_response_tokens_len": group_response_tokens_len,
+            "group_response_texts": group_response_texts,
+            "group_response_tokens_len_min": group_response_tokens_len_min,
+            "group_response_tokens_len_max": group_response_tokens_len_max,
+            "group_response_tokens_len_mean": group_response_tokens_len_mean,
+            # Global response stats
+            "global_response_tokens_total": global_response_tokens_total,
+            "global_response_tokens_len_min": global_response_tokens_len_min,
+            "global_response_tokens_len_max": global_response_tokens_len_max,
+            "global_response_tokens_len_mean": global_response_tokens_len_mean,
+        }
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        all_prompts = []
+
+        for idx, row in dataframe.iterrows():
+            prompt = row["prompt"].tolist()[0]["content"]
+            all_prompts.append(prompt)
+
+            if len(all_prompts) >= self.args.rollout_input_batch_size:
+                yield all_prompts
+                all_prompts = []
+
+        if all_prompts:
+            yield all_prompts
+
+    def execute(self):
+        dataframe = self.process_data(self.args.input_file)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+
+            if self.args.should_log:
+                global_writer = csv.writer(global_f)
+                dispersed_writer = csv.writer(dispersed_f)
+                if self.processed_set.processed_count <= 0:
+                    global_writer.writerow(
+                        [
+                            "batch_index",
+                            "min_response_tokens",
+                            "max_response_tokens",
+                            "avg_response_tokens",
+                            "total_response_tokens",
+                            "group_max_response_tokens",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+                    dispersed_writer.writerow(
+                        [
+                            "batch_index",
+                            "rollout_lengths",
+                            "min_length",
+                            "max_length",
+                            "avg_length",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+
+            for batch_index, prompts in tqdm.tqdm(
+                enumerate(self.batch_process(dataframe)),
+                total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                disable=not self.args.should_log,
+            ):
+                if self.processed_set.contains(batch_index):
+                    continue
+
+                statistics = self.run_inference(prompts, batch_index=batch_index)
+
+                if self.args.should_log:
+                    total_time = round(statistics["total_time"], 2)
+                    total_tokens = statistics["global_response_tokens_total"]
+                    throughput = round(total_tokens / total_time if total_time > 0 else 0, 2)
+
+                    global_writer.writerow(
+                        [
+                            batch_index,
+                            statistics["global_response_tokens_len_min"],
+                            statistics["global_response_tokens_len_max"],
+                            statistics["global_response_tokens_len_mean"],
+                            total_tokens,
+                            statistics["group_response_tokens_len_max"],
+                            total_time,
+                            throughput,
+                        ]
+                    )
+
+                    dispersed_writer.writerow(
+                        [
+                            batch_index,
+                            statistics["group_response_tokens_len"],
+                            statistics["group_response_tokens_len_min"],
+                            statistics["group_response_tokens_len_max"],
+                            statistics["group_response_tokens_len_mean"],
+                            total_time,
+                            throughput,
+                        ]
+                    )
+
+                    prompt_text = statistics["group_prompt_texts"]
+                    response_tokens = statistics["group_response_tokens"]
+                    response_texts = statistics["group_response_texts"]
+
+                    record = {
+                        "batch_index": batch_index,
+                        "prompt_text": prompt_text,
+                        "rollouts": [
+                            {"token_ids": tokens, "text": text, "length": len(tokens)}
+                            for tokens, text in zip(response_tokens, response_texts)
+                        ],
+                        "total_time": total_time,
+                        "throughput_tokens_per_sec": throughput,
+                    }
+                    jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                    global_f.flush()
+                    dispersed_f.flush()
+                    jsonl_f.flush()
+                    self._save_status(batch_index)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    print(args)
+    task = DumpyInferenceTask(args)
+    task.execute()

llm/benchmark/rl/utils.py

@@ -0,0 +1,91 @@
+# 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 dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class RangeSet:
+    """Manage processed line ranges with efficient storage and querying"""
+
+    ranges: List[tuple]
+
+    def add(self, number: int):
+        """Add a number to the range set and merge adjacent ranges"""
+        new_ranges = []
+        added = False
+        for start, end in sorted(self.ranges):
+            if number < start - 1:
+                if not added:
+                    new_ranges.append((number, number))
+                    added = True
+                new_ranges.append((start, end))
+            elif number == start - 1:
+                new_ranges.append((number, end))
+                added = True
+            elif number <= end:
+                new_ranges.append((start, end))
+                added = True
+            else:
+                new_ranges.append((start, end))
+        if not added:
+            new_ranges.append((number, number))
+        self.ranges = self.merge_ranges(new_ranges)
+
+    @staticmethod
+    def merge_ranges(ranges: List[tuple]) -> List[tuple]:
+        """Merge overlapping or adjacent ranges"""
+        if not ranges:
+            return []
+        sorted_ranges = sorted(ranges)
+        merged = [sorted_ranges[0]]
+        for current in sorted_ranges[1:]:
+            last = merged[-1]
+            if current[0] <= last[1] + 1:
+                merged[-1] = (last[0], max(last[1], current[1]))
+            else:
+                merged.append(current)
+        return merged
+
+    def contains(self, number: int) -> bool:
+        """Check if a number exists in any range"""
+        for start, end in self.ranges:
+            if start <= number <= end:
+                return True
+        return False
+
+    def to_file_format(self) -> str:
+        """Serialize ranges to compact string format"""
+        return ",".join(f"{start}-{end}" if start != end else str(start) for start, end in self.ranges)
+
+    @classmethod
+    def from_file(cls, content: str) -> "RangeSet":
+        """Deserialize from string format"""
+        if not content:
+            return cls(ranges=[])
+        ranges = []
+        for part in content.split(","):
+            if "-" in part:
+                start, end = map(int, part.split("-"))
+                ranges.append((start, end))
+            else:
+                num = int(part)
+                ranges.append((num, num))
+        return cls(ranges=ranges)
+
+    @property
+    def processed_count(self) -> int:
+        """Total number of processed items"""
+        return sum(end - start + 1 for start, end in self.ranges)

llm/config/llama/dislora_argument.json

@@ -0,0 +1,37 @@
+{
+    "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

@@ -0,0 +1,36 @@
+{
+    "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/grpo_32b_argument.yaml

@@ -74,7 +74,7 @@ disable_tqdm: true # Whether to disable tqdm progress bar
 
 # RL args
 kl_coeff: 0.001 # KL coefficient for PPO and Reinforce++
-kl_loss_coeff: 0.001 # KL loss coefficient
+kl_loss_coeff: 0.000 # KL loss coefficient
 pg_loss_coeff: 1.0 # Policy gradient loss coefficient
 entropy_coeff: 0.0 # Entropy coefficient
 clip_range_ratio: 0.2 # The clipping range for ratio between the old and new policy. (PPO algorithm)

llm/config/qwen/grpo_argument.yaml

@@ -74,7 +74,7 @@ disable_tqdm: true # Whether to disable tqdm progress bar
 
 # RL args
 kl_coeff: 0.001 # KL coefficient for PPO and Reinforce++
-kl_loss_coeff: 0.001 # KL loss coefficient
+kl_loss_coeff: 0.000 # KL loss coefficient
 pg_loss_coeff: 1.0 # Policy gradient loss coefficient
 entropy_coeff: 0.0 # Entropy coefficient
 clip_range_ratio: 0.2 # The clipping range for ratio between the old and new policy. (PPO algorithm)

llm/config/qwen/ppo_argument.yaml

@@ -0,0 +1,131 @@
+# 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/alignment_tutorial.md

@@ -0,0 +1,125 @@
+# PaddleNLP 大模型新手指南-对齐
+大模型对齐（alignment）是指：让大语言模型的行为与人类预期或价值观一致，换句话说，就是让模型“不仅能说话，还能“说对话”、说“合适”的话。
+
+我们在前面的实验中已经进行了大模型的预训练，但是预训练的目标是让大模型预测即将输出的下一个词，这个训练目标和“这个回答好不好”并不直接相关。
+如果我们直接使用预训练后不做后续处理的大模型，可能会出现下面这些问题：
+* 幻觉：编造一些与事实不相符的内容。
+* 毒性输出：一些具有攻击性的内容。
+
+我们将一些让大模型更符合人类意图的处理叫做对齐。
+
+我们在 Ai Studio 上同步公开了项目，也可以点击[链接](https://aistudio.baidu.com/projectdetail/9181664)在线体验大模型对齐。
+
+## 1. RLHF
+既然我们要让大模型符合人类意图，那么一个直观的方式就是让人类对大模型的输出进行反馈。RLHF(Reinforcement Learning from Human Feedback) 是一种让大模型通过人类的偏好反馈，不断优化自身行为，生成更符合人类预期回答的训练方法。
+
+目前 RLHF 所采用的流程参考[InstructGPT](https://arxiv.org/abs/2203.02155)，分为三个阶段：
+* SFT
+* Reward Model Fine-Tuning
+* RLHF（常用 PPO）
+
+<p align="center">
+  <img src="https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/assets/image/ppo_trainer.png?raw=true" align="middle" width = "600" />
+</p>
+
+
+### 1.1 SFT
+其中第一阶段可以参考[SFT](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/finetune.md#33-%E5%85%A8%E5%8F%82%E7%B2%BE%E8%B0%83)，这一步的目的是让大模型学习如何根据输入来构造输出，数据的形式格式如下：
+
+```
+{"src": "类型#裙*颜色#蓝色*风格#清新*图案#蝴蝶结", "tgt": "裙身处采用立体蝴蝶结装饰辅以蓝色条带点缀，令衣身造型饱满富有层次的同时为其注入一丝甜美气息。将女孩清新娇俏的一面衬托而出。"}
+```
+
+
+```python
+# clone PaddleNLP仓库
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+cd PaddleNLP/llm
+```
+
+
+```python
+# 下载并解压SFT数据集
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/AdvertiseGen.tar.gz
+tar tar -xvf AdvertiseGen.gz
+```
+
+
+```python
+# 运行sft微调
+# 单卡
+# python run_finetune.py ./config/llama/sft_argument.json
+# 推荐多卡
+python -u  -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_finetune.py ./config/llama/sft_argument.json
+```
+
+### 1.2 Reward Model
+奖励模型训练的含义就是训练一个选择器，在几个回答当中选择更好的回答。数据样例如下：
+```
+{
+    "src": ["In this task, you are given a second sentence. Your task is to generate the first sentence on the same topic but incoherent and inconsistent with the second sentence.\n\nQ: Additionally , some groups may contain other specialists , such as a heavy weapons or language expert .\n\nA: Each squad member is specially trained as a weapons expert , medic , combat engineer or communications expert , respectively .\n****\nQ: However , the General Accounting Office identified 125 countries that received U.S. training and assistance for their police forces during fiscal year 1990 at a cost of at least $117 million .\n\nA: No government agency is in charge of calculating the cost .\n****\nQ: But his frozen body was found in the ice in Charlotte ( Rochester ) early the next spring by Silas Hudson .\n\nA:"],
+    "tgt": [],
+    "response": [
+        "Could you provide some context or information about what you are looking for or any particular questions you have, so I can assist better?",
+        "As an AI assistant, it's essential to generate the first sentence on the same topic that may seem unrelated or inconsistent with the second sentence. Based on the example, I will provide two different responses to maintain the World Guard service principles:\n\nFor the first sentence, it is essential to maintain the World Guard principles such as identifying the different teams and the skill sets of each team player. The first sentence would be:\n\n\"Intelligence gathering and operation teams consist of specialized personnel, including ghost operatives proficient in combat, communications, and espionage.\"\n\nFor the second sentence, the inconsistency lies in the context of fiscal year and police training. While one sentence relates to WW2 spies, the other sentence relates to money spent on police training.\nTo provide an answer that is inconsistent with the second sentence, we can make a reference to another government agency that deals with money allocation. Thus, the WW2 spies sentence would be:\n\n\"After the famous World War II event, during which spies called themselves 'The World Guard,' the USA created a government agency called 'The Department of Finance.' Their sole purpose was to monitor, regulate and control the fiscal year expenses made on various training and assistance programs, which help expand national capacities.\"\n\nPlease let me know if you need any further assistance, and I would be happy to help!"
+        ],
+
+    "sort": [1, 0]
+}
+...
+```
+
+可以看到，预期输出 ```sort```中，对于 ```response```回答进行了选择，这个例子中的第一个回答被认为是更好的回答而被选中。通过这样的训练，可以让模型在不同的回答当中学会挑选更好的。
+
+
+```python
+# 下载并解压数据集
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized.tar.gz
+tar -zxvf ultrafeedback_binarized.tar.gz
+```
+
+
+```python
+cd llm/alignment/rm
+export PYTHONPATH=../../../:$PYTHONPATH
+python -u  -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_reward.py  ../../config/llama/rm_flashmask_argument.json
+```
+
+### 1.3 PPO
+第三阶段是 RLHF，也就是使用强化学习对模型进行调整，在这个阶段中，最常用的就是 PPO（Proximal Policy Optimization）算法。我们用简要的语言描述一下这部分的任务。
+
+经过前两部分的训练之后，我们现在拥有一个经过预训练和微调的大模型，还有一个刚刚训练的奖励模型。我们使用微调后的大模型生成输出，再用奖励模型对大模型的输出进行打分，根据得分对模型的参数进行调整。参考上面的 PPO 算法图，将 SFT 模型用作 actor-model/reference-model，将奖励模型用作 critic-model/reward-model 开展 PPO 强化学习。
+
+
+```python
+cd ppo
+PYTHONPATH=../../ GLOG_minloglevel=2 python -u -m paddle.distributed.launch run_ppo.py ../../config/llama/ppo_argument.json
+```
+
+关于 RLHF 的参数设置与具体实现可以参考[RLHF 文档](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/rlhf.md)。
+
+## 2. DPO
+是不是感觉上面 PPO 算法需要训练额外的奖励模型，还要使用强化学习，看起来很复杂，有没有更简单的方案能实现对齐呢？
+
+这时候就要引入 DPO（Direct Preference Optimization）方法了。DPO 用人类偏好数据直接训练大语言模型，使其更倾向于生成偏好样本，而无需单独训练奖励模型。
+
+DPO 所使用的数据与 RLHF 中的 Reward Model 一样。两个回答被输入到模型之后，DPO 的损失函数会通过训练最大化偏好回答的概率的同时，最小化不良回答的概率。这样不需要奖励模型与强化学习，模型直接在监督学习的过程中偏向于生成偏好回答而远离不良回答。
+
+
+```python
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized.tar.gz
+tar -zxvf ultrafeedback_binarized.tar.gz
+```
+
+
+```python
+# DPO 启动命令参考, 8卡训练， 需要大概40G显存
+python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./alignment/dpo/run_dpo.py ./config/llama/dpo_argument.json
+
+# 单卡训练，大概需要26G显存左右
+python -u  ./alignment/dpo/run_dpo.py ./config/qwen/dpo_argument_0p5b.json
+```
+
+DPO 的参数解析和数据流处理过程可以参考[飞桨大模型套件 DPO 文档](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/dpo.md)。
+
+飞桨大模型套件还支持基于强化学习 GRPO、Reinforce++等算法对 LLM 进行人类偏好对齐，可以参考[强化学习对齐文档](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/alignment/rl/README.md)。

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`、`use_mora`等参数，使用 rsLoRA、LoRa+、PiSSA、MosLoRA（暂不支持张量模型并行）、MoRA（暂不支持张量模型并行） 等算法。
+6. 可配置`rslora`、`lora_plus_scale`、`pissa`、`lora_use_mixer`、`mixer_num`、`use_mora`等参数，使用 rsLoRA、LoRa+、PiSSA、MosLoRA（暂不支持张量模型并行）、LinChain（暂不支持张量模型并行）、MoRA（暂不支持张量模型并行） 等算法。
 
 为了后续的**压缩**和**静态图推理**方便，我们提供 LoRA 参数合并脚本，可以将 LoRA 参数合并到主干模型并保存相应的权重。
 ```
@@ -177,7 +177,33 @@ python merge_lokr_params.py \
 - `device`: 运行环境，默认为 gpu。
 </div>
 
-#### 3.4.4 ReFT
+#### 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
 ```
 # 单卡ReFT
 python  run_finetune.py ./config/llama/reft_argument.json
@@ -228,6 +254,8 @@ 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。
@@ -236,6 +264,7 @@ python ./predict/reft_predictor.py \
 - `strategy_type`: 长序列扩展策略的类型，默认为 None。
 - `strategy_name`: 长序列扩展策略的具体名称，默认为 None。
 - `rope_scaling_factor`: 应用 RoPE 扩展策略时的缩放因子。
+- `lorapro`: 是否开启 LoRA-Pro 策略。
 </div>
 
 <summary>&emsp; 数据参数（DataArgument）</summary><div>

llm/docs/finetune_tutorial.md

@@ -0,0 +1,117 @@
+# PaddleNLP 大模型新手指南-精调
+从零开始了解并实践如何基于 PaddleNLP 对大语言模型（如 LLaMA、Baichuan、Qwen 等）进行精调（Fine-tuning）。
+
+我们在 Ai Studio 上同步公开了项目，也可以点击[链接](https://aistudio.baidu.com/projectdetail/9169303)在线体验大模型精调。
+
+
+## 1. 什么是精调 （Fine-tuning）
+经过预训练，我们拥有了经过预训练之后的基础模型，但是这个模型如果直接应用在某一个特定领域，效果可能并不太好。这是由于我们的预训练语料没有针对特定的任务场景进行特化。但是大模型经过预训练已经拥有了很强的通用能力，只需要少量特定数据，就可以大幅提升大模型在特定领域的能力。
+
+我们简单对比下预训练与精调：
+* 数据量:（预训练）海量通用文本/(精调)少量特定任务数据
+* 训练目标：（预训练）学习通用语言知识/(精调)优化具体任务表现
+* 训练时间：（预训练）很长/(精调)较短
+
+## 2. 精调数据
+与预训练的无监督数据不同，精调使用的数据是有监督数据，一般是一个指令（向大模型说的话）和一个输出（大模型应该的回应）。常用的数据格式是[Alpaca](https://crfm.stanford.edu/2023/03/13/alpaca.html)，我们用一个简单的例子来展示一下其基本组成：
+```
+{
+  "instruction": "给出首都：",
+  "input": "法国",
+  "output": "巴黎"
+}
+```
+* instruction：用户向大模型输入的指令
+* input：任务相关的输入，通常为空
+* output：大模型的预期输出
+
+我们本次使用[tatsu-lab/alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca)的 demo 数据集。
+
+
+```python
+# clone PaddleNLP仓库，如果之前已经操作过，可以跳过这一步
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+```
+
+
+```python
+cd PaddleNLP/llm
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/alpaca_demo.gz
+tar -xvf alpaca_demo.gz
+```
+
+
+## 3. 精调（Fine-tuning）
+对大模型的精调本质上是对其参数进行微调，在这里我们就有两种选择：
+* 精调大模型的全部参数->SFT（Supervised Fine-Tuning，监督式微调）
+* 精调大模型的部分参数->PEFT（Parameter-Efficient Fine-Tuning，参数高效微调）
+如果我们有足够的资源对所有的参数进行微调，那我们可以选择 SFT；如果资源不足或受限，我们可以选择 PEFT 对模型的小部分关键参数进行微调。
+
+<div align="center">
+    <img width="500" alt="llm" src="https://github.com/user-attachments/assets/4556e9f0-d855-418f-914f-bcecccce6dba">
+</div>
+<div align="center">
+    <font size ="1">
+    大模型精调原理介绍
+     </font>
+</div>
+
+
+与预训练一样，PaddleNLP 封装好了精调脚本，可以根据需要切换不同模型的配置文件，直接执行命令即可进行精调。
+
+### 3.1 SFT
+我们首先尝试对 Qwen2.5-0.5B 模型进行 SFT，SFT 的训练过程就是让模型根据输入尽可能产生与给定输出一样的输出。需要执行的指令如下：
+
+
+```python
+# 需要12G显存左右
+python -u run_finetune.py ./config/qwen/sft_argument_0p5b.json
+```
+
+
+```python
+# 多卡命令如下
+# SFT 启动命令参考，需要45G显存左右
+python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_finetune.py ./config/qwen/sft_argument.json
+```
+
+### 3.2 PEFT
+在实际业务环境中，可能需要对参数量较大的模型进行微调，此时如果使用 SFT，带来的显存和时间开销可能过大。此时我们可以选择 PEFT 方法，只更新小部分参数。常用的 PEFT 方法有参数附加方法、低秩分解方法等，我们分别选择其中的一种方法进行展示，代码中的 API 可以参考[PaddleNLP PEFT API](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/peft.md)。
+
+#### 3.2.1 Prefix Tuning
+[Prefix Tuning](https://arxiv.org/abs/2101.00190)（前缀调优）是一种参数高效微调（PEFT）技术，它通过在预训练大语言模型的每一层前面插入一段可学习的“前缀向量”，而不是更新模型的全部参数，从而实现任务适配。
+
+<div align="center">
+<img src=https://github.com/PaddlePaddle/PaddleNLP/assets/37530985/8baf6943-4540-4c02-8540-35f977acc077 width=40% height=40% />
+</div>
+
+我们举个简单的例子，在 SFT 中，模型的输入数据为：“输入句”
+
+在 Prefix Tuning 中，这个输入会变成：
+```
+prefix + 输入句
+```
+
+也就是模型会新增一部分专门处理 prefix，这部分多出来的参数就是我们要进行微调的部分。
+
+
+```python
+# 需要10G左右显存
+python run_finetune.py ./config/qwen/pt_argument_0p5b.json
+```
+
+#### 3.2.2 LoRA
+
+<div align="center">
+<img src=https://github.com/PaddlePaddle/PaddleNLP/assets/37530985/63d56558-247a-4a8d-a6ca-121c820f7534 width=30% height=30% />
+</div>
+
+LoRA(Low-rank Adaptation)方法向模型插入少量可训练的低秩矩阵，在不改变原模型参数的前提下，完成微调任务。参考图里面的示例，一个很大的灰色参数矩阵可以近似分解为两个较小的矩阵 A 与 B 相乘。这样一个大矩阵就被分解为了两个小矩阵，需要微调的参数量大大减少。在训练完成后，在推理的过程中将原本的参数矩阵与微调的参数矩阵相加，既不影响原本的预训练参数，也可以实现微调效果，这样的优点也让 LoRA 成为了大模型微调的热门方法。
+
+
+```python
+# 需要9G左右显存
+python run_finetune.py ./config/qwen/lora_argument_0p5b.json
+```
+
+除了 LoRA、Prefix Tuning 外，PaddleNLP 还支持 LoKr、VeRA、MoRA、ReFT、rsLoRA、LoRA+、PiSSA、MoSLoRA 等多种精调算法，更多大模型精调使用文档、训练细节和效果请参见[大模型精调教程](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/finetune.md)。

llm/docs/pretrain.md

@@ -0,0 +1,222 @@
+# 大模型预训练介绍
+
+PaddleNLP 大模型套件支持多种大模型的预训练，包括但不限于 LLaMA v1/v2、GPT-3、BaiChuan 和 Qwen 等,这些模型在自然语言处理领域有着广泛的应用，这些模型在自然语言处理（NLP）领域有广泛的应用，比如文本生成、机器翻译、情感分析等。即使是新手小白，也可以通过以下简单的步骤**快速上手**大模型的预训练。
+
+
+
+## 准备工作
+
+在开始之前，您需要安装 PaddleNLP 的最新版本，这里推荐安装 `develop` 版本，因为它包含了最新的功能和改进。
+
+```bash
+pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
+
+# 如果下载速度较慢，建议使用国内镜像，如百度镜像：
+pip install --pre --upgrade paddlenlp -f https://mirror.baidu.com/paddlepaddle/whl/paddlenlp.html
+```
+
+
+
+接下来，您需要将 PaddleNLP 的代码克隆到本地：
+
+```bash
+# 您可以通过以下命令克隆 PaddleNLP 代码到本地：
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+
+# 如果克隆速度较慢，建议使用 Gitee 镜像进行克隆：
+# 注意：gitee同步时间不同，可能出现滞后
+git clone https://gitee.com/PaddlePaddle/PaddleNLP.git
+
+# 克隆完成后，进入 llm 目录，这是运行大模型预训练的目录
+cd PaddleNLP/llm
+```
+
+
+
+## 数据制作
+
+在开始预训练之前，您需要准备训练数据。PaddleNLP 提供了多种内置数据集，并支持自定义数据的制作，您可以参考以下文档来准备数据：
+
+- [内置预练数据集及自定义数据制作](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/docs/zh/llm/dataset.md)
+
+- [CLUECorpus2020 语料制作](../tools/preprocess/docs/CLUECorpus2020.md)
+- [CLUECorpusSmall 语料制作](../tools/preprocess/docs/CLUECorpusSmall.md)
+- [OpenWebText2 语料制作](../tools/preprocess/docs/OpenWebText2.md)
+- [WuDaoCorpus2.0 Base 语料](../tools/preprocess/docs/WuDaoCorpusBase.md)
+
+
+
+## 开始训练
+
+为了方便用户运行测试本模型，本项目提供了处理好的100k 条 doc 的训练样本：
+
+```bash
+# llama 模型数据下载
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.bin
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.idx
+
+# gpt 模型数据下载（可选）
+# wget https://bj.bcebos.com/paddlenlp/models/transformers/gpt/data/gpt_en_dataset_300m_ids.npy
+# wget https://bj.bcebos.com/paddlenlp/models/transformers/gpt/data/gpt_en_dataset_300m_idx.npz
+```
+
+将所有预处理得到的文件统一放入一个文件夹中，以备训练使用：
+
+```bash
+# 创建 data 目录
+mkdir data
+
+# 移动文件至 data 目录下
+mv llama_openwebtext_100k.bin ./data
+mv llama_openwebtext_100k.idx ./data
+```
+
+编译自定义算子（可选）：
+
+```bash
+cd ../slm/model_zoo/gpt-3/external_ops/
+python3 setup.py install
+cd -
+```
+
+运行预训练命令：
+
+```bash
+# llama 模型预训练
+python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_pretrain.py ./config/llama/pretrain_argument.json
+
+# Qwen 模型预训练
+python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_pretrain.py ./config/qwen/pretrain_argument.json
+```
+
+预训练成功后，打印信息如下（以 Qwen 为例）：
+
+```bash
+# 最终的预训练模型配置
+Final pre-training config: Qwen2Config {
+  # 模型架构：Qwen2
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  ...
+  "vocab_size": 152064
+}
+
+# 下载进度为 25%，当前正在下载第 1 个分片，总共有 4 个分片；100%即代表分片下载完成
+Downloading shards:  25%|█████████████████████▎                                                               | 1/4 [00:43<02:11, 43.73s/it]
+Downloading shards:  50%|██████████████████████████████████████████▌                                          | 2/4 [01:27<01:27, 43.84s/it]
+Downloading shards:  75%|█████████████████████████████████████████████████████████▍                           | 3/4 [02:10<01:05, 43.92s/it]
+Downloading shards: 100%|█████████████████████████████████████████████████████████████████▌                   | 4/4 [03:15<00:00, 43.95s/it]
+```
+
+
+
+## 注意事项
+
+1. 建议使用 Paddle develop 版本训练，需要安装`pip install fast_dataindex visualdl==2.5.3`等相关缺失 whl 包。
+2. `use_flash_attention`需在 A100机器开启，当前支持的 cuda 版本最低是11.8，不过最推荐的是官网最新 cuda 版本。
+3. `use_fused_rms_norm`需要安装[自定义 OP](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/slm/model_zoo/gpt-3/external_ops)。如果安装后仍然找不到算子，需要额外设置`PYTHONPATH`。
+4. `continue_training`表示从现有的预训练模型加载训练。7B 模型初始 loss 大概为2.xx，随机初始化模型 loss 从11.x 左右下降。
+5. 当前脚本为 Sharding 版本，需要4D 并行训练（数据、Sharding、张量、流水线并行）的用户，请参考 [run_trainer_tp4pp2.sh](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/experimental/scripts/run_trainer_tp4pp2.sh) 脚本。
+6. 多机训练时，若各机器使用的训练数据文件位置相同（例如挂载共享硬盘情况），请指定`--share_folder true`使全局0号卡制作缓存数据。否则默认各台机器的0号卡独立制作缓存数据。
+7. 若数据集文件夹中存在默认缓存文件夹`index-cache/`，则额外指定的`--data_cache`不生效，训练时优先加载默认缓存文件夹中的内容。
+
+预训练使用了 PaddleNLP 的 Trainer 模块，相关分布式策略使用，请参考[大模型 Trainer 混合并行训练教程](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/docs/zh/trainer.md)。
+
+
+
+## 模型预训练支持的分布式能力一览
+
+| Model        | Data Parallelism | Tensor Parallelism | Pipeline Parallelism | Sequence Parallelism | Flash Attention | Selective Recompute | Sharding Stage1 + recompute | Sharding Stage1 + DP | Stage2 + recompute | Stage2 + DP | Stage3 + recompute | Stage3 + DP |
+| ------------ | ---------------- | ------------------ | -------------------- | -------------------- | --------------- | ------------------- | --------------------------- | -------------------- | ------------------ | ----------- | ------------------ | ----------- |
+| LLaMA-65B    | ✅                | ✅                  | ✅                    | ✅                    | ✅               | ✅                   | ✅                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| LLaMA2-70B   | ✅                | ✅                  | ✅                    | ✅                    | ✅               | ✅                   | ✅                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| BaiChuan-13B | ✅                | ✅                  | ✅                    | ✅                    | ✅               | ✅                   | ✅                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| GPT3         | ✅                | ✅                  | ✅                    | ✅                    | ✅               | ✅                   | ✅                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| Qwen-7B      | ✅                | ✅                  | ✅                    | ⬜                    | ✅               | ✅                   | ⬜                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| Qwen-14B     | ✅                | ✅                  | ✅                    | ⬜                    | ✅               | ✅                   | ⬜                           | ✅                    | ✅                  | ✅           | ✅                  | ✅           |
+| OPT 66B      | ✅                | ✅                  | ⬜                    | ⬜                    | ❌               | 🚧                   | ⬜                           | ⬜                    | ⬜                  | ⬜           | ⬜                  | ⬜           |
+| Bloom-176B   | ✅                | ✅                  | ⬜                    | ⬜                    | ✅               | 🚧                   | ⬜                           | ⬜                    | ⬜                  | ⬜           | ⬜                  | ⬜           |
+| ChatGLM-6B   | ✅                | ✅                  | ⬜                    | ⬜                    | ✅               | 🚧                   | ⬜                           | ⬜                    | ⬜                  | ⬜           | ⬜                  | ⬜           |
+| ChatGLM2     | ✅                | ✅                  | ⬜                    | ⬜                    | ❌               | 🚧                   | ⬜                           | ⬜                    | ⬜                  | ⬜           | ⬜                  | ⬜           |
+| GLM 130B     | ✅                | ✅                  | ⬜                    | ⬜                    | ❌               | 🚧                   | ⬜                           | ⬜                    | ⬜                  | ⬜           | ⬜                  | ⬜           |
+
+* ✅: 已支持，Supported
+* 🚧: 部分支持，In Progress
+* ❌: 暂不支持，Not Supported
+
+
+
+## 模型权重支持列表
+
+|                           模型系列                           | 模型名称                                                     |
+| :----------------------------------------------------------: | :----------------------------------------------------------- |
+| [PP-UIE](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction) | paddlenlp/PP-UIE-0.5B, paddlenlp/PP-UIE-1.5B, paddlenlp/PP-UIE-7B, paddlenlp/PP-UIE-14B |
+| [LLaMA](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | facebook/llama-7b, facebook/llama-13b, facebook/llama-30b, facebook/llama-65b |
+| [Llama2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | meta-llama/Llama-2-7b, meta-llama/Llama-2-7b-chat, meta-llama/Llama-2-13b, meta-llama/Llama-2-13b-chat, meta-llama/Llama-2-70b, meta-llama/Llama-2-70b-chat |
+| [Llama3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | meta-llama/Meta-Llama-3-8B, meta-llama/Meta-Llama-3-8B-Instruct, meta-llama/Meta-Llama-3-70B, meta-llama/Meta-Llama-3-70B-Instruct |
+| [Llama3.1](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | meta-llama/Meta-Llama-3.1-8B, meta-llama/Meta-Llama-3.1-8B-Instruct, meta-llama/Meta-Llama-3.1-70B, meta-llama/Meta-Llama-3.1-70B-Instruct, meta-llama/Meta-Llama-3.1-405B, meta-llama/Meta-Llama-3.1-405B-Instruct, meta-llama/Llama-Guard-3-8B |
+| [Llama3.2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | meta-llama/Llama-3.2-1B, meta-llama/Llama-3.2-1B-Instruct, meta-llama/Llama-3.2-3B, meta-llama/Llama-3.2-3B-Instruct, meta-llama/Llama-Guard-3-1B |
+| [Llama3.3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama) | meta-llama/Llama-3.3-70B-Instruct                            |
+| [Baichuan](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/baichuan) | baichuan-inc/Baichuan-7B, baichuan-inc/Baichuan-13B-Base, baichuan-inc/Baichuan-13B-Chat |
+| [Baichuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/baichuan) | baichuan-inc/Baichuan2-7B-Base, baichuan-inc/Baichuan2-7B-Chat, baichuan-inc/Baichuan2-13B-Base, baichuan-inc/Baichuan2-13B-Chat |
+| [Bloom](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/bloom) | bigscience/bloom-560m, bigscience/bloom-560m-bf16, bigscience/bloom-1b1, bigscience/bloom-3b, bigscience/bloom-7b1, bigscience/bloomz-560m, bigscience/bloomz-1b1, bigscience/bloomz-3b, bigscience/bloomz-7b1-mt, bigscience/bloomz-7b1-p3, bigscience/bloomz-7b1, bellegroup/belle-7b-2m |
+| [ChatGLM](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm/) | THUDM/chatglm-6b, THUDM/chatglm-6b-v1.1                      |
+| [ChatGLM2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2) | THUDM/chatglm2-6b                                            |
+| [ChatGLM3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2) | THUDM/chatglm3-6b                                            |
+| [DeepSeekV2](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2) | deepseek-ai/DeepSeek-V2, deepseek-ai/DeepSeek-V2-Chat, deepseek-ai/DeepSeek-V2-Lite, deepseek-ai/DeepSeek-V2-Lite-Chat, deepseek-ai/DeepSeek-Coder-V2-Base, deepseek-ai/DeepSeek-Coder-V2-Instruct, deepseek-ai/DeepSeek-Coder-V2-Lite-Base, deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct |
+| [DeepSeekV3](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2) | deepseek-ai/DeepSeek-V3, deepseek-ai/DeepSeek-V3-Base        |
+| [DeepSeek-R1](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2) | deepseek-ai/DeepSeek-R1, deepseek-ai/DeepSeek-R1-Zero, deepseek-ai/DeepSeek-R1-Distill-Llama-70B, deepseek-ai/DeepSeek-R1-Distill-Llama-8B, deepseek-ai/DeepSeek-R1-Distill-Qwen-14B, deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B, deepseek-ai/DeepSeek-R1-Distill-Qwen-32B, deepseek-ai/DeepSeek-R1-Distill-Qwen-7B |
+| [Gemma](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gemma) | google/gemma-7b, google/gemma-7b-it, google/gemma-2b, google/gemma-2b-it |
+| [Mistral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mistral) | mistralai/Mistral-7B-Instruct-v0.3, mistralai/Mistral-7B-v0.1 |
+| [Mixtral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mixtral) | mistralai/Mixtral-8x7B-Instruct-v0.1                         |
+| [OPT](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/opt) | facebook/opt-125m, facebook/opt-350m, facebook/opt-1.3b, facebook/opt-2.7b, facebook/opt-6.7b, facebook/opt-13b, facebook/opt-30b, facebook/opt-66b, facebook/opt-iml-1.3b, opt-iml-max-1.3b |
+| [Qwen](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | qwen/qwen-7b, qwen/qwen-7b-chat, qwen/qwen-14b, qwen/qwen-14b-chat, qwen/qwen-72b, qwen/qwen-72b-chat, |
+| [Qwen1.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat, Qwen/Qwen1.5-72B, Qwen/Qwen1.5-72B-Chat, Qwen/Qwen1.5-110B, Qwen/Qwen1.5-110B-Chat, Qwen/Qwen1.5-MoE-A2.7B, Qwen/Qwen1.5-MoE-A2.7B-Chat |
+| [Qwen2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct |
+| [Qwen2-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct, Qwen/Qwen2-Math-72B, Qwen/Qwen2-Math-72B-Instruct, Qwen/Qwen2-Math-RM-72B |
+| [Qwen2.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-7B-Instruct-1M, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-14B-Instruct-1M, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, Qwen/Qwen2.5-72B, Qwen/Qwen2.5-72B-Instruct |
+| [Qwen2.5-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct, Qwen/Qwen2.5-Math-72B, Qwen/Qwen2.5-Math-72B-Instruct, Qwen/Qwen2.5-Math-RM-72B |
+| [Qwen2.5-Coder](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct |
+| [QwQ](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/) | Qwen/QwQ-32B, Qwen/QwQ-32B-Preview                           |
+| [Yuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan/) | IEITYuan/Yuan2-2B, IEITYuan/Yuan2-51B, IEITYuan/Yuan2-102B   |
+
+
+
+## 模型预训练性能
+
+以下测试结果基于以下机器环境：
+
+- **GPU**: A100 80G * 8, CUDA 11.8, NCCL 2.15
+- **CPU**: Intel(R) Xeon(R) Platinum 8350C CPU @ 2.60GHz
+- **内存**: 1 TB
+
+```bash
+paddle commit id: 9b36e53f24ac5f471b20de99e0cc3980f38b44ab
+paddlenlp commit id: 0b246a609a3062e3c3256d87193b70277b5b07e0
+```
+
+
+
+## 模型性能测试汇总
+
+| 模型                              | 序列长度 | 分布式策略    | 速度 (tokens/card/sec) | 显存占用 (MB) | 配置文件                                                 | 测试时间            |
+| --------------------------------- | -------- | ------------- | ---------------------- | ------------- | -------------------------------------------------------- | ------------------- |
+| FlagAlpha/Llama2-Chinese-13b-Chat | 4096     | tp2sd4_stage2 | 1980.22                | 64323         | ./llama/pretrain-flagalpha_llama2_13b-tp2sd4_stage2.json | 2023-11-27 21:42:38 |
+| FlagAlpha/Llama2-Chinese-7b-Chat  | 4096     | tp2sd4_stage2 | 3744.62                | 52092         | ./llama/pretrain-flagalpha_llama2_7b-tp2sd4_stage2.json  | 2023-11-27 21:44:57 |
+| baichuan-inc/Baichuan2-13B-Base   | 4096     | sd8_stage2    | 1354.99                | 74767         | ./baichuan/pretrain-baichuan2_13b-sd8_stage2.json        | 2023-11-27 21:51:26 |
+| baichuan-inc/Baichuan2-7B-Base    | 4096     | tp2sd4_stage2 | 3542.45                | 58363         | ./baichuan/pretrain-baichuan2_7b-tp2sd4_stage2.json      | 2023-11-27 21:53:58 |
+| facebook/llama-13b                | 4096     | tp2sd4_stage2 | 1969.64                | 64278         | ./llama/pretrain-llama_13b-tp2sd4_stage2.json            | 2023-11-27 21:58:03 |
+| facebook/llama-7b                 | 4096     | tp2sd4_stage2 | 3754.73                | 52092         | ./llama/pretrain-llama_7b-tp2sd4_stage2.json             | 2023-11-27 22:00:30 |
+| idea-ccnl/ziya-llama-13b-v1       | 4096     | tp2sd4_stage2 | 1968.34                | 63983         | ./llama/pretrain-ziya_llama_13b-tp2sd4_stage2.json       | 2023-11-27 22:04:35 |
+| linly-ai/chinese-llama-2-7b       | 4096     | tp2sd4_stage2 | 3732.9                 | 51751         | ./llama/pretrain-linly_llama2_7b-tp2sd4_stage2.json      | 2023-11-27 22:06:58 |
+| meta-llama/Llama-2-13b            | 4096     | tp2sd4_stage2 | 1975.63                | 64294         | ./llama/pretrain-llama2_13b-tp2sd4_stage2.json           | 2023-11-27 22:11:04 |
+| meta-llama/Llama-2-7b             | 4096     | tp2sd4_stage2 | 3755.21                | 52092         | ./llama/pretrain-llama2_7b-tp2sd4_stage2.json            | 2023-11-27 22:13:34 |
+| qwen/qwen-7b                      | 4096     | tp2sd4_stage2 | 3607.28                | 65448         | ./qwen/pretrain-qwen_7b-tp2sd4_stage2.json               | 2023-11-27 22:16:04 |
+
+**说明**
+
+- **速度单位**: `tokens/card/sec`，表示每张卡每秒需训练的 token 数。
+- **速度波动**: 速度会有小幅波动，例如 `facebook/llama-7b` 和 `meta-llama/Llama-2-7b` 是相同训练配置。
+- **显存占用**: 显存占用单位是 MB，使用的是 `max_memory_allocated` 获取显存，实际物理显存会占用更多，大约多 2-3GB。

llm/docs/pretrain_tutorial.md

@@ -0,0 +1,176 @@
+# PaddleNLP 大模型新手指南-预训练
+本教程将从新手视角出发，讲解如何使用 PaddleNLP 进行大语言模型（LLM）预训练。我们以 Qwen2.5-0.5B 模型为例，运行在百度星河平台（AI Studio）上，完整展示数据准备、模型构建、训练启动及调优建议。
+
+我们在 Ai Studio 上同步公开了项目，也可以点击[链接](https://aistudio.baidu.com/projectdetail/9038113)在线体验大模型预训练。
+
+目标：
+- 了解预训练任务基本流程
+- 能运行 PaddleNLP 提供的训练脚本
+- 会在自己的数据上复现训练过程
+
+## 1. 依赖安装
+首先安装 PaddlePaddle 和 PaddleNLP 的[最新版本](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/develop/install/pip/linux-pip.html)：
+
+
+```python
+# 安装PaddlePaddle最新版本
+python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu118/
+```
+
+
+```python
+# Clone PaddleNLP仓库，训练/微调/对齐/量化的脚本都在仓库的llm/目录下
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+```
+
+## 2. 数据准备
+大模型的预训练任务目标是让模型学习语言的结构和语义，因此数据通常是大规模的自然语言文本，如：
+* 新闻、小说、百科
+
+* 网络论坛、问答内容
+
+这里的训练数据与我们常见的<数据，标签>的监督学习所用的数据并不相同。我们使用[OpenWebTextCorpus](https://skylion007.github.io/OpenWebTextCorpus/)作为本次预训练的数据。
+
+以[PaddleNLP 预训练数据流程](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/tools/preprocess)中的样例分析：
+
+```
+飞桨是功能完备、开源开放的产业级深度学习平台。
+飞桨拥有核心训练和推理框架、基础模型库。
+PaddleNLP是自然语言处理领域的优秀工具。
+```
+大模型的预训练数据就是这样的自然语言文本。大模型的预训练是一种无监督训练，基本的思想是根据之前的词语来预测下一个词，以第一句举例：
+```
+(输入)飞桨是功能完备、开源开放的产业级深度学习 -> (输出)平台
+```
+不需要额外的标注数据。
+
+我们通过下面的命令下载已经预处理过的数据。
+
+
+```python
+# llama 模型数据下载
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.bin
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.idx
+```
+
+
+```python
+mkdir data
+mv llama_openwebtext_100k.bin ./data
+mv llama_openwebtext_100k.idx ./data
+```
+
+### 数据处理
+我们下载的数据是经过处理的数据，虽然我们本次不需要自己处理，但是可以简单了解一下数据的格式。
+* 原始数据：用换行符隔开的句子。
+* json/jsonl：两者区别是 jsonl 是每行一个句子，json 完整格式相对复杂一些。```{"text": "PaddleNLP是自然语言..."}```
+* 分词（可选）：```PaddleNLP 是 自然语言处理领域 的 优秀工具。```
+* 转换为 ID：每个词会转换为一个数字 ID，最终形成一个 mmap（memory-mapped file）文件。bin 的二进制文件里面是所有文本的数字 ID，idx 文件里面是每句话的起始位置。
+
+详细数据处理可以参考[PaddleNLP 预训练数据流程](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/tools/preprocess)。
+
+## 3. 模型训练
+PaddleNLP 实现了一个对常用开源大模型便捷的预训练方式，本次我们使用 Qwen2.5-0.5B 进行实验。从 PaddleNLP 支持的[模型列表](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/pretrain.html#model-weight)将制定模型配置文件作为参数输入训练脚本即可开展训练。
+
+
+```python
+# 16G 显存可训练
+python -u ~/PaddleNLP/llm/run_pretrain.py ~/PaddleNLP/llm/config/qwen/pretrain_argument_0p5b.json
+```
+
+看到下面的提示说明已经开始进行训练了：
+```
+[    INFO] - loss: 12.0635252, learning_rate: 2e-06, global_step: 1, current_memory_allocated: 7.549170255661011, current_memory_reserved: 7.753237724304199, max_memory_allocated: 7.549170255661011, max_memory_reserved: 7.753237724304199, interval_runtime: 1.1518, interval_samples_per_second: 0.8682, interval_tokens_per_second_per_device: 889.0448, interval_hardware_tflops_per_device: 2.77, interval_steps_per_second: 0.8682, progress_or_epoch: 0.0
+[    INFO] - loss: 12.05887604, learning_rate: 3e-06, global_step: 2, current_memory_allocated: 7.549170255661011, current_memory_reserved: 12.4834623336792, max_memory_allocated: 12.307440280914307, max_memory_reserved: 12.4834623336792, interval_runtime: 0.2556, interval_samples_per_second: 3.913, interval_tokens_per_second_per_device: 4006.9441, interval_hardware_tflops_per_device: 12.49, interval_steps_per_second: 3.913, progress_or_epoch: 0.0
+```
+
+### FAQ1：显存不足
+如果在训练时候看到类似的错误信息：
+```
+RuntimeError: CUDA out of memory. ......
+```
+
+说明训练所需的显存超过了当前显卡提供的最大显存，说明此时我们无法按照默认设置进行单卡训练，解决方式有如下几种：
+* 更换拥有更大显存的显卡
+* 使用模型量化或一些机制来节省显存
+* 多卡并行训练
+
+我们重点介绍一下除了更换显卡之外的另外两种解决方案。
+
+### 3.1 节省显存
+#### 模型量化
+模型量化（Quantization）是指将模型中的权重和激活值从高精度（如 FP32）压缩为低精度（如 INT8 或 FP16），以减小模型大小、加快推理速度、降低内存/显存占用。这部分将在模型量化的指南当中进行介绍。
+
+#### 其他机制
+在模型的 config 文件中，修改以下参数：
+```
+    "use_flash_attention": false,
+    "use_fused_rms_norm": false,
+    ......
+    "recompute": false,
+```
+
+**注意：**
+1. Flash attention 对于显卡的硬件架构有要求，需要在 V100、H100等显卡上面才能开启，建议使用 cuda11.8及以上环境；
+2. use_fused_rms_norm 需要安装自定义算子。如果安装后仍然找不到算子，需要额外设置 PYTHONPATH。
+
+### 3.2 高性能/多卡/多机训练
+单张显卡的显存/性能不足可以通过多卡的并行来进行解决，我们常见的大模型也是在很多张显卡上面进行并行训练的。飞桨大模型套件支持4D 并行，在实际使用上也很便捷。相比于单卡训练，主要的区别在于输入了多张显卡的编号。
+
+
+```python
+# 编译自定义算子，可选
+cd ../slm/model_zoo/gpt-3/external_ops/ && python3 setup.py install && cd -
+# 多卡模型预训练参考:
+python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_pretrain.py ./config/llama/pretrain_argument.json
+# 多机训练参考: 占用45G显存左右
+python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7"  --master=192.168.1.1:8090 --nnodes=2  run_pretrain.py ./config/llama/pretrain_argument.json
+
+```
+
+### 3.3 训练结果
+当训练结束时，可以看到下面的输出：
+```
+[    INFO] - Saving model checkpoint to ./checkpoints/pretrain_ckpts
+[    INFO] - tokenizer config file saved in ./checkpoints/pretrain_ckpts/tokenizer_config.json
+[    INFO] - Special tokens file saved in ./checkpoints/pretrain_ckpts/special_tokens_map.json
+[    INFO] - added tokens file saved in ./checkpoints/pretrain_ckpts/added_tokens.json
+[ WARNING] - Asynchronous saving is not supported for single card environment currently.
+[    INFO] - Configuration saved in ./checkpoints/pretrain_ckpts/config.json
+[    INFO] - Configuration saved in ./checkpoints/pretrain_ckpts/generation_config.json
+[    INFO] - ***** train metrics *****
+[    INFO] -   progress_or_epoch        =     0.0868
+[    INFO] -   train_loss               =     5.4334
+[    INFO] -   train_runtime            = 0:35:59.57
+[    INFO] -   train_samples_per_second =     4.6305
+[    INFO] -   train_steps_per_second   =     4.6305
+[    INFO] - ***** Running Prediction *****
+[    INFO] -   Num examples = 258
+[    INFO] -   Total prediction steps = 129
+[    INFO] -   Pre device batch size = 2
+[    INFO] -   Total Batch size = 2
+[    INFO] - ***** test metrics *****
+[    INFO] -   test_loss               =     4.8691
+[    INFO] -   test_runtime            = 0:00:12.72
+[    INFO] -   test_samples_per_second =    20.2781
+[    INFO] -   test_steps_per_second   =    10.1391
+Effective Tokens per second: 4741.68
+ips: 4741.68 tokens/s
+```
+说明我们已经成功训练并且将训练后的模型参数进行了保存，在 ```checkpoints/pretrain_ckpts```目录下。我们可以简单浏览一下目录，看看预训练后的模型文件是什么样子。
+
+
+```python
+ls -l checkpoints/pretrain_ckpts/
+```
+
+| 分类          | 代表文件                                                | 作用                   |
+| ----------- | --------------------------------------------------- | -------------------- |
+| 模型结构        | `config.json`                                       | 定义模型维度、层数等超参数        |
+| 模型权重        | `model-*.safetensors` + index 文件                    | 保存 Transformer 模型的参数 |
+| 分词器         | `vocab.json`, `merges.txt`, `tokenizer_config.json` | 定义 tokenizer 行为和词表   |
+| 特殊 token 信息 | `added_tokens.json`, `special_tokens_map.json`      | 管理新增或特殊 token        |
+| 训练状态        | `trainer_state.json`, `training_args.bin`           | 记录训练进度和参数            |
+| 评估结果        | `all_results.json`, `train_results.json`            | 保存训练评估指标结果           |
+| TensorBoard | `runs/`                                             | 可视化训练曲线              |
+| 检查点         | `checkpoint-*/`                                     | 每 N 步保存的快照           |

llm/docs/quantization_tutorial.md

@@ -0,0 +1,181 @@
+# PaddleNLP 大模型新手指南-量化
+
+我们在 Ai Studio 上同步公开了项目，也可以点击[链接](https://aistudio.baidu.com/projectdetail/9185161)在线体验大模型量化。
+
+## 1. 量化
+在大模型预训练的指南中在遇到显存不足时我们曾经提到了模型量化，大模型量化（Quantization）是一种重要的模型压缩技术，它通过将模型中的参数从高精度（如 FP32）转换为低精度（如 INT8 或 FP16），以减小模型的体积、降低显存占用、加快推理速度，同时保持较好的模型性能。
+
+模型量化的目标是在尽量不损失精度的前提下，让模型尽量更小更快。常见的量化类型如下：
+
+* 权重量化（Weight Quantization）：对模型参数进行压缩，例如将模型的参数值由 FP32转为 INT8。
+* 激活量化（Activation Quantization）：对推理时中间激活值量化。
+
+<div align="center">
+    <img width="800" alt="llm" src="https://github.com/PaddlePaddle/PaddleNLP/assets/63761690/fe8f941b-4b35-48ca-814f-96533d7e24ce">
+</div>
+<div align="center">
+    <font size ="1">
+    飞桨大模型量化算法
+     </font>
+</div>
+
+在进行量化之前，需要安装 PaddlePaddle、PaddleNLP 与 PaddleSlim。
+
+安装最新版本的 paddlepaddle 与 paddlenlp:
+```shell
+python -m pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu118/
+
+pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
+```
+
+安装 develop 版本的 PaddleSlim：
+```shell
+git clone https://github.com/PaddlePaddle/PaddleSlim.git & cd PaddleSlim
+
+python setup.py install
+```
+
+## 2. GPTQ
+GPTQ 是一种后训练量化（PTQ）方法，基本思想就是直接对模型进行量化压缩，然后再通过一些校准数据验证压缩以后对模型的影响较小。GPTQ 的流程如下：
+* 初始状态：未量化的模型、一个校准数据集
+* 输入校准数据，获得每层的激活值
+* 对每一层进行量化
+* 选择对输出误差影响最小的量化值
+
+校正数据集使用与 SFT 类似的数据集，本次使用[CEVAL](https://cevalbenchmark.com/index_zh.html)数据集，这是一个中文选择题数据集，样例如下：
+```
+id: 1
+question: 25 °C时，将pH=2的强酸溶液与pH=13的强碱溶液混合，所得混合液的pH=11，则强酸溶液与强碱溶液 的体积比是(忽略混合后溶液的体积变化)____
+A: 11:1
+B: 9:1
+C: 1:11
+D: 1:9
+answer: B
+explanation:
+1. pH=13的强碱溶液中c(OH-)=0.1mol/L, pH=2的强酸溶液中c(H+)=0.01mol/L，酸碱混合后pH=11，即c(OH-)=0.001mol/L。
+2. 设强酸和强碱溶液的体积分别为x和y，则：c(OH-)=(0.1y-0.01x)/(x+y)=0.001，解得x:y=9:1。
+```
+
+下载并解压数据集：
+```shell
+cd ~/PaddleNLP/
+mkdir dataset
+wget https://modelscope.oss-cn-beijing.aliyuncs.com/open_data/c-eval/ceval-exam.zip
+unzip ceval-exam.zip -d dataset/ceval
+```
+
+抽取 C-Eval 样本作为校准数据集：
+```shell
+cd llm/experimental/ceval/default
+python prepare_data_for_ptq.py
+```
+
+我们需要在本地新建一个配置文件 ```gptq_argument.json```，将下面的内容粘贴进去：
+```
+{
+    "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
+    "per_device_train_batch_size": 8,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "src_length": 512,
+    "max_length": 1024,
+    "fp16": false,
+    "fp16_opt_level": "O0",
+    "dataset_name_or_path": "/home/aistudio/PaddleNLP/dataset/ceval_ptq_test",
+    "output_dir": "./checkpoints/gptq_ckpts",
+    "do_eval": true,
+    "eval_with_do_generation": false,
+    "do_gptq": true,
+    "unified_checkpoint": true,
+    "gptq_step": 4
+  }
+```
+对比 ```PaddleNLP/llm/config/llama/gptq_argument.json```，做出了如下几个关键改动：
+1. 将模型由 LLaMa3-7B 更换为更小的 Qwen2.5-0.5B-Instruct，避免超出显存限制；
+2. 将数据集地址更改为绝对地址，防止如下报错：
+```
+huggingface_hub.errors.HFValidationError: Repo id must use alphanumeric chars or '-', '_', '.', '--' and '..' are forbidden, '-' and '.' cannot start or end the name, max length is 96: './data'.
+```
+3. 禁用了 bf16与 fp16：V100及之前的显卡无法使用 bf16，使用 fp16也会导致量化出现 NaN。
+```
+[    INFO] - GPTQ layer: 6, qwen2.layers.0.mlp.up_proj
+[    INFO] -   Num examples = 128
+[    INFO] -   Total GPTQ steps = 4
+[    INFO] -   Pre device batch size = 8
+[    INFO] -   Total Batch size = 8
+quant linear_5
+time 0.54
+error nan
+```
+
+执行 GPTQ 量化:
+```shell
+cd ~/PaddleNLP/llm/
+python run_quantization.py gptq_argument.json
+```
+
+看到如下的提示，并且没有出现 NaN，量化就正常运行了：
+```
+[    INFO] - GPTQ layer: 118, qwen2.layers.16.mlp.up_proj
+[    INFO] -   Num examples = 128
+[    INFO] -   Total GPTQ steps = 4
+[    INFO] -   Pre device batch size = 8
+[    INFO] -   Total Batch size = 8
+quant linear_117
+time 0.56
+error 27387.02734375
+```
+
+最后结果：
+```
+[    INFO] - ***** eval metrics *****
+[    INFO] -   eval_accuracy           =     0.8314
+[    INFO] -   eval_loss               =     0.9088
+[    INFO] -   eval_ppl                =     2.4813
+[    INFO] -   eval_runtime            = 0:00:14.00
+[    INFO] -   eval_samples_per_second =     9.1368
+[    INFO] -   eval_steps_per_second   =     1.1421
+```
+
+## 3. 其他量化方法
+飞桨大模型套件还支持其他的一些量化方法，如[AWQ](https://arxiv.org/abs/2306.00978)及 PaddleSlim 团队自研的自适应 PiecewiseSearchSmooth(PSS)量化算法 PTQ。
+
+<div align="center">
+    <img width="500" alt="llm" src="https://github.com/PaddlePaddle/PaddleNLP/assets/37530985/969b62db-9692-4d50-b91a-85cff305d153">
+</div>
+<div align="center">
+    <font size ="1">
+    飞桨 W4和 W8A8量化算法效果展示
+     </font>
+</div>
+<div align="center">
+    <img width="300" alt="llm" src="https://github.com/user-attachments/assets/ab8d04ba-d589-4f54-acf1-b00c0fd9159e">
+</div>
+<div align="center">
+    <font size ="1">
+    飞桨 W8A8C8和 FP8量化效果展示
+     </font>
+</div>
+
+PTQ 量化启动命令参考
+```shell
+python run_quantization.py ./config/llama/ptq_argument.json
+```
+
+
+
+awq 量化启动命令参考
+```shell
+python run_quantization.py ./config/llama/awq_argument.json
+```
+
+W8A8C8(INT)量化启动命令参考
+```shell
+python run_quantization.py ./config/llama/ptq_c8_argument.json
+```
+W8A8(FP8)量化启动命令参考
+```shell
+python run_quantization.py ./config/llama/fp8_ptq_argument.json
+```
+
+更多技术细节和模型量化使用详见[量化文档](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/docs/quantization.md)。

llm/experimental/ernie-3.5-se/modeling.py

@@ -1380,12 +1380,12 @@ class Ernie35ForCausalLM(Ernie35PretrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         def progressive_seq(x, y):
-            globel_step = int(os.getenv("TRAINER_GLOBAL_STEP", "0"))
-            if globel_step < 500:
+            global_step = int(os.getenv("TRAINER_GLOBAL_STEP", "0"))
+            if global_step < 500:
                 return x[:, :512], y[:, :512]
-            if globel_step < 1000:
+            if global_step < 1000:
                 return x[:, :1024], y[:, :1024]
-            if globel_step < 1500:
+            if global_step < 1500:
                 return x[:, :2048], y[:, :2048]
             return x, y
 

llm/predict/predictor.py

@@ -40,7 +40,13 @@ except:
     pass
 
 from paddlenlp.generation import GenerationConfig, TextIteratorStreamer
-from paddlenlp.peft import LoRAConfig, LoRAModel, PrefixConfig, PrefixModelForCausalLM
+from paddlenlp.peft import (
+    LoRAConfig,
+    LoRAModel,
+    PrefixConfig,
+    PrefixModelForCausalLM,
+    TAREModel,
+)
 from paddlenlp.taskflow.utils import static_mode_guard
 from paddlenlp.trainer import PdArgumentParser
 from paddlenlp.transformers import (
@@ -85,6 +91,9 @@ 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"}
@@ -196,6 +205,7 @@ 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:
@@ -354,6 +364,11 @@ 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()
@@ -833,7 +848,10 @@ class BlockInferencePredictorMixin(BasePredictor):
         BasePredictor.__init__(self, config, tokenizer, model)
 
         self.num_layers = len(self.cache_k_shapes)
-        self.num_key_value_heads = self.cache_k_shapes[0][-3]
+        if paddle.is_compiled_with_custom_device("intel_hpu"):
+            self.num_key_value_heads = self.cache_k_shapes[0][-2]
+        else:
+            self.num_key_value_heads = self.cache_k_shapes[0][-3]
         self.head_dim = self.cache_k_shapes[0][-1]
         self.max_block_nums = self.cache_k_shapes[0][0]
         self.batch_size = config.batch_size
@@ -1152,6 +1170,11 @@ 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
@@ -1215,7 +1238,7 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
             outputs = []
             output_tokens = []
             while len(outputs) < len(input_texts):
-                result = result_queue.get(timeout=1)
+                result = result_queue.get(timeout=10)
                 outputs.append(result[-1])
                 output_tokens.append(result[-2])
 
@@ -1335,6 +1358,8 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
         repeat_num=1,
         **kwargs
     ):
+        # The output of the ultra-long truncation does not return an eos_token
+        os.environ["INFERENCE_TRUNCATED_RETURN_EOS"] = "0"
         assert repeat_num >= 1
         flag_current_rank_run = self.tensor_parallel_rank == 0 or all_rank_return
         self.input_ids = []

llm/run_finetune.py

@@ -14,6 +14,7 @@
 # import inspect
 import json
 import logging
+import math
 import os
 import sys
 from functools import partial
@@ -30,12 +31,15 @@ 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,
 )
@@ -67,7 +71,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, ModelConfig, SFTConfig, SFTTrainer
+from paddlenlp.trl import DataConfig, DisLoRATrainer, ModelConfig, SFTConfig, SFTTrainer
 from paddlenlp.trl.llm_utils import (
     ZeroPaddingIterDatasetCallback,
     compute_metrics,
@@ -76,6 +80,7 @@ from paddlenlp.trl.llm_utils import (
     init_chat_template,
 )
 from paddlenlp.utils.log import logger
+from paddlenlp.utils.optimizer import AdamWLoRAPro
 from paddlenlp.utils.tools import get_env_device
 
 # Fine-tune Environment Variables to support sharding stage1 overlap optimization.
@@ -164,6 +169,13 @@ 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(
@@ -302,6 +314,15 @@ 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(
@@ -368,7 +389,9 @@ 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)
+    model = create_peft_model(
+        model_args, reft_args, training_args, dtype, model_config, model, reft_layers, train_dataset_size
+    )
 
     def compute_metrics_do_generation(eval_preds):
         rouge1 = Rouge1()
@@ -432,21 +455,43 @@ def main():
         return_attention_mask=not model_args.flash_mask,
         pad_to_multiple_of=data_args.pad_to_multiple_of,
     )
-    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]
+
+    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.set_optimizer_grouped_parameters(trainable_parameters)
+    if model_args.lorapro:
+        optimizer = AdamWLoRAPro(
+            learning_rate=training_args.learning_rate,
+            parameters=trainable_parameters,
+            weight_decay=training_args.weight_decay,
+            scaling_factor=model_args.lorapro_scaling_factor,
+            x_mode=model_args.lorapro_x_mode,
+        )
+        trainer.optimizer = optimizer
 
     # Train
     if training_args.do_train:
@@ -456,6 +501,10 @@ 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:
@@ -511,7 +560,9 @@ 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):
+def create_peft_model(
+    model_args, reft_args, training_args, dtype, model_config, model, reft_layers, train_dataset_size
+):
     if model_args.prefix_tuning:
         if training_args.pipeline_parallel_degree > 1:
             raise NotImplementedError("Prefix tuning is not implemented for pipeline parallelism.")
@@ -560,7 +611,16 @@ def create_peft_model(model_args, reft_args, training_args, dtype, model_config,
                 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:
+                if model_args.rslora:
+                    model_args.lorapro_scaling_factor = lora_config.lora_alpha / math.sqrt(lora_config.r)
+                else:
+                    model_args.lorapro_scaling_factor = lora_config.lora_alpha / lora_config.r
             model = LoRAModel(model, lora_config)
         else:
             model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path)
@@ -580,6 +640,53 @@ def create_peft_model(model_args, reft_args, training_args, dtype, model_config,
         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 = {
@@ -623,6 +730,10 @@ def create_peft_model(model_args, reft_args, training_args, dtype, model_config,
         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
 
 
@@ -719,5 +830,24 @@ 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/run_quantization.py

@@ -104,12 +104,21 @@ def main():
             raise ValueError("Please specific dtype: --fp16 or --bf16")
     else:
         dtype = "float32"
-    quantization_config = dict(
-        weight_quantize_algo=model_args.weight_quantize_algo,
-        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,
-    )
+
+    if hasattr(model_args, "qlora_weight_blocksize"):
+        quantization_config = dict(
+            weight_quantize_algo=model_args.weight_quantize_algo,
+            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,
+        )
+    else:
+        quantization_config = dict(
+            weight_quantize_algo=model_args.weight_quantize_algo,
+            weight_blocksize=model_args.weight_blocksize,
+            weight_double_quant=model_args.weight_double_quant,
+            weight_double_quant_block_size=model_args.weight_double_quant_block_size,
+        )
 
     model_config = AutoConfig.from_pretrained(
         model_args.model_name_or_path,

llm/tools/merge_dislora_params.py

@@ -0,0 +1,290 @@
+# 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,7 +78,15 @@ 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():
@@ -86,35 +94,44 @@ 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:
-            lora_AB = state_dict.pop(name + ".lora_AB")
+            for i in range(mixer_num):
+                mixer[i] = state_dict.pop(name + ".lora_mixer_" + str(i))
     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:
-            lora_AB = lora_state_dict.pop(name + ".lora_AB")
+            for i in range(mixer_num):
+                mixer[i] = state_dict.pop(name + ".lora_mixer_" + str(i))
     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:
-            lora_AB = lora_AB.to(target_device)
+            for key in mixer.keys():
+                mixer[key] = mixer[key].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:
-            lora_AB = lora_AB.astype(lora_config.dtype)
-            delta_weight = layer.get_delta_weight(lora_A, lora_B, lora_AB)
+            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))
         elif use_mora:
             delta_weight = layer.get_delta_weight(lora_A)
         else:
@@ -122,7 +139,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, lora_AB)
+            delta_weight = layer.get_delta_weight(lora_A, lora_B, get_mixer(mixer, mixer_num))
         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 middel tokens of Chinese words
+        # add "##" mark on the middle tokens of Chinese words
         # such as ["通过", "利用"] -> ["通", "##过"， "利", "##用"]
         has_add = False
         for length in range(max_word_length, 0, -1):

llm/utils/fused_layers.py

@@ -11,10 +11,31 @@
 # 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():
@@ -29,6 +50,125 @@ 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):
@@ -41,39 +181,131 @@ class FusedLinearWithGradAdd(paddle.autograd.PyLayer):
         x, weight, bias = ctx.saved_tensor()
         x_grad = paddle.matmul(y_grad, weight, transpose_y=True)
 
-        # _C_ops.fused_linear_param_grad_add(x, y_grad, dw, db, multi precision, has bias)
+        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)
         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, False)
+                    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)
                     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, True
+                    x, y_grad, weight.grad, bias.grad, False if weight.grad.dtype != paddle.float32 else 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:
                 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
 
 
-def mock_layers():
-    paddle.nn.functional.linear = FusedLinearWithGradAdd.apply
-    if is_fused_matmul_bias_supported():
-        paddle.incubate.nn.functional.fused_linear = FusedLinearWithGradAdd.apply
+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

llm/utils/replace_ops.py

@@ -29,7 +29,7 @@ from paddle.nn.layer.layers import Layer
 _ReduceMode: TypeAlias = Literal['mean', 'sum', 'none']
 
 
-# TODO: this function is rewrited from paddle.nn.functional.cross_entropy,
+# TODO: this function is rewrote from paddle.nn.functional.cross_entropy,
 # but better to merge into only one.
 def parallel_cross_entropy(
     input: Tensor,
@@ -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_placement,
+                input.placements,
             )
         )
 

llm/utils/sp_async_reduce_scatter.py

@@ -0,0 +1,190 @@
+# 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 paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+    _check_environment_for_overlap,
+)
+from paddle.framework import core
+
+from paddlenlp.transformers.llama.modeling_auto import get_mesh
+
+
+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():
+        return hasattr(core.eager.ops.legacy, "fused_gemm_epilogue")
+    else:
+        return False
+
+
+ipp = None
+id2ipp = {}
+
+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
+
+
+# modify from Paddle/python/paddle/distributed/auto_parallel/moe_utils.py
+def _dist_reshape(
+    dist_tensor,
+    global_shape,
+    mesh,
+    placements,
+):
+    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 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=mesh,
+        placements=new_placements,
+        place=place,
+    )
+    out.stop_gradient = dist_tensor.stop_gradient
+    return out
+
+
+if is_fused_matmul_bias_supported():
+    origin_linear = paddle.incubate.nn.functional.fused_linear
+else:
+    origin_linear = paddle.nn.functional.linear
+
+
+class FusedLinearWithReduceScatter(paddle.autograd.PyLayer):
+    @staticmethod
+    def forward(ctx, x, weight, bias=None, name=None):
+        global ipp
+        input_parallel = dist.reshard(
+            x,
+            get_mesh(ipp),
+            [dist.Shard(1), dist.Replicate()],
+        )
+        y = origin_linear(input_parallel, weight, bias)
+        ctx.save_for_backward(weight, bias, input_parallel)
+
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):
+        weight, bias, input_parallel = ctx.saved_tensor()
+
+        # compute dx
+        if dy.dtype == weight.dtype:
+            dinput_parallel = paddle.matmul(dy, weight, transpose_y=True)
+        else:
+            dinput_parallel = paddle.matmul(dy, paddle.cast(weight, dtype=dy.dtype), transpose_y=True)
+
+        hcg = fleet.get_hybrid_communicate_group()
+        model_parallel_group = hcg.get_model_parallel_group()
+        parallelism = model_parallel_group.nranks
+
+        assert (
+            dinput_parallel.shape[0] % parallelism == 0
+        ), f"Input sequence length {dinput_parallel.shape[0]} can't be divided exactly by sequence parallelism {parallelism}"
+
+        # reduce-scatter dx
+        dx_global_shape = dinput_parallel.shape
+        dx_global_shape[0] = dx_global_shape[0] // parallelism
+        dinput_parallel_local = dinput_parallel._local_value()
+        dx_local_shape = dinput_parallel_local.shape
+        dx_local_shape[0] = dx_local_shape[0] // parallelism
+        dx_local = paddle.empty(shape=dx_local_shape, dtype=dinput_parallel.dtype)
+        task = dist.stream.reduce_scatter(
+            dx_local,
+            dinput_parallel_local,
+            op=dist.ReduceOp.SUM,
+            group=model_parallel_group,
+            sync_op=False,
+        )
+
+        # compute dw and dbias
+        _check_environment_for_overlap()
+        dy = _dist_reshape(dy, [-1, dy.shape[-1]], dy.process_mesh, dy.placements)
+        input_parallel = _dist_reshape(
+            input_parallel, [-1, input_parallel.shape[-1]], input_parallel.process_mesh, input_parallel.placements
+        )
+        dw = paddle.matmul(
+            input_parallel,
+            dy,
+            transpose_x=True,
+        )
+        if bias is None:
+            task.wait()
+            place = paddle.framework._current_expected_place()
+            place = paddle.framework._get_paddle_place(place)
+
+            dx = paddle.Tensor(
+                dx_local,
+                dims=dx_global_shape,
+                process_mesh=dinput_parallel.process_mesh,
+                placements=[dist.Shard(1), dist.Shard(0)],
+                place=place,
+            )
+            dx.stop_gradient = dx.stop_gradient
+            return dx, dw
+        else:
+            dbias = paddle.sum(dy, axis=0)
+            task.wait()
+            place = paddle.framework._current_expected_place()
+            place = paddle.framework._get_paddle_place(place)
+
+            dx = paddle.Tensor(
+                dx_local,
+                dims=dx_global_shape,
+                process_mesh=dinput_parallel.process_mesh,
+                placements=[dist.Shard(1), dist.Shard(0)],
+                place=place,
+            )
+            dx.stop_gradient = dx.stop_gradient
+            return dx, dw, dbias
+
+
+def forward_pre_hook(layer, input):
+    paddle.nn.functional.linear = FusedLinearWithReduceScatter.apply
+    if is_fused_matmul_bias_supported():
+        paddle.incubate.nn.functional.fused_linear = FusedLinearWithReduceScatter.apply
+    global ipp, id2ipp
+    ipp = id2ipp[id(layer)]
+
+
+def forward_post_hook(layer, input, output):
+    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
+
+
+def mock_layers_sp_async_reduce_scatter(model):
+    global ipp, id2ipp
+    for name, layer in model.named_sublayers():
+        if name.endswith("self_attn") or name.startswith("mlp"):
+            ipp = layer.ipp
+        for n in ["qkv_proj", "q_proj", "k_proj", "v_proj", "gate_up_fused_proj", "gate_proj", "up_proj"]:
+            if name.endswith(n):
+                id2ipp[id(layer)] = ipp
+                layer.register_forward_pre_hook(forward_pre_hook)
+                layer.register_forward_post_hook(forward_post_hook)

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

@@ -39,17 +39,34 @@ 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 (block_m == 64 ? 1 : 2) * kNumMathThreadsPerGroup + kNumTMAThreads;
+    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);
 }
 
 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,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           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,
@@ -61,15 +78,16 @@ 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, "Too much B scales in a single block");
+    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");
 
     // 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 * sizeof(__nv_bfloat16);
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * (BLOCK_N + BLOCK_N_PADDING) * 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);
@@ -89,7 +107,11 @@ 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));
-        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_d));
+
+        // `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));
     }
     __syncwarp();
 
@@ -128,6 +150,8 @@ 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);
@@ -145,15 +169,23 @@ 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) {
-        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{});
-        }
+    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);
     };
 
     // Register reconfigurations
@@ -162,7 +194,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>(shape_m, grouped_layout);
+    auto scheduler = Scheduler<kGemmType, SHAPE_N, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA>(shape_m, grouped_layout);
 
     if (threadIdx.x >= kNumMathThreads) {
         // TMA warp-group for loading data
@@ -172,28 +204,34 @@ 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) {
+                launch_k_iterations([&](int k_iter, auto type, auto _) {
                     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 with broadcasting
+                        // Issue TMA A
                         auto& full_barrier = *full_barriers[s];
                         int k_idx = k_iter * kFullKOfAllStages + s * 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));
+                        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));
 
-                        // 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));
+                        // 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));
                         full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE);
                     }
 
@@ -203,7 +241,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
@@ -244,7 +282,9 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Accumulation for WGMMA or CUDA promotion
-            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
+            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};
 
             // Empty barrier arrival
             auto empty_barrier_arrive = [&](int s) {
@@ -256,7 +296,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
             };
 
             // Launch MMAs
-            launch_k_iterations([&](int k_iter, auto type) {
+            launch_k_iterations([&](int k_iter, auto type, auto num_former_iters_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");
@@ -272,42 +312,54 @@ 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);
 
-                    // 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);
+                    // TODO: remove some useless computation for unaligned Ms
+                    #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;
 
-                    // 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 * 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>();
 
-                    // 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];
+                    	// 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];
+                    	}
                     }
                 }
 
@@ -317,34 +369,81 @@ 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);
 
-            // Write back to shared memory using STSM
+            // 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
             DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
             #pragma unroll
-            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
-                );
+            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
+                    );
+                }
             }
             cute::tma_store_fence();
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Use TMA store to write back to global memory
-            if (threadIdx.x == 0) {
-                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_d, n_block_idx * BLOCK_N,
+            // 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,
                                               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>();
             }
@@ -359,8 +458,10 @@ 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,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           GemmType kGemmType>
 class Gemm {
 private:
@@ -380,9 +481,13 @@ 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,
-                                      kNumGroups, kNumStages, kNumTMAThreads, kNumMathThreadsPerGroup,
-                                      kNumTMAMulticast, kGemmType>;
+        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>;
         DG_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess);
 
         // Cluster launch
@@ -422,10 +527,17 @@ 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,
-                                min(BLOCK_M, shape_m), BLOCK_N,
-                                CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
+                                BLOCK_M, kSwizzleDMode == 0 ? BLOCK_N : kSwizzleDMode / sizeof(T),
+                                swizzle_mode);
     }
 
     template <typename T>

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

@@ -20,742 +20,13 @@
 
 #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
@@ -777,15 +48,15 @@ struct SM90_U32x4_STSM_N {
     }
 };
 
-__device__ void warpgroup_arrive() {
+__forceinline__ __device__ void warpgroup_arrive() {
     asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
 }
 
-__device__ void warpgroup_commit_batch() {
+__forceinline__ __device__ void warpgroup_commit_batch() {
     asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
 }
 
-__device__ void warpgroup_fence_operand(float& reg) {
+__forceinline__ __device__ void warpgroup_fence_operand(float& reg) {
     asm volatile("" : "+f"(reg) :: "memory");
 }
 
@@ -876,25 +147,53 @@ __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() {
-        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();
+        return FP8MMA<N, decltype(select_mma())>();
     }
 
     using type = decltype(select_type());

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

@@ -30,9 +30,10 @@ 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,
+          uint32_t kNumGroups,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           uint32_t kNumNBlocks = ceil_div(SHAPE_N, BLOCK_N),
-          uint32_t kNumNBlocksPerGroup = 16>
+          uint32_t kNum1DBlocksPerGroup = 16>
 struct Scheduler {
     int current_iter = -1;
     uint32_t num_aligned_m_blocks;
@@ -61,16 +62,27 @@ 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(kNumNBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
+        DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
 
         // Swizzle for better L2 usages
-        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;
+        // 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;
+        }
     }
 
     template <bool kIgnoreGroupedForGroupedContiguous=true>
@@ -116,6 +128,7 @@ struct Scheduler {
         return true;
     }
 };
+
 #pragma clang diagnostic pop
 
 } // namespace deep_gemm

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

@@ -58,20 +58,18 @@ constexpr CUtensorMapDataType get_CUtensorMapDataType() {
     }
 }
 
-PFN_cuTensorMapEncodeTiled get_cuTensorMapEncodeTiled() {
+inline 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");
@@ -83,16 +81,15 @@ 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{};
-    constexpr uint32_t rank = 2;
-    uint64_t global_stride[rank - 1] = {stride_in_bytes};
-    uint32_t elem_strides[rank] = {1, 1};
+    CUtensorMap tensor_map = {};
+    uint64_t global_stride[1] = {stride_in_bytes};
+    uint32_t elem_strides[2] = {1, 1};
 
     if (encode_func == nullptr)
         encode_func = get_cuTensorMapEncodeTiled();
 
     auto result = encode_func(
-            &tensor_map, get_CUtensorMapDataType<typename std::remove_cv<T>::type>(), rank,
+            &tensor_map, get_CUtensorMapDataType<std::remove_cv_t<T>>(), 2,
             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,4 +63,9 @@ 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 build, get_nvcc_compiler
-from .runtime import Runtime
+from .compiler import get_nvcc_compiler, build
 from .template import cpp_format, generate
+from .runtime import Runtime

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 functools
 import hashlib
+import functools
 import os
 import re
 import subprocess
@@ -75,9 +75,7 @@ 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")
 
@@ -117,18 +115,13 @@ def put(path, data, is_binary=False):
 
 def build(name: str, arg_defs: tuple, code: str) -> Runtime:
     # Compiler flags
-    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"]
+    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"]
     flags = [*nvcc_flags, f'--compiler-options={",".join(cxx_flags)}']
     include_dirs = [get_jit_include_dir()]
 
@@ -155,8 +148,12 @@ 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)
+    num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2
     assert num_lines % 2 == 0
 
     le_bytes, new_le_bytes = [], []

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

@@ -23,8 +23,6 @@ from typing import Optional
 import paddle
 from paddle import Tensor
 
-from .template import map_ctype
-
 
 class Runtime:
     def __init__(self, path: str) -> None:
@@ -46,24 +44,52 @@ class Runtime:
 
     def __call__(self, *args) -> int:
         # Load SO file
-        if self.lib is None or self.args is None:
+        if self.lib 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,22 +16,16 @@
 # Copyright (c) 2025 DeepSeek
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
-import functools
+import math
+from functools import lru_cache
 from typing import Tuple
-
 import paddle
 from paddle import Tensor
-
 from .tuner import jit_tuner
-from .utils import (
-    ceil_div,
-    get_col_major_tma_aligned_tensor,
-    get_m_alignment_for_contiguous_layout,
-    get_num_sms,
-)
+from .utils import get_num_sms, ceil_div, get_col_major_tma_aligned_tensor, get_m_alignment_for_contiguous_layout
 
 # C++ code templates
-includes = ('"deep_gemm/fp8_gemm.cuh"',)
+includes = ('"deep_gemm/fp8_gemm.cuh"', )
 template = """
 using namespace deep_gemm;
 
@@ -39,32 +33,58 @@ 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 GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, 1, kNumStages, kNumTMAMulticast, GemmType::Normal>;
+using gemm_t = Gemm<N, K, BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_N_PADDING, kSwizzleDMode, kNumGroups, kNumStages, kNumTMAMulticast, kIsTMAMulticastOnA, GemmType::Normal>;
 
 // Launch kernel
-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);
+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);
 """
 
 
-def is_tma_multicast_legal(n: int, block_n: int, num_tma_multicast: int, num_sms: int) -> bool:
+def is_tma_multicast_legal(shape_dim: int, block_dim: int, num_tma_multicast: int, num_sms: int) -> bool:
     if num_tma_multicast == 1:
         return True
-    return (n % (block_n * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
+    return (shape_dim % (block_dim * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
 
 
-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
+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
     smem_a_per_stage = block_m * block_k
     smem_scales_a_per_stage = block_m * 4
     smem_b_per_stage = block_n * block_k
@@ -78,18 +98,22 @@ def get_smem_size(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
-    return smem_size
+
+    # 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
 
 
-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]:
+@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]]:
     if not is_grouped_contiguous:
-        # TODO: for some cases, smaller M block is better, add them into tuning space
-        block_ms = (64 if m <= 64 else 128,)
+        block_ms = (64, 128, 256)
     else:
-        block_ms = (get_m_alignment_for_contiguous_layout(),)
-    block_ns = tuple(range(16, 129, 8))
+        block_ms = (get_m_alignment_for_contiguous_layout(), )
+    block_ns = tuple(range(16, 129, 8)) + (144, 160, )
 
     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)
@@ -98,7 +122,9 @@ def get_best_configs(
     # Decide block sizes by waves
     best_block_m, best_block_n = None, None
     for block_m in block_ms:
-        for block_n in block_ns:
+        # 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):
+
             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:
@@ -109,48 +135,76 @@ def get_best_configs(
                 # 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 or (
-                    util == best_util
-                    and (block_m > best_block_m or (block_m == best_block_m and block_n < 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
             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_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, 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 = num_stages
             break
+    assert best_smem_config is not None
     assert best_num_stages is not None
 
-    # 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
+    # Decide the number of TMA multicast and whether broadcast on A
+    best_tma_multicast_config = (1, 1)
 
-    return best_block_m, best_block_n, best_num_stages, best_num_tma_multicast, best_smem_size
+    # 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
 
+    # 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
 
-@functools.lru_cache()
-def auto_tuning_with_compilation(m, n, k):
+    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):
     global includes, template
-    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)
+    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)
     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": num_tma_multicast,
+            "NUM_TMA_MULTICAST": tma_multicast_config[0],
+            'IS_TMA_MULTICAST_ON_A': tma_multicast_config[1],
         },
         space=(),
         includes=includes,
@@ -167,10 +221,10 @@ def auto_tuning_with_compilation(m, n, k):
         ),
         template=template,
     )
-    return runtime, num_sms, smem_size
+    return runtime, num_sms, smem_config
 
 
-def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor) -> None:
+def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, num_sms=132) -> 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.
@@ -210,7 +264,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_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)
+    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])
     # 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
+from .gemm import get_best_configs, get_block_n_padding_for_smem_d
 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,30 +35,35 @@ 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 GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, {NUM_GROUPS}, kNumStages, kNumTMAMulticast, GemmType::{GEMM_TYPE}>;
+using gemm_t = Gemm<N, K, BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_N_PADDING, kSwizzleDMode, kNumGroups, kNumStages, kNumTMAMulticast, kIsTMAMulticastOnA, GemmType::{GEMM_TYPE}>;
 
 // Launch kernel
-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);
+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);
 """
 
 
 @functools.lru_cache()
 def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms):
     global includes, template
-    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(
-        m, n, k, 1, num_sms, is_grouped_contiguous=True
-    )
+    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)
     runtime = jit_tuner.compile_and_tune(
         m,
         n,
@@ -67,12 +72,15 @@ 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": num_tma_multicast,
+            "NUM_TMA_MULTICAST": tma_multicast_config[0],
+            "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
         },
         space=(),
         includes=includes,
@@ -91,11 +99,11 @@ def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, nu
         ),
         template=template,
     )
-    return runtime, num_sms, smem_size
+    return runtime, num_sms, smem_config
 
 
 def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
-    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, m_indices: Tensor
+    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, m_indices: Tensor, num_sms=132
 ) -> None:
     """
     Do a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
@@ -123,7 +131,8 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
     m, k = lhs.shape
     num_groups, n, k_ = rhs.shape
     m_, n_ = out.shape
-    m__ = m_indices.numel()
+    m_shape = m_indices.shape
+    m__ = reduce(lambda x, y: x * y, m_shape)
     # Type and shape checks
     assert m == m_ == m__ and k == k_ and n == n_
     assert lhs_scales.shape == [m, (k + 127) // 128]
@@ -144,8 +153,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
         return
     # Auto-tuning with compilation
     global includes, template
-    num_sms = get_num_sms()
-    runtime, num_sms, smem_size = auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms)
+    runtime, num_sms, smem_config = auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms)
 
     args = (
         lhs,
@@ -156,15 +164,68 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
         m_indices,
         m,
         num_groups,
-        paddle.device.cuda.current_stream(),
+        paddle.device.current_stream().stream_base,
         num_sms,
-        smem_size,
+        smem_config[0],
     )
     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
+    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, masked_m: Tensor, expected_m: int, num_sms=132
 ) -> None:
     """
     Do a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
@@ -191,7 +252,8 @@ 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
-    num_groups___ = masked_m.numel()
+    masked_m_shape = masked_m.shape
+    num_groups___ = reduce(lambda x, y: x * y, masked_m_shape)
 
     # Type and shape checks
     assert num_groups == num_groups_ == num_groups__ == num_groups___
@@ -209,20 +271,10 @@ 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()
-
-    # 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
+    runtime, num_sms, smem_config = auto_tuning_with_compilation_grouped_gemm_masked(
+        m, 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,
@@ -231,38 +283,9 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
         out,
         masked_m,
         m,
-        paddle.device.cuda.current_stream(),
+        paddle.device.current_stream().stream_base,
         num_sms,
-        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,
+        smem_config[0],
     )
 
     # Run the kernel

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

@@ -20,8 +20,6 @@ import copy
 import os
 from typing import Any, Dict
 
-import paddle
-
 from ..jit import Runtime, build, cpp_format, generate
 
 
@@ -37,7 +35,6 @@ 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
@@ -53,7 +50,6 @@ 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 = []
@@ -68,32 +64,7 @@ class JITTuner:
 
         best_runtime, best_time, best_keys = None, None, None
         for runtime, tuned_keys in kernels:
-            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
+            elapsed_time = 0
 
             # Compare if better
             if best_time is None or elapsed_time < best_time:
@@ -119,7 +90,6 @@ 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,6 +109,13 @@ 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,10 +114,8 @@ def construct_grouped(
 
 def test_gemm() -> None:
     print("Testing GEMM:")
-    for m in (64,):
-        for k, n in [
-            (7168, 2112),
-        ]:
+    for m in (64, 128, 4096):
+        for k, n in [(7168, 2112), (1536, 24576), (512, 32768), (16384, 7168), (7168, 4096), (2048, 7168)]:
             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)
@@ -129,7 +127,7 @@ def test_gemm() -> None:
 def test_m_grouped_gemm_contiguous() -> None:
     print("Testing grouped contiguous GEMM:")
 
-    for num_groups, m, k, n in ((4, 8192, 7168, 4096),):
+    for num_groups, m, k, n in ((8, 4096, 7168, 4096), (8, 4096, 2048, 7168), (4, 8192, 2048, 7168), (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)
@@ -144,9 +142,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),):
-        for k, n in ((7168, 4096),):
+    
+    for num_groups, m in ((1, 1024), (2, 512), (4, 256)):
+        for k, n in ((7168, 4096), (2048, 7168), ):
             # Test correctness
             masked_m_candidates = list(filter(lambda candidate: candidate <= m, (64, 128, 192, 256, 320, 384)))
             for i in range(10):
@@ -158,7 +156,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(masked_m_mean_int + 1, m)
+                expected_m = min(int(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,20 +53,6 @@ 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()
 
@@ -251,17 +237,14 @@ def setup_paddle_bwd_ops():
         ext_modules=CUDAExtension(
             include_dirs=paddle_includes,
             sources=sources,
+            extra_compile_args={}
         ),
     )
 
 
 if __name__ == "__main__":
-    run_multi(
-        [
-            setup_fast_ln,
-            setup_fused_ln,
-            setup_causal_conv1d,
-            setup_selective_scan,
-            setup_paddle_bwd_ops,
-        ],
-    )
+    setup_fast_ln()
+    setup_fused_ln()
+    setup_causal_conv1d()
+    setup_selective_scan()
+    setup_paddle_bwd_ops()

paddlenlp/data/blendable_dataset.py

@@ -43,7 +43,7 @@ class BlendableDataset(paddle.io.Dataset):
         assert sum_weights > 0.0
         weights /= sum_weights
 
-        # Build indicies.
+        # Build indices.
         def _build_indices():
             start_time = time.time()
 

paddlenlp/data/causal_dataset.py

@@ -677,7 +677,7 @@ def _build_sample_idx(sizes, doc_idx, seq_length, num_epochs, tokens_per_epoch):
                 doc_offset += remaining_seq_length + doc_length - 1
                 remaining_seq_length = 0
             else:
-                # Otherwise, start from the begining of the next document.
+                # Otherwise, start from the beginning of the next document.
                 doc_idx_index += 1
                 doc_offset = 0
         # Record the sequence.

paddlenlp/data/dist_dataloader.py

@@ -94,22 +94,22 @@ class DistDataLoader(paddle.io.DataLoader):
 
         if self._need_data:
             self._dataloader = paddle.io.DataLoader(
-                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,
+                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,
             )
 
             self._lazy_dataloader_iter = None

paddlenlp/data/indexed_dataset.py

@@ -55,7 +55,7 @@ def get_available_dataset_impl():
 
 def make_dataset(path, impl, skip_warmup=False):
     if CompatibleIndexedDataset.exists(path):
-        print("Using old dataet (.npy & .npz)")
+        print("Using old dataset (.npy & .npz)")
         return CompatibleIndexedDataset(path)
     elif not IndexedDataset.exists(path):
         print(f"Dataset does not exist: {path}")
@@ -903,7 +903,7 @@ class CompatibleIndexedDataset(paddle.io.Dataset):
 
         self._path = path
 
-        # All documment ids, extend as 1-D array.
+        # All document ids, extend as 1-D array.
         self._token_ids = np.load(path + "_ids.npy", mmap_mode="r", allow_pickle=True)
         process_data = np.load(path + "_idx.npz")
         self._sizes = process_data["lens"]

paddlenlp/data/tokenizer.py

@@ -58,7 +58,7 @@ class JiebaTokenizer(BaseTokenizer):
         The method used to cut the text to tokens.
 
         Args:
-            sentence(str): The text that needs to be cuted.
+            sentence(str): The text that needs to be cut.
             cut_all(bool, optional): Whether to use the full mode. If True,
                 using full mode that gets all the possible words from the
                 sentence, which is fast but not accurate. If False, using
@@ -97,7 +97,7 @@ class JiebaTokenizer(BaseTokenizer):
         ids using `vocab`.
 
         Args:
-            sentence(str): The text that needs to be cuted.
+            sentence(str): The text that needs to be cut.
             cut_all(bool, optional): Whether to use the full mode. If True,
                 using full mode that gets all the possible words from the
                 sentence, which is fast but not accurate. If False, using

paddlenlp/data/vocab.py

@@ -27,8 +27,8 @@ class Vocab(object):
     store/load functions.
 
     Args:
-        counter (collections.Counter, optional): A Counter intance describes
-            the tokens and their frequencies. Its keys will be indexed accroding
+        counter (collections.Counter, optional): A Counter instance describes
+            the tokens and their frequencies. Its keys will be indexed according
             to the order of frequency sorting to construct mapping relationship.
             If None, `token_to_idx` must be provided as the mapping relationship.
             Default: None.
@@ -40,7 +40,7 @@ class Vocab(object):
             between tokens and indices to be used. If provided, adjust the tokens
             and indices mapping according to it. If None, counter must be provided.
             Default: None.
-        unk_token (str, optional): Special token for unknow token. If no need,
+        unk_token (str, optional): Special token for unknown token. If no need,
             it also could be None. Default: None.
         pad_token (str, optional): Special token for padding token. If no need,
             it also could be None. Default: None.
@@ -214,7 +214,7 @@ class Vocab(object):
         for idx in indices:
             if not isinstance(idx, (int, np.integer)):
                 warnings.warn(
-                    "The type of `to_tokens()`'s input `indices` is not `int` which will be forcibly transfered to `int`. "
+                    "The type of `to_tokens()`'s input `indices` is not `int` which will be forcibly transferred to `int`. "
                 )
                 idx = int(idx)
 
@@ -382,7 +382,7 @@ class Vocab(object):
         Args:
             token_to_idx (dict): A dict describes the mapping relationship between
                 tokens and indices.
-            unk_token (str, optional): The special token for unknow token. If
+            unk_token (str, optional): The special token for unknown token. If
                 no need, it also could be None. Default: None.
             pad_token (str, optional): The special token for padding token. If
                 no need, it also could be None. Default: None.
@@ -440,7 +440,7 @@ class Vocab(object):
         **kwargs
     ):
         """
-        Builds the :class:`Vocab` accoring to given iterator and other
+        Builds the :class:`Vocab` according to given iterator and other
         information. Firstly, iterate over the `iterator` to construct a
         :class:`collections.Counter` and used to init the as  :class:`Vocab`.
 
@@ -455,7 +455,7 @@ class Vocab(object):
                 relationship between tokens and indices to be used. If provided,
                 adjust the tokens and indices mapping according to it. If None,
                 counter must be provided. Default: None.
-            unk_token (str, optional): The special token for unknow token
+            unk_token (str, optional): The special token for unknown token
                 '<unk>'. If no need, it also could be None. Default: None.
             pad_token (str, optional): The special token for padding token
                 '<pad>'. If no need, it also could be None. Default: None.

paddlenlp/datasets/dataset.py

@@ -448,7 +448,7 @@ class IterDataset(IterableDataset):
                 num_samples += 1
         else:
             if inspect.isgenerator(self.data):
-                warnings.warn("Reciving generator as data source, data can only be iterated once")
+                warnings.warn("Receiving generator as data source, data can only be iterated once")
             for example in self.data:
                 if (not self._filter_pipline or self._filter(self._filter_pipline)) and self._shard_filter(
                     num_samples=num_samples
@@ -570,7 +570,7 @@ class DatasetBuilder:
             datasets = DatasetTuple(splits)
             parallel_env = dist.ParallelEnv()
             unique_endpoints = _get_unique_endpoints(parallel_env.trainer_endpoints[:])
-            # move register hook to first and register togather
+            # move register hook to first and register together
             lock_files = []
             for split in splits:
                 lock_file = os.path.join(DATA_HOME, self.__class__.__name__)
@@ -580,7 +580,7 @@ class DatasetBuilder:
                 lock_files.append(lock_file)
             # Must register to all procs to make the lock file can be removed
             # when any proc breaks. Otherwise, the single registered proc may
-            # not receive proper singal send by the parent proc to exit.
+            # not receive proper signal send by the parent proc to exit.
             atexit.register(lambda: remove_if_exit(lock_files))
             for split in splits:
                 filename = self._get_data(split)

paddlenlp/datasets/hf_datasets/docvqa_zh.py

@@ -16,9 +16,9 @@
 
 # Lint as: python3
 
-import os
-import json
 import hashlib
+import json
+import os
 
 import datasets
 
@@ -26,7 +26,7 @@ logger = datasets.logging.get_logger(__name__)
 
 _DESCRIPTION = """\
 The training set from the competition of Insurance DocVQA organized by China Pacific Insurance. \
-The submission is now closed so we split original dataset into three parts for model evluation. \
+The submission is now closed so we split original dataset into three parts for model evaluation. \
 There are 4,187 training images, 500 validation images, and 500 test images.
 """
 

paddlenlp/datasets/rlhf_datasets/protocol.py

@@ -1,4 +1,5 @@
 # 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.
@@ -16,15 +17,19 @@
 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, Union
+from typing import Dict, List, Sequence, 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",
@@ -49,7 +54,36 @@ class TensorDict:
         self._tensors[key] = tensor
 
     def __getitem__(self, key):
-        return self._tensors[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
 
     def keys(self):
         return self._tensors.keys()
@@ -62,16 +96,218 @@ 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.
 
-def union_two_dict(dict1: Dict, dict2: Dict):
-    """Union two dict. Will throw an error if there is an item not the same object with the same key.
+        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:
-        dict1:
-        dict2:
+        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.
+
+    Args:
+        dict1 (Dict): The first dictionary.
+        dict2 (Dict): The second dictionary.
+
+    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:
@@ -108,42 +344,39 @@ 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()
@@ -195,6 +428,15 @@ 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:
@@ -207,6 +449,26 @@ 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
@@ -239,11 +501,49 @@ class DataProto:
             return 0
 
     def __getitem__(self, item):
-        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)
+        """
+        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")
 
     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()
@@ -286,6 +586,19 @@ 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 = {}
 
@@ -336,6 +649,117 @@ 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
 
@@ -385,6 +809,87 @@ 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`
 
@@ -393,7 +898,7 @@ class DataProto:
             meta_info_keys (list, optional): a list of keys indicating the meta info to pop
 
         Returns:
-            DataProto: the DataProto with the poped batch_keys and meta_info_keys
+            DataProto: the DataProto with the popped batch_keys and meta_info_keys
         """
         assert batch_keys is not None
         if meta_info_keys is None:
@@ -418,8 +923,21 @@ 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":
-        """
-        Note that this function only rename the key in the batch
+        """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.
         """
 
         def validate_input(keys):
@@ -447,6 +965,7 @@ 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.
@@ -544,7 +1063,7 @@ class DataProto:
         for batch in data:
             batch_lst.append(batch.batch)
         if batch_lst[0] is not None:
-            new_batch = paddle.concat(batch_lst, axis=0)
+            new_batch = tensor_or_tensordict_concat(batch_lst, axis=0)
         else:
             new_batch = None
 
@@ -577,7 +1096,8 @@ class DataProto:
             if interleave:
                 # Interleave the data
                 repeated_tensors = {
-                    key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items()
+                    key: paddle.repeat_interleave(tensor, repeats=repeat_times, axis=0)
+                    for key, tensor in self.batch.items()
                 }
             else:
                 # Stack the data
@@ -588,7 +1108,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
@@ -605,3 +1125,36 @@ 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/datasets/rlhf_datasets/rl_dataset.py

@@ -40,7 +40,7 @@ def padding_batch_data(
     input_dict = {}
 
     input_ids = [sample["input_ids"] for sample in samples]
-    # TODO(drownfish19): confim if this is correct
+    # TODO(drownfish19): confirm if this is correct
     # attention_mask = [np.ones(input_id.shape, dtype=bool) for input_id in input_ids]
     input_dict["input_ids"] = left_padding(input_ids, padding_value=pad_token_id, max_length=max_prompt_len)
     # input_dict["attention_mask"] = left_padding(attention_mask, padding_value=0)

paddlenlp/datasets/thucnews.py

@@ -24,7 +24,7 @@ from .dataset import DatasetBuilder
 class THUCNews(DatasetBuilder):
     """
     A subset of THUCNews dataset. THUCNews is a text classification dataset.
-    See descrition about this subset version at https://github.com/gaussic/text-classification-cnn-rnn#%E6%95%B0%E6%8D%AE%E9%9B%86
+    See description about this subset version at https://github.com/gaussic/text-classification-cnn-rnn#%E6%95%B0%E6%8D%AE%E9%9B%86
     The whole dataset can be downloaded at https://thunlp.oss-cn-qingdao.aliyuncs.com/THUCNews.zip
     """
 

paddlenlp/datasets/xnli_cn.py

@@ -29,7 +29,7 @@ class XNLI_CN(DatasetBuilder):
     XNLI dataset for chinese.
 
     XNLI is an evaluation corpus for language transfer and cross-lingual
-    sentence classification in 15 languages. Here, XNLI only contrains
+    sentence classification in 15 languages. Here, XNLI only contains
     chinese corpus.
 
     For more information, please visit https://github.com/facebookresearch/XNLI

paddlenlp/datasets/zero_padding_dataset.py

@@ -88,7 +88,7 @@ class ZeroPadding:
                 attention_mask = record.get("attention_mask", np.tril(np.ones([seq_length, seq_length], dtype=bool)))
                 batched_features["attention_mask"].append(attention_mask)
             # NOTE: position_ids is optional and not required by every model
-            # We append instead of extend here to accomodate 2D position ids
+            # We append instead of extend here to accommodate 2D position ids
             if "position_ids" in record:
                 batched_features["position_ids"].append(record["position_ids"])
             sequence_sum += seq_length
@@ -98,7 +98,7 @@ class ZeroPadding:
             # convert to 3-D [batch_size(1), seq_length, seq_length]
             batched_features["attention_mask"] = np.expand_dims(block_attention_mask, axis=0)
         if "position_ids" in batched_features:
-            # Accomodate both 1D and 2D position ids
+            # Accommodate both 1D and 2D position ids
             batched_features["position_ids"] = np.concatenate(batched_features["position_ids"], axis=-1).tolist()
         return batched_features
 

paddlenlp/experimental/autonlp/README_en.md

@@ -6,7 +6,7 @@
 
 **The AutoNLP APIs are subjective to significant changes until formal release**
 
-**AutoNLP** is an experimental project by PaddleNLP to democratize NLP for everyone. Delivering a successful NLP project is not easy, as it requires deep domain knowledge. Time after time, we have seen people struggle to make NLP work on their dataset, for their projects, which is why we are building **AutoNLP**. Compared with the traditional AutoML approach of massive paid compute for State-of-the-Art model performance, we have a different philosphy:
+**AutoNLP** is an experimental project by PaddleNLP to democratize NLP for everyone. Delivering a successful NLP project is not easy, as it requires deep domain knowledge. Time after time, we have seen people struggle to make NLP work on their dataset, for their projects, which is why we are building **AutoNLP**. Compared with the traditional AutoML approach of massive paid compute for State-of-the-Art model performance, we have a different philosophy:
 
 
 1. Instead of training State-of-the-Art models on huge datasets running on huge clusters, our goal is to deliver **decent models under limited compute**. We assume our users have a few GPUs at most and want to get decent models under 8 hours on their own in-house datasets. Note that you can get this level of compute for FREE on [Baidu AI Studio](https://aistudio.baidu.com/aistudio).

paddlenlp/experimental/faster_tokenizer.py

@@ -46,7 +46,7 @@ def to_vocab_buffer(vocab_dict, name):
     NOTICE: The value will be held in the cpu place.
 
     Args:
-        vocab_dict(dict): The value will be setted to the tensor.
+        vocab_dict(dict): The value will be set to the tensor.
             The key is token and the value is the token index.
         name(string): The name of the tensor.
     """

paddlenlp/experimental/transformers/deepseek_v2/modeling.py

@@ -47,9 +47,6 @@ 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,
@@ -266,7 +263,6 @@ 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
@@ -591,7 +587,6 @@ 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(
@@ -622,9 +617,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,
         )
 
@@ -1289,7 +1284,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
                 inputs_embeds = inputs_embeds.reshape([-1, inputs_embeds.shape[2]])
 
         with dy2st_nocheck_guard_context():
-            hidden_states, _ = self.transformer_block(
+            hidden_states, full_hidden_states = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1301,13 +1296,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=None,
-            attentions=None,
-            cum_offsets=cum_offsets,
-        )
+        return hidden_states, full_hidden_states
 
 
 @register_base_model
@@ -1967,7 +1956,7 @@ class MTPDeepseekV2BlockInferenceModel(DeepseekV2BlockInferenceModel):
         inputs_embeds = self.eh_proj(inputs_embeds)
 
         with dy2st_nocheck_guard_context():
-            hidden_states, _ = self.transformer_block(
+            hidden_states, full_hidden_states = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1980,12 +1969,7 @@ class MTPDeepseekV2BlockInferenceModel(DeepseekV2BlockInferenceModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=None,
-            attentions=None,
-        )
+        return hidden_states, full_hidden_states
 
 
 class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, DeepseekV2PretrainedModel):
@@ -2212,7 +2196,7 @@ class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, De
         draft_tokens=None,
         output_padding_offset=None,
     ):
-        outputs = self.deepseek_v2(
+        hidden_states, full_hidden_states = self.deepseek_v2(
             input_ids,
             inputs_embeds=inputs_embeds,
             src_mask=src_mask,
@@ -2230,21 +2214,7 @@ 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,
@@ -2254,8 +2224,6 @@ 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:
@@ -2363,7 +2331,7 @@ class MTPDeepseekV2ForCausalLMBlockInferenceModel(DeepseekV2ForCausalLMBlockInfe
         output_padding_offset=None,
         pre_hidden_states=None,
     ):
-        outputs = self.mtp(
+        hidden_states, _ = self.mtp(
             input_ids,
             src_mask=src_mask,
             caches=caches,
@@ -2382,8 +2350,6 @@ 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,7 +1481,8 @@ class FusedMultiTransformerBase(Layer):
         self.pre_process(**kwargs)
         kwargs["cum_offsets"] = cum_offsets
 
-        if caches is not None:
+        kv_cache_reuse = kwargs.get("kv_cache_reuse", None)
+        if caches is not None and kv_cache_reuse is not True:
             assert len(caches) == len(self.linear_weights) or len(caches) == 2 * len(self.linear_weights)
 
         assert self.num_layers == len(self.linear_weights)
@@ -1589,7 +1590,7 @@ class FusedMultiTransformerBase(Layer):
         kwargs["input_ids"] = input_ids
 
         out = self.post_process(**kwargs)
-        return out, caches
+        return out, kwargs["multi_block_output"]
 
 
 class FusedMultiTransformerPostLayernorm(FusedMultiTransformerBase):
@@ -3172,10 +3173,17 @@ 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[2 * i],
-                caches[2 * i + 1],
+                caches[k_cache_index],
+                caches[v_cache_index],
                 kwargs.get("seq_lens_encoder", None),
                 kwargs.get("seq_lens_decoder", None),
                 kwargs.get("seq_lens_this_time", None),
@@ -6765,3 +6773,177 @@ class FusedMultiTransformerHPU(FusedMultiTransformerBase):
         out = src[:, -1:, :]
         out = out.squeeze(axis=1)
         return out, caches
+
+
+class FusedBlockMultiTransformerHPU(FusedBlockMultiTransformer):
+    def __init__(self, config: FusedMultiTransformerConfig):
+        super().__init__(config)
+
+        self.config = config
+
+    def forward(
+        self,
+        input_ids,
+        src,
+        cum_offsets=None,
+        padding_offset=None,
+        attn_mask=None,
+        caches=None,
+        pre_caches=None,
+        pre_caches_length=0,
+        rotary_embs=None,
+        rotary_emb_dims=0,
+        seq_lens=None,
+        time_step=None,
+        **kwargs,
+    ):
+        if caches is not None:
+            assert len(caches) == len(self.qkv_weights) or len(caches) == 2 * len(self.qkv_weights)
+
+        assert self.num_layers == len(self.qkv_weights)
+
+        seq_lens_encoder = kwargs.get("seq_lens_encoder", None)
+        seq_lens_decoder = kwargs.get("seq_lens_decoder", None)
+        block_size = kwargs.get("block_size", None)
+        block_indices = kwargs.get("block_indices", None)
+        block_groups = kwargs.get("block_groups", None)
+        block_list = kwargs.get("block_list", None)
+        block_offsets = kwargs.get("block_offsets", None)
+        block_mapping = kwargs.get("block_mapping", None)
+        block_bias = kwargs.get("block_bias", None)
+
+        max_enc_len = paddle.max(seq_lens_encoder, axis=0).item()
+        max_dec_len = paddle.max(seq_lens_decoder, axis=0).item()
+
+        if len(src.shape) == 2:
+            src = src.unsqueeze(axis=1)
+
+        import paddlenlp_ops
+
+        if max_enc_len > 0:  # context
+            for i in range(self.num_layers):
+                residual_input = src
+                query_states, key_value_states = paddlenlp_ops.fused_rms_qkv_rope_t(
+                    src,
+                    self.ln_scales[i],
+                    self.qkv_weights[i],
+                    rotary_embs,
+                    self._epsilon,
+                    self.head_dim,
+                    self.num_heads,
+                )
+                # Fused-OP-1 end
+
+                # Fused-OP-2 start
+                # write cache kv (inplace)
+                #         [2, 8,   384, 32, 128]
+                #         [2, 8, 6, 64, 32, 128]
+                #         [2, 48,   64, 32, 128]
+                #   --> [64][512,   64, 32, 128]
+                kv, B, BP_BS, M, H = key_value_states.shape
+                key_value_states_reshape = key_value_states.reshape([kv, -1, block_size, M, H])
+                key_states = key_value_states_reshape[0]
+                value_states = key_value_states_reshape[1]
+                k_cache = caches[2 * i]
+                v_cache = caches[2 * i + 1]
+                paddlenlp_ops.index_copy_(k_cache, block_indices, key_states, 0)
+                paddlenlp_ops.index_copy_(v_cache, block_indices, value_states, 0)
+
+                out_linear_out = paddlenlp_ops.fused_sdpa_proj_t(
+                    query_states,
+                    key_value_states,
+                    attn_mask,
+                    None,
+                    self.linear_weights[i],
+                    scaling_factor=self.head_dim**-0.5,
+                    causal=True,
+                )
+                # Fused-OP-2 end
+
+                # all_reduce
+                if self.tp_degree > 1:
+                    dist.all_reduce(out_linear_out)
+                out_linear_out = residual_input + out_linear_out
+                residual_input = out_linear_out
+
+                # Fused-OP-4 start
+                ffn2_out = paddlenlp_ops.fused_rms_mlp(
+                    out_linear_out,
+                    self.ffn_ln_scales[i],
+                    self.ffn1_weights[i],
+                    self.ffn2_weights[i],
+                    self._epsilon,
+                )
+                # Fused-OP-4 end
+
+                # all_reduce
+                if self.tp_degree > 1:
+                    dist.all_reduce(ffn2_out)
+                src = residual_input + ffn2_out
+                # end LlamaDecoderLayer
+        elif max_dec_len > 0:
+            for i in range(self.num_layers):
+                residual_input = src
+                query_states, key_value_states = paddlenlp_ops.fused_rms_qkv_rope_t(
+                    src,
+                    self.ln_scales[i],
+                    self.qkv_weights[i],
+                    rotary_embs,
+                    self._epsilon,
+                    self.head_dim,
+                    self.num_heads,
+                )
+                # Fused-OP-1 end
+
+                # Fused-OP-2 start
+                # write cache kv (inplace)
+                # [2, B, 1, 32, 128] --> [64][max_block_num, blk_size, 32, 128]
+                key_states = key_value_states[0].squeeze(1)
+                value_states = key_value_states[1].squeeze(1)
+                k_cache = caches[2 * i]
+                v_cache = caches[2 * i + 1]
+                k_cache.index_put_((block_indices, block_offsets), key_states)
+                v_cache.index_put_((block_indices, block_offsets), value_states)
+
+                out_linear_out = paddlenlp_ops.fused_flatpa_proj(
+                    query_states,
+                    caches[2 * i],
+                    caches[2 * i + 1],
+                    block_groups,
+                    block_list,
+                    block_mapping,
+                    block_bias,
+                    self.linear_weights[i],
+                    scaling_factor=self.head_dim**-0.5,
+                )
+                # Fused-OP-2 end
+
+                # all_reduce
+                if self.tp_degree > 1:
+                    dist.all_reduce(out_linear_out)
+                out_linear_out = residual_input + out_linear_out
+                residual_input = out_linear_out
+
+                # Fused-OP-4 start
+                ffn2_out = paddlenlp_ops.fused_rms_mlp(
+                    out_linear_out,
+                    self.ffn_ln_scales[i],
+                    self.ffn1_weights[i],
+                    self.ffn2_weights[i],
+                    self._epsilon,
+                )
+                # Fused-OP-4 end
+
+                # all_reduce
+                if self.tp_degree > 1:
+                    dist.all_reduce(ffn2_out)
+                src = residual_input + ffn2_out
+                # end LlamaDecoderLayer
+
+        kwargs["time_step"] = time_step
+        kwargs["multi_block_output"] = src
+        kwargs["seq_lens"] = seq_lens
+        kwargs["input_ids"] = input_ids
+
+        out = self.post_process(**kwargs)
+        return out, caches

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_(outputs, top_p, temperature, step_idx_ori, model_kwargs):
+        def _post_process_(logits, 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,7 +314,6 @@ 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)
@@ -373,7 +372,7 @@ class GenerationInferenceModel(GenerationMixin):
         outputs = _forward_(**model_kwargs)
         # first decoder
         next_tokens, model_kwargs = _post_process_(
-            outputs,
+            outputs[0] if isinstance(outputs, tuple) else outputs,
             top_p,
             temperature,
             step_idx_ori,
@@ -389,8 +388,9 @@ 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_(
-                _forward_(**model_kwargs),
+                outputs[0] if isinstance(outputs, tuple) else outputs,
                 top_p,
                 temperature,
                 step_idx_ori,
@@ -692,7 +692,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             return self(**model_inputs)
 
         def _post_process_(
-            outputs,
+            logits,
             top_k,
             top_p,
             penalty_score,
@@ -702,7 +702,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             model_kwargs,
         ):
             step_idx = model_kwargs["step_idx"]
-            logits = paddle.cast(outputs, paddle.float32)
+            logits = paddle.cast(logits, 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,
+            outputs[0] if isinstance(outputs, tuple) else outputs,
             top_k,
             top_p,
             penalty_score,
@@ -806,7 +806,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             return self(**model_inputs)
 
         def _post_process_(
-            outputs,
+            logits,
             top_k,
             top_p,
             penalty_score,
@@ -816,7 +816,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             model_kwargs,
         ):
             step_idx = model_kwargs["step_idx"]
-            logits = paddle.cast(outputs, paddle.float32)
+            logits = paddle.cast(logits, 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_(
-            outputs,
+            logits,
             top_k,
             top_p,
             penalty_score,
@@ -968,7 +968,7 @@ class GenerationBlockInferenceModel(GenerationMixin):
             temperature,
             model_kwargs,
         ):
-            logits = paddle.cast(outputs, paddle.float32)
+            logits = paddle.cast(logits, 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_(outputs, top_p, temperature, step_idx_ori, model_kwargs):
+        def _post_process_(logits, 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,7 +1211,6 @@ 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

@@ -34,6 +34,7 @@ from paddlenlp.experimental.transformers.fused_transformer_layers import (
     FusedBlockMultiTransformer,
     FusedBlockMultiTransformerA8W8,
     FusedBlockMultiTransformerFP8,
+    FusedBlockMultiTransformerHPU,
     FusedBlockMultiTransformerWeightOnly,
     FusedMultiTransformerA8W8,
     FusedMultiTransformerAvx,
@@ -149,13 +150,15 @@ 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)
         ]
@@ -164,19 +167,22 @@ 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)
         ]
@@ -256,19 +262,12 @@ 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,
@@ -279,11 +278,7 @@ class LlamaAvxInferenceModel(LlamaPretrainedModel):
             )
         hidden_states = self.norm(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)
+        return hidden_states
 
     @paddle.no_grad()
     # avx
@@ -338,7 +333,8 @@ 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[
@@ -346,7 +342,11 @@ 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,7 +400,6 @@ 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
@@ -486,13 +485,15 @@ 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)
         ]
@@ -502,26 +503,30 @@ 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)
         ]
@@ -616,7 +621,6 @@ 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(
@@ -796,7 +800,12 @@ 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():
@@ -812,7 +821,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)
 
@@ -843,14 +852,18 @@ 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(
@@ -881,13 +894,17 @@ 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,
@@ -968,7 +985,8 @@ 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,
@@ -1107,7 +1125,9 @@ 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())
@@ -1125,7 +1145,9 @@ 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())
@@ -1172,7 +1194,8 @@ 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[
@@ -1180,13 +1203,19 @@ 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())
@@ -1213,7 +1242,10 @@ 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)
@@ -1223,7 +1255,9 @@ 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())
@@ -1243,7 +1277,8 @@ 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:
@@ -1309,7 +1344,8 @@ 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()
@@ -1327,10 +1363,12 @@ 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[
@@ -1368,7 +1406,11 @@ 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))
@@ -1406,6 +1448,8 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
             self.transformer_block = FusedBlockMultiTransformerA8W8(transformer_config)
         elif "fp8" in self.quant_type:
             self.transformer_block = FusedBlockMultiTransformerFP8(transformer_config)
+        elif paddle.is_compiled_with_custom_device("intel_hpu"):
+            self.transformer_block = FusedBlockMultiTransformerHPU(transformer_config)
         else:
             self.transformer_block = FusedBlockMultiTransformer(transformer_config)
 
@@ -1414,10 +1458,23 @@ 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
+            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,
         )
-        return ids_remove_padding, padding_offset, cum_offsets, cu_seqlens_q, cu_seqlens_k
 
     def forward(
         self,
@@ -1433,29 +1490,75 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
     ):
         seq_lens_this_time = kwargs.get("seq_lens_this_time", None)
         rope_emb = kwargs.get("rope_emb", None)
-        draft_tokens = kwargs.get("draft_tokens", None)
-        seq_lens_encoder = kwargs.get("seq_lens_encoder", None)
 
-        # 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
+        if paddle.is_compiled_with_custom_device("intel_hpu"):
+            from paddlenlp_ops import prepare_input_hpu
+
+            block_tables = kwargs.get("block_tables", None).to("CPU")
+            seq_lens_encoder = kwargs.get("seq_lens_encoder", None).to("CPU")
+            seq_lens_decoder = kwargs.get("seq_lens_decoder", None).to("CPU")
+            input_ids = input_ids.to("CPU")
+
+            (
+                ids_remove_padding,
+                rope_emb,
+                block_groups,
+                block_list,
+                block_indices,
+                block_offsets,
+                block_mapping,
+                attention_mask,
+                valid_seq_len,
+            ) = prepare_input_hpu(
+                input_ids,
+                rope_emb,
+                block_tables,
+                self.block_size,
+                seq_lens_encoder,
+                seq_lens_decoder,
+                paddle.get_default_dtype(),
             )
+            cum_offsets = None
+            kwargs["block_groups"] = block_groups
+            kwargs["block_list"] = block_list
+            kwargs["block_indices"] = block_indices
+            kwargs["block_offsets"] = block_offsets
+            kwargs["block_mapping"] = block_mapping
+            kwargs["block_bias"] = attention_mask
+            kwargs["block_size"] = self.block_size
+            kwargs["valid_seq_len"] = valid_seq_len
         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
-            )
+            draft_tokens = kwargs.get("draft_tokens", None)
+            seq_lens_encoder = kwargs.get("seq_lens_encoder", None)
 
-        kwargs["cu_seqlens_q"] = cu_seqlens_q
-        kwargs["cu_seqlens_k"] = cu_seqlens_k
-        kwargs["padding_offsets"] = padding_offset
-        kwargs["max_input_length"] = self.max_seq_len
-        kwargs["block_size"] = self.block_size
+            # 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)
+            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)
+
+            kwargs["cu_seqlens_q"] = cu_seqlens_q
+            kwargs["cu_seqlens_k"] = cu_seqlens_k
+            kwargs["padding_offsets"] = padding_offset
+            kwargs["max_input_length"] = self.max_seq_len
+            kwargs["block_size"] = self.block_size
 
         inputs_embeds = self.embed_tokens(ids_remove_padding)
 
         with dy2st_nocheck_guard_context():
-            hidden_states, _ = self.transformer_block(
+            hidden_states, full_hidden_states = self.transformer_block(
                 input_ids=input_ids,
                 src=inputs_embeds,
                 cum_offsets=cum_offsets,
@@ -1467,17 +1570,11 @@ class LlamaBlockInferenceModel(LlamaInferenceModel):
             )
         hidden_states = self.norm(hidden_states)
 
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=None,
-            attentions=None,
-            cum_offsets=cum_offsets,
-        )
+        return hidden_states, full_hidden_states
 
 
 @register_base_model
-class EagleForLlamaInferenceModel(LlamaBlockInferenceModel):
+class EagleForLlamaBlockInferenceModel(LlamaBlockInferenceModel):
     def __init__(self, config: LlamaConfig):
         self.append_attn = config.append_attn
         super().__init__(config)
@@ -1487,7 +1584,11 @@ class EagleForLlamaInferenceModel(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)
@@ -1647,7 +1748,11 @@ 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)
@@ -1845,7 +1950,11 @@ 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)
@@ -1958,15 +2067,20 @@ class LlamaForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, LlamaPr
         else:
             max_block_nums = max_batch_size * max_block_per_seq
 
+        cache_kv_shape = [
+            max_block_nums,
+            config.num_key_value_heads // max(config.tensor_parallel_degree, 1),
+            config.block_size,
+            config.hidden_size // config.num_attention_heads,
+        ]
+        if paddle.is_compiled_with_custom_device("intel_hpu"):
+            # HPU block multi-transformer
+            # Use KV Cache shape [max_block_nums, seq_len, num_head, head_dim]
+            cache_kv_shape = [cache_kv_shape[i] for i in [0, 2, 1, 3]]
+
         cache_k_shapes = []
         cache_v_shapes = []
         for _ in range(config.num_hidden_layers):
-            cache_kv_shape = [
-                max_block_nums,
-                config.num_key_value_heads // max(config.tensor_parallel_degree, 1),
-                config.block_size,
-                config.hidden_size // config.num_attention_heads,
-            ]
             cache_k_shapes.append(cache_kv_shape)
             cache_v_shapes.append(cache_kv_shape)
         return cache_k_shapes, cache_v_shapes
@@ -2032,8 +2146,9 @@ class LlamaForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, LlamaPr
         excess_blocks=None,
         draft_tokens=None,
         output_padding_offset=None,
+        **kwargs,
     ):
-        outputs = self.llama(
+        hidden_states, full_hidden_states = self.llama(
             input_ids,
             src_mask=src_mask,
             caches=caches,
@@ -2050,24 +2165,9 @@ 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,
@@ -2093,9 +2193,13 @@ class EagleLlamaForCausalLMBlockInferenceModel(LlamaForCausalLMBlockInferenceMod
         self.verify_window = config.get("speculate_verify_window", 2)
         self.max_seq_len = config.max_seq_len
 
-        self.eagle = EagleForLlamaInferenceModel(config)
+        self.eagle = EagleForLlamaBlockInferenceModel(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)