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)$

README.md

@@ -33,17 +33,20 @@
 
 ## News 📢
 
-* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成 DeepSeek-R1-MTP 单机部署任务、提交高质量测评 blog，即可实时赢取奖金！💰💰💰
-报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+* **2025.04.29 PaddleNLP 现已支持 Qwen3 系列模型**: Qwen3 系列模型支持持两种思考模式，预训练约 36 万亿个 token、119 种语言和方言。包括六个 Dense 模型, Qwen3-32B、Qwen3-14B、Qwen3-8B、Qwen3-4B、Qwen3-1.7B 和 Qwen3-0.6B。两个 MoE 模型的权重：Qwen3-235B-A22B，Qwen3-30B-A3B。
 
 * **2025.03.12 [PaddleNLP v3.0 Beta4](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta4)**：全面支持 DeepSeek V3/R1/R1-Distill, 及 QwQ-32B 等热门思考模型。**DeepSeek V3/R1完整版支持 FP8、INT8、4-bit 量化推理，MTP 投机解码**。单机 FP8推理输出超**1000 tokens/s**; 4-bit 推理输出超**2100 tokens/s**! 发布新版推理部署镜像，热门模型[一键部署](https://paddlenlp.readthedocs.io/zh/latest/llm/server/docs/general_model_inference.html)。推理部署[使用文档](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/predict/index.html)全面更新，体验全面提升！自研下一代通用信息抽取模型 PP-UIE [全新发布](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction)，支持8K 长度信息抽取。新增大模型 Embedding 训练，支持 INF-CL 超大 batch size 训练。新增[MergeKit](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/mergekit.html)模型融合工具，缓解对齐代价。低资源训练全面优化，16G 小显存可以流畅训练。
 
-* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
 
 * **2025.02.10 PaddleNLP 现已支持 DeepSeek-R1系列模型，[在线使用](https://aistudio.baidu.com/projectdetail/8775758)**：依托全新的 PaddleNLP 3.0套件，DeepSeek-R1系列模型现已全面支持。凭借数据并行、数据分组切分并行、模型并行、流水线并行以及专家并行等一系列先进的分布式训练能力，结合 Paddle 框架独有的列稀疏注意力掩码表示技术——FlashMask 方法，DeepSeek-R1系列模型在训练过程中显著降低了显存消耗，同时取得了卓越的训练性能提升。
 
 <details><summary> <b>点击展开</b> </summary><div>
 
+* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成 DeepSeek-R1-MTP 单机部署任务、提交高质量测评 blog，即可实时赢取奖金！💰💰💰
+报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+
+* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
+
 * **2025.02.20 🔥🔥《PP-UIE 信息抽取智能引擎全新升级》** 强化零样本学习能力，支持极少甚至零标注数据实现高效冷启动与迁移学习，显著降低数据标注成本；具备处理长文本能力，支持 8192 个 Token 长度文档信息抽取，实现跨段落识别关键信息，形成完整理解；提供完整可定制化的训练和推理全流程，训练效率相较于 LLama-Factory 实现了1.8倍的提升。
 2月26日（周三）19：00为您深度解析全新 PP-UIE 技术方案及在部署方面的功能、优势与技巧。报名链接：https://www.wjx.top/vm/mBKC6pb.aspx?udsid=606418
 
@@ -119,6 +122,7 @@
 |           [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                                                                                                                                                                                                                                                                              |
+|           [Qwen3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/Qwen3-0.6B, Qwen/Qwen3-1.7B, Qwen/Qwen3-4B, Qwen/Qwen3-8B, Qwen/Qwen3-14B, Qwen/Qwen3-32B, Qwen/Qwen3-30B-A3B, Qwen/Qwen3-235B-A22B, Qwen/Qwen3-0.6B-Base, Qwen/Qwen3-1.7B-Base, Qwen/Qwen3-4B-Base, Qwen/Qwen3-8B-Base, Qwen/Qwen3-14B-Base, Qwen/Qwen3-30B-A3B-Base                                                                                                                    |
 |             [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                                                                                                                                                                                                                                                                                                                                    |
 

csrc/gpu/all_reduce.cu

@@ -72,12 +72,14 @@ void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
                           reinterpret_cast<half*>(out.data()), out.numel());
       break;
     }
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
     case phi::DataType::BFLOAT16: {
       fa->allreduce<nv_bfloat16>(
           stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
           reinterpret_cast<nv_bfloat16*>(out.data()), out.numel());
       break;
     }
+#endif
     default:
       throw std::runtime_error(
           "custom allreduce only supports float32, float16 and bfloat16");

csrc/gpu/all_reduce.cuh

@@ -98,7 +98,7 @@ DINLINE half& assign_add(half& a, half b) {
 }
 DINLINE float& assign_add(float& a, float b) { return a += b; }
 
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
 DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
 template <>
 DINLINE nv_bfloat16 downcast_s(float val) {

csrc/gpu/cpp_extensions.cu

@@ -238,11 +238,9 @@ std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& msg_queue_id,      // cpu
                  int64_t rank_id);
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& msg_queue_id, // cpu
                int64_t rank_id,
                bool wait_flag);
 

csrc/gpu/get_output.cc

@@ -20,6 +20,7 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define GET_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -27,14 +28,25 @@ struct msgdata {
 };
 
 void GetOutput(const paddle::Tensor& x,
-               const paddle::Tensor& msg_queue_id,
                int64_t rank_id,
                bool wait_flag) {
   if (rank_id > 0) return;
 
   static struct msgdata msg_rcv;
-  int queue_id_val = msg_queue_id.data<int>()[0];
-  static key_t key = ftok("./", queue_id_val);
+  int msg_queue_id = 1;
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+      std::string inference_msg_queue_id_env_str(
+          inference_msg_queue_id_env_p);
+      int inference_msg_queue_id_from_env =
+          std::stoi(inference_msg_queue_id_env_str);
+#ifdef GET_OUTPUT_DEBUG
+      std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                << inference_msg_queue_id_from_env << std::endl;
+#endif
+      msg_queue_id = inference_msg_queue_id_from_env;
+  }
+  static key_t key = ftok("./", msg_queue_id);
 
   static int msgid = msgget(key, IPC_CREAT | 0666);
 
@@ -62,7 +74,7 @@ void GetOutput(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(get_output)
-    .Inputs({"x", "msg_queue_id"})
+    .Inputs({"x"})
     .Attrs({"rank_id: int64_t",
             "wait_flag: bool"})
     .Outputs({"x_out"})

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/multi_head_latent_attention.cu

@@ -205,6 +205,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
   meta_data.batch_size = cum_offsets.dims()[0];
 
   switch (query.dtype()) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
     case paddle::DataType::BFLOAT16: {
       return MultiHeadLatentAttentionKernel<paddle::DataType::BFLOAT16>(
           meta_data,
@@ -253,6 +254,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           causal,
           speculate_decoder);
     }
+#endif
     case paddle::DataType::FLOAT16: {
       return MultiHeadLatentAttentionKernel<paddle::DataType::FLOAT16>(
           meta_data,

csrc/gpu/save_with_output_msg.cc

@@ -20,6 +20,7 @@
 #include "paddle/extension.h"
 
 #define MAX_BSZ 512
+// #define SAVE_WITH_OUTPUT_DEBUG
 
 struct msgdata {
     long mtype;
@@ -28,7 +29,6 @@ struct msgdata {
 
 void SaveOutMmsg(const paddle::Tensor& x,
                  const paddle::Tensor& not_need_stop, // cpu
-                 const paddle::Tensor& msg_queue_id,      // cpu
                  int64_t rank_id) {
     if (rank_id > 0) return;
     auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
@@ -36,8 +36,25 @@ void SaveOutMmsg(const paddle::Tensor& x,
     auto not_need_stop_data = not_need_stop.data<bool>()[0];
 
     static struct msgdata msg_sed;
-    int queue_id_val = msg_queue_id.data<int>()[0];
-    static key_t key = ftok("./", queue_id_val);
+    int msg_queue_id = 1;
+    if (const char* inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+        std::string inference_msg_queue_id_env_str(
+            inference_msg_queue_id_env_p);
+        int inference_msg_queue_id_from_env =
+            std::stoi(inference_msg_queue_id_env_str);
+        msg_queue_id = inference_msg_queue_id_from_env;
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+                  << inference_msg_queue_id_from_env << std::endl;
+#endif
+    } else {
+#ifdef SAVE_WITH_OUTPUT_DEBUG
+        std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                  << std::endl;
+#endif
+    }
+    static key_t key = ftok("./", msg_queue_id);
     static int msgid = msgget(key, IPC_CREAT | 0666);
 
     msg_sed.mtype = 1;
@@ -54,7 +71,7 @@ void SaveOutMmsg(const paddle::Tensor& x,
 }
 
 PD_BUILD_OP(save_output)
-    .Inputs({"x", "not_need_stop", "msg_queue_id"})
+    .Inputs({"x", "not_need_stop"})
     .Attrs({"rank_id: int64_t"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})

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);

csrc/setup_cuda.py

@@ -130,14 +130,11 @@ sources = [
     "./gpu/speculate_decoding_kernels/speculate_save_output.cc",
     "./gpu/speculate_decoding_kernels/speculate_get_output.cc",
     "./gpu/save_output_dygraph.cu",
-    "./gpu/cpp_extensions.cu",
     "./gpu/all_reduce.cu",
     "./gpu/quantization/per_token_group_quant.cu",
     "./gpu/quantization/per_tensor_quant_fp8.cu",
 ]
 sources += find_end_files("./gpu/speculate_decoding_kernels", ".cu")
-sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
-sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
 
 nvcc_compile_args = gencode_flags
 update_git_submodule()
@@ -174,6 +171,9 @@ if cc >= 80:
 
     sources += find_end_files("./gpu/append_attn", ".cu")
     sources += find_end_files("./gpu/append_attn/template_instantiation", ".cu")
+    sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
+    sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
+    sources += "./gpu/cpp_extensions.cu",
 
 
 fp8_auto_gen_directory = "gpu/cutlass_kernels/fp8_gemm_fused/autogen"

csrc/tools/build_wheel.sh

@@ -61,7 +61,7 @@ function generate_sm_version(){
         sm_versions=($SM_VERSION )
     elif [ "$ARCHITECTURE" = "all" ]; then
         if awk -v version="$cuda_version" 'BEGIN { exit !(version >= 12.0) }'; then
-          sm_versions=(70 75 80 80 86 89 90 )
+          sm_versions=(70 75 80 86 89 90 )
         else
           sm_versions=(70 75 80 86 89 ) 
         fi 
@@ -72,10 +72,12 @@ function generate_sm_version(){
 }
 
 function create_directories(){
-  mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops
-  touch $OPS_SRC_DIR/tmp/setup.py
-  touch $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
-  echo '# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+  for sm_version in "${sm_versions[@]}"; do
+    echo "create sm$sm_version"
+    mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops
+    touch $OPS_SRC_DIR/tmp/setup.py
+    touch $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
+    echo '# 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.
@@ -93,6 +95,7 @@ function create_directories(){
 
 import os
 from datetime import datetime
+import paddle
 
 from setuptools import find_packages, setup
 
@@ -109,14 +112,19 @@ def read(file: str):
         content = f.read().strip()
     return content
 
-
 def read_version():
     """
     read version and return content
     """
     __version__ = "3.0.0b4.post"
+
     formatted_date = datetime.now().date().strftime("%Y%m%d")
-    __version__ = __version__.replace(".post", ".post{}".format(formatted_date))
+    cuda_version = float(paddle.version.cuda())
+    sm_version=80
+    paddle_commit = paddle.__git_commit__[:7]
+    build_tag = "{}+cuda{}sm{}paddle{}".format(formatted_date, cuda_version, sm_version, paddle_commit)
+
+    __version__ = __version__.replace(".post", ".post{}".format(build_tag))
     
     return __version__
 
@@ -184,12 +192,9 @@ try:
 except ImportError:
     logger.WARNING(f"No {module_name} ")
 ' > $OPS_SRC_DIR/tmp/paddlenlp_ops/__init__.py
-
-  for sm_version in "${sm_versions[@]}"; do
-        echo "create sm$sm_version"
-        mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}
-        touch $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
-        echo '# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+    mkdir -p $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}
+    touch $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
+    echo '# 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.
@@ -210,6 +215,7 @@ try:
 except ImportError:
     logger.WARNING("No paddlenlp_ops_'${sm_version}' ops")
 ' > $OPS_SRC_DIR/tmp/paddlenlp_ops/sm${sm_version}/__init__.py
+    build_ops
   done
 }
 
@@ -228,11 +234,11 @@ function init() {
 }
 
 function build_ops() {
-    for sm_version in "${sm_versions[@]}"; do
-        echo "Building and installing for sm_version: $sm_version"
-        build_and_install_ops $sm_version
-    done
-    return 
+    echo "Building and installing for sm_version: $sm_version"
+    build_and_install_ops $sm_version
+    build_and_install_whl
+    unittest
+    cleanup
 }
 
 function copy_ops(){
@@ -269,6 +275,7 @@ function build_and_install_whl() {
   echo -e "${BLUE}[build]${NONE} building paddlenlp_ops wheel..."
   rm -rf ./dist
   cd ${TMP_DIR}
+  sed -i "s/sm_version=80/sm_version=${sm_version}/g" setup.py
   ${python} setup.py bdist_wheel --dist-dir ./$DIST_DIR
   if [ $? -ne 0 ]; then
     echo -e "${RED}[FAIL]${NONE} build paddlenlp_ops wheel failed !"
@@ -286,7 +293,8 @@ function build_and_install_whl() {
   fi
   echo -e "${BLUE}[install]${NONE} ${GREEN}paddlenlp_ops install success\n"
   cd ..
-  mv $DIST_DIR ../
+  mkdir -p ../$DIST_DIR
+  mv $DIST_DIR/* ../$DIST_DIR/
   cd ..
 }
 
@@ -321,10 +329,6 @@ trap 'abort' 0
 set -e
 
 init
-build_ops
-build_and_install_whl
-unittest
-cleanup
 
 # get Paddle version
 PADDLE_VERSION=`${python} -c "import paddle; print(paddle.version.full_version)"`

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/alignment/rl/README.md

@@ -0,0 +1 @@
+../../../../../llm/alignment/rl/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#模型支持)
 
 ## 🚀 快速开始 🚀
 
@@ -186,7 +168,7 @@ python run_finetune.py ./config/qwen/pt_argument.json
 
 ### 3. 对齐
 
-我们支持 DPO、KTO、RLHF 等偏好对齐策略。DPO、KTO 策略采用 zero_padding 策略，结合 FlashMask 策略，有效提升模型训练效率。
+我们支持 DPO、KTO、RL 等偏好对齐策略。DPO、KTO 策略采用 zero_padding 策略，结合 FlashMask 策略，有效提升模型训练效率。
 
 #### 3.1 DPO
 
@@ -272,8 +254,8 @@ python -u  ./alignment/dpo/run_dpo.py ./config/llama/dpo_lora_argument.json
 为了方便测试，我们也提供了偏好数据集可以直接使用：
 
 ```bash
-wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized_pointwise.tar.gz
-tar -zxvf ultrafeedback_binarized_pointwise.tar.gz
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/ultrafeedback_binarized_pointwise.tar
+tar -xvf ultrafeedback_binarized_pointwise.tar.gz
 ```
 
 ##### 全参 KTO
@@ -289,20 +271,18 @@ python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./alignment/
 python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./alignment/kto/run_kto.py ./config/llama/kto_lora_argument.json
 ```
 
-#### 3.3 RLHF
+#### 3.3 RL
 
-飞桨大模型套件提供了提供了基于强化学习 PPO 算法对 LLM 进行人类偏好对齐的代码及完整使用示例，支持**3D 分布式并行训练以及 rollout 阶段使用预测优化进行生成加速**。详细使用教程详见[RLHF 文档](./docs/rlhf.md)。
+飞桨大模型套件提供了提供了基于强化学习 GRPO、Reinforce++、PPO 等 算法对 LLM 进行人类偏好对齐的代码及完整使用示例，支持**3D 分布式并行训练以及 rollout 阶段使用预测优化进行生成加速**。详细使用教程详见[RL 文档](./alignment/rl/README.md)。
 
 ### 4. 模型融合
 PadlleNLP 支持多种模型融合方法，包括**Linear、Slerp、Ties、DARE、DELLA**，并支持模型参数稀疏化方法与模型融合算法的灵活组合使用。
 ```shell
 # 模型融合启动命令参考
-# cd PaddleNLP/llm/tools
-python mergekit.py \
-    --tensor_type np \
-    --n_process 2 \
+python -u  -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" ./llm/tools/mergekit.py \
+    --tensor_type pd \
     --merge_method linear \
-    --model_path_list ../checkpoints/model1 ../checkpoints/model \
+    --model_path_str "../checkpoints/model1,../checkpoints/model2" \
     --output_path ../checkpoints/model_merge
 
 ```

llm/alignment/ppo/README.md

@@ -1,141 +0,0 @@
-# PPO && GRPO
-
-PPO（Proximal Policy Optimization，近端策略优化）是一种强化学习算法，旨在通过优化策略来最大化累积奖励。PPO 算法结合了 Policy Gradient 和‌TRPO 的优点，通过使用随机梯度上升优化一个“替代”目标函数，实现小批量更新，而不是每个数据样本只进行一次梯度更新。
-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），以此作为优化策略的依据。
-
-以下是详细的使用文档和示例：
-
-## 环境依赖
-
-* 训练环境：
-1. 参考 Paddle 官网安装 PaddlePaddle-GPU
-2. clone 并安装 PaddleNLP
-```shell
-git clone https://github.com/PaddlePaddle/PaddleNLP.git
-```
-3. 安装 paddlenlp_ops，参考 PaddleNLP/csrc 进行安装（必需）
-```shell
-cd your_PaddleNLP_path/csrc
-python setup_cuda.py install
-```
-
-## 支持模型
-
-|   模型系列    | 模型名称                                                                                                                                                                                                                                                                      |
-|:-------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-|   Llama3.1    | 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    | 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                                                                                                                                                          |
-|    Qwen1.5    | 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                          |
-|     Qwen2     | 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   | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct                                                                                                                                                                          |
-|    Qwen2.5    | 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-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, |
-| Qwen2.5-Math  | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct                                                                                                                                                                  |
-| Qwen2.5-Coder | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                              |
-
-## 数据协议
-
-### 字段说明
-
-- src (list(str)): 用户对话内容，可能会包含 markup 内容，如 [<search-res>]；
-- tgt (list(str)): 除了最后一轮的系统多轮回复内容，以对话轮次排列，可能会包含 markup 内容，如 [<search>]；注意：len(tgt)==len(src)-1
-
-### 数据示例
-
-```json
-{
-    "src": [
-        "需要你帮我写几个有创意的广告语来打开市场。",
-        "目标用户是年轻人，追求时尚、个性和自我。"
-    ],
-    "tgt": [
-        "当然！我很乐意帮助你创作几个有创意的广告语来推广你的新洗发露。请告诉我一些关于你的产品的特点，目标受众以及你希望传达的核心信息，我会根据这些信息为你提供几个创意的广告语。"
-    ]
-}
-```
-
-
-### PPO & GRPO 数据准备
-
-```
-wget https://paddlenlp.bj.bcebos.com/datasets/examples/ppo-kk.tgz && tar zxf ppo-kk.tgz
-```
-
-## 训练
-
-### 训练配置
-
-我们采用的配置文件在放置在`llm/config/llama/ppo_argument.json`和`llm/config/llama/grpo_argument.json`中，同时我们提供了详细参数释义如下：
-
-- `train_task_config`: 训练数据 config, 请以`config/task_ppo.json`为例
-- `eval_task_config`: 评估数据 config, 请以`config/task_ppo.json`为例
-- `ptx_task_config`: SFT 辅助数据, 请以`config/task_sft.json`为例，默认为""
-- `actor_model_name_or_path`: PPO 中 actor-model 和 reference-model 模型本地的模型路径
-- `reward_model_name_or_path`: PPO 中 reward-model 和 critic-model 模型本地的模型路径
-- `use_fusemt`: 是否通过 FustMT 加速生成，默认为 True
-- `use_flash_attention`: 是否启用 FlashAttention-2，默认为 False
-- `output_dir`: 模型参数保存目录
-- `max_seq_len`: 输入数据的最大长度，默认为 4096
-- `max_dec_len`: 最大生成长度
-- `min_dec_len`: 最小生成长度
-- `top_p`: 生成解码超参数
-- `temperature`: 生成解码超参数
-- `repetition_penalty`: 生成解码超参数
-- `rollout_n`: 生成 response 的数量
-- `min_learning_rate`: Actor 模型的最小学习率
-- `critic_learning_rate`: Critic 模型的最小学习率
-- `recompute`: Actor 模型是否使用重计算策略，开启后可节省训练显存
-- `critic_recompute`: Critic 模型是否使用重计算策略，开启后可节省训练显存
-- `recompute_granularity` Actor 模型的重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
-- `critic_recompute_granularity` Critic 模型重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
-- `warmup_ratio`: Actor 模型用于从 0 到 `learning_rate` 的线性 warmup 的总训练步骤的比例
-- `critic_warmup_ratio`: Critic 模型用于从 0 到 `critic_learning_rate` 的线性 warmup 的总训练步骤的比例
-- `lr_scheduler_type`: Actor 模型要使用的学习率调度策略。 (`str`, 可选, 默认为 `"linear"`)
-- `critic_lr_scheduler_type`: Critic 模型要使用的学习率调度策略。 (`str`, 可选, 默认为 `"linear"`)
-- `weight_decay`: Actor 模型除了所有 bias 和 LayerNorm 权重之外，应用于所有层的权重衰减数值。（`float`，可选，默认为 0.0）
-- `critic_weight_decay`: Critic 模型除了所有 bias 和 LayerNorm 权重之外，应用于所有层的权重衰减数值。（`float`，可选，默认为 0.0）
-- `max_prompt_len`: 生成样本时的最大生成长度， max_length 调大会增加生成时间，并且增加显存占用。注意：
-max_dec_len + max_prompt_len 应当小于 max_seq_len。
-- `per_device_prompt_batch_size`: PPO 生成样本时的批处理大小，同 micro batch size，即满足 global_batch_size = dp（data parallel）* sharding * micro batch size。batch_size 调大会增加生成时间，并且增加显存占用
-- `per_device_train_batch_size`: 训练 batch 大小, 当前为了优化性能设为1，请避免更改
-- `per_device_eval_batch_size`: 评估 batch 大小。
-- `max_steps`: 总的训练步数
-- `eval_steps`: 模型评估的间隔步数
-- `max_evaluate_steps`: 模型单次评估的最大步数
-- `logging_steps`: 训练日志打印的间隔步数
-- `save_steps`: 模型参数保存的间隔步数
-- `weight_decay`: 权重衰减数值
-- `do_train`: 是否进行训练任务
-- `do_eval`: 是否进行评估任务
-- `fp16`: 使用 float16 精度进行模型训练和推理。
-- `bf16`: 使用 bfloat16 精度进行模型训练和推理。
-- `fp16_opt_level`: float16 精度训练模式，`O2`表示纯 float16 训练
-
-
-<!-- ### PPO 训练命令
-
-```shell
-python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7"  run_ppo.py llm/config/llama/ppo_argument.json
-``` -->
-
-### GRPO 训练命令
-```shell
-cd your_PaddleNLP_path/llm/alignment/ppo
-```
-
-```shell
-# 启动 reward server
-python reward_server.py
-```
-
-```shell
-export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
-export PYTHONPATH=your_PaddleNLP_path/llm:$PYTHONPATH
-python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_ppo.py ../../config/qwen/grpo_argument.json
-# python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_ppo.py ../../config/llama/grpo_argument.json
-```
-
-### 在线监控
-在`grpo_argument.json`中设置的输出目录为`"logging_dir": "vdl_log"`, 可以通过以下命令查看训练过程
-```shell
-visualdl --logdir vdl_log --host 0.0.0.0
-```

llm/alignment/rl/README.md

@@ -0,0 +1,214 @@
+# 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），以此作为优化策略的依据。
+REINFORCE++ 是经典 REINFORCE 算法的改进版本，通过融合 PPO 的关键优化技术并移除 Critic Model，实现了更加简洁高效的策略优化。相比于传统的 REINFORCE，REINFORCE++ 在 Token-Level KL 惩罚、PPO-Clip、优势标准化、奖励裁剪和奖励标准化等关键技术上进行了改进，从而提高了训练过程的效率和稳定性。
+
+以下是详细的使用文档和示例：
+
+## 环境依赖
+
+* 训练环境：
+1. 参考 [Paddle 官网](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/develop/install/pip/linux-pip.html)安装 PaddlePaddle-GPU, 要求 PaddlePaddle>=3.0
+2. clone 并安装 PaddleNLP
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git
+```
+3. 安装 paddlenlp_ops 推理算子，参考 PaddleNLP/csrc 进行安装（必需）
+```shell
+cd your_PaddleNLP_path/csrc
+python setup_cuda.py install
+```
+4. 安装 fused_ln 和 fast_ln 训练算子，参考 PaddleNLP/slm/model_zoo/gpt-3/external_ops (必须)
+```shell
+cd your_PaddleNLP_path/slm/model_zoo/gpt-3/external_ops
+python setup.py install
+```
+
+## 支持模型
+
+|   模型系列    | 模型名称                                                                                                                                                                                                                                                                      |
+|:-------------:|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+|    Qwen1.5    | 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                          |
+|     Qwen2     | 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   | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct                                                                                                                                                                          |
+|    Qwen2.5    | 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-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, |
+| Qwen2.5-Math  | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct                                                                                                                                                                  |
+| Qwen2.5-Coder | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                              |
+
+## 数据协议
+
+### 字段说明
+
+- src (list(str)): 经过 chat_template 处理后的 prompt 输入；或者根据需要自己拼接构造 prompt；
+- tgt (list(str)): 标签内容；
+
+### 数据示例
+
+```json
+{
+    "src": ["<|im_start|>system\nYou are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and<answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>.  Now the user asks you to solve a logical reasoning problem. After thinking, when you finally reach a conclusion, clearly state the identity of each character within <answer> </answer> tags. i.e., <answer> (1) Zoey is a knight\n(2) ... </answer>.\n<|im_end|>\n<|im_start|>user\nA very special island is inhabited only by knights and knaves. Knights always tell the truth, and knaves always lie. You meet 3 inhabitants: Michael, Zoey, and Ethan. Michael was heard saying, \"Ethan is a knight if and only if Michael is a knight\". \"Zoey is a knight or Ethan is a knight,\" Zoey mentioned. Ethan asserted: \"Michael is a knave if and only if Zoey is a knave\". So who is a knight and who is a knave?\n<|im_end|>\n<|im_start|>assistant\n<think>"],
+    "tgt": ["(1) Michael is a knight\n(2) Zoey is a knight\n(3) Ethan is a knight"]
+}
+```
+
+
+### GRPO & REINFORCE++ 数据准备
+我们提供了一版使用 `Qwen/Qwen2.5-7B-Instruct-1M` 的`chat template`预处理后的[KK 数据集](https://hf-mirror.com/datasets/K-and-K/knights-and-knaves)。
+```
+wget https://paddlenlp.bj.bcebos.com/datasets/examples/ppo-kk.tgz && tar zxf ppo-kk.tgz
+```
+
+## 训练
+
+### GRPO && REINFORCE++ 训练配置
+
+我们采用的配置文件放置在`llm/config/qwen/grpo_argument.yaml`中，同时我们提供了详细参数释义如下：
+- `rl_algorithm`: 使用的强化学习算法，支持`grpo`、`reinforce_plus_plus`
+- `actor_model_name_or_path`: actor-model 和 reference-model 模型本地的模型路径
+- `reward_model_name_or_path`: reward 模型的名称或本地路径
+- `use_rm_server`: 是否使用 reward model server，设置为`False`时需要提供`reward_model_name_or_path`
+- `reward_server`: reward model server 的 URL 地址, 比如`http://127.0.0.1:8731`
+- `logging_dir`: 日志保存的文件夹
+- `logging_steps`: 训练日志打印的间隔步数
+- `output_dir`: 模型参数保存目录
+- `report_to`: 训练可视化工具，支持 "all"、"wandb"、"tensorboard"、"visualdl"、"none"
+- `wandb_http_proxy`: 连接 wandb 使用的 HTTP 代理
+- `run_name`: 实验名称
+- `train_datasets`: 训练集路径
+- `eval_datasets`: 验证集路径
+- `prompt_key`: 数据集中 query 对应的字段名
+- `response_key`: 数据集中 response 对应的字段名
+- `dataloader_drop_last`: dataloader 是否丢弃最后不完整的 batch
+- `balance_batch`: 该参数用于指定是否在数据并行场景下，对批次内的 token 数量进行均衡分配。若设置为 True，系统将尝试在不同并行设备间平衡 token 的分布；若设置为 False（默认值），则不进行此类均衡操作。
+- `use_remove_padding`: 此参数决定是否在训练过程中去除输入数据中的 padding 部分。启用该选项（设置为 True）可有效提高训练过程中有效 token 的占比，从而提升训练效率；若设置为 False（默认值），则保留输入数据中的 padding。
+- `tensor_parallel_degree`: 张量并行度
+- `sequence_parallel`: 是否启用序列并行
+- `sharding_parallel_degree`: sharding 并行度
+- `sharding`: 分片策略，支持 "stage1" 或 "stage2"
+- `sharding_parallel_config`: sharding 并行配置
+- `pipeline_parallel_degree`: 流水线并行度
+- `virtual_pp_degree`: 虚拟流水线并行度
+- `max_prompt_len`: 生成样本时的最大生成长度， max_length 调大会增加生成时间，并且增加显存占用。注意：
+max_dec_len + max_prompt_len 应当小于 max_seq_len。
+- `max_dec_len`: 最大生成长度
+- `min_dec_len`: 最小生成长度
+- `top_p`: 生成解码超参数
+- `temperature`: 生成解码超参数
+- `repetition_penalty`: 生成解码超参数
+- `rollout_max_num_seqs`: 单次推理可以处理的最大序列数
+- `rollout_quant_type`: 量化类型，例如 "weight_only_int8"
+- `seed`: 随机种子
+- `global_batch_size`: 一次（一个 step）推理（rollout)采样的 prompt 数量
+- `global_mini_batch_size`: actor model 更新一次参数训练的 prompt 数量
+- `rollout_n`: 一个 prompt 采样的 response 数量
+- `update_iters`: 同一批数据训练次数
+- `per_device_logprob_batch_size`: 计算 log_probs 时，一个 batch 的样本数量
+- `per_device_reward_batch_size`: critic model 计算 loss 与反向传播时，一个 batch 的的样本数量
+- `per_device_value_batch_size`: critic model 前向计算 values 时，一个 batch 的的样本数量
+- `per_device_train_batch_size`: actor model 计算 loss 与反向传播时，一个 batch 的样本数量
+- `num_train_epochs`: 训练的 epoch 数
+- `max_length`: 训练时的最大长度，应大于 `max_prompt_len` 和 `max_dec_len` 之和
+- `learning_rate`: 学习率
+- `lr_scheduler_type`: Actor 模型要使用的学习率调度策略。 (`str`, 可选, 默认为`linear`)
+- `weight_decay`: AdamW 优化器的权重衰减
+- `adam_beta1`: AdamW 优化器的 beta1
+- `adam_beta2`: AdamW 优化器的 beta2
+- `adam_epsilon`: AdamW 优化器的 epsilon
+- `max_grad_norm`: 梯度裁剪的最大值
+- `max_steps`: 总的训练步数
+- `save_steps`: 模型参数保存的间隔步数
+- `ignore_save_lr_and_optim`: 是否忽略保存学习率和优化器状态
+- `kl_coeff`: KL 惩罚系数
+- `kl_loss_coeff`: KL Loss 系数
+- `pg_loss_coeff`: 策略梯度损失系数
+- `entropy_coeff`: entropy loss 系数
+- `clip_range_ratio`: PPO-Clip 裁剪阈值
+- `clip_range_ratio_low`: PPO-Clip 裁剪下限阈值
+- `clip_range_ratio_high`: PPO-Clip 裁剪上限阈值
+- `clip_range_score`: reward 的剪切范围，reward 会被限制在 [-clip_range_score, clip_range_score] 范围内
+- `clip_range_value`: value 模型输出的剪切范围，value 会被限制在 [-clip_range_value, clip_range_value] 范围内
+- `normalize_reward`: 是否使用 reward 标准化
+- `normalize_advantage`: 是否使用 advantage 标准化
+- `use_fp32_compute`: 是否使用 fp32 来计算 log_prob、reward、advantage 和 loss
+- `do_eval`: 是否进行评估
+- `per_device_eval_batch_size`: 估 batch 大小
+- `evaluation_strategy`: 评估策略，例如 `steps`
+- `eval_steps`: 模型评估的间隔步数
+- `use_flash_attention`: 是否启用 FlashAttention-2，默认为 False
+- `use_fused_rms_norm`: 是否使用融合的 RMSNorm 算子，需安装 fused_ln
+- `recompute`: Actor 模型是否使用重计算策略，开启后可节省训练显存
+- `recompute_granularity`: Actor 模型的重计算的粒度，可选项为`core_attn`和`full`. `core_attn`速度快但是显存占用，`full`速度慢但是显存占用低
+- `bf16`: 使用 bfloat16 精度进行模型训练和推理。
+- `fp16_opt_level`: float16 精度训练模式，`O2`表示纯 float16 训练
+- `amp_custom_black_list`: 自定义 AMP 黑名单
+- `amp_custom_white_list`: 自定义 AMP 白名单
+
+
+
+### GRPO 训练命令
+```shell
+cd your_PaddleNLP_path/llm/alignment/rl
+```
+
+```shell
+# 启动 reward server
+python reward_server.py
+```
+
+```shell
+export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
+export PYTHONPATH=your_PaddleNLP_path/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_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+python -u -m paddle.distributed.launch --devices "0,1,2,3" run_rl.py ../../config/qwen/grpo_argument.yaml
+
+# QWEN32B 2k prompt + 30k response 9台8x80G 显卡训练命令如下：
+# python -u -m paddle.distributed.launch --devices "0,1,2,3,4,5,6,7" run_rl.py ../../config/qwen/grpo_32b_argument.yaml
+```
+我们提供根据上述脚本可复现的[wandb 日志](https://api.wandb.ai/links/junyu/5jiulhem)。
+
+
+### Reinforce++ 训练命令
+```shell
+cd your_PaddleNLP_path/llm/alignment/rl
+```
+
+```shell
+# 启动 reward server
+python reward_server.py
+```
+
+```shell
+export PYTHONPATH=your_PaddleNLP_path/:$PYTHONPATH
+export PYTHONPATH=your_PaddleNLP_path/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_mla_use_tensorcore=0
+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"`, 可以通过以下命令查看训练过程
+```shell
+visualdl --logdir vdl_log --host 0.0.0.0
+```
+
+也支持 wandb 等多种监控，可设置`"logging_dir": "wandb"`，需要提前安装好 wandb 依赖并登录。

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/reward/fake_reward_server.py

@@ -0,0 +1,112 @@
+# 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.
+
+"""Launch Reward HTTP Server."""
+
+import argparse
+import json
+import logging
+import threading
+import traceback
+from typing import List
+
+import uvicorn
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+
+class Request(BaseModel):
+    """The request for RM server."""
+
+    src: List[str]
+    tgt: List[str]
+    response: List[str]
+
+
+class Response(BaseModel):
+    """The response for RM server."""
+
+    error_code: int = 0
+    error_msg: str = "Success"
+    score: List[float] = None
+
+
+def compute_score(
+    solution_str: str, ground_truth: str, query=None, format_reward: int = 1, answer_reward: float = 1.0
+):
+    score = float(1.0)
+    print(
+        f"==============================================================={ground_truth}=========================================================================="
+    )
+    print(f"score {score}, solution_str\n", solution_str)
+    print(
+        "================================================================================================================================================="
+    )
+    return score
+
+
+def setup_args():
+    """Setup inerance server arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8731)
+    parser.add_argument("--log_file", type=str, default="rm_server.log")
+    args = parser.parse_args()
+    return args
+
+
+def server(args):
+    """Launch RM server."""
+    app = FastAPI()
+    lock = threading.Lock()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        filename=args.log_file,
+        filemode="w",
+        format="%(asctime)s - %(message)s",
+    )
+
+    @app.post("/")
+    async def _server(request: Request) -> Response:
+        lock.acquire()
+        logging.info(f"Request: {request}")
+        try:
+            all_result = []
+            if len(request.tgt) != len(request.response) or len(request.tgt) != len(request.src):
+                raise ValueError("The length of response, tgt, and src should be equal.")
+            for i in range(len(request.response)):
+                reward = compute_score(request.response[i], request.tgt[i], request.src[i])
+                all_result.append(reward)
+            output = {
+                "error_code": 0,
+                "error_msg": "Success",
+                "score": all_result,
+            }
+        except Exception as err:
+            logging.error(f"Server error: when process {request}\n{traceback.format_stack()}")
+            output = {
+                "error_code": 500,
+                "error_msg": f"{err}",
+                "score": [0] * len(request.tgt),
+            }
+        logging.info(f"Response: {json.dumps(output, indent=2, ensure_ascii=False)}")
+        lock.release()
+        return output
+
+    uvicorn.run(app, host="0.0.0.0", port=args.port)
+
+
+if __name__ == "__main__":
+    args = setup_args()
+    server(args)

llm/alignment/rl/run_rl.py

@@ -23,6 +23,7 @@ import paddle
 from paddle.distributed import fleet
 
 from paddlenlp.datasets.rlhf_datasets import RLHFDataset, collate_fn
+from paddlenlp.generation import GenerationConfig
 from paddlenlp.rl.models.score_model import AutoModelForScore
 from paddlenlp.rl.trainer.ppo_trainer import PPOTrainer
 from paddlenlp.rl.utils.config_utils import (
@@ -31,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,
@@ -41,6 +42,7 @@ from paddlenlp.trainer import (
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForTokenClassification,
     AutoTokenizer,
     PretrainedConfig,
 )
@@ -80,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
@@ -100,11 +103,9 @@ def create_actor_models(
         actor_model_config.set_attn_func = True
         actor_model_config.max_position_embeddings = data_args.max_length
         actor_model_config.use_sparse_head_and_loss_fn = False
-        actor_model_config.fused_linear = model_args.fused_linear
-        actor_model_config.use_fused_rms_norm = training_args.use_fused_rms_norm
         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(
@@ -114,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
 
@@ -136,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,
@@ -172,8 +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
-        reward_model_config.fused_linear = model_args.fused_linear
-        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:
@@ -213,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:
@@ -264,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
 
@@ -273,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,
@@ -325,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
@@ -350,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"))
 
@@ -358,9 +362,22 @@ 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:
+        generation_config = GenerationConfig.from_pretrained(model_args.actor_model_name_or_path)
+    except:
+        logger.warning("Can't find generation config, so it will not use generation_config field in the model config")
+        generation_config = None
+
     trainer = PPOTrainer(
         actor_model=actor_model,
         reference_model=reference_model,
@@ -378,10 +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,6 +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
+    ):
+        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/llama/grpo_argument.yaml

@@ -11,7 +11,7 @@ reward_server: "http://127.0.0.1:8731" # The address of the reward model server
 logging_dir: grpo-logs # Directory for logging
 logging_steps: 1 # Number of steps between logging
 output_dir: "qwen2.5-7b-kk-dataset-grpo/checkpoints" # Directory for output ckpts
-report_to: "wandb" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
+report_to: "visualdl" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
 wandb_http_proxy: "http://127.0.0.1:8962" # HTTP proxy for wandb
 run_name: "qwen2.5-7b-kk-dataset-grpo" # Name of the run
 
@@ -22,12 +22,13 @@ 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: "stage2" # Sharding strategy, e.g., "stage1" or "stage2"
+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
@@ -39,24 +40,23 @@ min_dec_len: 32 # Minimum length of the response
 top_p: 1.0 # Top-p sampling parameter
 temperature: 0.7 # Temperature parameter for sampling
 repetition_penalty: 1.0 # Repetition penalty parameter
-# rollout_use_dynamic_insert: 1 # Whether to use dynamic insert for rollout
-# rollout_continue_batching_batch_size: 32 # Base batch size for dynamic insert
-quant_type: "" # Quantization type, e.g., "weight_only_int8"
+rollout_max_num_seqs: 32 # 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: 2 # Global batch size for training
-mini_batch_size: 2 # Mini-batch size for training
+global_batch_size: 4 # Global batch size for training
+global_gen_batch_size: -1 # Global generation batch size for dynamic sampling
+global_mini_batch_size: -1 # Mini-batch size for training
 rollout_n: 8 # Number of rollouts
 update_iters: 1 # Number of training iterations for rollout samples
-per_device_rollout_batch_size: 1 # Rollout batch size per device
 per_device_logprob_batch_size: 8 # Log probability batch size per device
 per_device_reward_batch_size: 8 # Reward batch size per device
 per_device_value_batch_size: 8 # Value batch size per device
 per_device_train_batch_size: 8 # Training batch size per device
 # gradient_accumulation_steps: 1 # Gradient accumulation steps (auto-calculated)
-num_train_epochs: 3 # Number of training epochs
+num_train_epochs: 6 # Number of training epochs
 max_length: 4608 # Maximum length for training, should be larger than max_prompt_len + max_dec_len
 learning_rate: 5e-7 # Learning rate for training
 lr_scheduler_type: "constant" # Learning rate scheduler type
@@ -65,7 +65,7 @@ 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: 3600 # Maximum number of training steps
+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)
@@ -73,7 +73,7 @@ disable_tqdm: true # Whether to disable tqdm progress bar
 
 # RL args
 kl_coeff: 0.0 # KL coefficient
-kl_loss_coeff: 0.0 # KL loss coefficient
+kl_loss_coeff: 0.001 # 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)
@@ -84,12 +84,11 @@ 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: 5.0 # The clipping range for the output of the value model. The value is clipped into [-clip_range_value, clip_range_value].
-normalize_reward: true # Whether to normalize reward
-normalize_advantage: true # Whether to normalize advantage
+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
-per_device_sample_batch_size: 1 # Sample batch size per device for dynamic sampling
 max_gen_batches: 2 # Maximum number of generation batches for dynamic sampling
-use_fp32_compute: false # Whether to use fp32 to compute xx_log_prob,rewards, advantages and loss
+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
@@ -99,19 +98,18 @@ eval_steps: 20 # 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
-use_fused_rope: true # Whether to use fused rope operations
-use_fused_head_and_loss_fn: false # Whether to use fused head and loss function
-use_fused_linear: false # Whether to use fused linear operations, which needs to install fused_ln in slm/model_zoo/gpt-3/external_ops
-fused_linear: false # Whether to use fused_gemm_epilogue
+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: true # Whether to enable gradient checkpointing for memory optimization
 recompute_use_reentrant: true # 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
+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

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

@@ -0,0 +1,119 @@
+# RL algorithms
+rl_algorithm: "grpo" # The reinforcement learning algorithm used, supported: "ppo", "grpo", "reinforce_plus_plus"
+
+# models
+actor_model_name_or_path: "Qwen/Qwen2.5-32B-Instruct" # The name or path of the actor model
+reward_model_name_or_path: "" # The name or path of the reward model
+use_rm_server: true # Whether to use the reward model server
+reward_server: "http://127.0.0.1:8731" # The address of the reward model server
+
+# logging
+logging_dir: grpo-logs # Directory for logging
+logging_steps: 1 # Number of steps between logging
+output_dir: "qwen2.5-32b-kk-dataset-grpo/checkpoints" # Directory for output ckpts
+report_to: "visualdl" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
+wandb_http_proxy: "http://127.0.0.1:8962" # HTTP proxy for wandb
+run_name: "qwen2.5-32b-kk-dataset-grpo" # Name of the run
+
+# data
+train_datasets: "ppo-kk/34567ppl/train.jsonl" # Path to the training dataset
+eval_datasets: "ppo-kk/5ppl/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
+dataloader_shuffle: true # Whether to shuffle the train dataset
+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: 4 # 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: 2048 # Maximum length of the prompt, exceeding which will be automatically truncated
+max_dec_len: 30720 # 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: 24 # 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: 64 # Global batch size for training
+global_gen_batch_size: -1 # Global generation batch size for dynamic sampling
+global_mini_batch_size: 16 # Mini-batch size for training
+rollout_n: 8 # Number of rollouts
+update_iters: 1 # Number of training iterations for rollout samples
+per_device_logprob_batch_size: 2 # 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 batch size per device
+# gradient_accumulation_steps: 1 # Gradient accumulation steps (auto-calculated)
+num_train_epochs: -1 # Number of training epochs
+max_length: 32768 # Maximum length for training, should be larger than max_prompt_len + max_dec_len
+learning_rate: 1e-7 # Learning rate for training
+lr_scheduler_type: "constant" # Learning rate scheduler type
+weight_decay: 1e-2 # Weight decay for the AdamW optimizer
+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: 1000 # Maximum number of training steps
+save_steps: 500 # 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
+
+# RL args
+kl_coeff: 0.001 # KL coefficient for PPO and Reinforce++
+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)
+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: 5.0 # 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: 8 # Evaluation batch size per device
+evaluation_strategy: "steps" # Evaluation strategy, e.g., "steps"
+eval_steps: 100 # 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: true # Whether to enable gradient checkpointing for memory optimization
+recompute_use_reentrant: true # 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: true # Whether to offload optimizer to pinned memory
+
+# benchmark args
+skip_profile_timer: false # Whether to skip profiling timer

llm/config/qwen/grpo_argument.yaml

@@ -11,7 +11,7 @@ reward_server: "http://127.0.0.1:8731" # The address of the reward model server
 logging_dir: grpo-logs # Directory for logging
 logging_steps: 1 # Number of steps between logging
 output_dir: "qwen2.5-7b-kk-dataset-grpo/checkpoints" # Directory for output ckpts
-report_to: "wandb" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
+report_to: "visualdl" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
 wandb_http_proxy: "http://127.0.0.1:8962" # HTTP proxy for wandb
 run_name: "qwen2.5-7b-kk-dataset-grpo" # Name of the run
 
@@ -21,13 +21,15 @@ eval_datasets: "ppo-kk/5ppl/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
+dataloader_shuffle: false # Whether to shuffle the train dataset
 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: "stage2" # Sharding strategy, e.g., "stage1" or "stage2"
+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
@@ -39,25 +41,23 @@ min_dec_len: 32 # Minimum length of the response
 top_p: 1.0 # Top-p sampling parameter
 temperature: 0.7 # Temperature parameter for sampling
 repetition_penalty: 1.0 # Repetition penalty parameter
-# rollout_use_dynamic_insert: 1 # Whether to use dynamic insert for rollout
-# rollout_continue_batching_batch_size: 32 # Base batch size for dynamic insert
-quant_type: "" # Quantization type, e.g., "weight_only_int8"
+rollout_max_num_seqs: 32 # 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: 2 # Global batch size for training
+global_batch_size: 8 # Global batch size for training
 global_gen_batch_size: -1 # Global generation batch size for dynamic sampling
-mini_batch_size: 2 # Mini-batch size for training
+global_mini_batch_size: -1 # Mini-batch size for training
 rollout_n: 8 # Number of rollouts
 update_iters: 1 # Number of training iterations for rollout samples
-per_device_rollout_batch_size: 1 # Rollout batch size per device
 per_device_logprob_batch_size: 8 # Log probability batch size per device
 per_device_reward_batch_size: 8 # Reward batch size per device
 per_device_value_batch_size: 8 # Value batch size per device
 per_device_train_batch_size: 8 # Training batch size per device
 # gradient_accumulation_steps: 1 # Gradient accumulation steps (auto-calculated)
-num_train_epochs: 3 # Number of training epochs
+num_train_epochs: 6 # Number of training epochs
 max_length: 4608 # Maximum length for training, should be larger than max_prompt_len + max_dec_len
 learning_rate: 5e-7 # Learning rate for training
 lr_scheduler_type: "constant" # Learning rate scheduler type
@@ -73,8 +73,8 @@ ignore_save_lr_and_optim: true # Whether to ignore saving learning rate and opti
 disable_tqdm: true # Whether to disable tqdm progress bar
 
 # RL args
-kl_coeff: 0.0 # KL coefficient
-kl_loss_coeff: 0.0 # KL loss coefficient
+kl_coeff: 0.001 # KL coefficient for PPO and Reinforce++
+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)
@@ -85,11 +85,11 @@ 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: 5.0 # The clipping range for the output of the value model. The value is clipped into [-clip_range_value, clip_range_value].
-normalize_reward: true # Whether to normalize reward
-normalize_advantage: true # Whether to normalize advantage
+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: false # Whether to use fp32 to compute xx_log_prob,rewards, advantages and loss
+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
@@ -100,21 +100,20 @@ eval_steps: 20 # Number of steps between evaluations
 # device memory optimization args
 use_flash_attention: true # Whether to use fused attention operations
 use_fused_rms_norm: false # Whether to use fused RMS norm operations, which needs to install fused_ln in slm/model_zoo/gpt-3/external_ops
-use_fused_rope: true # Whether to use fused rope operations
-use_fused_head_and_loss_fn: false # Whether to use fused head and loss function
-use_fused_linear: false # Whether to use fused linear operations, which needs to install fused_ln in slm/model_zoo/gpt-3/external_ops
-fused_linear: false # Whether to use fused_gemm_epilogue
+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: true # Whether to enable gradient checkpointing for memory optimization
 recompute_use_reentrant: true # 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
+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 timer
+skip_profile_timer: false # Whether to skip profiling timer

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/config/qwen/reinforce_plus_plus_argument.yaml

@@ -0,0 +1,119 @@
+# RL algorithms
+rl_algorithm: "reinforce_plus_plus" # The reinforcement learning algorithm used, supported: "ppo", "grpo", "reinforce_plus_plus"
+
+# models
+actor_model_name_or_path: "Qwen/Qwen2.5-7B-Instruct-1M" # The name or path of the actor model
+reward_model_name_or_path: "" # The name or path of the reward model
+use_rm_server: true # Whether to use the reward model server
+reward_server: "http://127.0.0.1:8731" # The address of the reward model server
+
+# logging
+logging_dir: grpo-logs # Directory for logging
+logging_steps: 1 # Number of steps between logging
+output_dir: "qwen2.5-7b-kk-dataset-grpo/checkpoints" # Directory for output ckpts
+report_to: "visualdl" # Supported reporting options: "all", "wandb", "tensorboard", "visualdl"(default), "none"
+wandb_http_proxy: "http://127.0.0.1:8962" # HTTP proxy for wandb
+run_name: "qwen2.5-7b-kk-dataset-grpo" # Name of the run
+
+# data
+train_datasets: "ppo-kk/34567ppl/train.jsonl" # Path to the training dataset
+eval_datasets: "ppo-kk/5ppl/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
+dataloader_shuffle: false # Whether to shuffle the train dataset
+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: 512 # Maximum length of the prompt, exceeding which will be automatically truncated
+max_dec_len: 4096 # Maximum length of the response
+min_dec_len: 32 # Minimum length of the response
+top_p: 1.0 # Top-p sampling parameter
+temperature: 0.7 # Temperature parameter for sampling
+repetition_penalty: 1.0 # Repetition penalty parameter
+rollout_max_num_seqs: 32 # 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: 8 # Global batch size for training
+global_gen_batch_size: -1 # Global generation batch size for dynamic sampling
+global_mini_batch_size: -1 # Mini-batch size for training
+rollout_n: 8 # Number of rollouts
+update_iters: 1 # Number of training iterations for rollout samples
+per_device_logprob_batch_size: 8 # Log probability batch size per device
+per_device_reward_batch_size: 8 # Reward batch size per device
+per_device_value_batch_size: 8 # Value batch size per device
+per_device_train_batch_size: 8 # Training batch size per device
+# gradient_accumulation_steps: 1 # Gradient accumulation steps (auto-calculated)
+num_train_epochs: 6 # Number of training epochs
+max_length: 4608 # Maximum length for training, should be larger than max_prompt_len + max_dec_len
+learning_rate: 5e-7 # Learning rate for training
+lr_scheduler_type: "constant" # Learning rate scheduler type
+weight_decay: 1e-2 # Weight decay for the AdamW optimizer
+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: 3600 # 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
+
+# RL args
+kl_coeff: 0.001 # KL coefficient for PPO and Reinforce++
+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)
+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: 5.0 # 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: 32 # Evaluation batch size per device
+evaluation_strategy: "steps" # Evaluation strategy, e.g., "steps"
+eval_steps: 20 # Number of steps between evaluations
+
+# device memory optimization args
+use_flash_attention: true # Whether to use fused attention operations
+use_fused_rms_norm: false # 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: true # Whether to enable gradient checkpointing for memory optimization
+recompute_use_reentrant: true # 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 timer

llm/devices/intel_hpu/tests/README.md

@@ -0,0 +1,20 @@
+#setup custom tpc library
+unset FLAGS_selected_intel_hpus
+export GC_KERNEL_PATH=/workspace/pdpd_automation/repo/PaddleCustomDevice/backends/intel_hpu/build/libcustom_tpc_perf_lib.so:/usr/lib/habanalabs/libtpc_kernels.so
+
+#test cmdline example
+#PR test cases
+python e2e-test-run.py --context pr --data /data/ckpt/ --filter stable --device intel_hpu --junit test_result.xml --platform gaudi2d
+#BAT test cases
+python e2e-test-run.py --context bat --data /data/ckpt/ --filter stable --device intel_hpu --junit test_result.xml --platform gaudi2d
+#smoke test cases
+python e2e-test-run.py --context sanity --data /data/ckpt/ --filter stable --device intel_hpu --junit test_result.xml --platform gaudi2d
+python e2e-test-run.py --context sanity --data /data/ckpt/ --filter stable --device intel_hpu:2 --junit test_result.xml --platform gaudi2d
+
+#static graph mode Inference
+export PYTHONPATH=$PYTHONPATH:/workspace/pdpd_automation/repo/PaddleNLP/
+export FLAGS_intel_hpu_execution_queue_size=10
+#export relative static mode file
+python export_model.py --model_name_or_path /data/ckpt/meta-llama/Llama-2-7b-chat/ --inference_model --output_path ./inference --dtype bfloat16 --device intel_hpu
+#run static mode inference
+python e2e-test-run.py --context sanity --data /data/ckpt/ --filter stable --device intel_hpu:2 --mode static --junit test_result.xml --platform gaudi2d

llm/devices/intel_hpu/tests/config/llm.py

@@ -0,0 +1,79 @@
+#   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 itertools
+
+
+def create_variants(
+    scr_length: list, max_length: list, total_max_length: list, batch_size: list, decode_strategy: list
+):
+    return [
+        dict(zip(["src_length", "max_length", "total_max_length", "batch_size", "decode_strategy"], v))
+        for v in itertools.product(scr_length, max_length, total_max_length, batch_size, decode_strategy)
+    ]
+
+
+def add_testcase(
+    dict_lst: dict,
+    model_full_name,
+    scr_length: list,
+    max_length: list,
+    total_max_length: list,
+    batch_size: list,
+    decode_strategy: list,
+):
+    case_dict = dict_lst.setdefault(model_full_name, {})
+    case_dict["variants"] = {
+        "inference": create_variants(scr_length, max_length, total_max_length, batch_size, decode_strategy)
+    }
+    case_dict["output_file"] = f"{model_full_name.lower().split('/')[-1]}.json"
+
+
+test_case_lst = {}
+skip_case_lst = {}
+
+for i in ["bat", "pr", "sanity", "full"]:
+    test_case_lst.setdefault(i, {})
+    skip_case_lst.setdefault(i, {})
+    add_testcase(
+        test_case_lst[i],
+        "meta-llama/Llama-2-7b-chat",
+        ["128"],
+        ["128"],
+        ["256"],
+        ["1", "16"],
+        ["sampling", "greedy_search", "beam_search"],
+    )
+
+# testcase list + model name + scr_length(list) + max_length(list) + total_max_length (list)  + batch_size(list) + decode_strategy (list)
+add_testcase(
+    test_case_lst["sanity"],
+    "meta-llama/Llama-2-13b-chat",
+    ["128"],
+    ["128"],
+    ["256"],
+    ["1", "16"],
+    ["greedy_search", "greedy_search", "beam_search"],
+)
+add_testcase(
+    test_case_lst["sanity"], "meta-llama/Llama-2-70b-chat", ["128"], ["128"], ["256"], ["1", "16"], ["greedy_search"]
+)
+
+# when filter passwdown 'stable' will load this list
+# this list for the unstable test case to skip
+skip_case_lst["sanity"]["stable"] = []
+
+# when filter passwdown 'unstable' will load this list
+# this list for the stable test case to skip
+skip_case_lst["sanity"]["unstable"] = []

llm/devices/intel_hpu/tests/e2e-test-run.py

@@ -0,0 +1,217 @@
+#!/usr/bin/python3
+
+#   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 os
+import subprocess
+
+from util import get_model_path
+
+realwd = os.path.dirname(os.path.realpath(__file__))
+paddlenlp_path = os.path.realpath(f"{realwd}/../../")
+
+data_path = os.environ.get("DATA_DIR", "/data")
+
+os.environ.setdefault("GAUDI2_CI", "1")
+
+cmd_args = {}
+parser = argparse.ArgumentParser(
+    description="help scriopt for pdpd e2e run test on intel hpu",
+    add_help=True,
+    formatter_class=argparse.RawTextHelpFormatter,
+)
+parser.add_argument(
+    "--context",
+    choices=["pr", "bat", "sanity", "full"],
+    help="which test suites to be used: pr(PR testing), bat (BAT testing), sanity (Smoke testing), full (full testing); default: bat",
+    default="bat",
+)
+parser.add_argument("--data", type=str, help="data folder which should include huggingface folder", default=data_path)
+parser.add_argument(
+    "--filter",
+    choices=["stable", "unstable", "all"],
+    help="filter test case list: stable/unstable/all",
+    default="all",
+)
+parser.add_argument("--device", type=str, help="device name", default="intel_hpu")
+parser.add_argument("--mode", type=str, help="it should be one of [dynamic, static]", default="dynamic")
+parser.add_argument("--junit", type=str, help="junit result file")
+parser.add_argument("--platform", type=str, help="platform name")
+
+cmd_args.update(vars(parser.parse_args()))
+
+if cmd_args["junit"]:
+    libpath = os.path.dirname(os.path.dirname(realwd))
+    if os.path.exists(f"{libpath}/junitxml.py"):
+        import sys
+
+        sys.path.append(libpath)
+    from junitxml import jTestCase, jTestSuite
+
+cmd_args["platform"] = "" if cmd_args["platform"] is None else cmd_args["platform"].lower()
+
+script_path = os.path.dirname(os.path.realpath(__file__))
+
+if os.path.exists(cmd_args["data"]) is False:
+    print("data path not exist, please check for parameter: data / environment DATA_DIR")
+    exit(2)
+
+data_path = f"{cmd_args['data']}"
+os.environ.setdefault("DATA_HOME", data_path)
+if os.path.exists(data_path) is False:
+    print(f"Couldn't find mode data path, please confirm this folder under the {cmd_args['data']} folder")
+    exit(2)
+
+
+def case_command(command, test_case=None, test_suite=None):
+    output = ""
+    output += f"Command: {command}\n"
+    proc = subprocess.Popen(command, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    try:
+        line = proc.stdout.readline().decode()
+        while len(line) > 0:
+            print(f"{line[:-1]}")  # line include '\n' at last
+            output += line
+            line = proc.stdout.readline().decode()
+    except Exception as e:
+        if test_case:
+            test_case.setFail("Command abnormally")
+        print(e)
+    finally:
+        proc.communicate()
+    if proc.returncode != 0 and test_case:
+        test_case.setFail(f"Command return code:{proc.returncode}")
+    if test_case:
+        test_case.AddOutput(output)
+    if test_suite:
+        ts.addCase(test_case)
+
+    return proc.returncode, output
+
+
+def run_e2e_test_case(
+    model,
+    model_case_name,
+    src_length,
+    max_length,
+    total_max_length,
+    batch_size,
+    decode_strategy,
+    output_file,
+    skip_lst=None,
+    test_suite=None,
+    cmd_env="",
+):
+    # mode option: dynamic or static
+    mode_opt = ""
+    mode_param = ""
+    if cmd_args["mode"] and cmd_args["mode"] in ["dynamic", "static"]:
+        mode_opt = f"--mode {cmd_args['mode']}"
+        mode_param = cmd_args["mode"]
+    else:
+        mode_param = "dynamic"
+
+    # device option
+    device_opt = "intel_hpu"
+    if cmd_args["device"]:
+        device_opt = cmd_args["device"]
+
+    # gaudi2d platform only support bfloat16
+    float_opt = "--dtype bfloat16" if cmd_args["platform"] == "gaudi2d" else ""
+
+    testcase = None
+    testcase_name = (
+        f"{model_case_name}-{src_length}-{max_length}-{total_max_length}-{batch_size}-{decode_strategy}-{mode_param}"
+    )
+    if test_suite:
+        testcase = jTestCase(testcase_name)
+    else:
+        testcase = None
+
+    cmd_line = f"python {paddlenlp_path}/predict/predictor.py --model_name_or_path {model} --inference_model --device {device_opt} {mode_opt} {float_opt} "
+    cmd_line += f"--src_length {src_length} --max_length {max_length} --total_max_length {total_max_length} --batch_size {batch_size} --decode_strategy {decode_strategy} --output_file result/{testcase_name}.json"
+    _env = os.environ.copy()
+    for opt in cmd_env.split():
+        _env.setdefault(opt.split("=")[0], opt.split("=")[1])
+    print(f"RUN shell CMD: {cmd_env} {cmd_line}")
+    ret, output = case_command(cmd_line, testcase, test_suite)
+
+    return ret, output
+
+
+case_dict_lst = {}
+
+skip_lst = []
+
+from config.llm import skip_case_lst, test_case_lst
+
+case_dict_lst = test_case_lst[cmd_args["context"]]
+skip_lst = skip_case_lst.get(cmd_args["context"], {}).get(cmd_args["filter"], [])
+
+ts = None
+if cmd_args["junit"]:
+    ts = jTestSuite("E2E Test")
+    ts.setPlatform(cmd_args["platform"])
+
+total_case_num = 0
+pass_case_num = 0
+fail_case_num = 0
+
+for model_name, case_dict in case_dict_lst.items():
+    model_case_name = model_name.split("/")[-1]
+    model_path_or_name = get_model_path(model_name)
+    variants = case_dict.get("variants", dict()).get("inference", [])
+    mode_param = "dynamic"
+    if cmd_args["mode"] and cmd_args["mode"] in ["dynamic", "static"]:
+        mode_param = cmd_args["mode"]
+
+    for variant_dict in variants:
+        case_tag = f"{model_name} src:{variant_dict['src_length']}-max_length:{variant_dict['max_length']}-total_max_length:{variant_dict['total_max_length']}-bs:{variant_dict['batch_size']}-decode_strategy:{variant_dict['decode_strategy']}-mode:{mode_param}"
+        ret_test, _ = run_e2e_test_case(
+            model_path_or_name,
+            model_case_name,
+            variant_dict["src_length"],
+            variant_dict["max_length"],
+            variant_dict["total_max_length"],
+            variant_dict["batch_size"],
+            variant_dict["decode_strategy"],
+            case_dict["output_file"],
+            skip_lst,
+            ts,
+        )
+        if "skip" == _:
+            pass
+        else:
+            total_case_num = total_case_num + 1
+            if ret_test == 0:
+                pass_case_num = pass_case_num + 1
+                print(f"\033[0;32mtest case {total_case_num} : {case_tag} pass \033[0m")
+            else:
+                fail_case_num = fail_case_num + 1
+                print(f"\033[0;31mtest case {total_case_num} : {case_tag} fail \033[0m")
+
+if cmd_args["junit"]:
+    with open(cmd_args["junit"], "w+") as f:
+        ts.toString()
+        f.write(ts.toString())
+
+print("...............................Summary.......................................")
+print(f"\033[0;37mE2E total {total_case_num} test case running \033[0m")
+print(f"\033[0;32mE2E total {pass_case_num} test case pass \033[0m")
+if fail_case_num != 0:
+    print(f"\033[0;31mE2E total {fail_case_num} test case fail \033[0m")
+
+exit(fail_case_num)

llm/devices/intel_hpu/tests/junitxml.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+
+# 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 platform
+import re
+import subprocess
+from enum import Enum
+
+from lxml import etree as et
+
+
+class jResult(Enum):
+    PASS = 0
+    FAIL = 1
+    SKIP = 2
+    ERROR = 3
+    DISABLE = 4
+
+
+def remove_invalidel_xml_characters(src, output, start=10):
+    # "invalid XML character" will return ""
+    numerical = int(src, start)
+    if (
+        numerical in (0xB, 0xC, 0xFFFE, 0xFFFF)
+        or 0x0 <= numerical <= 0x8
+        or 0xE <= numerical <= 0x1F
+        or 0xD800 <= numerical <= 0xDFFF
+    ):
+        return ""
+    return output
+
+
+class jXMLBase:
+    def __init__(self, tag: str):
+        self._attr_dict = dict()
+        self._tag = tag
+        self._text = ""
+
+    def setName(self, name):
+        self._attr_dict["name"] = name
+
+    def covertET(self) -> et:
+        elem = et.Element(self._tag)
+        for k, v in self._attr_dict.items():
+            elem.set(k, str(v))
+        if len(self._text) > 0:
+            # to deal all non-ascii characters for XML
+            self._text = re.sub(
+                r"&#(\d+);?", lambda ch: remove_invalidel_xml_characters(ch.group(1), ch.group(0)), self._text
+            )
+            self._text = re.sub(
+                r"&#[xX]([0-9a-fA-F]+);?",
+                lambda ch: remove_invalidel_xml_characters(ch.group(1), ch.group(0), start=16),
+                self._text,
+            )
+            self._text = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1F\uD800-\uDFFF\uFFFE\uFFFF]").sub("", self._text)
+            elem.text = self._text
+        return elem
+
+    def toString(self, pretty_print=False):
+        return et.tostring(self.covertET(), encoding="utf-8", pretty_print=pretty_print).decode()
+
+
+class jTestCase(jXMLBase):
+    result = jResult.PASS
+    _child_elem = None
+    _output = None
+
+    def __init__(self, name: str):
+        super().__init__("testcase")
+        self.setName(name)
+        self._child_elem = list()
+
+    def setClass(self, classname: str):
+        self._attr_dict["classname"] = classname
+
+    def setTime(self, time):
+        self._attr_dict["time"] = time
+
+    def AddOutput(self, output: str):
+        if self._output is None:
+            self._output = jXMLBase("system-out")
+            self._child_elem.append(self._output)
+        if len(output) >= 1:
+            self._output._text += output
+
+    def setFail(self, message: str):
+        self.result = jResult.FAIL
+        el = jXMLBase("failure")
+        el._attr_dict["message"] = message
+        el._attr_dict["type"] = "failure"
+        self._child_elem.append(el)
+
+    def setSkip(self, message: str):
+        self.result = jResult.SKIP
+        self._attr_dict["skipped"] = 1
+        self.AddOutput("")
+        el = jXMLBase("skipped")
+        el._attr_dict["type"] = message
+        self._child_elem.append(el)
+
+    def covertET(self) -> et:
+        elem = super().covertET()
+        for jxml in self._child_elem:
+            elem.append(jxml.covertET())
+        return elem
+
+
+class jTestSuite(jXMLBase):
+    _case_lst = None
+
+    def __init__(self, name: str):
+        super().__init__("testsuite")
+        self.setName(name)
+        self._case_lst = list()
+        for i in ["tested", "passed", "errors", "failed", "skipped", "disabled"]:
+            self._attr_dict[i] = 0
+        self._attr_dict["kernel"] = subprocess.getoutput("uname -r")
+        self._attr_dict["os"] = "wsl" if "microsoft" in platform.platform().lower() else "linux"
+
+    def print_attr_dict(self):
+        for k, v in self._attr_dict.items():
+            print(f"{k} : {v}")
+
+    def setPlatform(self, platform):
+        self._attr_dict["platform"] = platform
+
+    def addCase(self, case: jTestCase):
+        if case.result == jResult.FAIL:
+            self._attr_dict["failed"] += 1
+        elif case.result == jResult.ERROR:
+            self._attr_dict["errors"] += 1
+        elif case.result == jResult.SKIP:
+            self._attr_dict["skipped"] += 1
+        elif case.result == jResult.DISABLE:
+            self._attr_dict["disabled"] += 1
+        else:
+            self._attr_dict["passed"] += 1
+        self._case_lst.append(case)
+
+    def covertET(self) -> et:
+        self._attr_dict["tested"] = len(self._case_lst)
+        elem = super().covertET()
+        for jxml in self._case_lst:
+            elem.append(jxml.covertET())
+        return elem

llm/devices/intel_hpu/tests/util.py

@@ -0,0 +1,56 @@
+#!/usr/bin/python3
+
+#   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 json
+import os
+import subprocess
+
+
+def get_model_path(model_path_or_name):
+    if os.path.exists(model_path_or_name):
+        return model_path_or_name
+    else:
+        data_home = os.environ.get("DATA_HOME", "/data")
+        model_name_path = f"{data_home}/{model_path_or_name}"
+        if os.path.exists(model_name_path):
+            return model_name_path
+        else:
+            return None
+
+
+def get_model_path_hf(model_path_or_name):
+    if os.path.exists(model_path_or_name):
+        return model_path_or_name
+    else:
+        hf_home = os.environ.get("HF_HOME", f'{os.environ["HOME"]}/.cache/huggingface')
+        model_name_path = f"{hf_home}/hub/models--{model_path_or_name.replace('/', '--')}"
+        if os.path.exists(model_name_path):
+            model_uuid = subprocess.getoutput(f"cat {model_name_path}/refs/main")
+            return f"{model_name_path}/snapshots/{model_uuid}"
+        else:
+            return None
+
+
+def model_config_add_type(model_name, mt):
+    model_path = get_model_path(model_name)
+    if model_path is None:
+        return
+    with open(f"{model_path}/config.json", "r") as f:
+        config_dict = json.load(f)
+    if "model_type" not in config_dict:
+        config_dict["model_type"] = mt
+        with open(f"{model_path}/config.json", "w+") as f:
+            json.dump(config_dict, f)

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/predict/best_practices.md

@@ -44,3 +44,6 @@ PaddleNLP 提供了多种环境变量，用于优化推理性能和资源使用
 
 **自定义算子通用 优化**
 - `DYNAMIC_INFERENCE_MODE`：动态图推理时自定义算子是否采用 pybind 调用方式，默认为 True。
+
+**其他**
+- `INFERENCE_MSG_QUEUE_ID`：多实例消息队列 id，默认为 进程 id。

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
@@ -848,7 +866,8 @@ class BlockInferencePredictorMixin(BasePredictor):
 
         self.pre_cache_length = 0
 
-        self.msg_queue_id = os.getpid()
+        msg_queue_id_str = os.getenv("INFERENCE_MSG_QUEUE_ID", str(os.getpid()))
+        os.environ["INFERENCE_MSG_QUEUE_ID"] = msg_queue_id_str
 
         if config.export_precache:
             pre_cache_npy = np.load(config.prefix_path)
@@ -948,7 +967,6 @@ class BlockInferencePredictorMixin(BasePredictor):
         )
         self.model_inputs["bad_tokens"] = paddle.to_tensor([-1], dtype="int64")
         self.model_inputs["is_block_step"] = paddle.full(shape=[config.batch_size], fill_value=False, dtype="bool")
-        self.model_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=self.msg_queue_id, dtype="int32").cpu()
 
         # bloom model needs src_mask and tgt_mask!
         if "bloom" in self.architectures:
@@ -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
@@ -1185,7 +1208,7 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
 
         read_res_process = mp.Process(
             target=read_res_func,
-            args=[self.model_name_or_path, tensor_queue, result_queue, done_event, self.model_inputs["msg_queue_id"]],
+            args=[self.model_name_or_path, tensor_queue, result_queue, done_event],
         )
         if self.tensor_parallel_rank == 0:
             read_res_process.start()
@@ -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])
 
@@ -1304,14 +1327,14 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
     def insert_task(self, pos, task_id, repeat_num):
         query_id = task_id // repeat_num
         length = len(self.input_ids[query_id])
-        logger.debug(f"Insert task {task_id} while query id is {query_id} inserting pos {pos}")
+        # logger.debug(f"Insert task {task_id} while query id is {query_id} inserting pos {pos}")
         self.model_inputs["input_ids"][pos, 0] = self.model_inputs["all_token_ids"][task_id, 0]
         self.model_inputs["seq_lens_this_time"][pos] = 1
         self.model_inputs["seq_lens_decoder"][pos] = length
         self.model_inputs["stop_flags"][pos] = False
         self.model_inputs["result_id"][pos][0] = task_id
         self.model_inputs["step_idx"][pos, 0] = 1
-        self.model_inputs["pre_ids"][pos][0] = self.input_ids[query_id][-1]
+        self.model_inputs["pre_ids"][pos][0] = np.array(self.input_ids[query_id][-1])
         self.model_inputs["pre_ids"][pos][1:] = -1
         self.model_inputs["not_need_stop"][0] = True
 
@@ -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 = []
@@ -1477,7 +1502,6 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
                     task_queue,
                     result_queue,
                     done_event,
-                    self.model_inputs["msg_queue_id"],
                     len(self.input_ids),
                     detokenize,
                 ],
@@ -1568,7 +1592,7 @@ class DygraphBlockInferencePredictor(BlockInferencePredictorMixin):
         else:
             if flag_current_rank_run:
                 output_tokens = self.model_inputs["all_token_ids"].numpy()
-                output_tokens[output_tokens < 0] = self.tokenizer.pad_token_id
+                output_tokens[output_tokens == -1] = self.tokenizer.eos_token_id
                 if detokenize:
                     outputs = self.tokenizer.batch_decode(
                         output_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False
@@ -1716,7 +1740,7 @@ class StaticGraphBlockInferencePredictor(BlockInferencePredictorMixin):
 
         read_res_process = mp.Process(
             target=read_res_func,
-            args=[self.model_name_or_path, tensor_queue, result_queue, done_event, self.model_inputs["msg_queue_id"]],
+            args=[self.model_name_or_path, tensor_queue, result_queue, done_event],
         )
         if self.tensor_parallel_rank == 0:
             read_res_process.start()

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/server/dockerfiles/Dockerfile_serving_xpu_p800

@@ -0,0 +1,42 @@
+FROM registry.baidubce.com/paddlepaddle/fastdeploy:llm-base-gcc12.3-cuda12.4-cudnn9-nccl2.15.5
+
+ENV http_proxy=http://agent.baidu.com:8891
+ENV https_proxy=http://agent.baidu.com:8891
+ENV no_proxy=localhost,bj.bcebos.com,su.bcebos.com,pypi.tuna.tsinghua.edu.cn,paddle-ci.gz.bcebos.com
+
+WORKDIR /opt/output/
+
+# 安装 paddlepaddle & paddlenlp & paddlenlp_ops
+RUN pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple \
+    && python3 -m pip install -U --no-deps --force-reinstall https://paddle-whl.bj.bcebos.com/nightly/xpu-p800/paddlepaddle-xpu/paddlepaddle_xpu-3.0.0.dev20250426-cp310-cp310-linux_x86_64.whl \
+    && python3 -m pip install --no-cache-dir sentencepiece pycryptodome tritonclient[all]==2.41.1 \
+    && apt update && apt install net-tools \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# clone paddle & paddlenlp 源码（代码版本应与上述安装版本对齐）
+RUN mkdir -p /opt/source/ && cd /opt/source/ \
+    && git clone https://github.com/PaddlePaddle/Paddle.git \
+    && python3 -m pip install --no-cache-dir -r Paddle/python/requirements.txt \
+    && git clone https://github.com/PaddlePaddle/PaddleNLP.git && cd PaddleNLP && git checkout 323ffec4fad74b5df9b8f6c2ec37d6855f6e62ca && cd ..\
+    && python3 -m pip install --no-cache-dir -r PaddleNLP/requirements.txt \
+    && python3 -m pip install --no-cache-dir -r PaddleNLP/llm/server/server/requirements.txt
+
+# 安装 paddlenlp_ops
+RUN cd /opt/source/ && wget https://paddlenlp.bj.bcebos.com/xpu_release/xpu_release/paddlenlp_ops_xpu.tar.gz \
+    && mkdir -p paddlenlp_ops_xpu && tar zxf paddlenlp_ops_xpu.tar.gz -C paddlenlp_ops_xpu \
+    && cd paddlenlp_ops_xpu/src && python3 setup.py install \
+    && rm -f paddlenlp_ops_xpu.tar.gz
+
+RUN cp /opt/source/PaddleNLP/llm/server/server/scripts/start_server.sh start_server \
+    && chmod +x start_server \
+    && mv start_server /usr/local/bin/ \
+    && cp /opt/source/PaddleNLP/llm/server/server/scripts/stop_server.sh stop_server \
+    && chmod +x stop_server \
+    && mv stop_server /usr/local/bin/ 
+    
+# RUN cp /opt/source/PaddleNLP/llm/server/server/server/download_model.py download_model.py
+
+ENV PYTHONPATH="/opt/source/PaddleNLP/llm/server/server:/opt/source/PaddleNLP"
+
+ENV http_proxy=""
+ENV https_proxy=""

llm/server/server/server/engine/engine.py

@@ -24,8 +24,7 @@ from multiprocessing import shared_memory
 
 import numpy as np
 from server.engine.resource_manager import ResourceManager
-from server.engine.task_queue_manager import (TaskQueueManager,
-                                              launch_queue_service)
+from server.engine.task_queue_manager import TaskQueueManager, launch_queue_service
 from server.engine.token_processor import TokenProcessor, WarmUpTokenProcessor
 from server.utils import model_server_logger
 
@@ -34,10 +33,11 @@ class Engine(object):
     """
     Engine Class
     """
+
     def __init__(self, cfg, token_processor):
         self.cfg = cfg
         # Master node only
-        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv('POD_0_IP', '127.0.0.1'):
+        if self.cfg.nnode == 1 or self.cfg.host_ip == os.getenv("POD_0_IP", "127.0.0.1"):
             self.queue_service = self._start_tasks_queue_service()
         self.tasks_queue = TaskQueueManager(mp_num=self.cfg.mp_num, port=self.cfg.infer_port)
         self.resource_manager = ResourceManager(self.cfg)
@@ -53,8 +53,10 @@ class Engine(object):
         initialize engine and start sub services
         """
         assert not self.is_started, "The engine is already started.!"
+        msg_queue_id_str = os.getenv("INFERENCE_MSG_QUEUE_ID", str(os.getpid()))
+        os.environ["INFERENCE_MSG_QUEUE_ID"] = msg_queue_id_str
         start_time = time.time()
-        
+
         self.token_processor.tasks_queue = self.tasks_queue
         self.infer_proc = self._start_infer_service()
         model_server_logger.info("Waiting infer processes ready...")
@@ -80,9 +82,10 @@ class Engine(object):
         """
         # get eos_token_id
         from server.data.processor import DataProcessor
+
         eos_token_ids = DataProcessor().get_eos_tokens()
 
-       # construct test tasks
+        # construct test tasks
         res_task = []
         for j in range(2 * self.cfg.max_batch_size):
             data = {
@@ -90,7 +93,7 @@ class Engine(object):
                 "req_id": j,
                 "max_dec_len": self.cfg.dec_len_limit,
                 "min_dec_len": int(self.cfg.dec_len_limit * 0.5) + 1,
-                "eos_token_ids": eos_token_ids
+                "eos_token_ids": eos_token_ids,
             }
             res_task.append(data)
         for j in range(2 * self.cfg.max_prefill_batch):
@@ -99,7 +102,7 @@ class Engine(object):
                 "req_id": j + 2 * self.cfg.max_batch_size,
                 "max_dec_len": 1,
                 "min_dec_len": 1,
-                "eos_token_ids": eos_token_ids
+                "eos_token_ids": eos_token_ids,
             }
             res_task.append(data)
 
@@ -130,8 +133,9 @@ class Engine(object):
 
         available_batch = np.sum(self.resource_manager.stop_flags)
         if len(tasks) > available_batch:
-            model_server_logger.error("Inserting batch:{} exceeds the available batch:{}.".format(
-                len(tasks), available_batch))
+            model_server_logger.error(
+                "Inserting batch:{} exceeds the available batch:{}.".format(len(tasks), available_batch)
+            )
             model_server_logger.error("The exceeded part will be ignored!")
             tasks = tasks[:available_batch]
 
@@ -140,21 +144,23 @@ class Engine(object):
             input_token_num = len(tasks[i]["input_ids"])
             if input_token_num >= self.cfg.max_seq_len - 1:
                 model_server_logger.warning(f"{req_id}: Input length:{input_token_num}, exceed the limits.")
-                tasks[i]["input_ids"] = tasks[i]["input_ids"][:self.cfg.max_seq_len - 1]
+                tasks[i]["input_ids"] = tasks[i]["input_ids"][: self.cfg.max_seq_len - 1]
             if "seq_len" in tasks[i] and "max_dec_len" not in tasks[i]:
                 tasks[i]["max_dec_len"] = tasks[i]["seq_len"]
 
             # max_dec_len + input_token_num > MAX_SEQ_LEN
             if input_token_num + tasks[i]["max_dec_len"] > self.cfg.max_seq_len:
                 tasks[i]["max_dec_len"] = self.cfg.max_seq_len - input_token_num
-                model_server_logger.warning("Force max_dec_len to be {} for req_id={}.".format(
-                    tasks[i]["max_dec_len"], tasks[i]["req_id"]))
+                model_server_logger.warning(
+                    "Force max_dec_len to be {} for req_id={}.".format(tasks[i]["max_dec_len"], tasks[i]["req_id"])
+                )
 
             # min_dec_len + input_token_num > MAX_SEQ_LEN
             if input_token_num + tasks[i]["min_dec_len"] > self.cfg.max_seq_len:
                 tasks[i]["min_dec_len"] = self.cfg.max_seq_len - input_token_num
-                model_server_logger.warning("Force min_dec_len to be {} for req_id={}.".format(
-                    tasks[i]["min_dec_len"], tasks[i]["req_id"]))
+                model_server_logger.warning(
+                    "Force min_dec_len to be {} for req_id={}.".format(tasks[i]["min_dec_len"], tasks[i]["req_id"])
+                )
 
         tasks = self.resource_manager.allocate_resources_for_new_tasks(tasks)
         if not tasks:
@@ -292,9 +298,7 @@ class Engine(object):
         self.shm_flag_ready = shared_memory.SharedMemory(
             create=True, size=flag_array.nbytes, name=self.cfg.get_unique_name("shm_flag_infer_ready")
         )
-        self.flag_ready_array = np.ndarray(
-            flag_array.shape, dtype=flag_array.dtype, buffer=self.shm_flag_ready.buf
-        )
+        self.flag_ready_array = np.ndarray(flag_array.shape, dtype=flag_array.dtype, buffer=self.shm_flag_ready.buf)
         self.flag_ready_array[:] = 0
 
         # broadcast flag for engine
@@ -324,7 +328,8 @@ class Engine(object):
             tmp = shared_memory.SharedMemory(
                 create=False,
                 size=has_block_step_flag_array.nbytes,
-                name=self.cfg.get_unique_name("shm_flag_has_block_step"))
+                name=self.cfg.get_unique_name("shm_flag_has_block_step"),
+            )
             tmp.close()
             tmp.unlink()
         except:
@@ -332,11 +337,13 @@ class Engine(object):
         self.shm_flag_has_block_step = shared_memory.SharedMemory(
             create=True,
             size=has_block_step_flag_array.nbytes,
-            name=self.cfg.get_unique_name("shm_flag_has_block_step"))
+            name=self.cfg.get_unique_name("shm_flag_has_block_step"),
+        )
         self.flag_has_block_step_array = np.ndarray(
             has_block_step_flag_array.shape,
             dtype=has_block_step_flag_array.dtype,
-            buffer=self.shm_flag_has_block_step.buf)
+            buffer=self.shm_flag_has_block_step.buf,
+        )
         self.flag_has_block_step_array[:] = 0
 
     def _exit_sub_services(self):
@@ -362,8 +369,9 @@ class Engine(object):
         if p.is_alive():
             model_server_logger.info("start tasks queue service successfully")
         else:
-            error_msg = "Failed to start tasks queue service, please check " \
-                        "the log/task_queue_manager.log for details"
+            error_msg = (
+                "Failed to start tasks queue service, please check " "the log/task_queue_manager.log for details"
+            )
             model_server_logger.info(error_msg)
             raise Exception(error_msg)
         return p
@@ -380,14 +388,16 @@ class Engine(object):
         pd_cmd = "python3 -m paddle.distributed.launch "
         py_script = os.path.join(current_dir_path, "infer.py")
 
-        arguments = (f" --nnodes {str(self.cfg.nnode)}"
-                    f" --devices {self.cfg.device_ids} {py_script} --model_dir {self.cfg.model_dir}"
-                    f" --max_batch_size {self.cfg.max_batch_size} --max_seq_len {self.cfg.max_seq_len}"
-                    f" --max_dec_len {self.cfg.max_dec_len}"
-                    f" --max_block_num {self.cfg.total_block_num} --block_size {self.cfg.block_size}"
-                    f" --use_cache_kv_int8 {self.cfg.use_cache_kv_int8}"
-                    f" --enc_dec_block_num {self.cfg.enc_dec_block_num}"
-                    f" --block_ratio {self.cfg.block_ratio} --dtype {self.cfg.dtype}")
+        arguments = (
+            f" --nnodes {str(self.cfg.nnode)}"
+            f" --devices {self.cfg.device_ids} {py_script} --model_dir {self.cfg.model_dir}"
+            f" --max_batch_size {self.cfg.max_batch_size} --max_seq_len {self.cfg.max_seq_len}"
+            f" --max_dec_len {self.cfg.max_dec_len}"
+            f" --max_block_num {self.cfg.total_block_num} --block_size {self.cfg.block_size}"
+            f" --use_cache_kv_int8 {self.cfg.use_cache_kv_int8}"
+            f" --enc_dec_block_num {self.cfg.enc_dec_block_num}"
+            f" --block_ratio {self.cfg.block_ratio} --dtype {self.cfg.dtype}"
+        )
         if self.cfg.nnode > 1:
             pd_cmd = pd_cmd + f" --ips {self.cfg.ips}"
         pd_cmd = pd_cmd + arguments + " >log/launch_infer.log 2>&1"

llm/server/server/server/engine/infer.py

@@ -279,7 +279,6 @@ class ModelRunner:
         self.share_inputs["input_ids"] = paddle.full(
             shape=[self.args.max_batch_size, self.args.max_seq_len], fill_value=self.pad_token_id, dtype="int64"
         )
-        self.share_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=1, dtype="int32").cpu()
         self.share_inputs["top_p"] = paddle.full(
             shape=[self.args.max_batch_size, 1], fill_value=self.top_p, dtype="float32"
         )
@@ -743,7 +742,7 @@ class InferenceEngine(object):
             config.set_xpu_device_id(device_id)
             xpu_config = paddle.inference.XpuConfig()
             xpu_config.device_id = device_id
-            xpu_config.l3_size = 0 
+            xpu_config.l3_size = 0
             xpu_config.l3_autotune_size = 0
             config.set_xpu_config(xpu_config)
             config.switch_ir_optim(True)

llm/server/server/server/engine/token_processor.py

@@ -47,7 +47,6 @@ class TokenProcessor(object):
         self.tokens_counter = Counter()
 
         self.is_speculate_decoding = self.cfg.get_speculate_config().speculate_method != "None"
-        self.msg_queue_id = paddle.full(shape=[1], fill_value=1, dtype="int32")
         if self.is_speculate_decoding:
             self.output_tokens = paddle.full(
                 shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
@@ -97,7 +96,7 @@ class TokenProcessor(object):
                 if self.is_speculate_decoding:
                     speculate_get_output(self.output_tokens, rank_id, is_blocking)
                 else:
-                    get_output(self.output_tokens, self.msg_queue_id, rank_id, is_blocking)
+                    get_output(self.output_tokens, rank_id, is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue
@@ -281,7 +280,7 @@ class WarmUpTokenProcessor(TokenProcessor):
                 if self.is_speculate_decoding:
                     speculate_get_output(self.output_tokens, rank_id, self._is_blocking)
                 else:
-                    get_output(self.output_tokens, self.msg_queue_id, rank_id, self._is_blocking)
+                    get_output(self.output_tokens, rank_id, self._is_blocking)
 
                 if self.output_tokens[0, 0] == -2:
                     continue

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)