update penalty_decay

* patch: fix tsc error

* chore: resolve ts-ignore with type-safe version code

.github/workflows/docker.yml

@@ -0,0 +1,171 @@
+name: Publish Docker Image
+on: [push]
+
+concurrency:
+  group: ${{ github.ref }}-${{ github.workflow }}
+  cancel-in-progress: true
+
+jobs:
+  docker_build:
+    name: Build ${{ matrix.arch }} Image
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        include:
+          - arch: amd64
+            name: amd64
+        #   - arch: arm64
+        #     name: arm64
+
+    steps:
+      - name: Free up disk spaces
+        run: |
+          sudo rm -rf /usr/share/dotnet || true
+          sudo rm -rf /opt/ghc || true
+          sudo rm -rf "/usr/local/share/boost" || true
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY" || true
+
+      - name: Get lowercase string for the repository name
+        id: lowercase-repo-name
+        uses: ASzc/change-string-case-action@v2
+        with:
+          string: ${{ github.event.repository.name }}
+
+      - name: Checkout base
+        uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Cache Docker layers
+        uses: actions/cache@v2
+        with:
+          path: /tmp/.buildx-cache
+          key: ${{ github.ref }}-${{ matrix.arch }}
+          restore-keys: |
+            ${{ github.ref }}-${{ matrix.arch }}
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+        with:
+          platforms: linux/${{ matrix.arch }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Docker login
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_PASSWORD }}
+
+      - name: Get commit SHA
+        id: vars
+        run: echo "::set-output name=sha_short::$(git rev-parse --short HEAD)"
+
+      - name: Build and export
+        id: build
+        if: github.ref == 'refs/heads/master'
+        uses: docker/build-push-action@v3
+        with:
+          push: true
+          platforms: linux/${{ matrix.arch }}
+          tags: ${{ secrets.DOCKER_USERNAME }}/${{ steps.lowercase-repo-name.outputs.lowercase }}:${{ matrix.name }}-latest
+          build-args: |
+            SHA=${{ steps.vars.outputs.sha_short }}
+          outputs: type=image,push=true
+          cache-from: type=local,src=/tmp/.buildx-cache
+          cache-to: type=local,dest=/tmp/.buildx-cache
+
+      - name: Replace tag without `v`
+        if: startsWith(github.ref, 'refs/tags/')
+        uses: actions/github-script@v1
+        id: version
+        with:
+          script: |
+            return context.payload.ref.replace(/\/?refs\/tags\/v/, '')
+          result-encoding: string
+
+      - name: Build release and export
+        id: build_rel
+        if: startsWith(github.ref, 'refs/tags/')
+        uses: docker/build-push-action@v3
+        with:
+          push: true
+          platforms: linux/${{ matrix.arch }}
+          tags: ${{ secrets.DOCKER_USERNAME }}/${{ steps.lowercase-repo-name.outputs.lowercase }}:${{ matrix.name }}-${{steps.version.outputs.result}}
+          build-args: |
+            SHA=${{ steps.version.outputs.result }}
+          outputs: type=image,push=true
+          cache-from: type=local,src=/tmp/.buildx-cache
+          cache-to: type=local,dest=/tmp/.buildx-cache
+
+      - name: Save digest
+        if: github.ref == 'refs/heads/master'
+        run: echo ${{ steps.build.outputs.digest }} > /tmp/digest.txt
+
+      - name: Save release digest
+        if: startsWith(github.ref, 'refs/tags/')
+        run: echo ${{ steps.build_rel.outputs.digest }} > /tmp/digest.txt
+
+      - name: Upload artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: digest_${{ matrix.name }}
+          path: /tmp/digest.txt
+
+  manifests:
+    name: Build manifests
+    needs: [docker_build]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Get lowercase string for the repository name
+        id: lowercase-repo-name
+        uses: ASzc/change-string-case-action@v2
+        with:
+          string: ${{ github.event.repository.name }}
+
+      - name: Checkout base
+        uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      # https://github.com/docker/setup-qemu-action
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+
+      # https://github.com/docker/setup-buildx-action
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+        with:
+          config-inline: |
+            [worker.oci]
+              max-parallelism = 1
+
+      - name: Download artifact
+        uses: actions/download-artifact@v3
+        with:
+          path: /tmp/images/
+
+      - name: Docker login
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_PASSWORD }}
+
+      - name: Replace tag without `v`
+        if: startsWith(github.ref, 'refs/tags/')
+        uses: actions/github-script@v1
+        id: version
+        with:
+          script: |
+            return context.payload.ref.replace(/\/?refs\/tags\/v/, '')
+          result-encoding: string
+
+      - name: Merge and push manifest on master branch
+        if: github.ref == 'refs/heads/master'
+        run: python scripts/merge_manifest.py "${{ secrets.DOCKER_USERNAME }}/${{ steps.lowercase-repo-name.outputs.lowercase }}"
+
+      - name: Merge and push manifest on release
+        if: startsWith(github.ref, 'refs/tags/')
+        run: python scripts/merge_manifest.py "${{ secrets.DOCKER_USERNAME }}/${{ steps.lowercase-repo-name.outputs.lowercase }}" ${{steps.version.outputs.result}}

.github/workflows/pre-release.yml

@@ -0,0 +1,114 @@
+name: pre-release
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - master
+    paths:
+      - "backend-python/**"
+    tags-ignore:
+      - "v*"
+
+jobs:
+  windows:
+    runs-on: windows-2022
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: master
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.20.5"
+      - uses: actions/setup-python@v5
+        id: cp310
+        with:
+          python-version: "3.10"
+      - uses: crazy-max/ghaction-chocolatey@v3
+        with:
+          args: install upx
+      - run: |
+          Start-BitsTransfer https://github.com/josStorer/ai00_rwkv_server/releases/latest/download/webgpu_server_windows_x86_64.exe ./backend-rust/webgpu_server.exe
+          Start-BitsTransfer https://github.com/josStorer/web-rwkv-converter/releases/latest/download/web-rwkv-converter_windows_x86_64.exe ./backend-rust/web-rwkv-converter.exe
+          Start-BitsTransfer https://github.com/josStorer/LibreHardwareMonitor.Console/releases/latest/download/LibreHardwareMonitor.Console.zip ./LibreHardwareMonitor.Console.zip
+          Expand-Archive ./LibreHardwareMonitor.Console.zip -DestinationPath ./components/LibreHardwareMonitor.Console
+          Start-BitsTransfer https://www.python.org/ftp/python/3.10.11/python-3.10.11-embed-amd64.zip ./python-3.10.11-embed-amd64.zip
+          Expand-Archive ./python-3.10.11-embed-amd64.zip -DestinationPath ./py310
+          $content=Get-Content "./py310/python310._pth"; $content | ForEach-Object {if ($_.ReadCount -eq 3) {"Lib\\site-packages"} else {$_}} | Set-Content ./py310/python310._pth
+          ./py310/python ./backend-python/get-pip.py
+          ./py310/python -m pip install Cython==3.0.4
+          Copy-Item -Path "${{ steps.cp310.outputs.python-path }}/../include" -Destination "py310/include" -Recurse
+          Copy-Item -Path "${{ steps.cp310.outputs.python-path }}/../libs" -Destination "py310/libs" -Recurse
+          ./py310/python -m pip install cyac==1.9
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
+          del ./backend-python/rwkv_pip/cpp/librwkv.dylib
+          del ./backend-python/rwkv_pip/cpp/librwkv.so
+          (Get-Content -Path ./backend-golang/app.go) -replace "//go:custom_build windows ", "" | Set-Content -Path ./backend-golang/app.go
+          (Get-Content -Path ./backend-golang/utils.go) -replace "//go:custom_build windows ", "" | Set-Content -Path ./backend-golang/utils.go
+          make
+          Rename-Item -Path "build/bin/RWKV-Runner.exe" -NewName "RWKV-Runner_windows_x64.exe"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: RWKV-Runner_windows_x64.exe
+          path: build/bin/RWKV-Runner_windows_x64.exe
+
+  linux:
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: master
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.20.5"
+      - run: |
+          wget https://github.com/josStorer/ai00_rwkv_server/releases/latest/download/webgpu_server_linux_x86_64 -O ./backend-rust/webgpu_server
+          wget https://github.com/josStorer/web-rwkv-converter/releases/latest/download/web-rwkv-converter_linux_x86_64 -O ./backend-rust/web-rwkv-converter
+          sudo apt-get update
+          sudo apt-get install upx
+          sudo apt-get install build-essential libgtk-3-dev libwebkit2gtk-4.0-dev libasound2-dev
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
+          rm ./backend-python/rwkv_pip/wkv_cuda.pyd
+          rm ./backend-python/rwkv_pip/rwkv5.pyd
+          rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/get-pip.py
+          rm ./backend-python/rwkv_pip/cpp/librwkv.dylib
+          rm ./backend-python/rwkv_pip/cpp/rwkv.dll
+          rm ./backend-python/rwkv_pip/webgpu/web_rwkv_py.cp310-win_amd64.pyd
+          make
+          mv build/bin/RWKV-Runner build/bin/RWKV-Runner_linux_x64
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: RWKV-Runner_linux_x64
+          path: build/bin/RWKV-Runner_linux_x64
+
+  macos:
+    runs-on: macos-13
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: master
+      - uses: actions/setup-go@v5
+        with:
+          go-version: "1.20.5"
+      - run: |
+          wget https://github.com/josStorer/ai00_rwkv_server/releases/latest/download/webgpu_server_darwin_aarch64 -O ./backend-rust/webgpu_server
+          wget https://github.com/josStorer/web-rwkv-converter/releases/latest/download/web-rwkv-converter_darwin_aarch64 -O ./backend-rust/web-rwkv-converter
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
+          rm ./backend-python/rwkv_pip/wkv_cuda.pyd
+          rm ./backend-python/rwkv_pip/rwkv5.pyd
+          rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/get-pip.py
+          rm ./backend-python/rwkv_pip/cpp/rwkv.dll
+          rm ./backend-python/rwkv_pip/cpp/librwkv.so
+          rm ./backend-python/rwkv_pip/webgpu/web_rwkv_py.cp310-win_amd64.pyd
+          make
+          cp build/darwin/Readme_Install.txt build/bin/Readme_Install.txt
+          cp build/bin/RWKV-Runner.app/Contents/MacOS/RWKV-Runner build/bin/RWKV-Runner_darwin_universal
+          cd build/bin && zip -r RWKV-Runner_macos_universal.zip RWKV-Runner.app Readme_Install.txt
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: RWKV-Runner_macos_universal.zip
+          path: build/bin/RWKV-Runner_macos_universal.zip

.github/workflows/release.yml

@@ -14,11 +14,11 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
       - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           ref: master
 
-      - uses: jossef/action-set-json-field@v2.1
+      - uses: jossef/action-set-json-field@v2.2
         with:
           file: manifest.json
           field: version
@@ -38,17 +38,17 @@ jobs:
     runs-on: windows-2022
     needs: create-draft
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           ref: master
-      - uses: actions/setup-go@v4
+      - uses: actions/setup-go@v5
         with:
-          go-version: '1.20.5'
-      - uses: actions/setup-python@v4
+          go-version: "1.20.5"
+      - uses: actions/setup-python@v5
         id: cp310
         with:
-          python-version: '3.10'
-      - uses: crazy-max/ghaction-chocolatey@v2
+          python-version: "3.10"
+      - uses: crazy-max/ghaction-chocolatey@v3
         with:
           args: install upx
       - run: |
@@ -64,7 +64,7 @@ jobs:
           Copy-Item -Path "${{ steps.cp310.outputs.python-path }}/../include" -Destination "py310/include" -Recurse
           Copy-Item -Path "${{ steps.cp310.outputs.python-path }}/../libs" -Destination "py310/libs" -Recurse
           ./py310/python -m pip install cyac==1.9
-          go install github.com/wailsapp/wails/v2/cmd/wails@latest
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
           del ./backend-python/rwkv_pip/cpp/librwkv.dylib
           del ./backend-python/rwkv_pip/cpp/librwkv.so
           (Get-Content -Path ./backend-golang/app.go) -replace "//go:custom_build windows ", "" | Set-Content -Path ./backend-golang/app.go
@@ -78,23 +78,22 @@ jobs:
     runs-on: ubuntu-20.04
     needs: create-draft
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           ref: master
-      - uses: actions/setup-go@v4
+      - uses: actions/setup-go@v5
         with:
-          go-version: '1.20.5'
+          go-version: "1.20.5"
       - run: |
           wget https://github.com/josStorer/ai00_rwkv_server/releases/latest/download/webgpu_server_linux_x86_64 -O ./backend-rust/webgpu_server
           wget https://github.com/josStorer/web-rwkv-converter/releases/latest/download/web-rwkv-converter_linux_x86_64 -O ./backend-rust/web-rwkv-converter
           sudo apt-get update
           sudo apt-get install upx
           sudo apt-get install build-essential libgtk-3-dev libwebkit2gtk-4.0-dev libasound2-dev
-          go install github.com/wailsapp/wails/v2/cmd/wails@latest
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
-          rm ./backend-python/rwkv_pip/beta/wkv_cuda.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/librwkv.dylib
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
@@ -108,20 +107,19 @@ jobs:
     runs-on: macos-13
     needs: create-draft
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           ref: master
-      - uses: actions/setup-go@v4
+      - uses: actions/setup-go@v5
         with:
-          go-version: '1.20.5'
+          go-version: "1.20.5"
       - run: |
           wget https://github.com/josStorer/ai00_rwkv_server/releases/latest/download/webgpu_server_darwin_aarch64 -O ./backend-rust/webgpu_server
           wget https://github.com/josStorer/web-rwkv-converter/releases/latest/download/web-rwkv-converter_darwin_aarch64 -O ./backend-rust/web-rwkv-converter
-          go install github.com/wailsapp/wails/v2/cmd/wails@latest
+          go install github.com/wailsapp/wails/v2/cmd/wails@v2.8.0
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
-          rm ./backend-python/rwkv_pip/beta/wkv_cuda.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
           rm ./backend-python/rwkv_pip/cpp/librwkv.so
@@ -137,5 +135,5 @@ jobs:
     runs-on: ubuntu-22.04
     needs: [ windows, linux, macos ]
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - run: gh release edit ${{github.ref_name}} --draft=false

.gitignore

@@ -19,7 +19,6 @@ __pycache__
 /cmd-helper.bat
 /install-py-dep.bat
 /backend-python/wkv_cuda
-/backend-python/rwkv*
 *.exe
 *.old
 .DS_Store

CURRENT_CHANGE.md

@@ -1,10 +1,26 @@
-## Changes
+## v1.8.4
 
-- rwkv5 lora finetune support (https://github.com/JL-er/RWKV-v5-lora)
-- improve memory usage and speed of convert_safetensors.py
-- webgpu(python) state cache support
-- improve state cache performance (especially for rwkv.cpp)
-- chore
+- fix f05a4a, __init__.py is not embedded
+
+## v1.8.3
+
+### Deprecations
+
+- rwkv-beta is deprecated
+
+### Upgrades
+
+- bump webgpu(python) (https://github.com/cryscan/web-rwkv-py)
+- sync https://github.com/JL-er/RWKV-PEFT (LoRA)
+
+### Improvements
+
+- improve default LoRA fine-tune params
+
+### Fixes
+
+- fix #342, #345: cannot import name 'packaging' from 'pkg_resources'
+- fix the huge error prompt that pops up when running in webgpu mode
 
 ## Install
 

Dockerfile

@@ -0,0 +1,55 @@
+FROM node:21-slim AS frontend
+
+RUN echo "registry=https://registry.npmmirror.com/" > ~/.npmrc
+
+WORKDIR /app
+
+COPY manifest.json manifest.json
+COPY frontend frontend
+
+WORKDIR /app/frontend
+
+RUN npm ci
+RUN npm run build
+
+FROM nvidia/cuda:11.6.1-devel-ubuntu20.04 AS runtime
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt update && \
+    apt install -yq git curl wget build-essential ninja-build aria2 jq software-properties-common
+
+RUN add-apt-repository -y ppa:deadsnakes/ppa && \
+    add-apt-repository -y ppa:ubuntu-toolchain-r/test && \
+    apt install -y g++-11 python3.10 python3.10-distutils python3.10-dev && \
+    curl -sS http://mirrors.aliyun.com/pypi/get-pip.py | python3.10
+
+RUN python3.10 -m pip install cmake
+
+FROM runtime AS librwkv
+
+WORKDIR /app
+
+RUN git clone https://github.com/RWKV/rwkv.cpp.git && \
+    cd rwkv.cpp && \
+    git submodule update --init --recursive && \
+    mkdir -p build && \
+    cd build && \
+    cmake -G Ninja .. && \
+    cmake --build .
+
+FROM runtime AS final
+
+WORKDIR /app
+
+COPY ./backend-python/requirements.txt ./backend-python/requirements.txt
+
+RUN python3.10 -m pip install --quiet -r ./backend-python/requirements.txt
+
+COPY . .
+COPY --from=frontend /app/frontend/dist /app/frontend/dist
+COPY --from=librwkv /app/rwkv.cpp/build/librwkv.so /app/backend-python/rwkv_pip/cpp/librwkv.so
+
+EXPOSE 27777
+
+CMD ["python3.10", "./backend-python/main.py", "--port", "27777", "--host", "0.0.0.0", "--webui"]

Makefile

@@ -8,7 +8,8 @@ endif
 
 build-windows:
 	@echo ---- build for windows
-	wails build -upx -ldflags '-s -w -extldflags "-static"' -platform windows/amd64
+	wails build -ldflags '-s -w -extldflags "-static"' -platform windows/amd64
+	upx -9 --lzma ./build/bin/RWKV-Runner.exe
 
 build-macos:
 	@echo ---- build for macos
@@ -16,7 +17,8 @@ build-macos:
 
 build-linux:
 	@echo ---- build for linux
-	wails build -upx -ldflags '-s -w' -platform linux/amd64
+	wails build -ldflags '-s -w' -platform linux/amd64
+	upx -9 --lzma ./build/bin/RWKV-Runner
 
 build-web:
 	@echo ---- build for web

README.md

@@ -1,5 +1,5 @@
 <p align="center">
-    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/d24834b0-265d-45f5-93c0-fac1e19562af">
+    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/65c46133-7506-4b54-b64f-fe49f188afa7">
 </p>
 
 <h1 align="center">RWKV Runner</h1>
@@ -12,6 +12,7 @@ compatible with the OpenAI API, which means that every ChatGPT client is an RWKV
 
 [![license][license-image]][license-url]
 [![release][release-image]][release-url]
+[![py-version][py-version-image]][py-version-url]
 
 English | [简体中文](README_ZH.md) | [日本語](README_JA.md)
 
@@ -31,6 +32,10 @@ English | [简体中文](README_ZH.md) | [日本語](README_JA.md)
 
 [release-url]: https://github.com/josStorer/RWKV-Runner/releases/latest
 
+[py-version-image]: https://img.shields.io/pypi/pyversions/fastapi.svg
+
+[py-version-url]: https://github.com/josStorer/RWKV-Runner/tree/master/backend-python
+
 [download-url]: https://github.com/josStorer/RWKV-Runner/releases
 
 [Windows-image]: https://img.shields.io/badge/-Windows-blue?logo=windows
@@ -89,7 +94,8 @@ English | [简体中文](README_ZH.md) | [日本語](README_JA.md)
 - Built-in model conversion tool.
 - Built-in download management and remote model inspection.
 - Built-in one-click LoRA Finetune. (Windows Only)
-- Can also be used as an OpenAI ChatGPT and GPT-Playground client. (Fill in the API URL and API Key in Settings page)
+- Can also be used as an OpenAI ChatGPT, GPT-Playground, Ollama and more clients. (Fill in the API URL and API Key in
+  Settings page)
 - Multilingual localization.
 - Theme switching.
 - Automatic updates.
@@ -231,22 +237,24 @@ computer keyboard as MIDI input.
 - ChatRWKV: https://github.com/BlinkDL/ChatRWKV
 - RWKV-LM: https://github.com/BlinkDL/RWKV-LM
 - RWKV-LM-LoRA: https://github.com/Blealtan/RWKV-LM-LoRA
+- RWKV-v5-lora: https://github.com/JL-er/RWKV-v5-lora
 - MIDI-LLM-tokenizer: https://github.com/briansemrau/MIDI-LLM-tokenizer
 - ai00_rwkv_server: https://github.com/cgisky1980/ai00_rwkv_server
 - rwkv.cpp: https://github.com/saharNooby/rwkv.cpp
 - web-rwkv-py: https://github.com/cryscan/web-rwkv-py
+- web-rwkv: https://github.com/cryscan/web-rwkv
 
 ## Preview
 
 ### Homepage
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/c9b9cdd0-63f9-4319-9f74-5bf5d7df5a67)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/c1923ed8-22f7-48b4-a274-e215e27a8e01)
 
 ### Chat
 
 ![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/80009872-528f-4932-aeb2-f724fa892e7c)
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/e98c9038-3323-47b0-8edb-d639fafd37b2)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/de8d3fa7-c31f-4941-a22b-ded785427ac0)
 
 ### Completion
 

README_JA.md

@@ -1,5 +1,5 @@
 <p align="center">
-    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/d24834b0-265d-45f5-93c0-fac1e19562af">
+    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/65c46133-7506-4b54-b64f-fe49f188afa7">
 </p>
 
 <h1 align="center">RWKV Runner</h1>
@@ -12,6 +12,7 @@
 
 [![license][license-image]][license-url]
 [![release][release-image]][release-url]
+[![py-version][py-version-image]][py-version-url]
 
 [English](README.md) | [简体中文](README_ZH.md) | 日本語
 
@@ -31,6 +32,10 @@
 
 [release-url]: https://github.com/josStorer/RWKV-Runner/releases/latest
 
+[py-version-image]: https://img.shields.io/pypi/pyversions/fastapi.svg
+
+[py-version-url]: https://github.com/josStorer/RWKV-Runner/tree/master/backend-python
+
 [download-url]: https://github.com/josStorer/RWKV-Runner/releases
 
 [Windows-image]: https://img.shields.io/badge/-Windows-blue?logo=windows
@@ -84,8 +89,8 @@
 - 内蔵モデル変換ツール
 - ダウンロード管理とリモートモデル検査機能内蔵
 - 内蔵のLoRA微調整機能を搭載しています (Windowsのみ)
-- このプログラムは、OpenAI ChatGPTとGPT Playgroundのクライアントとしても使用できます（設定ページで `API URL` と `API Key`
-  を入力してください）
+- このプログラムは、OpenAI ChatGPT、GPT Playground、Ollama などのクライアントとしても使用できます（設定ページで `API URL`
+  と `API Key` を入力してください）
 - 多言語ローカライズ
 - テーマ切り替え
 - 自動アップデート
@@ -228,22 +233,24 @@ MIDIキーボードをお持ちでない場合、`Virtual Midi Controller 3 LE`
 - ChatRWKV: https://github.com/BlinkDL/ChatRWKV
 - RWKV-LM: https://github.com/BlinkDL/RWKV-LM
 - RWKV-LM-LoRA: https://github.com/Blealtan/RWKV-LM-LoRA
+- RWKV-v5-lora: https://github.com/JL-er/RWKV-v5-lora
 - MIDI-LLM-tokenizer: https://github.com/briansemrau/MIDI-LLM-tokenizer
 - ai00_rwkv_server: https://github.com/cgisky1980/ai00_rwkv_server
 - rwkv.cpp: https://github.com/saharNooby/rwkv.cpp
 - web-rwkv-py: https://github.com/cryscan/web-rwkv-py
+- web-rwkv: https://github.com/cryscan/web-rwkv
 
 ## Preview
 
 ### ホームページ
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/c9b9cdd0-63f9-4319-9f74-5bf5d7df5a67)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/c1923ed8-22f7-48b4-a274-e215e27a8e01)
 
 ### チャット
 
 ![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/80009872-528f-4932-aeb2-f724fa892e7c)
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/e98c9038-3323-47b0-8edb-d639fafd37b2)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/de8d3fa7-c31f-4941-a22b-ded785427ac0)
 
 ### 補完
 

README_ZH.md

@@ -1,5 +1,5 @@
 <p align="center">
-    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/d24834b0-265d-45f5-93c0-fac1e19562af">
+    <img src="https://github.com/josStorer/RWKV-Runner/assets/13366013/65c46133-7506-4b54-b64f-fe49f188afa7">
 </p>
 
 <h1 align="center">RWKV Runner</h1>
@@ -11,6 +11,7 @@ API兼容的接口，这意味着一切ChatGPT客户端都是RWKV客户端。
 
 [![license][license-image]][license-url]
 [![release][release-image]][release-url]
+[![py-version][py-version-image]][py-version-url]
 
 [English](README.md) | 简体中文 | [日本語](README_JA.md)
 
@@ -30,6 +31,10 @@ API兼容的接口，这意味着一切ChatGPT客户端都是RWKV客户端。
 
 [release-url]: https://github.com/josStorer/RWKV-Runner/releases/latest
 
+[py-version-image]: https://img.shields.io/pypi/pyversions/fastapi.svg
+
+[py-version-url]: https://github.com/josStorer/RWKV-Runner/tree/master/backend-python
+
 [download-url]: https://github.com/josStorer/RWKV-Runner/releases
 
 [Windows-image]: https://img.shields.io/badge/-Windows-blue?logo=windows
@@ -78,7 +83,7 @@ API兼容的接口，这意味着一切ChatGPT客户端都是RWKV客户端。
 - 内置模型转换工具
 - 内置下载管理和远程模型检视
 - 内置一键LoRA微调 (仅限Windows)
-- 也可用作 OpenAI ChatGPT 和 GPT Playground 客户端 (在设置内填写API URL和API Key)
+- 也可用作 OpenAI ChatGPT, GPT Playground, Ollama 等服务的客户端 (在设置内填写API URL和API Key)
 - 多语言本地化
 - 主题切换
 - 自动更新
@@ -210,22 +215,24 @@ for i in np.argsort(embeddings_cos_sim)[::-1]:
 - ChatRWKV: https://github.com/BlinkDL/ChatRWKV
 - RWKV-LM: https://github.com/BlinkDL/RWKV-LM
 - RWKV-LM-LoRA: https://github.com/Blealtan/RWKV-LM-LoRA
+- RWKV-v5-lora: https://github.com/JL-er/RWKV-v5-lora
 - MIDI-LLM-tokenizer: https://github.com/briansemrau/MIDI-LLM-tokenizer
 - ai00_rwkv_server: https://github.com/cgisky1980/ai00_rwkv_server
 - rwkv.cpp: https://github.com/saharNooby/rwkv.cpp
 - web-rwkv-py: https://github.com/cryscan/web-rwkv-py
+- web-rwkv: https://github.com/cryscan/web-rwkv
 
 ## Preview
 
 ### 主页
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/cd82674e-3ee3-4175-bd9c-a11d45437327)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/3265b11a-ab19-4e19-bfea-fc687f59aaf9)
 
 ### 聊天
 
 ![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/9570e73b-dca2-4316-9e92-09961f3c48c4)
 
-![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/54bb0e2b-cdc4-4ea0-8d16-9beaf57c232c)
+![image](https://github.com/josStorer/RWKV-Runner/assets/13366013/162fce43-8568-4850-a6af-ab60af988da6)
 
 ### 续写
 

backend-golang/app.go

@@ -1,15 +1,22 @@
 package backend_golang
 
 import (
+	"archive/zip"
 	"bufio"
+	"bytes"
 	"context"
 	"errors"
 	"io"
+	"log"
+	"net"
 	"net/http"
+	"net/http/httputil"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"syscall"
 	"time"
 
@@ -23,6 +30,8 @@ type App struct {
 	ctx           context.Context
 	HasConfigData bool
 	ConfigData    map[string]any
+	Dev           bool
+	proxyPort     int
 	exDir         string
 	cmdPrefix     string
 }
@@ -32,6 +41,63 @@ func NewApp() *App {
 	return &App{}
 }
 
+func (a *App) newFetchProxy() {
+	go func() {
+		handler := func(w http.ResponseWriter, r *http.Request) {
+			if r.Method == "OPTIONS" {
+				w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+				w.Header().Set("Access-Control-Allow-Headers", "*")
+				w.Header().Set("Access-Control-Allow-Origin", "*")
+				return
+			}
+			proxy := &httputil.ReverseProxy{
+				ModifyResponse: func(res *http.Response) error {
+					res.Header.Set("Access-Control-Allow-Origin", "*")
+					return nil
+				},
+				Director: func(req *http.Request) {
+					realTarget := req.Header.Get("Real-Target")
+					if realTarget != "" {
+						realTarget, err := url.PathUnescape(realTarget)
+						if err != nil {
+							log.Printf("Error decoding target URL: %v\n", err)
+							return
+						}
+						target, err := url.Parse(realTarget)
+						if err != nil {
+							log.Printf("Error parsing target URL: %v\n", err)
+							return
+						}
+						req.Header.Set("Accept", "*/*")
+						req.Header.Del("Origin")
+						req.Header.Del("Referer")
+						req.Header.Del("Real-Target")
+						req.Header.Del("Sec-Fetch-Dest")
+						req.Header.Del("Sec-Fetch-Mode")
+						req.Header.Del("Sec-Fetch-Site")
+						req.URL.Scheme = target.Scheme
+						req.URL.Host = target.Host
+						req.URL.Path = target.Path
+						req.URL.RawQuery = url.PathEscape(target.RawQuery)
+						log.Println("Proxying to", realTarget)
+					} else {
+						log.Println("Real-Target header is missing")
+					}
+				},
+			}
+			proxy.ServeHTTP(w, r)
+		}
+		http.HandleFunc("/", handler)
+		listener, err := net.Listen("tcp", "127.0.0.1:0")
+		if err != nil {
+			return
+		}
+		a.proxyPort = listener.Addr().(*net.TCPAddr).Port
+
+		http.Serve(listener, nil)
+	}()
+}
+
 // startup is called when the app starts. The context is saved
 // so we can call the runtime methods
 func (a *App) OnStartup(ctx context.Context) {
@@ -39,20 +105,31 @@ func (a *App) OnStartup(ctx context.Context) {
 	a.exDir = ""
 	a.cmdPrefix = ""
 
-	if runtime.GOOS == "darwin" {
-		ex, _ := os.Executable()
-		a.exDir = filepath.Dir(ex) + "/../../../"
-		a.cmdPrefix = "cd " + a.exDir + " && "
+	ex, err := os.Executable()
+	if err == nil {
+		if runtime.GOOS == "darwin" {
+			a.exDir = filepath.Dir(ex) + "/../../../"
+			a.cmdPrefix = "cd " + a.exDir + " && "
+		} else {
+			a.exDir = filepath.Dir(ex) + "/"
+			a.cmdPrefix = "cd " + a.exDir + " && "
+		}
+		if a.Dev {
+			a.exDir = ""
+		} else {
+			os.Chdir(a.exDir)
+		}
 	}
 
 	os.Chmod(a.exDir+"backend-rust/webgpu_server", 0777)
 	os.Chmod(a.exDir+"backend-rust/web-rwkv-converter", 0777)
 	os.Mkdir(a.exDir+"models", os.ModePerm)
 	os.Mkdir(a.exDir+"lora-models", os.ModePerm)
+	os.Mkdir(a.exDir+"state-models", os.ModePerm)
 	os.Mkdir(a.exDir+"finetune/json2binidx_tool/data", os.ModePerm)
-	trainLogPath := a.exDir + "lora-models/train_log.txt"
+	trainLogPath := "lora-models/train_log.txt"
 	if !a.FileExists(trainLogPath) {
-		f, err := os.Create(trainLogPath)
+		f, err := os.Create(a.exDir + trainLogPath)
 		if err == nil {
 			f.Close()
 		}
@@ -62,6 +139,7 @@ func (a *App) OnStartup(ctx context.Context) {
 	a.midiLoop()
 	a.watchFs()
 	a.monitorHardware()
+	a.newFetchProxy()
 }
 
 func (a *App) OnBeforeClose(ctx context.Context) bool {
@@ -74,8 +152,9 @@ func (a *App) OnBeforeClose(ctx context.Context) bool {
 func (a *App) watchFs() {
 	watcher, err := fsnotify.NewWatcher()
 	if err == nil {
-		watcher.Add(a.exDir + "./lora-models")
 		watcher.Add(a.exDir + "./models")
+		watcher.Add(a.exDir + "./lora-models")
+		watcher.Add(a.exDir + "./state-models")
 		go func() {
 			for {
 				select {
@@ -149,6 +228,7 @@ func (a *App) UpdateApp(url string) (broken bool, err error) {
 	ticker := time.NewTicker(250 * time.Millisecond)
 	defer ticker.Stop()
 
+	// update progress
 	go func() {
 		for {
 			<-ticker.C
@@ -168,13 +248,35 @@ func (a *App) UpdateApp(url string) (broken bool, err error) {
 			}
 		}
 	}()
-	err = selfupdate.Apply(pr, selfupdate.Options{})
+
+	var updateFile io.Reader = pr
+	// extract macos binary from zip
+	if strings.HasSuffix(url, ".zip") && runtime.GOOS == "darwin" {
+		zipBytes, err := io.ReadAll(pr)
+		if err != nil {
+			return false, err
+		}
+		archive, err := zip.NewReader(bytes.NewReader(zipBytes), int64(len(zipBytes)))
+		if err != nil {
+			return false, err
+		}
+		file, err := archive.Open("RWKV-Runner.app/Contents/MacOS/RWKV-Runner")
+		if err != nil {
+			return false, err
+		}
+		defer file.Close()
+		updateFile = file
+	}
+
+	// apply update
+	err = selfupdate.Apply(updateFile, selfupdate.Options{})
 	if err != nil {
 		if rerr := selfupdate.RollbackError(err); rerr != nil {
 			return true, rerr
 		}
 		return false, err
 	}
+	// restart app
 	if runtime.GOOS == "windows" {
 		name, err := os.Executable()
 		if err != nil {
@@ -202,3 +304,7 @@ func (a *App) RestartApp() error {
 func (a *App) GetPlatform() string {
 	return runtime.GOOS
 }
+
+func (a *App) GetProxyPort() int {
+	return a.proxyPort
+}

backend-golang/download.go

@@ -10,7 +10,11 @@ import (
 )
 
 func (a *App) DownloadFile(path string, url string) error {
-	_, err := grab.Get(a.exDir+path, url)
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return err
+	}
+	_, err = grab.Get(absPath, url)
 	if err != nil {
 		return err
 	}
@@ -88,11 +92,15 @@ func (a *App) ContinueDownload(url string) {
 }
 
 func (a *App) AddToDownloadList(path string, url string) {
-	if !existsInDownloadList(a.exDir+path, url) {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return
+	}
+	if !existsInDownloadList(absPath, url) {
 		downloadList = append(downloadList, &DownloadStatus{
 			resp:        nil,
 			Name:        filepath.Base(path),
-			Path:        a.exDir + path,
+			Path:        absPath,
 			Url:         url,
 			Downloading: false,
 		})

backend-golang/file.go

@@ -14,27 +14,55 @@ import (
 	wruntime "github.com/wailsapp/wails/v2/pkg/runtime"
 )
 
+func (a *App) GetAbsPath(path string) (string, error) {
+	var absPath string
+	var err error
+	if filepath.IsAbs(path) {
+		absPath = filepath.Clean(path)
+	} else {
+		absPath, err = filepath.Abs(filepath.Join(a.exDir, path))
+		if err != nil {
+			return "", err
+		}
+	}
+	absPath = strings.ReplaceAll(absPath, "/", string(os.PathSeparator))
+	println("GetAbsPath:", absPath)
+	return absPath, nil
+}
+
 func (a *App) SaveFile(path string, savedContent []byte) error {
-	if err := os.WriteFile(a.exDir+path, savedContent, 0644); err != nil {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return err
+	}
+	if err := os.WriteFile(absPath, savedContent, 0644); err != nil {
 		return err
 	}
 	return nil
 }
 
-func (a *App) SaveJson(fileName string, jsonData any) error {
+func (a *App) SaveJson(path string, jsonData any) error {
 	text, err := json.MarshalIndent(jsonData, "", "  ")
 	if err != nil {
 		return err
 	}
 
-	if err := os.WriteFile(a.exDir+fileName, text, 0644); err != nil {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return err
+	}
+	if err := os.WriteFile(absPath, text, 0644); err != nil {
 		return err
 	}
 	return nil
 }
 
-func (a *App) ReadJson(fileName string) (any, error) {
-	file, err := os.ReadFile(a.exDir + fileName)
+func (a *App) ReadJson(path string) (any, error) {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return nil, err
+	}
+	file, err := os.ReadFile(absPath)
 	if err != nil {
 		return nil, err
 	}
@@ -48,8 +76,12 @@ func (a *App) ReadJson(fileName string) (any, error) {
 	return data, nil
 }
 
-func (a *App) FileExists(fileName string) bool {
-	_, err := os.Stat(a.exDir + fileName)
+func (a *App) FileExists(path string) bool {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return false
+	}
+	_, err = os.Stat(absPath)
 	return err == nil
 }
 
@@ -60,8 +92,12 @@ type FileInfo struct {
 	ModTime string `json:"modTime"`
 }
 
-func (a *App) ReadFileInfo(fileName string) (*FileInfo, error) {
-	info, err := os.Stat(a.exDir + fileName)
+func (a *App) ReadFileInfo(path string) (*FileInfo, error) {
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return nil, err
+	}
+	info, err := os.Stat(absPath)
 	if err != nil {
 		return nil, err
 	}
@@ -74,7 +110,11 @@ func (a *App) ReadFileInfo(fileName string) (*FileInfo, error) {
 }
 
 func (a *App) ListDirFiles(dirPath string) ([]FileInfo, error) {
-	files, err := os.ReadDir(a.exDir + dirPath)
+	absDirPath, err := a.GetAbsPath(dirPath)
+	if err != nil {
+		return nil, err
+	}
+	files, err := os.ReadDir(absDirPath)
 	if err != nil {
 		return nil, err
 	}
@@ -96,7 +136,11 @@ func (a *App) ListDirFiles(dirPath string) ([]FileInfo, error) {
 }
 
 func (a *App) DeleteFile(path string) error {
-	err := os.Remove(a.exDir + path)
+	absPath, err := a.GetAbsPath(path)
+	if err != nil {
+		return err
+	}
+	err = os.Remove(absPath)
 	if err != nil {
 		return err
 	}
@@ -104,18 +148,27 @@ func (a *App) DeleteFile(path string) error {
 }
 
 func (a *App) CopyFile(src string, dst string) error {
-	sourceFile, err := os.Open(a.exDir + src)
+	absSrc, err := a.GetAbsPath(src)
+	if err != nil {
+		return err
+	}
+	absDst, err := a.GetAbsPath(dst)
+	if err != nil {
+		return err
+	}
+
+	sourceFile, err := os.Open(absSrc)
 	if err != nil {
 		return err
 	}
 	defer sourceFile.Close()
 
-	err = os.MkdirAll(a.exDir+dst[:strings.LastIndex(dst, "/")], 0755)
+	err = os.MkdirAll(filepath.Dir(absDst), 0755)
 	if err != nil {
 		return err
 	}
 
-	destFile, err := os.Create(a.exDir + dst)
+	destFile, err := os.Create(absDst)
 	if err != nil {
 		return err
 	}
@@ -166,14 +219,8 @@ func (a *App) OpenOpenFileDialog(filterPattern string) (string, error) {
 	return path, nil
 }
 
-func (a *App) OpenFileFolder(path string, relative bool) error {
-	var absPath string
-	var err error
-	if relative {
-		absPath, err = filepath.Abs(a.exDir + path)
-	} else {
-		absPath, err = filepath.Abs(path)
-	}
+func (a *App) OpenFileFolder(path string) error {
+	absPath, err := a.GetAbsPath(path)
 	if err != nil {
 		return err
 	}

backend-golang/rwkv.go

@@ -1,3 +1,4 @@
+// Considering some whitespace and multilingual support, the functions in rwkv.go should always be executed with cwd as RWKV-Runner, and never use a.GetAbsPath() here.
 package backend_golang
 
 import (
@@ -11,19 +12,23 @@ import (
 )
 
 func (a *App) StartServer(python string, port int, host string, webui bool, rwkvBeta bool, rwkvcpp bool, webgpu bool) (string, error) {
-	var err error
+	execFile := "./backend-python/main.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
 	if err != nil {
 		return "", err
 	}
-	args := []string{python, "./backend-python/main.py"}
+	args := []string{python, execFile}
 	if webui {
 		args = append(args, "--webui")
 	}
 	if rwkvBeta {
-		args = append(args, "--rwkv-beta")
+		// args = append(args, "--rwkv-beta")
 	}
 	if rwkvcpp {
 		args = append(args, "--rwkv.cpp")
@@ -36,41 +41,77 @@ func (a *App) StartServer(python string, port int, host string, webui bool, rwkv
 }
 
 func (a *App) StartWebGPUServer(port int, host string) (string, error) {
-	args := []string{"./backend-rust/webgpu_server"}
+	var execFile string
+	execFiles := []string{"./backend-rust/webgpu_server", "./backend-rust/webgpu_server.exe"}
+	for _, file := range execFiles {
+		_, err := os.Stat(file)
+		if err == nil {
+			execFile = file
+			break
+		}
+	}
+	if execFile == "" {
+		return "", errors.New(execFiles[0] + " not found")
+	}
+	args := []string{execFile}
 	args = append(args, "--port", strconv.Itoa(port), "--ip", host)
 	return Cmd(args...)
 }
 
 func (a *App) ConvertModel(python string, modelPath string, strategy string, outPath string) (string, error) {
-	var err error
+	execFile := "./backend-python/convert_model.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
 	if err != nil {
 		return "", err
 	}
-	return Cmd(python, "./backend-python/convert_model.py", "--in", modelPath, "--out", outPath, "--strategy", strategy)
+	return Cmd(python, execFile, "--in", modelPath, "--out", outPath, "--strategy", strategy)
 }
 
 func (a *App) ConvertSafetensors(modelPath string, outPath string) (string, error) {
-	args := []string{"./backend-rust/web-rwkv-converter"}
+	var execFile string
+	execFiles := []string{"./backend-rust/web-rwkv-converter", "./backend-rust/web-rwkv-converter.exe"}
+	for _, file := range execFiles {
+		_, err := os.Stat(file)
+		if err == nil {
+			execFile = file
+			break
+		}
+	}
+	if execFile == "" {
+		return "", errors.New(execFiles[0] + " not found")
+	}
+	args := []string{execFile}
 	args = append(args, "--input", modelPath, "--output", outPath)
 	return Cmd(args...)
 }
 
 func (a *App) ConvertSafetensorsWithPython(python string, modelPath string, outPath string) (string, error) {
-	var err error
+	execFile := "./backend-python/convert_safetensors.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
 	if err != nil {
 		return "", err
 	}
-	return Cmd(python, "./backend-python/convert_safetensors.py", "--input", modelPath, "--output", outPath)
+	return Cmd(python, execFile, "--input", modelPath, "--output", outPath)
 }
 
 func (a *App) ConvertGGML(python string, modelPath string, outPath string, Q51 bool) (string, error) {
-	var err error
+	execFile := "./backend-python/convert_pytorch_to_ggml.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
@@ -81,11 +122,15 @@ func (a *App) ConvertGGML(python string, modelPath string, outPath string, Q51 b
 	if Q51 {
 		dataType = "Q5_1"
 	}
-	return Cmd(python, "./backend-python/convert_pytorch_to_ggml.py", modelPath, outPath, dataType)
+	return Cmd(python, execFile, modelPath, outPath, dataType)
 }
 
 func (a *App) ConvertData(python string, input string, outputPrefix string, vocab string) (string, error) {
-	var err error
+	execFile := "./finetune/json2binidx_tool/tools/preprocess_data.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
@@ -129,19 +174,23 @@ func (a *App) ConvertData(python string, input string, outputPrefix string, voca
 		return "", err
 	}
 
-	return Cmd(python, "./finetune/json2binidx_tool/tools/preprocess_data.py", "--input", input, "--output-prefix", outputPrefix, "--vocab", vocab,
+	return Cmd(python, execFile, "--input", input, "--output-prefix", outputPrefix, "--vocab", vocab,
 		"--tokenizer-type", tokenizerType, "--dataset-impl", "mmap", "--append-eod")
 }
 
 func (a *App) MergeLora(python string, useGpu bool, loraAlpha int, baseModel string, loraPath string, outputPath string) (string, error) {
-	var err error
+	execFile := "./finetune/lora/merge_lora.py"
+	_, err := os.Stat(execFile)
+	if err != nil {
+		return "", err
+	}
 	if python == "" {
 		python, err = GetPython()
 	}
 	if err != nil {
 		return "", err
 	}
-	args := []string{python, "./finetune/lora/merge_lora.py"}
+	args := []string{python, execFile}
 	if useGpu {
 		args = append(args, "--use-gpu")
 	}
@@ -157,17 +206,21 @@ func (a *App) DepCheck(python string) error {
 	if err != nil {
 		return err
 	}
-	out, err := exec.Command(python, a.exDir+"./backend-python/dep_check.py").CombinedOutput()
+	out, err := exec.Command(python, a.exDir+"backend-python/dep_check.py").CombinedOutput()
 	if err != nil {
-		return errors.New("DepCheck Error: " + string(out))
+		return errors.New("DepCheck Error: " + string(out) + " GError: " + err.Error())
 	}
 	return nil
 }
 
 func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
 	var err error
+	torchWhlUrl := "torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 --index-url https://download.pytorch.org/whl/cu117"
 	if python == "" {
 		python, err = GetPython()
+		if cnMirror && python == "py310/python.exe" {
+			torchWhlUrl = "https://mirrors.aliyun.com/pytorch-wheels/cu117/torch-1.13.1+cu117-cp310-cp310-win_amd64.whl"
+		}
 		if runtime.GOOS == "windows" {
 			python = `"%CD%/` + python + `"`
 		}
@@ -178,14 +231,14 @@ func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
 
 	if runtime.GOOS == "windows" {
 		ChangeFileLine("./py310/python310._pth", 3, "Lib\\site-packages")
-		installScript := python + " ./backend-python/get-pip.py -i https://pypi.tuna.tsinghua.edu.cn/simple --no-warn-script-location\n" +
-			python + " -m pip install torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 --index-url https://download.pytorch.org/whl/cu117 --no-warn-script-location\n" +
-			python + " -m pip install -r ./backend-python/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple --no-warn-script-location\n" +
+		installScript := python + " ./backend-python/get-pip.py -i https://mirrors.aliyun.com/pypi/simple --no-warn-script-location\n" +
+			python + " -m pip install " + torchWhlUrl + " --no-warn-script-location\n" +
+			python + " -m pip install -r ./backend-python/requirements.txt -i https://mirrors.aliyun.com/pypi/simple --no-warn-script-location\n" +
 			"exit"
 		if !cnMirror {
-			installScript = strings.Replace(installScript, " -i https://pypi.tuna.tsinghua.edu.cn/simple", "", -1)
+			installScript = strings.Replace(installScript, " -i https://mirrors.aliyun.com/pypi/simple", "", -1)
 		}
-		err = os.WriteFile("./install-py-dep.bat", []byte(installScript), 0644)
+		err = os.WriteFile(a.exDir+"install-py-dep.bat", []byte(installScript), 0644)
 		if err != nil {
 			return "", err
 		}
@@ -193,7 +246,7 @@ func (a *App) InstallPyDep(python string, cnMirror bool) (string, error) {
 	}
 
 	if cnMirror {
-		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt", "-i", "https://pypi.tuna.tsinghua.edu.cn/simple")
+		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt", "-i", "https://mirrors.aliyun.com/pypi/simple")
 	} else {
 		return Cmd(python, "-m", "pip", "install", "-r", "./backend-python/requirements_without_cyac.txt")
 	}

backend-golang/utils.go

@@ -23,14 +23,19 @@ func CmdHelper(hideWindow bool, args ...string) (*exec.Cmd, error) {
 	if runtime.GOOS != "windows" {
 		return nil, errors.New("unsupported OS")
 	}
-	filename := "./cmd-helper.bat"
-	_, err := os.Stat(filename)
+	ex, err := os.Executable()
 	if err != nil {
-		if err := os.WriteFile(filename, []byte("start %*"), 0644); err != nil {
+		return nil, err
+	}
+	exDir := filepath.Dir(ex) + "/"
+	path := exDir + "cmd-helper.bat"
+	_, err = os.Stat(path)
+	if err != nil {
+		if err := os.WriteFile(path, []byte("start %*"), 0644); err != nil {
 			return nil, err
 		}
 	}
-	cmdHelper, err := filepath.Abs(filename)
+	cmdHelper, err := filepath.Abs(path)
 	if err != nil {
 		return nil, err
 	}
@@ -86,16 +91,18 @@ func Cmd(args ...string) (string, error) {
 }
 
 func CopyEmbed(efs embed.FS) error {
-	prefix := ""
+	ex, err := os.Executable()
+	if err != nil {
+		return err
+	}
+	var prefix string
 	if runtime.GOOS == "darwin" {
-		ex, err := os.Executable()
-		if err != nil {
-			return err
-		}
 		prefix = filepath.Dir(ex) + "/../../../"
+	} else {
+		prefix = filepath.Dir(ex) + "/"
 	}
 
-	err := fs.WalkDir(efs, ".", func(path string, d fs.DirEntry, err error) error {
+	err = fs.WalkDir(efs, ".", func(path string, d fs.DirEntry, err error) error {
 		if d.IsDir() {
 			return nil
 		}
@@ -136,13 +143,19 @@ func CopyEmbed(efs embed.FS) error {
 func GetPython() (string, error) {
 	switch platform := runtime.GOOS; platform {
 	case "windows":
-		_, err := os.Stat("py310/python.exe")
+		ex, err := os.Executable()
 		if err != nil {
-			_, err := os.Stat("python-3.10.11-embed-amd64.zip")
+			return "", err
+		}
+		exDir := filepath.Dir(ex) + "/"
+		pyexe := exDir + "py310/python.exe"
+		_, err = os.Stat(pyexe)
+		if err != nil {
+			_, err := os.Stat(exDir + "python-3.10.11-embed-amd64.zip")
 			if err != nil {
 				return "", errors.New("python zip not found")
 			} else {
-				err := Unzip("python-3.10.11-embed-amd64.zip", "py310")
+				err := Unzip(exDir+"python-3.10.11-embed-amd64.zip", exDir+"py310")
 				if err != nil {
 					return "", errors.New("failed to unzip python")
 				} else {

backend-golang/wsl_windows.go

@@ -9,7 +9,6 @@ import (
 	"io"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"strings"
 	"time"
 
@@ -133,26 +132,20 @@ func (a *App) WslStop() error {
 }
 
 func (a *App) WslIsEnabled() error {
-	ex, err := os.Executable()
-	if err != nil {
-		return err
-	}
-	exDir := filepath.Dir(ex)
-
-	data, err := os.ReadFile(exDir + "/wsl.state")
+	data, err := os.ReadFile(a.exDir + "wsl.state")
 	if err == nil {
 		if strings.Contains(string(data), "Enabled") {
 			return nil
 		}
 	}
 
-	cmd := `-Command (Get-WindowsOptionalFeature -Online -FeatureName Microsoft-Windows-Subsystem-Linux).State | Out-File -Encoding utf8 -FilePath ` + exDir + "/wsl.state"
-	_, err = su.ShellExecute(su.RUNAS, "powershell", cmd, exDir)
+	cmd := `-Command (Get-WindowsOptionalFeature -Online -FeatureName VirtualMachinePlatform).State | Out-File -Encoding utf8 -FilePath ` + a.exDir + "wsl.state"
+	_, err = su.ShellExecute(su.RUNAS, "powershell", cmd, a.exDir)
 	if err != nil {
 		return err
 	}
 	time.Sleep(2 * time.Second)
-	data, err = os.ReadFile(exDir + "/wsl.state")
+	data, err = os.ReadFile(a.exDir + "wsl.state")
 	if err != nil {
 		return err
 	}
@@ -164,13 +157,13 @@ func (a *App) WslIsEnabled() error {
 }
 
 func (a *App) WslEnable(forceMode bool) error {
-	cmd := `/online /enable-feature /featurename:Microsoft-Windows-Subsystem-Linux`
+	cmd := `/online /enable-feature /featurename:VirtualMachinePlatform`
 	_, err := su.ShellExecute(su.RUNAS, "dism", cmd, `C:\`)
 	if err != nil {
 		return err
 	}
 	if forceMode {
-		os.WriteFile("./wsl.state", []byte("Enabled"), 0644)
+		os.WriteFile(a.exDir+"wsl.state", []byte("Enabled"), 0644)
 	}
 	return nil
 }

backend-python/convert_safetensors.py

@@ -54,19 +54,21 @@ def convert_file(pt_filename: str, sf_filename: str, rename={}, transpose_names=
                     loaded[k].unsqueeze(1).repeat(1, n_emb // loaded[k].shape[0])
                 )
 
-    for k in kk:
-        new_k = rename_key(rename, k).lower()
-        v = loaded[k].half()
-        del loaded[k]
-        for transpose_name in transpose_names:
-            if transpose_name in k:
-                v = v.transpose(0, 1)
-        print(f"{new_k}\t{v.shape}\t{v.dtype}")
-        loaded[new_k] = {
-            "dtype": str(v.dtype).split(".")[-1],
-            "shape": v.shape,
-            "data": v.numpy().tobytes(),
-        }
+    with torch.no_grad():
+        for k in kk:
+            new_k = rename_key(rename, k).lower()
+            v = loaded[k].half()
+            del loaded[k]
+            for transpose_name in transpose_names:
+                if transpose_name in new_k:
+                    dims = len(v.shape)
+                    v = v.transpose(dims - 2, dims - 1)
+            print(f"{new_k}\t{v.shape}\t{v.dtype}")
+            loaded[new_k] = {
+                "dtype": str(v.dtype).split(".")[-1],
+                "shape": v.shape,
+                "data": v.numpy().tobytes(),
+            }
 
     dirname = os.path.dirname(sf_filename)
     os.makedirs(dirname, exist_ok=True)
@@ -100,6 +102,8 @@ if __name__ == "__main__":
                 "time_mix_w2",
                 "time_decay_w1",
                 "time_decay_w2",
+                "time_state",
+                "lora.0",
             ],
         )
         print(f"Saved to {args.output}")

backend-python/dep_check.py

@@ -1,3 +1,8 @@
+import setuptools
+
+if setuptools.__version__ >= "70.0.0":
+    raise ImportError("setuptools>=70.0.0 is not supported")
+
 import multipart
 import fitz
 import safetensors
@@ -7,7 +12,6 @@ import lm_dataformat
 import ftfy
 import tqdm
 import tiktoken
-import GPUtil
 
 import torch
 import rwkv

backend-python/global_var.py

@@ -5,6 +5,7 @@ Model = "model"
 Model_Status = "model_status"
 Model_Config = "model_config"
 Deploy_Mode = "deploy_mode"
+Midi_Vocab_Config_Type = "midi_vocab_config_type"
 
 
 class ModelStatus(Enum):
@@ -13,11 +14,17 @@ class ModelStatus(Enum):
     Working = 3
 
 
+class MidiVocabConfig(Enum):
+    Default = auto()
+    Piano = auto()
+
+
 def init():
     global GLOBALS
     GLOBALS = {}
     set(Model_Status, ModelStatus.Offline)
     set(Deploy_Mode, False)
+    set(Midi_Vocab_Config_Type, MidiVocabConfig.Default)
 
 
 def set(key, value):

backend-python/main.py

@@ -27,11 +27,6 @@ def get_args(args: Union[Sequence[str], None] = None):
         action="store_true",
         help="whether to enable WebUI (default: False)",
     )
-    group.add_argument(
-        "--rwkv-beta",
-        action="store_true",
-        help="whether to use rwkv-beta (default: False)",
-    )
     group.add_argument(
         "--rwkv.cpp",
         action="store_true",

backend-python/requirements.txt

@@ -1,9 +1,10 @@
 torch
 torchvision
 torchaudio
-rwkv==0.8.22
+setuptools==69.5.1
+rwkv==0.8.26
 langchain==0.0.322
-fastapi==0.104.0
+fastapi==0.109.1
 uvicorn==0.23.2
 sse-starlette==1.6.5
 pydantic==2.4.2
@@ -19,7 +20,7 @@ midi2audio==0.1.1
 mido==1.3.0
 safetensors==0.4.0
 PyMuPDF==1.23.5
-python-multipart==0.0.6
+python-multipart==0.0.7
 Cython==3.0.4
 cyac==1.9
-torch_directml==0.1.13.1.dev230413
+torch-directml==0.1.13.1.dev230413

backend-python/requirements_without_cyac.txt

@@ -1,9 +1,10 @@
 torch
 torchvision
 torchaudio
-rwkv==0.8.22
+setuptools==69.5.1
+rwkv==0.8.26
 langchain==0.0.322
-fastapi==0.104.0
+fastapi==0.109.1
 uvicorn==0.23.2
 sse-starlette==1.6.5
 pydantic==2.4.2
@@ -19,5 +20,5 @@ midi2audio==0.1.1
 mido==1.3.0
 safetensors==0.4.0
 PyMuPDF==1.23.5
-python-multipart==0.0.6
+python-multipart==0.0.7
 Cython==3.0.4

backend-python/routes/completion.py

@@ -4,6 +4,7 @@ from threading import Lock
 from typing import List, Union
 from enum import Enum
 import base64
+import time
 
 from fastapi import APIRouter, Request, status, HTTPException
 from sse_starlette.sse import EventSourceResponse
@@ -53,8 +54,11 @@ class ChatCompletionBody(ModelConfigBody):
     assistant_name: Union[str, None] = Field(
         None, description="Internal assistant name", min_length=1
     )
+    system_name: Union[str, None] = Field(
+        None, description="Internal system name", min_length=1
+    )
     presystem: bool = Field(
-        True, description="Whether to insert default system prompt at the beginning"
+        False, description="Whether to insert default system prompt at the beginning"
     )
 
     model_config = {
@@ -68,12 +72,13 @@ class ChatCompletionBody(ModelConfigBody):
                 "stop": None,
                 "user_name": None,
                 "assistant_name": None,
+                "system_name": None,
                 "presystem": True,
                 "max_tokens": 1000,
-                "temperature": 1.2,
-                "top_p": 0.5,
-                "presence_penalty": 0.4,
-                "frequency_penalty": 0.4,
+                "temperature": 1,
+                "top_p": 0.3,
+                "presence_penalty": 0,
+                "frequency_penalty": 1,
             }
         }
     }
@@ -94,10 +99,10 @@ class CompletionBody(ModelConfigBody):
                 "stream": False,
                 "stop": None,
                 "max_tokens": 100,
-                "temperature": 1.2,
-                "top_p": 0.5,
-                "presence_penalty": 0.4,
-                "frequency_penalty": 0.4,
+                "temperature": 1,
+                "top_p": 0.3,
+                "presence_penalty": 0,
+                "frequency_penalty": 1,
             }
         }
     }
@@ -144,39 +149,57 @@ async def eval_rwkv(
                 return
             set_rwkv_config(model, global_var.get(global_var.Model_Config))
             set_rwkv_config(model, body)
+            print(get_rwkv_config(model))
 
             response, prompt_tokens, completion_tokens = "", 0, 0
+            completion_start_time = None
             for response, delta, prompt_tokens, completion_tokens in model.generate(
                 prompt,
                 stop=stop,
             ):
+                if not completion_start_time:
+                    completion_start_time = time.time()
                 if await request.is_disconnected():
                     break
                 if stream:
                     yield json.dumps(
                         {
-                            "object": "chat.completion.chunk"
-                            if chat_mode
-                            else "text_completion",
+                            "object": (
+                                "chat.completion.chunk"
+                                if chat_mode
+                                else "text_completion"
+                            ),
                             # "response": response,
                             "model": model.name,
+                            "id": "chatcmpl-123",
+                            "system_fingerprint": "fp_44709d6fcb",
                             "choices": [
-                                {
-                                    "delta": {"content": delta},
-                                    "index": 0,
-                                    "finish_reason": None,
-                                }
-                                if chat_mode
-                                else {
-                                    "text": delta,
-                                    "index": 0,
-                                    "finish_reason": None,
-                                }
+                                (
+                                    {
+                                        "delta": {"role":Role.Assistant.value,"content": delta},
+                                        "index": 0,
+                                        "finish_reason": None,
+                                        "logprobs":None
+                                    }
+                                    if chat_mode
+                                    else {
+                                        "text": delta,
+                                        "index": 0,
+                                        "finish_reason": None,
+                                    }
+                                )
                             ],
                         }
                     )
             # torch_gc()
             requests_num = requests_num - 1
+            completion_end_time = time.time()
+            completion_interval = completion_end_time - completion_start_time
+            tps = 0
+            if completion_interval > 0:
+                tps = completion_tokens / completion_interval
+            print(f"Generation TPS: {tps:.2f}")
+
             if await request.is_disconnected():
                 print(f"{request.client} Stop Waiting")
                 quick_log(
@@ -193,23 +216,28 @@ async def eval_rwkv(
             if stream:
                 yield json.dumps(
                     {
-                        "object": "chat.completion.chunk"
-                        if chat_mode
-                        else "text_completion",
+                        "object": (
+                            "chat.completion.chunk" if chat_mode else "text_completion"
+                        ),
                         # "response": response,
                         "model": model.name,
+                        "id": "chatcmpl-123",
+                        "system_fingerprint": "fp_44709d6fcb",
                         "choices": [
-                            {
-                                "delta": {},
-                                "index": 0,
-                                "finish_reason": "stop",
-                            }
-                            if chat_mode
-                            else {
-                                "text": "",
-                                "index": 0,
-                                "finish_reason": "stop",
-                            }
+                            (
+                                {
+                                    "delta": {},
+                                    "index": 0,
+                                    "logprobs": None,
+                                    "finish_reason": "stop",
+                                }
+                                if chat_mode
+                                else {
+                                    "text": "",
+                                    "index": 0,
+                                    "finish_reason": "stop",
+                                }
+                            )
                         ],
                     }
                 )
@@ -225,38 +253,29 @@ async def eval_rwkv(
                         "total_tokens": prompt_tokens + completion_tokens,
                     },
                     "choices": [
-                        {
-                            "message": {
-                                "role": Role.Assistant.value,
-                                "content": response,
-                            },
-                            "index": 0,
-                            "finish_reason": "stop",
-                        }
-                        if chat_mode
-                        else {
-                            "text": response,
-                            "index": 0,
-                            "finish_reason": "stop",
-                        }
+                        (
+                            {
+                                "message": {
+                                    "role": Role.Assistant.value,
+                                    "content": response,
+                                },
+                                "index": 0,
+                                "finish_reason": "stop",
+                            }
+                            if chat_mode
+                            else {
+                                "text": response,
+                                "index": 0,
+                                "finish_reason": "stop",
+                            }
+                        )
                     ],
                 }
 
 
-@router.post("/v1/chat/completions", tags=["Completions"])
-@router.post("/chat/completions", tags=["Completions"])
-async def chat_completions(body: ChatCompletionBody, request: Request):
-    model: TextRWKV = global_var.get(global_var.Model)
-    if model is None:
-        raise HTTPException(status.HTTP_400_BAD_REQUEST, "model not loaded")
-
-    if body.messages is None or body.messages == []:
-        raise HTTPException(status.HTTP_400_BAD_REQUEST, "messages not found")
-
-    interface = model.interface
-    user = model.user if body.user_name is None else body.user_name
-    bot = model.bot if body.assistant_name is None else body.assistant_name
-
+def chat_template_old(
+    model: TextRWKV, body: ChatCompletionBody, interface: str, user: str, bot: str
+):
     is_raven = model.rwkv_type == RWKVType.Raven
 
     completion_text: str = ""
@@ -325,6 +344,53 @@ The following is a coherent verbose detailed conversation between a girl named {
         completion_text += append_message + "\n\n"
     completion_text += f"{bot}{interface}"
 
+    return completion_text
+
+
+def chat_template(
+    model: TextRWKV, body: ChatCompletionBody, interface: str, user: str, bot: str
+):
+    completion_text: str = ""
+    if body.presystem:
+        completion_text = (
+            f"{user}{interface} hi\n\n{bot}{interface} Hi. "
+            + "I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.\n\n"
+        )
+
+    system = "System" if body.system_name is None else body.system_name
+    for message in body.messages:
+        append_message: str = ""
+        if message.role == Role.User:
+            append_message = f"{user}{interface} " + message.content
+        elif message.role == Role.Assistant:
+            append_message = f"{bot}{interface} " + message.content
+        elif message.role == Role.System:
+            append_message = f"{system}{interface} " + message.content
+        completion_text += append_message + "\n\n"
+    completion_text += f"{bot}{interface}"
+
+    return completion_text
+
+
+@router.post("/v1/chat/completions", tags=["Completions"])
+@router.post("/chat/completions", tags=["Completions"])
+async def chat_completions(body: ChatCompletionBody, request: Request):
+    model: TextRWKV = global_var.get(global_var.Model)
+    if model is None:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, "model not loaded")
+
+    if body.messages is None or body.messages == []:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, "messages not found")
+
+    interface = model.interface
+    user = model.user if body.user_name is None else body.user_name
+    bot = model.bot if body.assistant_name is None else body.assistant_name
+
+    if model.version < 5:
+        completion_text = chat_template_old(model, body, interface, user, bot)
+    else:
+        completion_text = chat_template(model, body, interface, user, bot)
+
     user_code = model.pipeline.decode([model.pipeline.encode(user)[0]])
     bot_code = model.pipeline.decode([model.pipeline.encode(bot)[0]])
     if type(body.stop) == str:
@@ -333,9 +399,9 @@ The following is a coherent verbose detailed conversation between a girl named {
         body.stop.append(f"\n\n{user_code}")
         body.stop.append(f"\n\n{bot_code}")
     elif body.stop is None:
-        body.stop = default_stop
-    if not body.presystem:
-        body.stop.append("\n\n")
+        body.stop = default_stop + [f"\n\n{user_code}", f"\n\n{bot_code}"]
+    # if not body.presystem:
+    #     body.stop.append("\n\n")
 
     if body.stream:
         return EventSourceResponse(

backend-python/routes/config.py

@@ -86,32 +86,58 @@ def switch_model(body: SwitchModelBody, response: Response, request: Request):
 
     if body.deploy:
         global_var.set(global_var.Deploy_Mode, True)
-    if global_var.get(global_var.Model_Config) is None:
-        global_var.set(
-            global_var.Model_Config, get_rwkv_config(global_var.get(global_var.Model))
-        )
+
+    saved_model_config = global_var.get(global_var.Model_Config)
+    init_model_config = get_rwkv_config(global_var.get(global_var.Model))
+    if saved_model_config is not None:
+        merge_model(init_model_config, saved_model_config)
+    global_var.set(global_var.Model_Config, init_model_config)
     global_var.set(global_var.Model_Status, global_var.ModelStatus.Working)
 
     return "success"
 
 
+def merge_model(to_model: BaseModel, from_model: BaseModel):
+    from_model_fields = [x for x in from_model.dict().keys()]
+    to_model_fields = [x for x in to_model.dict().keys()]
+
+    for field_name in from_model_fields:
+        if field_name in to_model_fields:
+            from_value = getattr(from_model, field_name)
+
+            if from_value is not None:
+                setattr(to_model, field_name, from_value)
+
+
 @router.post("/update-config", tags=["Configs"])
 def update_config(body: ModelConfigBody):
     """
     Will not update the model config immediately, but set it when completion called to avoid modifications during generation
     """
 
-    print(body)
-    global_var.set(global_var.Model_Config, body)
+    model_config = global_var.get(global_var.Model_Config)
+    if model_config is None:
+        model_config = ModelConfigBody()
+        global_var.set(global_var.Model_Config, model_config)
+    merge_model(model_config, body)
+    exception = load_rwkv_state(
+        global_var.get(global_var.Model), model_config.state, True
+    )
+    if exception is not None:
+        raise exception
+    print("Updated Model Config:", model_config)
 
     return "success"
 
 
 @router.get("/status", tags=["Configs"])
 def status():
-    import GPUtil
+    try:
+        import GPUtil
 
-    gpus = GPUtil.getGPUs()
+        gpus = GPUtil.getGPUs()
+    except:
+        gpus = []
     if len(gpus) == 0:
         device_name = "CPU"
     else:

backend-python/routes/midi.py

@@ -23,7 +23,11 @@ class TextToMidiBody(BaseModel):
 
 @router.post("/text-to-midi", tags=["MIDI"])
 def text_to_midi(body: TextToMidiBody):
-    vocab_config = "backend-python/utils/midi_vocab_config.json"
+    vocab_config_type = global_var.get(global_var.Midi_Vocab_Config_Type)
+    if vocab_config_type == global_var.MidiVocabConfig.Piano:
+        vocab_config = "backend-python/utils/vocab_config_piano.json"
+    else:
+        vocab_config = "backend-python/utils/midi_vocab_config.json"
     cfg = VocabConfig.from_json(vocab_config)
     mid = convert_str_to_midi(cfg, body.text.strip())
     mid_data = io.BytesIO()
@@ -35,7 +39,11 @@ def text_to_midi(body: TextToMidiBody):
 
 @router.post("/midi-to-text", tags=["MIDI"])
 async def midi_to_text(file_data: UploadFile):
-    vocab_config = "backend-python/utils/midi_vocab_config.json"
+    vocab_config_type = global_var.get(global_var.Midi_Vocab_Config_Type)
+    if vocab_config_type == global_var.MidiVocabConfig.Piano:
+        vocab_config = "backend-python/utils/vocab_config_piano.json"
+    else:
+        vocab_config = "backend-python/utils/midi_vocab_config.json"
     cfg = VocabConfig.from_json(vocab_config)
     filter_config = "backend-python/utils/midi_filter_config.json"
     filter_cfg = FilterConfig.from_json(filter_config)
@@ -69,7 +77,11 @@ def txt_to_midi(body: TxtToMidiBody):
     if not body.midi_path.startswith("midi/"):
         raise HTTPException(status.HTTP_400_BAD_REQUEST, "bad output path")
 
-    vocab_config = "backend-python/utils/midi_vocab_config.json"
+    vocab_config_type = global_var.get(global_var.Midi_Vocab_Config_Type)
+    if vocab_config_type == global_var.MidiVocabConfig.Piano:
+        vocab_config = "backend-python/utils/vocab_config_piano.json"
+    else:
+        vocab_config = "backend-python/utils/midi_vocab_config.json"
     cfg = VocabConfig.from_json(vocab_config)
     with open(body.txt_path, "r") as f:
         text = f.read()

backend-python/routes/state_cache.py

@@ -76,6 +76,33 @@ class AddStateBody(BaseModel):
     logits: Any
 
 
+def copy_tensor_to_cpu(tensors):
+    import torch
+    import numpy as np
+
+    devices: List[torch.device] = []
+    copied: Union[Any, None] = None
+
+    tensors_type = type(tensors)
+    if tensors_type == list:
+        if hasattr(tensors[0], "device"):  # torch state
+            devices = [tensor.device for tensor in tensors]
+            copied = [tensor.cpu() for tensor in tensors]
+        else:  # WebGPU logits
+            copied = tensors
+    elif tensors_type == torch.Tensor:  # torch logits
+        devices = [tensors.device]
+        copied = tensors.cpu()
+    elif tensors_type == np.ndarray:  # rwkv.cpp
+        copied = tensors
+    else:  # WebGPU state
+        model = global_var.get(global_var.Model)
+        if model:
+            copied = model.model.model.back_state()
+
+    return copied, devices
+
+
 # @router.post("/add-state", tags=["State Cache"])
 def add_state(body: AddStateBody):
     global trie, dtrie, loop_del_trie_id
@@ -91,23 +118,24 @@ def add_state(body: AddStateBody):
 
     try:
         devices: List[torch.device] = []
+        logits_device: Union[torch.device, None] = None
         state: Union[Any, None] = None
+        logits: Union[Any, None] = None
 
         if body.state is not None:
-            if type(body.state) == list and hasattr(body.state[0], "device"):  # torch
-                devices = [tensor.device for tensor in body.state]
-                state = [tensor.cpu() for tensor in body.state]
-            elif type(body.state) == np.ndarray:  # rwkv.cpp
-                state = body.state
-            else:  # WebGPU
-                state = body.state.back()
+            state, devices = copy_tensor_to_cpu(body.state)
+        if body.logits is not None:
+            logits, logits_devices = copy_tensor_to_cpu(body.logits)
+            if len(logits_devices) > 0:
+                logits_device = logits_devices[0]
 
         id: int = trie.insert(body.prompt)
         dtrie[id] = {
             "tokens": body.tokens,
             "state": state,
-            "logits": body.logits,
+            "logits": logits,
             "devices": devices,
+            "logits_device": logits_device,
         }
 
         if len(trie) >= max_trie_len:
@@ -125,6 +153,7 @@ def add_state(body: AddStateBody):
         )
         return "success"
     except Exception as e:
+        print(e)  # should not happen
         raise HTTPException(
             status.HTTP_400_BAD_REQUEST, f"insert failed, bad prompt.\n{e}"
         )
@@ -149,6 +178,19 @@ def reset_state():
     return "success"
 
 
+def force_reset_state():
+    global trie, dtrie
+
+    if trie is None:
+        return
+
+    import cyac
+
+    trie = cyac.Trie()
+    dtrie = {}
+    gc.collect()
+
+
 class LongestPrefixStateBody(BaseModel):
     prompt: str
 
@@ -192,18 +234,38 @@ def longest_prefix_state(body: LongestPrefixStateBody, request: Request):
     if id != -1:
         prompt: str = trie[id]
         v = dtrie[id]
+        tokens: List[Union[str, int]] = copy.deepcopy(v["tokens"])
         devices: List[torch.device] = v["devices"]
+        logits_device: Union[torch.device, None] = v["logits_device"]
         state: Union[Any, None] = v["state"]
+        logits: Union[Any, None] = v["logits"]
 
-        if type(state) == list and hasattr(state[0], "device"):  # torch
-            state = [tensor.to(devices[i]) for i, tensor in enumerate(state)]
+        state_type = type(state)
+        if state_type == list and hasattr(state[0], "device"):  # torch
+            state = [
+                (
+                    tensor.to(devices[i])
+                    if devices[i] != torch.device("cpu")
+                    else tensor.clone()
+                )
+                for i, tensor in enumerate(state)
+            ]
+            logits = (
+                logits.to(logits_device)
+                if logits_device != torch.device("cpu")
+                else logits.clone()
+            )
+        elif state_type == np.ndarray:  # rwkv.cpp
+            logits = np.copy(logits)
+        else:  # WebGPU
+            logits = np.copy(logits)
 
         quick_log(request, body, "Hit:\n" + prompt)
         return {
             "prompt": prompt,
-            "tokens": v["tokens"],
+            "tokens": tokens,
             "state": state,
-            "logits": v["logits"],
+            "logits": logits,
         }
     else:
         return {"prompt": "", "tokens": [], "state": None, "logits": None}

backend-python/rwkv_pip/beta/cuda/att_one.cu

@@ -1,124 +0,0 @@
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include "element_wise.h"
-#include "util.h"
-
-// Equivalent Python code:
-// ww = t_first + k
-// p = torch.maximum(pp, ww)
-// e1 = torch.exp(pp - p)
-// e2 = torch.exp(ww - p)
-// wkv = ((e1 * aa + e2 * v) / (e1 * bb + e2)).to(dtype=x.dtype)
-// ww = t_decay + pp
-// p = torch.maximum(ww, k)
-// e1 = torch.exp(ww - p)
-// e2 = torch.exp(k - p)
-// t1 = e1 * aa + e2 * v
-// t2 = e1 * bb + e2
-// r = r * wkv
-// return t1, t2, p, r
-struct WkvForwardOne {
-  const float *t_first;
-  const float *k;
-  const float *pp;
-  const float *aa;
-  const float *bb;
-  const float *t_decay;
-  const float *v;
-  /* out */ float *t1;
-  /* out */ float *t2;
-  /* out */ float *p;
-  /* in & out */ half *r;
-
-  __device__ void operator()(int i) const {
-    float ww = t_first[i] + k[i];
-    float pp_ = pp[i];
-    float p_ = (pp_ > ww) ? pp_ : ww;
-    float e1 = expf(pp_ - p_);
-    float e2 = expf(ww - p_);
-    float aa_ = aa[i];
-    float bb_ = bb[i];
-    float v_ = v[i];
-    r[i] = __hmul(r[i], __float2half(((e1 * aa_ + e2 * v_) / (e1 * bb_ + e2))));
-    ww = t_decay[i] + pp_;
-    float k_ = k[i];
-    p_ = (ww > k_) ? ww : k_;
-    e1 = expf(ww - p_);
-    e2 = expf(k_ - p_);
-    t1[i] = e1 * aa_ + e2 * v_;
-    t2[i] = e1 * bb_ + e2;
-    p[i] = p_;
-  }
-};
-
-/*
-   Equivalent Python code:
-   kx = xx * k_mix + sx * (1 - k_mix)
-   vx = xx * v_mix + sx * (1 - v_mix)
-   rx = xx * r_mix + sx * (1 - r_mix)
-*/
-
-struct Mix {
-  const half *xx;
-  const half *sx;
-  const half *k_mix;
-  const half *v_mix;
-  const half *r_mix;
-  /* out */ half *kx;
-  /* out */ half *vx;
-  /* out */ half *rx;
-
-  __device__ void operator()(int i) const {
-    half xx_ = xx[i];
-    half sx_ = sx[i];
-    half k_mix_ = k_mix[i];
-    half v_mix_ = v_mix[i];
-    half r_mix_ = r_mix[i];
-    kx[i] = __hadd(__hmul(xx_, k_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), k_mix_)));
-    vx[i] = __hadd(__hmul(xx_, v_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), v_mix_)));
-    rx[i] = __hadd(__hmul(xx_, r_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), r_mix_)));
-  }
-};
-
-using torch::Tensor;
-
-void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c);
-
-Tensor att_one(Tensor x, Tensor ln_w, Tensor ln_b, Tensor sx, Tensor k_mix,
-             Tensor v_mix, Tensor r_mix, Tensor kw,
-             /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, Tensor rw,
-             /* imm */ Tensor rx, Tensor ow, Tensor t_first,
-             /* imm */ Tensor k, Tensor pp, Tensor ww, Tensor aa, Tensor bb,
-             Tensor t_decay, /* imm */ Tensor v, /* in & out */ Tensor r,
-             /* out */ Tensor x_plus_out, /* out */ Tensor t1,
-             /* out */ Tensor t2, /* out */ Tensor p) {
-  Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b);
-  element_wise(Mix{data_ptr<half>(xx), data_ptr<half>(sx),
-                   data_ptr<half>(k_mix), data_ptr<half>(v_mix),
-                   data_ptr<half>(r_mix), data_ptr<half>(kx),
-                   data_ptr<half>(vx), data_ptr<half>(rx)},
-               x.numel());
-
-  gemm_fp16_cublas_tensor(kx, kw, k);
-  gemm_fp16_cublas_tensor(vx, vw, v);
-  gemm_fp16_cublas_tensor(rx, rw, r);
-  at::sigmoid_(r);
-
-  element_wise(WkvForwardOne{data_ptr<float>(t_first), data_ptr<float>(k),
-                             data_ptr<float>(pp), data_ptr<float>(aa),
-                             data_ptr<float>(bb), data_ptr<float>(t_decay),
-                             data_ptr<float>(v), data_ptr<float>(t1),
-                             data_ptr<float>(t2), data_ptr<float>(p),
-                             data_ptr<half>(r)},
-               x.numel());
-
-  gemm_fp16_cublas_tensor(r, ow, x_plus_out);
-  x_plus_out += x;
-  return xx;
-}

backend-python/rwkv_pip/beta/cuda/att_one_v5.cu

@@ -1,109 +0,0 @@
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include "element_wise.h"
-#include "util.h"
-
-// Equivalent Python code:
-// s1 = t_first * a + s
-// s2 = a + t_decay * s
-struct Fused1 {
-  const float *t_first;
-  const float *t_decay;
-  const float *a;
-  const float *s;
-  const int32_t inner_size;
-  /* out */ float *s1;
-  /* out */ float *s2;
-
-  __device__ void operator()(int i) const {
-    const int j = i / inner_size;
-    s1[i] = t_first[j] * a[i] + s[i];
-    s2[i] = a[i] + t_decay[j] * s[i];
-  }
-};
-
-/*
-   Equivalent Python code:
-   kx = xx * k_mix + sx * (1 - k_mix)
-   vx = xx * v_mix + sx * (1 - v_mix)
-   rx = xx * r_mix + sx * (1 - r_mix)
-*/
-
-struct Mix {
-  const half *xx;
-  const half *sx;
-  const half *k_mix;
-  const half *v_mix;
-  const half *r_mix;
-  /* out */ half *kx;
-  /* out */ half *vx;
-  /* out */ half *rx;
-
-  __device__ void operator()(int i) const {
-    half xx_ = xx[i];
-    half sx_ = sx[i];
-    half k_mix_ = k_mix[i];
-    half v_mix_ = v_mix[i];
-    half r_mix_ = r_mix[i];
-    kx[i] = __hadd(__hmul(xx_, k_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), k_mix_)));
-    vx[i] = __hadd(__hmul(xx_, v_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), v_mix_)));
-    rx[i] = __hadd(__hmul(xx_, r_mix_),
-                   __hmul(sx_, __hsub(__float2half(1), r_mix_)));
-  }
-};
-
-using torch::Tensor;
-
-void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c);
-
-Tensor att_one_v5(Tensor x, Tensor sx, Tensor s, Tensor ln_w, Tensor ln_b,
-                  Tensor lx_w, Tensor lx_b, Tensor k_mix, Tensor v_mix,
-                  Tensor r_mix, Tensor kw,
-                  /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx,
-                  Tensor rw,
-                  /* imm */ Tensor rx, Tensor ow, Tensor t_first,
-                  /* imm */ Tensor k, Tensor t_decay, /* imm */ Tensor v,
-                  /* imm */ Tensor r, /* imm */ Tensor s1,
-                  /* out */ Tensor x_plus_out, /* out */ Tensor s2) {
-  Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b);
-  element_wise(Mix{data_ptr<half>(xx), data_ptr<half>(sx),
-                   data_ptr<half>(k_mix), data_ptr<half>(v_mix),
-                   data_ptr<half>(r_mix), data_ptr<half>(kx),
-                   data_ptr<half>(vx), data_ptr<half>(rx)},
-               x.numel());
-
-  int H = t_decay.size(0);
-  int S = x.size(-1) / H;
-  gemm_fp16_cublas_tensor(rx, rw, r);
-  r = at::reshape(r, {H, 1, S});
-  gemm_fp16_cublas_tensor(kx, kw, k);
-  k = at::reshape(k, {H, S, 1});
-  gemm_fp16_cublas_tensor(vx, vw, v);
-  v = at::reshape(v, {H, 1, S});
-
-  {
-    Tensor a = at::matmul(k, v);
-
-    // s1 = t_first * a + s
-    // s2 = a + t_decay * s
-    element_wise(Fused1{data_ptr<float>(t_first), data_ptr<float>(t_decay),
-                        data_ptr<float>(a), data_ptr<float>(s),
-                        static_cast<int32_t>(a.size(1) * a.size(2)),
-                        data_ptr<float>(s1), data_ptr<float>(s2)},
-                 a.numel());
-  }
-
-  Tensor out = at::matmul(r, s1);
-  out = at::flatten(out);
-  out = at::squeeze(at::group_norm(at::unsqueeze(out, 0), H, lx_w, lx_b), 0);
-  out = at::_cast_Half(out);
-
-  gemm_fp16_cublas_tensor(out, ow, x_plus_out);
-  x_plus_out += x;
-  return xx;
-}

backend-python/rwkv_pip/beta/cuda/att_seq.cu

@@ -1,178 +0,0 @@
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include "util.h"
-#include "element_wise.h"
-
-using torch::Tensor;
-
-void gemm_fp16_cublas(const void *a, const void *b, void *c, int m,
-                      int n, int k, bool output_fp32);
-
-// based on `kernel_wkv_forward`, fusing more operations
-__global__ void kernel_wkv_forward_new(
-    const int B, const int T, const int C, const float *__restrict__ const _w,
-    const float *__restrict__ const _u, const float *__restrict__ const _k,
-    const float *__restrict__ const _v, const half *__restrict__ const r,
-    half *__restrict__ const _y, float *__restrict__ const _aa,
-    float *__restrict__ const _bb, float *__restrict__ const _pp) {
-  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
-  const int _b = idx / C;
-  const int _c = idx % C;
-  const int _offset = _b * T * C + _c;
-  const int _state_offset = _b * C + _c;
-
-  float u = _u[_c];
-  float w = _w[_c];
-  const float *__restrict__ const k = _k + _offset;
-  const float *__restrict__ const v = _v + _offset;
-  half *__restrict__ const y = _y + _offset;
-
-  float aa = _aa[_state_offset];
-  float bb = _bb[_state_offset];
-  float pp = _pp[_state_offset];
-  for (int i = 0; i < T; i++) {
-    const int ii = i * C;
-    const float kk = k[ii];
-    const float vv = v[ii];
-    float ww = u + kk;
-    float p = max(pp, ww);
-    float e1 = exp(pp - p);
-    float e2 = exp(ww - p);
-    y[ii] = __float2half((e1 * aa + e2 * vv) / (e1 * bb + e2));
-    ww = w + pp;
-    p = max(ww, kk);
-    e1 = exp(ww - p);
-    e2 = exp(kk - p);
-    aa = e1 * aa + e2 * vv;
-    bb = e1 * bb + e2;
-    pp = p;
-  }
-  _aa[_state_offset] = aa;
-  _bb[_state_offset] = bb;
-  _pp[_state_offset] = pp;
-}
-
-void cuda_wkv_forward_new(int B, int T, int C, float *w, float *u, float *k,
-                          float *v, half *r, half *y, float *aa, float *bb,
-                          float *pp) {
-  dim3 threadsPerBlock(min(C, 32));
-  assert(B * C % threadsPerBlock.x == 0);
-  dim3 numBlocks(B * C / threadsPerBlock.x);
-  kernel_wkv_forward_new<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, r,
-                                                         y, aa, bb, pp);
-}
-
-__global__ void _att_mix(const half *xx, const half *sx, const half *k_mix,
-                        const half *v_mix, const half *r_mix,
-                        const int outer_size, const int inner_size, half *kx,
-                        half *vx, half *rx) {
-  for (int idx2 = blockIdx.x * blockDim.x + threadIdx.x; idx2 < inner_size;
-       idx2 += blockDim.x * gridDim.x) {
-    half k_mix_ = k_mix[idx2];
-    half v_mix_ = v_mix[idx2];
-    half r_mix_ = r_mix[idx2];
-    for (int row = 0; row < outer_size; ++row) {
-      int idx1 = row * inner_size + idx2;
-      half xx_ = xx[idx1];
-      half sx_ = sx[idx1];
-      kx[idx1] = __hadd(__hmul(xx_, k_mix_),
-                        __hmul(sx_, __hsub(__float2half(1), k_mix_)));
-      vx[idx1] = __hadd(__hmul(xx_, v_mix_),
-                        __hmul(sx_, __hsub(__float2half(1), v_mix_)));
-      rx[idx1] = __hadd(__hmul(xx_, r_mix_),
-                        __hmul(sx_, __hsub(__float2half(1), r_mix_)));
-    }
-  }
-}
-
-void att_mix(const half *xx, const half *sx, const half *k_mix,
-            const half *v_mix, const half *r_mix, const int outer_size,
-            const int inner_size, half *kx, half *vx, half *rx) {
-  // 256 is good enough on most GPUs
-  const int32_t BLOCK_SIZE = 256;
-  assert(inner_size % BLOCK_SIZE == 0);
-  _att_mix<<<inner_size / BLOCK_SIZE, BLOCK_SIZE>>>(
-      xx, sx, k_mix, v_mix, r_mix, outer_size, inner_size, kx, vx, rx);
-}
-
-struct InplaceSigmoid {
-  __device__ __forceinline__ half operator()(int i) const {
-    ptr[i] = __float2half(1.0 / (1.0 + exp(-__half2float(ptr[i]))));
-  }
-  half *ptr;
-};
-
-struct InplaceMul {
-  __device__ __forceinline__ half operator()(int i) const {
-    y[i] = __hmul(x[i], y[i]);
-  }
-  half *y;
-  half *x;
-};
-
-/*
-   Equivalent Python code:
-
-   xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b)
-   sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))
-   kx = xx * k_mix + sx * (1 - k_mix)
-   vx = xx * v_mix + sx * (1 - v_mix)
-   rx = xx * r_mix + sx * (1 - r_mix)
-
-   r = torch.sigmoid(gemm(rx, rw))
-   k = gemm(kx, kw, output_dtype=torch.float32)
-   v = gemm(vx, vw, output_dtype=torch.float32)
-
-   T = x.shape[0]
-   for t in range(T):
-       kk = k[t]
-       vv = v[t]
-       ww = t_first + kk
-       p = torch.maximum(pp, ww)
-       e1 = torch.exp(pp - p)
-       e2 = torch.exp(ww - p)
-       sx[t] = ((e1 * aa + e2 * vv) / (e1 * bb + e2)).to(dtype=x.dtype)
-       ww = t_decay + pp
-       p = torch.maximum(ww, kk)
-       e1 = torch.exp(ww - p)
-       e2 = torch.exp(kk - p)
-       aa = e1 * aa + e2 * vv
-       bb = e1 * bb + e2
-       pp = p
-   out = gemm(r * sx, ow)
-   return x + out, xx[-1,:], aa, bb, pp
-*/
-Tensor att_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor v_mix, Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               Tensor ow, Tensor t_first, Tensor pp, Tensor aa, Tensor bb,
-               Tensor t_decay, /* imm */ Tensor buf, /* out */ Tensor x_plus_out) {
-  Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b);
-  sx = at::cat({sx.unsqueeze(0), xx.slice(0, 0, -1)}, 0);
-  char* buf_ptr = (char*)buf.data_ptr();
-  half* kx = (half*)buf_ptr;
-  half* vx = kx + x.numel();
-  half* rx = vx + x.numel();
-  half* wkv_y = rx + x.numel();
-  att_mix(data_ptr<half>(xx), data_ptr<half>(sx), data_ptr<half>(k_mix),
-         data_ptr<half>(v_mix), data_ptr<half>(r_mix), xx.size(0), xx.size(1),
-         kx, vx, rx);
-  float* k = reinterpret_cast<float*>(wkv_y + x.numel());
-  float* v = k + x.size(0) * kw.size(1);
-  half* r = reinterpret_cast<half*>(v + x.size(0) * vw.size(1));
-
-  gemm_fp16_cublas(kx, kw.data_ptr(), k, x.size(0), kw.size(1), kw.size(0), true);
-  gemm_fp16_cublas(vx, vw.data_ptr(), v, x.size(0), vw.size(1), vw.size(0), true);
-  gemm_fp16_cublas(rx, rw.data_ptr(), r, x.size(0), rw.size(1), rw.size(0), false);
-  element_wise(InplaceSigmoid{r}, x.size(0) * rw.size(1));
-  cuda_wkv_forward_new(1, x.size(0), x.size(1), data_ptr<float>(t_decay),
-                       data_ptr<float>(t_first), k, v, r,
-                       wkv_y, data_ptr<float>(aa),
-                       data_ptr<float>(bb), data_ptr<float>(pp));
-  element_wise(InplaceMul{wkv_y, r}, x.numel());
-  gemm_fp16_cublas(wkv_y, ow.data_ptr(), x_plus_out.data_ptr(), x.size(0), ow.size(1), ow.size(0), false);
-  x_plus_out += x;
-  return xx;
-}

backend-python/rwkv_pip/beta/cuda/element_wise.h

@@ -1,21 +0,0 @@
-#include <cassert>
-#include <cstddef>
-#include <cstdint>
-
-template <typename Func> __global__ void _element_wise(Func func, int n) {
-  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n;
-       i += blockDim.x * gridDim.x) {
-    func(i);
-  }
-}
-
-// NOTE: packed data type (e.g. float4) is a overkill for current sizes
-// (4096 in 7B model and 768 in 0.1B model),
-// and is not faster than the plain float version.
-template <typename Func>
-void element_wise(Func func, int n) {
-  // 256 is good enough on most GPUs
-  const int32_t BLOCK_SIZE = 256;
-  assert(n % BLOCK_SIZE == 0);
-  _element_wise<<<n / BLOCK_SIZE, BLOCK_SIZE>>>(func, n);
-}

backend-python/rwkv_pip/beta/cuda/ffn.cu

@@ -1,165 +0,0 @@
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include "element_wise.h"
-#include "util.h"
-
-using torch::Tensor;
-
-void gemm_fp16_cublas(const void *a, const void *b, void *c, int ori_m,
-                      int ori_n, int ori_k, bool output_fp32);
-
-__global__ void _ffn_seq_mix(const half *xx, const half *sx, const half *k_mix,
-                             const half *r_mix, const int outer_size,
-                             const int inner_size, half *kx, half *rx) {
-  for (int idx2 = blockIdx.x * blockDim.x + threadIdx.x; idx2 < inner_size;
-       idx2 += blockDim.x * gridDim.x) {
-    half k_mix_ = k_mix[idx2];
-    half r_mix_ = r_mix[idx2];
-    for (int row = 0; row < outer_size; ++row) {
-      int idx1 = row * inner_size + idx2;
-      half xx_ = xx[idx1];
-      half sx_ = sx[idx1];
-      kx[idx1] = __hadd(__hmul(xx_, k_mix_),
-                        __hmul(sx_, __hsub(__float2half(1), k_mix_)));
-      rx[idx1] = __hadd(__hmul(xx_, r_mix_),
-                        __hmul(sx_, __hsub(__float2half(1), r_mix_)));
-    }
-  }
-}
-
-void ffn_seq_mix(const half *xx, const half *sx, const half *k_mix,
-                 const half *r_mix, const int outer_size, const int inner_size,
-                 half *kx, half *rx) {
-  // 256 is good enough on most GPUs
-  const int32_t BLOCK_SIZE = 256;
-  assert(inner_size % BLOCK_SIZE == 0);
-  _ffn_seq_mix<<<inner_size / BLOCK_SIZE, BLOCK_SIZE>>>(
-      xx, sx, k_mix, r_mix, outer_size, inner_size, kx, rx);
-}
-
-struct InplaceSigmoid {
-  __device__ __forceinline__ void operator()(int i) const {
-    ptr[i] = __float2half(1.0 / (1.0 + exp(-__half2float(ptr[i]))));
-  }
-  half *ptr;
-};
-
-struct InplaceReLUAndSquare {
-  __device__ __forceinline__ void operator()(int i) const {
-    // __hmax is not defined in old cuda
-    if (__hgt(ptr[i], __float2half(0))) {
-      ptr[i] = __hmul(ptr[i], ptr[i]);
-    } else {
-      ptr[i] = __float2half(0);
-    }
-  }
-  half *ptr;
-};
-
-struct InplaceFma {
-  __device__ __forceinline__ void operator()(int i) const {
-    a[i] = __hfma(a[i], b[i], c[i]);
-  }
-  half *a;
-  const half *b;
-  const half *c;
-};
-
-/*
-   Equivalent Python code:
-
-   xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b)
-   sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))
-   kx = xx * k_mix + sx * (1 - k_mix)
-   rx = xx * r_mix + sx * (1 - r_mix)
-
-   r = torch.sigmoid(gemm(rx, rw))
-   vx = torch.square(torch.relu(gemm(kx, kw)))
-   out = r * gemm(vx, vw)
-   return x + out, xx[-1,:]
-*/
-Tensor ffn_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               /* imm */ Tensor buf,
-               /* out */ Tensor x_plus_out) {
-  Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b);
-  sx = at::cat({sx.unsqueeze(0), xx.slice(0, 0, -1)}, 0);
-  char *buf_ptr = (char *)buf.data_ptr();
-  half *kx = (half *)buf_ptr;
-  half *rx = kx + x.numel();
-  half *vx = rx + x.numel();
-  half *r = vx + x.size(0) * kw.size(1);
-  ffn_seq_mix(data_ptr<half>(xx), data_ptr<half>(sx), data_ptr<half>(k_mix),
-              data_ptr<half>(r_mix), xx.size(0), xx.size(1), kx, rx);
-
-  gemm_fp16_cublas(rx, rw.data_ptr(), r, x.size(0), rw.size(1), x.size(1),
-                   false);
-  element_wise(InplaceSigmoid{r}, x.size(0) * rw.size(1));
-  gemm_fp16_cublas(kx, kw.data_ptr(), vx, x.size(0), kw.size(1), x.size(1),
-                   false);
-  element_wise(InplaceReLUAndSquare{vx}, x.size(0) * kw.size(1));
-  gemm_fp16_cublas(vx, vw.data_ptr(), x_plus_out.data_ptr(), x.size(0),
-                   vw.size(1), vw.size(0), false);
-  element_wise(InplaceFma{data_ptr<half>(x_plus_out), r, data_ptr<half>(x)},
-               x_plus_out.numel());
-  return xx;
-}
-
-struct FfnOneMix {
-  __device__ __forceinline__ void operator()(int idx) {
-    half k_mix_ = k_mix[idx];
-    half r_mix_ = r_mix[idx];
-    half xx_ = xx[idx];
-    half sx_ = sx[idx];
-    kx[idx] = __hadd(__hmul(xx_, k_mix_),
-                     __hmul(sx_, __hsub(__float2half(1), k_mix_)));
-    rx[idx] = __hadd(__hmul(xx_, r_mix_),
-                     __hmul(sx_, __hsub(__float2half(1), r_mix_)));
-  }
-  half *k_mix;
-  half *r_mix;
-  half *xx;
-  half *sx;
-  half *kx;
-  half *rx;
-};
-
-/*
-  Equivalent Python code:
-
-  xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b)
-  kx = xx * k_mix + sx * (1 - k_mix)
-  rx = xx * r_mix + sx * (1 - r_mix)
-
-  r = torch.sigmoid(gemm(rx, rw))
-  vx = torch.square(torch.relu(gemm(kx, kw)))
-  out = r * gemm(vx, vw)
-  return x + out, xx
-*/
-Tensor ffn_one(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               /* imm */ Tensor buf,
-               /* out */ Tensor x_plus_out) {
-  Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b);
-  char *buf_ptr = (char *)buf.data_ptr();
-  half *kx = (half *)buf_ptr;
-  half *rx = kx + x.numel();
-  half *vx = rx + x.numel();
-  half *r = vx + x.size(0) * kw.size(1);
-  element_wise(FfnOneMix{data_ptr<half>(k_mix), data_ptr<half>(r_mix),
-                         data_ptr<half>(xx), data_ptr<half>(sx), kx, rx},
-               x.numel());
-  // vector * matrix, so m = 1
-  gemm_fp16_cublas(rx, rw.data_ptr(), r, 1, rw.size(1), rw.size(0), false);
-  element_wise(InplaceSigmoid{r}, rw.size(1));
-  gemm_fp16_cublas(kx, kw.data_ptr(), vx, 1, kw.size(1), kw.size(0), false);
-  element_wise(InplaceReLUAndSquare{vx}, kw.size(1));
-  gemm_fp16_cublas(vx, vw.data_ptr(), x_plus_out.data_ptr(), 1, vw.size(1),
-                   vw.size(0), false);
-  element_wise(InplaceFma{data_ptr<half>(x_plus_out), r, data_ptr<half>(x)},
-               x_plus_out.numel());
-  return xx;
-}

backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp

@@ -1,128 +0,0 @@
-#include <cublas_v2.h>
-#include <cuda.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#define CUBLAS_CHECK(condition)                                                \
-  for (cublasStatus_t _cublas_check_status = (condition);                      \
-       _cublas_check_status != CUBLAS_STATUS_SUCCESS;)                         \
-    throw std::runtime_error("cuBLAS error " +                                 \
-                             std::to_string(_cublas_check_status) + " at " +   \
-                             std::to_string(__LINE__));
-
-#define CUDA_CHECK(condition)                                                  \
-  for (cudaError_t _cuda_check_status = (condition);                           \
-       _cuda_check_status != cudaSuccess;)                                     \
-    throw std::runtime_error(                                                  \
-        "CUDA error " + std::string(cudaGetErrorString(_cuda_check_status)) +  \
-        " at " + std::to_string(__LINE__));
-
-cublasHandle_t get_cublas_handle() {
-  static cublasHandle_t cublas_handle = []() {
-    cublasHandle_t handle = nullptr;
-    CUBLAS_CHECK(cublasCreate(&handle));
-#if CUDA_VERSION < 11000
-    CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
-#else
-    CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
-#endif // CUDA_VERSION < 11000
-    return handle;
-  }();
-  return cublas_handle;
-}
-
-/*
-  NOTE: blas gemm is column-major by default, but we need row-major output.
-  The data of row-major, transposed matrix is exactly the same as the
-  column-major, non-transposed matrix, and C = A * B ---> C^T = B^T * A^T
- */
-void gemm_fp16_cublas(const void *a, const void *b, void *c, int ori_m,
-                      int ori_n, int ori_k, bool output_fp32) {
-  const auto cuda_data_type = CUDA_R_16F;
-  const auto cuda_c_data_type = output_fp32 ? CUDA_R_32F : CUDA_R_16F;
-  const auto compute_type = CUDA_R_32F;
-  const float sp_alpha = 1.f;
-  // use CUBLAS_OP_N. see the notes above
-  const cublasOperation_t cublas_trans_a = CUBLAS_OP_N;
-  const cublasOperation_t cublas_trans_b = CUBLAS_OP_N;
-  // m = (B^T).size(0) = B.size(1) = n;
-  const int cublas_m = ori_n;
-  const int cublas_k = ori_k;
-  // comptiable with rwkv one mode, where 1-D tensor * 2-D tensor
-  // const int n = a.dense_dim() == 1 ? 1 : a.size(0);
-  const int cublas_n = ori_m;
-  const int cublas_lda = cublas_m;
-  const int cublas_ldb = cublas_k;
-  const int cublas_ldc = cublas_m;
-  cublasHandle_t cublas_handle = get_cublas_handle();
-
-#if CUDA_VERSION >= 11000
-  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT;
-#else
-  cublasGemmAlgo_t algo = CUBLAS_GEMM_DFALT_TENSOR_OP;
-#endif
-  const float sp_beta = 0.f;
-  CUBLAS_CHECK(cublasGemmEx(
-      cublas_handle, cublas_trans_a, cublas_trans_b, cublas_m, cublas_n,
-      cublas_k, &sp_alpha, b, cuda_data_type, cublas_lda,
-      a, cuda_data_type, cublas_ldb, &sp_beta, c,
-      cuda_c_data_type, cublas_ldc, compute_type, algo));
-}
-
-/*
-  NOTE: blas gemm is column-major by default, but we need row-major output.
-  The data of row-major, transposed matrix is exactly the same as the
-  column-major, non-transposed matrix, and C = A * B ---> C^T = B^T * A^T
- */
-void gemm_fp16_cublas_tensor(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
-  if (a.sizes().size() == 1) {
-    assert(b.sizes().size() == 2);
-    a = at::unsqueeze(a, 0);
-  }
-  const auto cuda_data_type = CUDA_R_16F;
-  const auto cuda_c_data_type =
-      c.dtype() == torch::kFloat32 ? CUDA_R_32F : CUDA_R_16F;
-  const auto compute_type = CUDA_R_32F;
-  const float sp_alpha = 1.f;
-  // swap a and b, and use CUBLAS_OP_N. see the notes above
-  std::swap(a, b);
-  const cublasOperation_t cublas_trans_a = CUBLAS_OP_N;
-  const cublasOperation_t cublas_trans_b = CUBLAS_OP_N;
-  // m = (B^T).size(0) = B.size(1), and = A.size(1) after swap,
-  // negative axis is used because of the existence of batch matmul.
-  const int m = a.size(-1);
-  const int k = a.size(-2);
-  const int n = b.size(-2);
-  const int cublas_lda = m;
-  const int cublas_ldb = k;
-  const int cublas_ldc = m;
-  cublasHandle_t cublas_handle = get_cublas_handle();
-
-#if CUDA_VERSION >= 11000
-  cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT;
-#else
-  cublasGemmAlgo_t algo = CUBLAS_GEMM_DFALT_TENSOR_OP;
-#endif
-  const float sp_beta = 0.f;
-  if (a.sizes().size() == 2 && b.sizes().size() == 2) {
-    CUBLAS_CHECK(cublasGemmEx(
-        cublas_handle, cublas_trans_a, cublas_trans_b, m, n, k, &sp_alpha,
-        a.data_ptr(), cuda_data_type, cublas_lda, b.data_ptr(), cuda_data_type,
-        cublas_ldb, &sp_beta, c.data_ptr(), cuda_c_data_type, cublas_ldc,
-        compute_type, algo));
-  } else {
-    // batch matmul
-    assert(a.sizes().size() == 3 && b.sizes().size() == 3);
-
-    const long long int cublas_stride_a = m * k;
-    const long long int cublas_stride_b = k * n;
-    const long long int cublas_stride_c = m * n;
-    CUBLAS_CHECK(cublasGemmStridedBatchedEx(
-        cublas_handle, cublas_trans_a, cublas_trans_b, m,
-        n, k, &sp_alpha, a.data_ptr(), cuda_data_type, cublas_lda,
-        cublas_stride_a, b.data_ptr(), cuda_data_type, cublas_ldb, cublas_stride_b,
-        &sp_beta, c.data_ptr(), cuda_c_data_type, cublas_ldc, cublas_stride_c,
-        a.size(0), compute_type, algo));
-  }
-}

backend-python/rwkv_pip/beta/cuda/operators.cu

@@ -1,246 +0,0 @@
-#include <stdio.h>
-#include <assert.h>
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-#define MIN_VALUE (-1e38)
-typedef at::Half fp16;
-__half *cast(fp16 *ptr) {
-    return reinterpret_cast<__half *>(ptr);
-}
-
-template <typename F>
-__global__ void kernel_wkv_forward(const int B, const int T, const int C,
-                               const float *__restrict__ const _w, const float *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
-                               F *__restrict__ const _y, float *__restrict__ const _aa, float *__restrict__ const _bb, float *__restrict__ const _pp) {
-    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    const int _b = idx / C;
-    const int _c = idx % C;
-    const int _offset = _b * T * C + _c;
-    const int _state_offset = _b * C + _c;
-
-    float u = _u[_c];
-    float w = _w[_c];
-    const F *__restrict__ const k = _k + _offset;
-    const F *__restrict__ const v = _v + _offset;
-    F *__restrict__ const y = _y + _offset;
-
-    float aa = _aa[_state_offset];
-    float bb = _bb[_state_offset];
-    float pp = _pp[_state_offset];
-    for (int i = 0; i < T; i++) {
-        const int ii = i * C;
-        const float kk = float(k[ii]);
-        const float vv = float(v[ii]);
-        float ww = u + kk;
-        float p = max(pp, ww);
-        float e1 = exp(pp - p);
-        float e2 = exp(ww - p);
-        y[ii] = F((e1 * aa + e2 * vv) / (e1 * bb + e2));
-        ww = w + pp;
-        p = max(ww, kk);
-        e1 = exp(ww - p);
-        e2 = exp(kk - p);
-        aa = e1 * aa + e2 * vv;
-        bb = e1 * bb + e2;
-        pp = p;
-    }
-    _aa[_state_offset] = aa;
-    _bb[_state_offset] = bb;
-    _pp[_state_offset] = pp;
-}
-
-template <typename F>
-void cuda_wkv_forward(int B, int T, int C, float *w, float *u, F *k, F *v, F *y, float *aa, float *bb, float *pp) {
-    dim3 threadsPerBlock( min(C, 32) );
-    assert(B * C % threadsPerBlock.x == 0);
-    dim3 numBlocks(B * C / threadsPerBlock.x);
-    kernel_wkv_forward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y, aa, bb, pp);
-}
-
-template void cuda_wkv_forward<fp16>(
-    int B, int T, int C,
-    float *w, float *u, fp16 *k, fp16 *v, fp16 *y,
-    float *aa, float *bb, float *pp);
-template void cuda_wkv_forward<float>(
-    int B, int T, int C,
-    float *w, float *u, float *k, float *v, float *y,
-    float *aa, float *bb, float *pp);
-
-__global__ void kernel_mm_seq_fp32i8(
-    const int B, const int N, const int M,
-    const float *__restrict__ const x, const int x_stride,
-    const uint8_t *__restrict__ const w, const int w_stride,
-    const float *__restrict__ const mx,
-    const float *__restrict__ const rx,
-    const float *__restrict__ const my,
-    const float *__restrict__ const ry,
-    float *__restrict__ const y, const int y_stride) {
-
-    const int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const int k = blockIdx.y * blockDim.y + threadIdx.y;
-
-    if (i < B && k < M) {
-        float y_local = 0;
-        for (int j = 0; j < N; ++j) {
-            y_local += x[i * x_stride + j] * (
-                (float(w[j * w_stride + k]) + 0.5f)
-                * rx[k] * ry[j] + mx[k] + my[j]
-            );
-        }
-        y[i * y_stride + k] = y_local;
-    }
-}
-
-template <typename F>
-void cuda_mm8_seq(int B, int N, int M,
-                  F *x, int x_stride,
-                  uint8_t *w, int w_stride,
-                  F *mx, F *rx,
-                  F *my, F *ry,
-                  F *y, int y_stride);
-
-template <>
-void cuda_mm8_seq<float>(int B, int N, int M,
-                         float *x, int x_stride,
-                         uint8_t *w, int w_stride,
-                         float *mx, float *rx,
-                         float *my, float *ry,
-                         float *y, int y_stride) {
-    dim3 blockSize(1, 128);
-    dim3 gridSize((B + blockSize.x - 1) / blockSize.x, (M + blockSize.y - 1) / blockSize.y);
-    kernel_mm_seq_fp32i8<<<gridSize, blockSize>>>(
-        B, N, M, x, x_stride, w, w_stride,
-        mx, rx, my, ry, y, y_stride);
-}
-
-__global__ void kernel_mm_seq_fp16i8(
-    const int B, const int N, const int M,
-    const __half *__restrict__ const x, const int x_stride,
-    const uint8_t *__restrict__ const w, const int w_stride,
-    const __half *__restrict__ const mx,
-    const __half *__restrict__ const rx,
-    const __half *__restrict__ const my,
-    const __half *__restrict__ const ry,
-    __half *__restrict__ const y, const int y_stride) {
-
-    const int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const int k = blockIdx.y * blockDim.y + threadIdx.y;
-
-    if (i < B && k < M) {
-        float y_local = 0;
-        for (int j = 0; j < N; ++j) {
-            y_local += __half2float(x[i * x_stride + j]) * (
-                (float(w[j * w_stride + k]) + 0.5f)
-                * __half2float(rx[k]) * __half2float(ry[j])
-                + __half2float(mx[k]) + __half2float(my[j])
-            );
-        }
-        y[i * y_stride + k] = __float2half(y_local);
-    }
-}
-
-template <>
-void cuda_mm8_seq<fp16>(int B, int N, int M,
-                        fp16 *x, int x_stride,
-                        uint8_t *w, int w_stride,
-                        fp16 *mx, fp16 *rx,
-                        fp16 *my, fp16 *ry,
-                        fp16 *y, int y_stride) {
-    dim3 blockSize(1, 128);
-    dim3 gridSize((B + blockSize.x - 1) / blockSize.x, (M + blockSize.y - 1) / blockSize.y);
-    kernel_mm_seq_fp16i8<<<gridSize, blockSize>>>(
-        B, N, M, cast(x), x_stride, w, w_stride,
-        cast(mx), cast(rx), cast(my), cast(ry), cast(y), y_stride);
-}
-
-#define MM8_ONE_JSPLIT 24
-#define MM8_ONE_TILE 1024
-
-__global__ void kernel_mm_one_fp32i8(
-    const int N, const int M,
-    const float *__restrict__ const x,
-    const uint8_t *__restrict__ const w, const int w_stride,
-    const float *__restrict__ const mx,
-    const float *__restrict__ const rx,
-    const float *__restrict__ const my,
-    const float *__restrict__ const ry,
-    float *__restrict__ const y) {
-
-    const int k = blockIdx.y * blockDim.y + threadIdx.y;
-    const int j0 = min(N, blockIdx.x * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT));
-    const int j1 = min(N, (blockIdx.x + 1) * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT));
-
-    if (k < M) {
-        float y_local = 0;
-        for (int j = j0; j < j1; ++j) {
-            y_local += x[j] * (
-                (float(w[j * w_stride + k]) + 0.5f)
-                * rx[k] * ry[j] + mx[k] + my[j]
-            );
-        }
-        atomicAdd(&y[k], y_local);
-    }
-}
-
-template <typename F>
-void cuda_mm8_one(int N, int M,
-                  F *x,
-                  uint8_t *w, int w_stride,
-                  F *mx, F *rx,
-                  F *my, F *ry,
-                  float *y);
-
-template <>
-void cuda_mm8_one<float>(int N, int M,
-                        float *x,
-                        uint8_t *w, int w_stride,
-                        float *mx, float *rx,
-                        float *my, float *ry,
-                        float *y) {
-    dim3 blockSize(1, MM8_ONE_TILE);
-    dim3 gridSize(MM8_ONE_JSPLIT, (M + blockSize.y - 1) / blockSize.y);
-    kernel_mm_one_fp32i8<<<gridSize, blockSize>>>(
-        N, M, x, w, w_stride,
-        mx, rx, my, ry, y);
-}
-
-__global__ void kernel_mm_one_fp16i8(
-    const int N, const int M,
-    const __half *__restrict__ const x,
-    const uint8_t *__restrict__ const w, const int w_stride,
-    const __half *__restrict__ const mx,
-    const __half *__restrict__ const rx,
-    const __half *__restrict__ const my,
-    const __half *__restrict__ const ry,
-    float *__restrict__ const y) {
-
-    const int k = blockIdx.y * blockDim.y + threadIdx.y;
-    const int j0 = min(N, blockIdx.x * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT));
-    const int j1 = min(N, (blockIdx.x + 1) * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT));
-
-    if (k < M) {
-        float y_local = 0;
-        for (int j = j0; j < j1; ++j) {
-            y_local += __half2float(x[j]) * (
-                (float(w[j * w_stride + k]) + 0.5f)
-                * __half2float(rx[k]) * __half2float(ry[j])
-                + __half2float(mx[k]) + __half2float(my[j])
-            );
-        }
-        atomicAdd(&y[k], y_local);
-    }
-}
-
-template <>
-void cuda_mm8_one<fp16>(int N, int M,
-                        fp16 *x,
-                        uint8_t *w, int w_stride,
-                        fp16 *mx, fp16 *rx,
-                        fp16 *my, fp16 *ry,
-                        float *y) {
-    dim3 blockSize(1, MM8_ONE_TILE);
-    dim3 gridSize(MM8_ONE_JSPLIT, (M + blockSize.y - 1) / blockSize.y);
-    kernel_mm_one_fp16i8<<<gridSize, blockSize>>>(
-        N, M, cast(x), w, w_stride,
-        cast(mx), cast(rx), cast(my), cast(ry), y);
-}

backend-python/rwkv_pip/beta/cuda/util.h

@@ -1,7 +0,0 @@
-#include "ATen/ATen.h"
-#include <cuda_fp16.h>
-
-template <typename T> T *data_ptr(torch::Tensor x) { return x.data_ptr<T>(); }
-template <> inline half *data_ptr(torch::Tensor x) {
-  return reinterpret_cast<half *>(x.data_ptr<at::Half>());
-}

backend-python/rwkv_pip/beta/cuda/wrapper.cpp

@@ -1,181 +0,0 @@
-#include <torch/extension.h>
-#include "ATen/ATen.h"
-#include <iostream>
-#include <c10/cuda/CUDAGuard.h>
-
-typedef at::Half fp16;
-
-template <typename F>
-void cuda_wkv_forward(int B, int T, int C,
-                      float *w, float *u, F *k, F *v, F *y,
-                      float *aa, float *bb, float *pp);
-template <typename F>
-void cuda_mm8_seq(int B, int N, int M,
-                  F *x, int x_stride,
-                  uint8_t *w, int w_stride,
-                  F *mx, F *rx,
-                  F *my, F *ry,
-                  F *y, int y_stride);
-template <typename F>
-void cuda_mm8_one(int N, int M,
-                  F *x,
-                  uint8_t *w, int w_stride,
-                  F *mx, F *rx,
-                  F *my, F *ry,
-                  float *y);
-
-void wkv_forward(int64_t B, int64_t T, int64_t C,
-                 torch::Tensor &w, torch::Tensor &u,
-                 torch::Tensor &k, torch::Tensor &v, torch::Tensor &y,
-                 torch::Tensor &aa, torch::Tensor &bb, torch::Tensor &pp) {
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(w));
-    switch (k.scalar_type()) {
-    case c10::ScalarType::Half:
-        cuda_wkv_forward(B, T, C,
-                         w.data_ptr<float>(), u.data_ptr<float>(),
-                         k.data_ptr<fp16>(), v.data_ptr<fp16>(), y.data_ptr<fp16>(),
-                         aa.data_ptr<float>(), bb.data_ptr<float>(), pp.data_ptr<float>());
-        break;
-    case c10::ScalarType::Float:
-        cuda_wkv_forward(B, T, C,
-                         w.data_ptr<float>(), u.data_ptr<float>(),
-                         k.data_ptr<float>(), v.data_ptr<float>(), y.data_ptr<float>(),
-                         aa.data_ptr<float>(), bb.data_ptr<float>(), pp.data_ptr<float>());
-        break;
-    default:
-        assert(false && "Only FP16 and FP32 are currently supported");
-    }
-}
-
-void mm8_seq(int64_t B, int64_t N, int64_t M,
-             torch::Tensor &x, torch::Tensor &w,
-             torch::Tensor &mx, torch::Tensor &rx,
-             torch::Tensor &my, torch::Tensor &ry,
-             torch::Tensor &y) {
-    assert(x.stride(1) == 1);
-    assert(w.stride(1) == 1);
-    assert(mx.stride(0) == 1 && rx.stride(0) == 1);
-    assert(my.stride(0) == 1 && ry.stride(0) == 1);
-    assert(y.stride(1) == 1);
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(w));
-    switch (x.scalar_type()) {
-    case c10::ScalarType::Half:
-        cuda_mm8_seq(
-            B, N, M,
-            x.data_ptr<fp16>(), x.stride(0),
-            w.data_ptr<uint8_t>(), w.stride(0),
-            mx.data_ptr<fp16>(), rx.data_ptr<fp16>(),
-            my.data_ptr<fp16>(), ry.data_ptr<fp16>(),
-            y.data_ptr<fp16>(), y.stride(0));
-        break;
-    case c10::ScalarType::Float:
-        cuda_mm8_seq(
-            B, N, M,
-            x.data_ptr<float>(), x.stride(0),
-            w.data_ptr<uint8_t>(), w.stride(0),
-            mx.data_ptr<float>(), rx.data_ptr<float>(),
-            my.data_ptr<float>(), ry.data_ptr<float>(),
-            y.data_ptr<float>(), y.stride(0));
-        break;
-    default:
-        assert(false && "Only FP16 and FP32 are currently supported");
-    }
-}
-void mm8_one(int64_t N, int64_t M,
-             torch::Tensor &x, torch::Tensor &w,
-             torch::Tensor &mx, torch::Tensor &rx,
-             torch::Tensor &my, torch::Tensor &ry,
-             torch::Tensor &y) {
-    assert(x.stride(0) == 1);
-    assert(w.stride(1) == 1);
-    assert(mx.stride(0) == 1 && rx.stride(0) == 1);
-    assert(my.stride(0) == 1 && ry.stride(0) == 1);
-    assert(y.stride(0) == 1);
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(w));
-    switch (x.scalar_type()) {
-    case c10::ScalarType::Half:
-        cuda_mm8_one(
-            N, M,
-            x.data_ptr<fp16>(),
-            w.data_ptr<uint8_t>(), w.stride(0),
-            mx.data_ptr<fp16>(), rx.data_ptr<fp16>(),
-            my.data_ptr<fp16>(), ry.data_ptr<fp16>(),
-            y.data_ptr<float>());
-        break;
-    case c10::ScalarType::Float:
-        cuda_mm8_one(
-            N, M,
-            x.data_ptr<float>(),
-            w.data_ptr<uint8_t>(), w.stride(0),
-            mx.data_ptr<float>(), rx.data_ptr<float>(),
-            my.data_ptr<float>(), ry.data_ptr<float>(),
-            y.data_ptr<float>());
-        break;
-    default:
-        assert(false && "Only FP16 and FP32 are currently supported");
-    }
-}
-
-using torch::Tensor;
-
-#ifndef DISABLE_CUBLAS_GEMM
-void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c);
-#endif
-
-Tensor att_one(Tensor x, Tensor ln_w, Tensor ln_b, Tensor sx, Tensor k_mix,
-             Tensor v_mix, Tensor r_mix, Tensor kw,
-             /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, Tensor rw,
-             /* imm */ Tensor rx, Tensor ow, Tensor t_first,
-             /* imm */ Tensor k, Tensor pp, Tensor ww, Tensor aa, Tensor bb,
-             Tensor t_decay, /* imm */ Tensor v, /* in & out */ Tensor r,
-             /* out */ Tensor x_plus_out, /* out */ Tensor t1,
-             /* out */ Tensor t2, /* out */ Tensor p);
-
-Tensor att_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor v_mix, Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               Tensor ow, Tensor t_first, Tensor pp, Tensor aa, Tensor bb,
-               Tensor t_decay, /* imm */ Tensor buf, /* out */ Tensor x_plus_out);
-
-Tensor att_one_v5(Tensor x, Tensor sx, Tensor s, Tensor ln_w, Tensor ln_b,
-                  Tensor lx_w, Tensor lx_b, Tensor k_mix, Tensor v_mix,
-                  Tensor r_mix, Tensor kw,
-                  /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx,
-                  Tensor rw,
-                  /* imm */ Tensor rx, Tensor ow, Tensor t_first,
-                  /* imm */ Tensor k, Tensor t_decay, /* imm */ Tensor v,
-                  /* imm */ Tensor r, /* imm */ Tensor s1,
-                  /* out */ Tensor x_plus_out, /* out */ Tensor s2);
-
-Tensor ffn_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               /* imm */ Tensor buf,
-               /* out */ Tensor x_plus_out);
-
-Tensor ffn_one(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix,
-               Tensor r_mix, Tensor kw, Tensor vw, Tensor rw,
-               /* imm */ Tensor buf,
-               /* out */ Tensor x_plus_out);
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("wkv_forward", &wkv_forward, "wkv forward");
-    m.def("mm8_seq", &mm8_seq, "mm8 seq");
-    m.def("mm8_one", &mm8_one, "mm8 one");
-    m.def("gemm_fp16_cublas", &gemm_fp16_cublas_tensor, "gemv fp16 cublas");
-    m.def("att_one", &att_one, "att one");
-    m.def("att_one_v5", &att_one_v5, "att one v5");
-    m.def("att_seq", &att_seq, "att seq");
-    m.def("ffn_seq", &ffn_seq, "ffn seq");
-    m.def("ffn_one", &ffn_one, "ffn one");
-}
-
-TORCH_LIBRARY(rwkv, m) {
-    m.def("wkv_forward", wkv_forward);
-    m.def("mm8_seq", mm8_seq);
-    m.def("mm8_one", mm8_one);
-    m.def("gemm_fp16_cublas", gemm_fp16_cublas_tensor);
-    m.def("att_one", att_one);
-    m.def("att_one_v5", &att_one_v5);
-    m.def("att_seq", att_seq);
-    m.def("ffn_seq", ffn_seq);
-    m.def("ffn_one", ffn_one);
-}

backend-python/rwkv_pip/cpp/model.py

@@ -9,6 +9,9 @@ class RWKV:
         self.model = rwkv_cpp_model.RWKVModel(self.library, model_path)
         self.w = {}  # fake weight
         self.w["emb.weight"] = [0] * self.model.n_vocab
+        self.version = (
+            self.model.arch_version_major + self.model.arch_version_minor / 10
+        )
 
     def forward(self, tokens: List[int], state: Union[Any, None] = None):
         return self.model.eval_sequence_in_chunks(tokens, state, use_numpy=True)

backend-python/rwkv_pip/cpp/rwkv_cpp_model.py

@@ -52,9 +52,14 @@ class RWKVModel:
         if 'gpu_layers_count' in kwargs:
             gpu_layer_count = kwargs['gpu_layers_count']
 
-        assert os.path.isfile(model_path), f'{model_path} is not a file'
-        assert thread_count > 0, 'Thread count must be > 0'
-        assert gpu_layer_count >= 0, 'GPU layer count must be >= 0'
+        if not os.path.isfile(model_path):
+            raise ValueError(f'{model_path} is not a file')
+        
+        if not (thread_count > 0):
+            raise ValueError('Thread count must be > 0')  
+
+        if not (gpu_layer_count >= 0):
+            raise ValueError('GPU layer count must be >= 0')
 
         self._library: rwkv_cpp_shared_library.RWKVSharedLibrary = shared_library
 
@@ -84,10 +89,19 @@ class RWKVModel:
             Count of layers to offload onto the GPU, must be >= 0.
         """
 
-        assert layer_count >= 0, 'Layer count must be >= 0'
+        if not (layer_count >= 0):
+            raise ValueError('Layer count must be >= 0')
 
         return self._library.rwkv_gpu_offload_layers(self._ctx, layer_count)
 
+    @property
+    def arch_version_major(self) -> int:
+        return self._library.rwkv_get_arch_version_major(self._ctx)
+
+    @property
+    def arch_version_minor(self) -> int:
+        return self._library.rwkv_get_arch_version_minor(self._ctx)
+
     @property
     def n_vocab(self) -> int:
         return self._library.rwkv_get_n_vocab(self._ctx)
@@ -133,7 +147,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -207,7 +222,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -281,7 +297,8 @@ class RWKVModel:
             Logits vector of shape (n_vocab); state for the next step.
         """
 
-        assert self._valid, 'Model was freed'
+        if not self._valid:
+            raise ValueError('Model was freed')
 
         use_numpy = self._detect_numpy_usage([state_in, state_out, logits_out], use_numpy)
 
@@ -320,7 +337,8 @@ class RWKVModel:
         The object must not be used anymore after calling this method.
         """
 
-        assert self._valid, 'Already freed'
+        if not self._valid:
+            raise ValueError('Already freed')
 
         self._valid = False
 
@@ -344,16 +362,25 @@ class RWKVModel:
     def _validate_tensor(self, tensor: NumpyArrayOrPyTorchTensor, name: str, size: int) -> None:
         if self._is_pytorch_tensor(tensor):
             tensor: torch.Tensor = tensor
-            assert tensor.device == torch.device('cpu'), f'{name} is not on CPU'
-            assert tensor.dtype == torch.float32, f'{name} is not of type float32'
-            assert tensor.shape == (size,), f'{name} has invalid shape {tensor.shape}, expected ({size})'
-            assert tensor.is_contiguous(), f'{name} is not contiguous'
+            
+            if tensor.device != torch.device('cpu'):
+                raise ValueError(f'{name} is not on CPU')
+            if tensor.dtype != torch.float32:
+                raise ValueError(f'{name} is not of type float32')
+            if tensor.shape != (size,):
+                raise ValueError(f'{name} has invalid shape {tensor.shape}, expected ({size})')
+            if not tensor.is_contiguous():
+                raise ValueError(f'{name} is not contiguous')
         else:
             import numpy as np
             tensor: np.ndarray = tensor
-            assert tensor.dtype == np.float32, f'{name} is not of type float32'
-            assert tensor.shape == (size,), f'{name} has invalid shape {tensor.shape}, expected ({size})'
-            assert tensor.data.contiguous, f'{name} is not contiguous'
+
+            if tensor.dtype != np.float32:
+                raise ValueError(f'{name} is not of type float32')
+            if tensor.shape != (size,):
+                raise ValueError(f'{name} has invalid shape {tensor.shape}, expected ({size})')
+            if not tensor.data.contiguous:
+                raise ValueError(f'{name} is not contiguous')
 
     def _get_data_ptr(self, tensor: NumpyArrayOrPyTorchTensor):
         if self._is_pytorch_tensor(tensor):

backend-python/rwkv_pip/cpp/rwkv_cpp_shared_library.py

@@ -6,21 +6,22 @@ import platform
 from typing import Optional, List, Tuple, Callable
 
 QUANTIZED_FORMAT_NAMES: Tuple[str, str, str, str, str] = (
-    'Q4_0',
-    'Q4_1',
-    'Q5_0',
-    'Q5_1',
-    'Q8_0'
+    "Q4_0",
+    "Q4_1",
+    "Q5_0",
+    "Q5_1",
+    "Q8_0",
 )
 
 P_FLOAT = ctypes.POINTER(ctypes.c_float)
 P_INT = ctypes.POINTER(ctypes.c_int32)
 
-class RWKVContext:
 
+class RWKVContext:
     def __init__(self, ptr: ctypes.pointer) -> None:
         self.ptr: ctypes.pointer = ptr
 
+
 class RWKVSharedLibrary:
     """
     Python wrapper around rwkv.cpp shared library.
@@ -39,7 +40,7 @@ class RWKVSharedLibrary:
         #  When Python is greater than 3.8, we need to reprocess the custom dll
         #  according to the documentation to prevent loading failure errors.
         #  https://docs.python.org/3/whatsnew/3.8.html#ctypes
-        if platform.system().lower() == 'windows':
+        if platform.system().lower() == "windows":
             self.library = ctypes.CDLL(shared_library_path, winmode=0)
         else:
             self.library = ctypes.cdll.LoadLibrary(shared_library_path)
@@ -47,39 +48,48 @@ class RWKVSharedLibrary:
         self.library.rwkv_init_from_file.argtypes = [ctypes.c_char_p, ctypes.c_uint32]
         self.library.rwkv_init_from_file.restype = ctypes.c_void_p
 
-        self.library.rwkv_gpu_offload_layers.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
+        self.library.rwkv_gpu_offload_layers.argtypes = [
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+        ]
         self.library.rwkv_gpu_offload_layers.restype = ctypes.c_bool
 
         self.library.rwkv_eval.argtypes = [
-            ctypes.c_void_p, # ctx
-            ctypes.c_int32, # token
-            P_FLOAT, # state_in
-            P_FLOAT, # state_out
-            P_FLOAT  # logits_out
+            ctypes.c_void_p,  # ctx
+            ctypes.c_int32,  # token
+            P_FLOAT,  # state_in
+            P_FLOAT,  # state_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval.restype = ctypes.c_bool
 
         self.library.rwkv_eval_sequence.argtypes = [
-            ctypes.c_void_p, # ctx
-            P_INT, # tokens
-            ctypes.c_size_t, # token count
-            P_FLOAT, # state_in
-            P_FLOAT, # state_out
-            P_FLOAT  # logits_out
+            ctypes.c_void_p,  # ctx
+            P_INT,  # tokens
+            ctypes.c_size_t,  # token count
+            P_FLOAT,  # state_in
+            P_FLOAT,  # state_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval_sequence.restype = ctypes.c_bool
 
         self.library.rwkv_eval_sequence_in_chunks.argtypes = [
-            ctypes.c_void_p, # ctx
-            P_INT, # tokens
-            ctypes.c_size_t, # token count
-            ctypes.c_size_t, # chunk size
-            P_FLOAT, # state_in
-            P_FLOAT, # state_out
-            P_FLOAT  # logits_out
+            ctypes.c_void_p,  # ctx
+            P_INT,  # tokens
+            ctypes.c_size_t,  # token count
+            ctypes.c_size_t,  # chunk size
+            P_FLOAT,  # state_in
+            P_FLOAT,  # state_out
+            P_FLOAT,  # logits_out
         ]
         self.library.rwkv_eval_sequence_in_chunks.restype = ctypes.c_bool
 
+        self.library.rwkv_get_arch_version_major.argtypes = [ctypes.c_void_p]
+        self.library.rwkv_get_arch_version_major.restype = ctypes.c_uint32
+
+        self.library.rwkv_get_arch_version_minor.argtypes = [ctypes.c_void_p]
+        self.library.rwkv_get_arch_version_minor.restype = ctypes.c_uint32
+
         self.library.rwkv_get_n_vocab.argtypes = [ctypes.c_void_p]
         self.library.rwkv_get_n_vocab.restype = ctypes.c_size_t
 
@@ -101,7 +111,11 @@ class RWKVSharedLibrary:
         self.library.rwkv_free.argtypes = [ctypes.c_void_p]
         self.library.rwkv_free.restype = None
 
-        self.library.rwkv_quantize_model_file.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+        self.library.rwkv_quantize_model_file.argtypes = [
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+        ]
         self.library.rwkv_quantize_model_file.restype = ctypes.c_bool
 
         self.library.rwkv_get_system_info_string.argtypes = []
@@ -109,7 +123,9 @@ class RWKVSharedLibrary:
 
         self.nullptr = ctypes.cast(0, ctypes.c_void_p)
 
-    def rwkv_init_from_file(self, model_file_path: str, thread_count: int) -> RWKVContext:
+    def rwkv_init_from_file(
+        self, model_file_path: str, thread_count: int
+    ) -> RWKVContext:
         """
         Loads the model from a file and prepares it for inference.
         Throws an exception in case of any error. Error messages would be printed to stderr.
@@ -122,9 +138,12 @@ class RWKVSharedLibrary:
             Count of threads to use, must be positive.
         """
 
-        ptr = self.library.rwkv_init_from_file(model_file_path.encode('utf-8'), ctypes.c_uint32(thread_count))
+        ptr = self.library.rwkv_init_from_file(
+            model_file_path.encode("utf-8"), ctypes.c_uint32(thread_count)
+        )
 
-        assert ptr is not None, 'rwkv_init_from_file failed, check stderr'
+        if ptr is None:
+            raise ValueError("rwkv_init_from_file failed, check stderr")
 
         return RWKVContext(ptr)
 
@@ -145,17 +164,20 @@ class RWKVSharedLibrary:
             Count of layers to offload onto the GPU, must be >= 0.
         """
 
-        assert layer_count >= 0, 'Layer count must be >= 0'
+        if not (layer_count >= 0):
+            raise ValueError("Layer count must be >= 0")
 
-        return self.library.rwkv_gpu_offload_layers(ctx.ptr, ctypes.c_uint32(layer_count))
+        return self.library.rwkv_gpu_offload_layers(
+            ctx.ptr, ctypes.c_uint32(layer_count)
+        )
 
     def rwkv_eval(
-            self,
-            ctx: RWKVContext,
-            token: int,
-            state_in_address: Optional[int],
-            state_out_address: int,
-            logits_out_address: int
+        self,
+        ctx: RWKVContext,
+        token: int,
+        state_in_address: Optional[int],
+        state_out_address: int,
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a single token.
@@ -176,21 +198,22 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval(
+        if not self.library.rwkv_eval(
             ctx.ptr,
             ctypes.c_int32(token),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
-        ), 'rwkv_eval failed, check stderr'
+            ctypes.cast(logits_out_address, P_FLOAT),
+        ):
+            raise ValueError("rwkv_eval failed, check stderr")
 
     def rwkv_eval_sequence(
-            self,
-            ctx: RWKVContext,
-            tokens: List[int],
-            state_in_address: Optional[int],
-            state_out_address: int,
-            logits_out_address: int
+        self,
+        ctx: RWKVContext,
+        tokens: List[int],
+        state_in_address: Optional[int],
+        state_out_address: int,
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a sequence of tokens.
@@ -223,23 +246,24 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval_sequence(
+        if not self.library.rwkv_eval_sequence(
             ctx.ptr,
             ctypes.cast((ctypes.c_int32 * len(tokens))(*tokens), P_INT),
             ctypes.c_size_t(len(tokens)),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
-        ), 'rwkv_eval_sequence failed, check stderr'
+            ctypes.cast(logits_out_address, P_FLOAT),
+        ):
+            raise ValueError("rwkv_eval_sequence failed, check stderr")
 
     def rwkv_eval_sequence_in_chunks(
-            self,
-            ctx: RWKVContext,
-            tokens: List[int],
-            chunk_size: int,
-            state_in_address: Optional[int],
-            state_out_address: int,
-            logits_out_address: int
+        self,
+        ctx: RWKVContext,
+        tokens: List[int],
+        chunk_size: int,
+        state_in_address: Optional[int],
+        state_out_address: int,
+        logits_out_address: int,
     ) -> None:
         """
         Evaluates the model for a sequence of tokens using `rwkv_eval_sequence`, splitting a potentially long sequence into fixed-length chunks.
@@ -269,15 +293,40 @@ class RWKVSharedLibrary:
             Address of the first element of a FP32 buffer of size rwkv_get_logits_buffer_element_count. This buffer will be written to.
         """
 
-        assert self.library.rwkv_eval_sequence_in_chunks(
+        if not self.library.rwkv_eval_sequence_in_chunks(
             ctx.ptr,
             ctypes.cast((ctypes.c_int32 * len(tokens))(*tokens), P_INT),
             ctypes.c_size_t(len(tokens)),
             ctypes.c_size_t(chunk_size),
             ctypes.cast(0 if state_in_address is None else state_in_address, P_FLOAT),
             ctypes.cast(state_out_address, P_FLOAT),
-            ctypes.cast(logits_out_address, P_FLOAT)
-        ), 'rwkv_eval_sequence_in_chunks failed, check stderr'
+            ctypes.cast(logits_out_address, P_FLOAT),
+        ):
+            raise ValueError("rwkv_eval_sequence_in_chunks failed, check stderr")
+
+    def rwkv_get_arch_version_major(self, ctx: RWKVContext) -> int:
+        """
+        Returns the major version used by the given model.
+
+        Parameters
+        ----------
+        ctx : RWKVContext
+            RWKV context obtained from rwkv_init_from_file.
+        """
+
+        return self.library.rwkv_get_arch_version_major(ctx.ptr)
+
+    def rwkv_get_arch_version_minor(self, ctx: RWKVContext) -> int:
+        """
+        Returns the minor version used by the given model.
+
+        Parameters
+        ----------
+        ctx : RWKVContext
+            RWKV context obtained from rwkv_init_from_file.
+        """
+
+        return self.library.rwkv_get_arch_version_minor(ctx.ptr)
 
     def rwkv_get_n_vocab(self, ctx: RWKVContext) -> int:
         """
@@ -358,7 +407,9 @@ class RWKVSharedLibrary:
 
         ctx.ptr = self.nullptr
 
-    def rwkv_quantize_model_file(self, model_file_path_in: str, model_file_path_out: str, format_name: str) -> None:
+    def rwkv_quantize_model_file(
+        self, model_file_path_in: str, model_file_path_out: str, format_name: str
+    ) -> None:
         """
         Quantizes FP32 or FP16 model to one of INT4 formats.
         Throws an exception in case of any error. Error messages would be printed to stderr.
@@ -373,20 +424,25 @@ class RWKVSharedLibrary:
             One of QUANTIZED_FORMAT_NAMES.
         """
 
-        assert format_name in QUANTIZED_FORMAT_NAMES, f'Unknown format name {format_name}, use one of {QUANTIZED_FORMAT_NAMES}'
+        if format_name not in QUANTIZED_FORMAT_NAMES:
+            raise ValueError(
+                f"Unknown format name {format_name}, use one of {QUANTIZED_FORMAT_NAMES}"
+            )
 
-        assert self.library.rwkv_quantize_model_file(
-            model_file_path_in.encode('utf-8'),
-            model_file_path_out.encode('utf-8'),
-            format_name.encode('utf-8')
-        ), 'rwkv_quantize_model_file failed, check stderr'
+        if not self.library.rwkv_quantize_model_file(
+            model_file_path_in.encode("utf-8"),
+            model_file_path_out.encode("utf-8"),
+            format_name.encode("utf-8"),
+        ):
+            raise ValueError("rwkv_quantize_model_file failed, check stderr")
 
     def rwkv_get_system_info_string(self) -> str:
         """
         Returns system information string.
         """
 
-        return self.library.rwkv_get_system_info_string().decode('utf-8')
+        return self.library.rwkv_get_system_info_string().decode("utf-8")
+
 
 def load_rwkv_shared_library() -> RWKVSharedLibrary:
     """
@@ -396,27 +452,27 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
 
     file_name: str
 
-    if 'win32' in sys.platform or 'cygwin' in sys.platform:
-        file_name = 'rwkv.dll'
-    elif 'darwin' in sys.platform:
-        file_name = 'librwkv.dylib'
+    if "win32" in sys.platform or "cygwin" in sys.platform:
+        file_name = "rwkv.dll"
+    elif "darwin" in sys.platform:
+        file_name = "librwkv.dylib"
     else:
-        file_name = 'librwkv.so'
+        file_name = "librwkv.so"
 
     # Possible sub-paths to the library relative to the repo dir.
     child_paths: List[Callable[[pathlib.Path], pathlib.Path]] = [
         # No lookup for Debug config here.
         # I assume that if a user wants to debug the library,
         # they will be able to find the library and set the exact path explicitly.
-        lambda p: p / 'backend-python' / 'rwkv_pip' / 'cpp' / file_name,
-        lambda p: p / 'bin' / 'Release' / file_name,
-        lambda p: p / 'bin' / file_name,
+        lambda p: p / "backend-python" / "rwkv_pip" / "cpp" / file_name,
+        lambda p: p / "bin" / "Release" / file_name,
+        lambda p: p / "bin" / file_name,
         # Some people prefer to build in the "build" subdirectory.
-        lambda p: p / 'build' / 'bin' / 'Release' / file_name,
-        lambda p: p / 'build' / 'bin' / file_name,
-        lambda p: p / 'build' / file_name,
+        lambda p: p / "build" / "bin" / "Release" / file_name,
+        lambda p: p / "build" / "bin" / file_name,
+        lambda p: p / "build" / file_name,
         # Fallback.
-        lambda p: p / file_name
+        lambda p: p / file_name,
     ]
 
     working_dir: pathlib.Path = pathlib.Path(os.path.abspath(os.getcwd()))
@@ -430,7 +486,7 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
         # .
         working_dir,
         # Repo dir relative to this Python file.
-        pathlib.Path(os.path.abspath(__file__)).parent.parent.parent
+        pathlib.Path(os.path.abspath(__file__)).parent.parent.parent,
     ]
 
     for parent_path in parent_paths:
@@ -440,5 +496,7 @@ def load_rwkv_shared_library() -> RWKVSharedLibrary:
             if os.path.isfile(full_path):
                 return RWKVSharedLibrary(str(full_path))
 
-    assert False, (f'Failed to find {file_name} automatically; '
-                   f'you need to find the library and create RWKVSharedLibrary specifying the path to it')
+    raise ValueError(
+        f"Failed to find {file_name} automatically; "
+        f"you need to find the library and create RWKVSharedLibrary specifying the path to it"
+    )

backend-python/rwkv_pip/model.py

@@ -488,14 +488,19 @@ class RWKV(MyModule):
             print_need_newline = False
 
             REAL_TIME_FIRST = False
+            args.time_state = False
             for x in list(w.keys()):
                 if ".time_faaaa" in x:
                     REAL_TIME_FIRST = True
+                if ".time_state" in x:
+                    args.time_state = True
             if REAL_TIME_FIRST:
                 w = {
-                    k.replace(".time_faaaa", ".time_first")
-                    if ".time_faaaa" in k
-                    else k: v
+                    (
+                        k.replace(".time_faaaa", ".time_first")
+                        if ".time_faaaa" in k
+                        else k
+                    ): v
                     for k, v in w.items()
                 }
                 self.w = w
@@ -552,7 +557,12 @@ class RWKV(MyModule):
                     elif ".ln_x" in x:  # need fp32 for group_norm
                         w[x] = w[x].float()
                     else:
-                        if (len(w[x].shape) == 2) and ("emb" not in x):
+                        if (
+                            (len(w[x].shape) == 2)
+                            and ("emb" not in x)
+                            and ("_w1" not in x)
+                            and ("_w2" not in x)
+                        ):
                             if WTYPE != torch.uint8:
                                 w[x] = w[x].to(dtype=WTYPE)
                             else:
@@ -626,10 +636,12 @@ class RWKV(MyModule):
                         torch.cuda.empty_cache()
 
                 shape = [i for i in w[x].shape if i != 1]
-                if len(shape) > 1:
-                    shape = f" {str(shape[0]).rjust(5)} {str(shape[1]).rjust(5)}"
+                if len(shape) > 2:
+                    shape = f" {str(shape[0]).rjust(5)} {str(shape[1]).rjust(5)} {str(shape[2]).rjust(5)}"
+                elif len(shape) > 1:
+                    shape = f" {str(shape[0]).rjust(5)} {str(shape[1]).rjust(5)}      "
                 else:
-                    shape = f" {str(shape[0]).rjust(5)}      "
+                    shape = f" {str(shape[0]).rjust(5)}            "
                 if layer_id == 0 or layer_id >= args.n_layer - 1:
                     if print_need_newline:
                         prxxx("\n", end="")
@@ -2103,16 +2115,25 @@ class RWKV(MyModule):
                         state[i * 3 + 0] = torch.zeros(
                             args.n_embd, dtype=atype, requires_grad=False, device=dev
                         ).contiguous()
-                        state[i * 3 + 1] = torch.zeros(
-                            (
-                                args.n_head,
-                                args.n_att // args.n_head,
-                                args.n_att // args.n_head,
-                            ),
-                            dtype=torch.float,
-                            requires_grad=False,
-                            device=dev,
-                        ).contiguous()
+                        if args.time_state:
+                            state[i * 3 + 1] = (
+                                w[f"blocks.{i}.att.time_state"]
+                                .transpose(1, 2)
+                                .to(dtype=torch.float, device=dev)
+                                .requires_grad_(False)
+                                .contiguous()
+                            )
+                        else:
+                            state[i * 3 + 1] = torch.zeros(
+                                (
+                                    args.n_head,
+                                    args.n_att // args.n_head,
+                                    args.n_att // args.n_head,
+                                ),
+                                dtype=torch.float,
+                                requires_grad=False,
+                                device=dev,
+                            ).contiguous()
                         state[i * 3 + 2] = torch.zeros(
                             args.n_embd, dtype=atype, requires_grad=False, device=dev
                         ).contiguous()

backend-python/rwkv_pip/utils.py

@@ -34,6 +34,25 @@ class PIPELINE_ARGS:
         )
 
 
+class ABC_TOKENIZER:
+    def __init__(self):
+        self.pad_token_id = 0
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+
+    def encode(self, text):
+        ids = [ord(c) for c in text]
+        return ids
+
+    def decode(self, ids):
+        txt = "".join(
+            chr(idx) if idx > self.eos_token_id else ""
+            for idx in ids
+            if idx != self.eos_token_id
+        )
+        return txt
+
+
 class PIPELINE:
     def __init__(self, model, WORD_NAME: str):
         self.model = model
@@ -48,6 +67,8 @@ class PIPELINE:
             self.tokenizer = TRIE_TOKENIZER(
                 os.path.dirname(os.path.abspath(__file__)) + "/rwkv_vocab_v20230424.txt"
             )
+        elif WORD_NAME == "abc_tokenizer":
+            self.tokenizer = ABC_TOKENIZER()
         else:
             if WORD_NAME.endswith(".txt"):
                 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -150,10 +171,17 @@ class PIPELINE:
             all_tokens += [token]
             for xxx in occurrence:
                 occurrence[xxx] *= args.alpha_decay
+
+            ttt = self.decode([token])
+            www = 1
+            if ttt in " \t0123456789":
+                www = 0
+            # elif ttt in '\r\n,.;?!"\':+-*/=#@$%^&_`~|<>\\()[]{}，。；“”：？！（）【】':
+            #     www = 0.5
             if token not in occurrence:
-                occurrence[token] = 1
+                occurrence[token] = www
             else:
-                occurrence[token] += 1
+                occurrence[token] += www
             # print(occurrence) # debug
 
             # output

backend-python/rwkv_pip/webgpu/model.py

@@ -13,19 +13,38 @@ except ModuleNotFoundError:
 
 class RWKV:
     def __init__(self, model_path: str, strategy: str = None):
-        self.model = wrp.v5.Model(
-            model_path,
-            turbo=True,
-            quant=32 if "i8" in strategy else None,
-            quant_nf4=26 if "i4" in strategy else None,
+        layer = (
+            int(s.lstrip("layer"))
+            for s in strategy.split()
+            for s in s.split(",")
+            if s.startswith("layer")
         )
+
+        chunk_size = (
+            int(s.lstrip("chunk"))
+            for s in strategy.split()
+            for s in s.split(",")
+            if s.startswith("chunk")
+        )
+        self.token_chunk_size = next(chunk_size, 32)
+
+        args = {
+            "path": model_path,
+            "quant": next(layer, 31) if "i8" in strategy else 0,
+            "quant_nf4": next(layer, 26) if "i4" in strategy else 0,
+        }
+        self.model = wrp.Model(**args)
+        self.info = self.model.info()
         self.w = {}  # fake weight
-        self.w["emb.weight"] = [0] * wrp.peek_info(model_path).num_vocab
+        self.w["emb.weight"] = [0] * self.info.num_vocab
+        self.version = str(self.info.version).lower()
+        self.version = float(self.version.lower().replace("v", ""))
 
     def forward(self, tokens: List[int], state: Union[Any, None] = None):
-        if type(state).__name__ == "BackedState":  # memory state
-            gpu_state = wrp.v5.ModelState(self.model, 1)
-            gpu_state.load(state)
-        else:
-            gpu_state = state
-        return wrp.v5.run_one(self.model, tokens, gpu_state)
+        if state is None:
+            self.model.clear_state()
+        elif type(state).__name__ == "State_Cpu":
+            self.model.load_state(state)
+        logits = self.model.run(tokens, self.token_chunk_size)
+        ret_state = "State_Gpu"
+        return logits, ret_state

backend-python/utils/rwkv.py

@@ -4,18 +4,14 @@ import os
 import pathlib
 import copy
 import re
-from typing import Dict, Iterable, List, Tuple, Union, Type
+import time
+from typing import Dict, Iterable, List, Tuple, Union, Type, Callable
 from utils.log import quick_log
-from fastapi import HTTPException
+from fastapi import HTTPException, status
 from pydantic import BaseModel, Field
 from routes import state_cache
 import global_var
 
-
-END_OF_TEXT = 0
-END_OF_LINE_DOUBLE = 535
-
-
 os.environ["TORCH_EXTENSIONS_DIR"] = f"{pathlib.Path(__file__).parent.parent.resolve()}"
 
 
@@ -28,7 +24,11 @@ class RWKVType(Enum):
 
 class AbstractRWKV(ABC):
     def __init__(self, model, pipeline):
+        self.EOS_ID = 0
+
         self.name = "rwkv"
+        self.model_path = ""
+        self.version = 4
         self.model = model
         self.pipeline = pipeline
         self.model_state = None
@@ -42,6 +42,10 @@ class AbstractRWKV(ABC):
         self.top_k = 0
         self.penalty_alpha_presence = 0
         self.penalty_alpha_frequency = 1
+        self.penalty_decay = 0.99
+        self.global_penalty = False
+        self.state_path = ""
+        self.state_tuned = None
 
     @abstractmethod
     def adjust_occurrence(self, occurrence: Dict, token: int):
@@ -235,7 +239,10 @@ class AbstractRWKV(ABC):
         except HTTPException:
             pass
         if cache is None or cache["prompt"] == "" or cache["state"] is None:
-            self.model_state = None
+            if self.state_path:
+                self.model_state = copy.deepcopy(self.state_tuned)
+            else:
+                self.model_state = None
             self.model_tokens = []
         else:
             delta_prompt = prompt[len(cache["prompt"]) :]
@@ -245,9 +252,16 @@ class AbstractRWKV(ABC):
 
         prompt_token_len = 0
         if delta_prompt != "":
+            prompt_start_time = time.time()
             logits, prompt_token_len = self.run_rnn(
                 self.fix_tokens(self.pipeline.encode(delta_prompt))
             )
+            prompt_end_time = time.time()
+            prompt_interval = prompt_end_time - prompt_start_time
+            tps = 0
+            if prompt_interval > 0:
+                tps = prompt_token_len / prompt_interval
+            print(f"Prompt Prefill TPS: {tps:.2f}", end=" ", flush=True)
             try:
                 state_cache.add_state(
                     state_cache.AddStateBody(
@@ -274,7 +288,18 @@ class AbstractRWKV(ABC):
                 logits, temperature=self.temperature, top_p=self.top_p, top_k=self.top_k
             )
 
-            if token == END_OF_TEXT:
+            if token == self.EOS_ID:
+                try:
+                    state_cache.add_state(
+                        state_cache.AddStateBody(
+                            prompt=prompt + response,
+                            tokens=self.model_tokens,
+                            state=self.model_state,
+                            logits=logits,
+                        )
+                    )
+                except HTTPException:
+                    pass
                 yield response, "", prompt_token_len, completion_token_len
                 break
 
@@ -305,22 +330,25 @@ class AbstractRWKV(ABC):
                             yield response, "", prompt_token_len, completion_token_len
                             break
                     elif type(stop) == list:
-                        stop_exist_regex = "|".join(stop)
-                        matched = re.search(stop_exist_regex, response)
-                        if matched:
-                            try:
-                                state_cache.add_state(
-                                    state_cache.AddStateBody(
-                                        prompt=prompt + response,
-                                        tokens=self.model_tokens,
-                                        state=self.model_state,
-                                        logits=logits,
+                        exit_flag = False
+                        for s in stop:
+                            if s in response:
+                                try:
+                                    state_cache.add_state(
+                                        state_cache.AddStateBody(
+                                            prompt=prompt + response,
+                                            tokens=self.model_tokens,
+                                            state=self.model_state,
+                                            logits=logits,
+                                        )
                                     )
-                                )
-                            except HTTPException:
-                                pass
-                            response = response.split(matched.group())[0]
-                            yield response, "", prompt_token_len, completion_token_len
+                                except HTTPException:
+                                    pass
+                                exit_flag = True
+                                response = response.split(s)[0]
+                                yield response, "", prompt_token_len, completion_token_len
+                                break
+                        if exit_flag:
                             break
                 out_last = begin + i + 1
                 if i == self.max_tokens_per_generation - 1:
@@ -363,18 +391,24 @@ class TextRWKV(AbstractRWKV):
             self.bot = "Assistant"
             self.END_OF_LINE = 11
 
-        self.AVOID_REPEAT_TOKENS = []
+        self.AVOID_REPEAT_TOKENS = set()
         AVOID_REPEAT = "，：？！"
         for i in AVOID_REPEAT:
             dd = self.pipeline.encode(i)
             assert len(dd) == 1
-            self.AVOID_REPEAT_TOKENS += dd
+            self.AVOID_REPEAT_TOKENS.add(dd[0])
+        self.AVOID_PENALTY_TOKENS = set()
+        AVOID_PENALTY = '\n,.:?!，。：？！"“”<>[]{}/\\|;；~`@#$%^&*()_+-=0123456789 '
+        for i in AVOID_PENALTY:
+            dd = self.pipeline.encode(i)
+            if len(dd) == 1:
+                self.AVOID_PENALTY_TOKENS.add(dd[0])
 
         self.__preload()
 
     def adjust_occurrence(self, occurrence: Dict, token: int):
         for xxx in occurrence:
-            occurrence[xxx] *= 0.996
+            occurrence[xxx] *= self.penalty_decay
         if token not in occurrence:
             occurrence[token] = 1
         else:
@@ -382,26 +416,24 @@ class TextRWKV(AbstractRWKV):
 
     def adjust_forward_logits(self, logits: List[float], occurrence: Dict, i: int):
         for n in occurrence:
+            # if n not in self.AVOID_PENALTY_TOKENS:
             logits[n] -= (
                 self.penalty_alpha_presence
                 + occurrence[n] * self.penalty_alpha_frequency
             )
 
-        if i == 0:
+        # set global_penalty to False to get the same generated results as the official RWKV Gradio
+        if self.global_penalty and i == 0:
             for token in self.model_tokens:
                 token = int(token)
-                for xxx in occurrence:
-                    occurrence[xxx] *= 0.996
-                if token not in occurrence:
-                    occurrence[token] = 1
-                else:
-                    occurrence[token] += 1
+                if token not in self.AVOID_PENALTY_TOKENS:
+                    self.adjust_occurrence(occurrence, token)
 
     # Model only saw '\n\n' as [187, 187] before, but the tokenizer outputs [535] for it at the end
     def fix_tokens(self, tokens) -> List[int]:
         if self.rwkv_type == RWKVType.World:
             return tokens
-        if len(tokens) > 0 and tokens[-1] == END_OF_LINE_DOUBLE:
+        if len(tokens) > 0 and tokens[-1] == 535:
             tokens = tokens[:-1] + [self.END_OF_LINE, self.END_OF_LINE]
         return tokens
 
@@ -459,7 +491,7 @@ The following is a coherent verbose detailed conversation between a girl named {
             pass
 
 
-class MusicRWKV(AbstractRWKV):
+class MusicMidiRWKV(AbstractRWKV):
     def __init__(self, model, pipeline):
         super().__init__(model, pipeline)
 
@@ -501,8 +533,47 @@ class MusicRWKV(AbstractRWKV):
         return " " + delta
 
 
+class MusicAbcRWKV(AbstractRWKV):
+    def __init__(self, model, pipeline):
+        super().__init__(model, pipeline)
+
+        self.EOS_ID = 3
+
+        self.max_tokens_per_generation = 500
+        self.temperature = 1
+        self.top_p = 0.8
+        self.top_k = 8
+
+        self.rwkv_type = RWKVType.Music
+
+    def adjust_occurrence(self, occurrence: Dict, token: int):
+        pass
+
+    def adjust_forward_logits(self, logits: List[float], occurrence: Dict, i: int):
+        pass
+
+    def fix_tokens(self, tokens) -> List[int]:
+        return tokens
+
+    def run_rnn(
+        self, _tokens: List[str], newline_adj: int = 0
+    ) -> Tuple[List[float], int]:
+        tokens = [int(x) for x in _tokens]
+        token_len = len(tokens)
+        self.model_tokens += tokens
+        out, self.model_state = self.model.forward(tokens, self.model_state)
+        return out, token_len
+
+    def delta_postprocess(self, delta: str) -> str:
+        return delta
+
+
 def get_tokenizer(tokenizer_len: int):
     tokenizer_dir = f"{pathlib.Path(__file__).parent.parent.resolve()}/rwkv_pip/"
+    if tokenizer_len < 2176:
+        return "abc_tokenizer"
+    if tokenizer_len < 20096:
+        return tokenizer_dir + "tokenizer-midipiano.json"
     if tokenizer_len < 50277:
         return tokenizer_dir + "tokenizer-midi.json"
     elif tokenizer_len < 65536:
@@ -511,21 +582,49 @@ def get_tokenizer(tokenizer_len: int):
         return "rwkv_vocab_v20230424"
 
 
+def get_model_path(model_path: str) -> str:
+    if os.path.isabs(model_path):
+        return model_path
+
+    working_dir: pathlib.Path = pathlib.Path(os.path.abspath(os.getcwd()))
+
+    parent_paths: List[pathlib.Path] = [
+        working_dir,  # [cwd](RWKV-Runner)/models/xxx
+        working_dir.parent,  # [cwd](backend-python)/../models/xxx
+        pathlib.Path(
+            os.path.abspath(__file__)
+        ).parent.parent,  # backend-python/models/xxx
+        pathlib.Path(
+            os.path.abspath(__file__)
+        ).parent.parent.parent,  # RWKV-Runner/models/xxx
+    ]
+
+    child_paths: List[Callable[[pathlib.Path], pathlib.Path]] = [
+        lambda p: p / model_path,
+        lambda p: p / "build" / "bin" / model_path,  # for dev
+    ]
+
+    for parent_path in parent_paths:
+        for child_path in child_paths:
+            full_path: pathlib.Path = child_path(parent_path)
+
+            if os.path.isfile(full_path):
+                return str(full_path)
+
+    return model_path
+
+
 def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV:
-    rwkv_beta = global_var.get(global_var.Args).rwkv_beta
+    model_path = get_model_path(model)
+
     rwkv_cpp = getattr(global_var.get(global_var.Args), "rwkv.cpp")
     webgpu = global_var.get(global_var.Args).webgpu
 
-    if "midi" in model.lower() or "abc" in model.lower():
+    if "midi" in model_path.lower() or "abc" in model_path.lower():
         os.environ["RWKV_RESCALE_LAYER"] = "999"
 
     # dynamic import to make RWKV_CUDA_ON work
-    if rwkv_beta:
-        print("Using rwkv-beta")
-        from rwkv_pip.beta.model import (
-            RWKV as Model,
-        )
-    elif rwkv_cpp:
+    if rwkv_cpp:
         print("Using rwkv.cpp, strategy is ignored")
         from rwkv_pip.cpp.model import (
             RWKV as Model,
@@ -541,8 +640,8 @@ def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV
         )
     from rwkv_pip.utils import PIPELINE
 
-    filename, _ = os.path.splitext(os.path.basename(model))
-    model = Model(model, strategy)
+    filename, _ = os.path.splitext(os.path.basename(model_path))
+    model = Model(model_path, strategy)
     if not tokenizer:
         tokenizer = get_tokenizer(len(model.w["emb.weight"]))
     pipeline = PIPELINE(model, tokenizer)
@@ -550,39 +649,152 @@ def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV
     rwkv_map: dict[str, Type[AbstractRWKV]] = {
         "20B_tokenizer": TextRWKV,
         "rwkv_vocab_v20230424": TextRWKV,
-        "tokenizer-midi": MusicRWKV,
+        "tokenizer-midi": MusicMidiRWKV,
+        "tokenizer-midipiano": MusicMidiRWKV,
+        "abc_tokenizer": MusicAbcRWKV,
     }
     tokenizer_name = os.path.splitext(os.path.basename(tokenizer))[0]
+    global_var.set(
+        global_var.Midi_Vocab_Config_Type,
+        (
+            global_var.MidiVocabConfig.Piano
+            if tokenizer_name == "tokenizer-midipiano"
+            else global_var.MidiVocabConfig.Default
+        ),
+    )
     rwkv: AbstractRWKV
     if tokenizer_name in rwkv_map:
         rwkv = rwkv_map[tokenizer_name](model, pipeline)
     else:
-        rwkv = TextRWKV(model, pipeline)
+        tokenizer_name = tokenizer_name.lower()
+        if "music" in tokenizer_name or "midi" in tokenizer_name:
+            rwkv = MusicMidiRWKV(model, pipeline)
+        elif "abc" in tokenizer_name:
+            rwkv = MusicAbcRWKV(model, pipeline)
+        else:
+            rwkv = TextRWKV(model, pipeline)
     rwkv.name = filename
+    rwkv.model_path = model_path
+    rwkv.version = model.version
 
     return rwkv
 
 
 class ModelConfigBody(BaseModel):
     max_tokens: int = Field(default=None, gt=0, le=102400)
-    temperature: float = Field(default=None, ge=0, le=2)
+    temperature: float = Field(default=None, ge=0, le=3)
     top_p: float = Field(default=None, ge=0, le=1)
     presence_penalty: float = Field(default=None, ge=-2, le=2)
     frequency_penalty: float = Field(default=None, ge=-2, le=2)
+    penalty_decay: float = Field(default=None, ge=0.99, le=0.999)
+    top_k: int = Field(default=None, ge=0, le=25)
+    global_penalty: bool = Field(
+        default=None,
+        description="When generating a response, whether to include the submitted prompt as a penalty factor. By turning this off, you will get the same generated results as official RWKV Gradio. If you find duplicate results in the generated results, turning this on can help avoid generating duplicates.",
+    )
+    state: str = Field(default=None, description="state-tuned file path")
 
     model_config = {
         "json_schema_extra": {
             "example": {
                 "max_tokens": 1000,
-                "temperature": 1.2,
-                "top_p": 0.5,
-                "presence_penalty": 0.4,
-                "frequency_penalty": 0.4,
+                "temperature": 1,
+                "top_p": 0.3,
+                "presence_penalty": 0,
+                "frequency_penalty": 1,
+                "penalty_decay": 0.996,
+                "global_penalty": False,
+                "state": "",
             }
         }
     }
 
 
+def load_rwkv_state(
+    model: AbstractRWKV, state_path: str, print_log: bool = True
+) -> HTTPException:
+    if model:
+        if state_path:
+            if model.model_path.endswith(".pth") and state_path.endswith(".pth"):
+                import torch
+
+                state_path = get_model_path(state_path)
+                if model.state_path == state_path:
+                    return
+
+                if not os.path.isfile(state_path):
+                    return HTTPException(
+                        status.HTTP_400_BAD_REQUEST, "state file not found"
+                    )
+
+                try:
+                    state_raw = torch.load(state_path, map_location="cpu")
+                except Exception as e:
+                    print(e)
+                    return HTTPException(
+                        status.HTTP_400_BAD_REQUEST, "state file failed to load"
+                    )
+                state_raw_shape = next(iter(state_raw.values())).shape
+
+                args = model.model.args
+                if (
+                    len(state_raw) != args.n_layer
+                    or state_raw_shape[0] * state_raw_shape[1] != args.n_embd
+                ):
+                    if model.state_path:
+                        pass
+                    elif print_log:
+                        print("state failed to load")
+                    return HTTPException(
+                        status.HTTP_400_BAD_REQUEST, "state shape mismatch"
+                    )
+
+                strategy = model.model.strategy
+                model.state_tuned = [None] * args.n_layer * 3
+
+                for i in range(args.n_layer):
+                    dd = strategy[i]
+                    dev = dd.device
+                    atype = dd.atype
+                    model.state_tuned[i * 3 + 0] = torch.zeros(
+                        args.n_embd, dtype=atype, requires_grad=False, device=dev
+                    ).contiguous()
+                    model.state_tuned[i * 3 + 1] = (
+                        state_raw[f"blocks.{i}.att.time_state"]
+                        .transpose(1, 2)
+                        .to(dtype=torch.float, device=dev)
+                        .requires_grad_(False)
+                        .contiguous()
+                    )
+                    model.state_tuned[i * 3 + 2] = torch.zeros(
+                        args.n_embd, dtype=atype, requires_grad=False, device=dev
+                    ).contiguous()
+
+                state_cache.force_reset_state()
+                model.state_path = state_path
+                if print_log:
+                    print("state loaded")
+            else:
+                if model.state_path:
+                    pass
+                elif print_log:
+                    print("state failed to load")
+                return HTTPException(
+                    status.HTTP_400_BAD_REQUEST,
+                    "file format of the model or state model not supported",
+                )
+        else:
+            if state_path == "" and model.state_path != "":
+                state_cache.force_reset_state()
+                model.state_path = ""
+                model.state_tuned = None  # TODO cached
+                if print_log:
+                    print("state unloaded")
+    else:
+        if print_log:
+            print("state not loaded")
+
+
 def set_rwkv_config(model: AbstractRWKV, body: ModelConfigBody):
     if body.max_tokens is not None:
         model.max_tokens_per_generation = body.max_tokens
@@ -597,6 +809,14 @@ def set_rwkv_config(model: AbstractRWKV, body: ModelConfigBody):
         model.penalty_alpha_presence = body.presence_penalty
     if body.frequency_penalty is not None:
         model.penalty_alpha_frequency = body.frequency_penalty
+    if body.penalty_decay is not None:
+        model.penalty_decay = body.penalty_decay
+    if body.top_k is not None:
+        model.top_k = body.top_k
+    if body.global_penalty is not None:
+        model.global_penalty = body.global_penalty
+    if body.state is not None:
+        load_rwkv_state(model, body.state, False)
 
 
 def get_rwkv_config(model: AbstractRWKV) -> ModelConfigBody:
@@ -606,4 +826,8 @@ def get_rwkv_config(model: AbstractRWKV) -> ModelConfigBody:
         top_p=model.top_p,
         presence_penalty=model.penalty_alpha_presence,
         frequency_penalty=model.penalty_alpha_frequency,
+        penalty_decay=model.penalty_decay,
+        top_k=model.top_k,
+        global_penalty=model.global_penalty,
+        state=model.state_path,
     )

backend-python/utils/torch.py

@@ -19,9 +19,12 @@ def set_torch():
 
 
 def torch_gc():
-    import torch
+    try:
+        import torch
 
-    if torch.cuda.is_available():
-        with torch.cuda.device(0):
-            torch.cuda.empty_cache()
-            torch.cuda.ipc_collect()
+        if torch.cuda.is_available():
+            with torch.cuda.device(0):
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+    except:
+        pass  # prevent 'torch' has no attribute 'cuda' error, so user can use CPU or WebGPU

backend-python/utils/vocab_config_piano.json

@@ -0,0 +1,279 @@
+{
+    "note_events": 128,
+    "wait_events": 125,
+    "max_wait_time": 1000,
+    "velocity_events": 128,
+    "velocity_bins": 16,
+    "velocity_exp": 0.33,
+    "do_token_sorting": true,
+    "unrolled_tokens": false,
+    "decode_end_held_note_delay": 5.0,
+    "decode_fix_repeated_notes": true,
+    "bin_instrument_names": [
+        "piano"
+    ],
+    "ch10_instrument_bin_name": "",
+    "program_name_to_bin_name": {
+        "Acoustic Grand Piano": "piano",
+        "Bright Acoustic Piano": "piano",
+        "Electric Grand Piano": "piano",
+        "Honky-tonk Piano": "piano",
+        "Electric Piano 1 (Rhodes Piano)": "piano",
+        "Electric Piano 2 (Chorused Piano)": "piano",
+        "Harpsichord": "piano",
+        "Clavinet": "piano",
+        "Celesta": "",
+        "Glockenspiel": "",
+        "Music Box": "",
+        "Vibraphone": "",
+        "Marimba": "",
+        "Xylophone": "",
+        "Tubular Bells": "",
+        "Dulcimer (Santur)": "",
+        "Drawbar Organ (Hammond)": "",
+        "Percussive Organ": "piano",
+        "Rock Organ": "piano",
+        "Church Organ": "piano",
+        "Reed Organ": "piano",
+        "Accordion (French)": "piano",
+        "Harmonica": "piano",
+        "Tango Accordion (Band neon)": "piano",
+        "Acoustic Guitar (nylon)": "",
+        "Acoustic Guitar (steel)": "",
+        "Electric Guitar (jazz)": "",
+        "Electric Guitar (clean)": "",
+        "Electric Guitar (muted)": "",
+        "Overdriven Guitar": "",
+        "Distortion Guitar": "",
+        "Guitar harmonics": "",
+        "Acoustic Bass": "",
+        "Electric Bass (fingered)": "",
+        "Electric Bass (picked)": "",
+        "Fretless Bass": "",
+        "Slap Bass 1": "",
+        "Slap Bass 2": "",
+        "Synth Bass 1": "",
+        "Synth Bass 2": "",
+        "Violin": "",
+        "Viola": "",
+        "Cello": "",
+        "Contrabass": "",
+        "Tremolo Strings": "",
+        "Pizzicato Strings": "",
+        "Orchestral Harp": "",
+        "Timpani": "",
+        "String Ensemble 1 (strings)": "",
+        "String Ensemble 2 (slow strings)": "",
+        "SynthStrings 1": "",
+        "SynthStrings 2": "",
+        "Choir Aahs": "",
+        "Voice Oohs": "",
+        "Synth Voice": "",
+        "Orchestra Hit": "",
+        "Trumpet": "",
+        "Trombone": "",
+        "Tuba": "",
+        "Muted Trumpet": "",
+        "French Horn": "",
+        "Brass Section": "",
+        "SynthBrass 1": "",
+        "SynthBrass 2": "",
+        "Soprano Sax": "",
+        "Alto Sax": "",
+        "Tenor Sax": "",
+        "Baritone Sax": "",
+        "Oboe": "",
+        "English Horn": "",
+        "Bassoon": "",
+        "Clarinet": "",
+        "Piccolo": "",
+        "Flute": "",
+        "Recorder": "",
+        "Pan Flute": "",
+        "Blown Bottle": "",
+        "Shakuhachi": "",
+        "Whistle": "",
+        "Ocarina": "",
+        "Lead 1 (square wave)": "",
+        "Lead 2 (sawtooth wave)": "",
+        "Lead 3 (calliope)": "",
+        "Lead 4 (chiffer)": "",
+        "Lead 5 (charang)": "",
+        "Lead 6 (voice solo)": "",
+        "Lead 7 (fifths)": "",
+        "Lead 8 (bass + lead)": "",
+        "Pad 1 (new age Fantasia)": "",
+        "Pad 2 (warm)": "",
+        "Pad 3 (polysynth)": "",
+        "Pad 4 (choir space voice)": "",
+        "Pad 5 (bowed glass)": "",
+        "Pad 6 (metallic pro)": "",
+        "Pad 7 (halo)": "",
+        "Pad 8 (sweep)": "",
+        "FX 1 (rain)": "",
+        "FX 2 (soundtrack)": "",
+        "FX 3 (crystal)": "",
+        "FX 4 (atmosphere)": "",
+        "FX 5 (brightness)": "",
+        "FX 6 (goblins)": "",
+        "FX 7 (echoes, drops)": "",
+        "FX 8 (sci-fi, star theme)": "",
+        "Sitar": "",
+        "Banjo": "",
+        "Shamisen": "",
+        "Koto": "",
+        "Kalimba": "",
+        "Bag pipe": "",
+        "Fiddle": "",
+        "Shanai": "",
+        "Tinkle Bell": "",
+        "Agogo": "",
+        "Steel Drums": "",
+        "Woodblock": "",
+        "Taiko Drum": "",
+        "Melodic Tom": "",
+        "Synth Drum": "",
+        "Reverse Cymbal": "",
+        "Guitar Fret Noise": "",
+        "Breath Noise": "",
+        "Seashore": "",
+        "Bird Tweet": "",
+        "Telephone Ring": "",
+        "Helicopter": "",
+        "Applause": "",
+        "Gunshot": ""
+    },
+    "bin_name_to_program_name": {
+        "piano": "Acoustic Grand Piano"
+    },
+    "instrument_names": {
+        "0": "Acoustic Grand Piano",
+        "1": "Bright Acoustic Piano",
+        "2": "Electric Grand Piano",
+        "3": "Honky-tonk Piano",
+        "4": "Electric Piano 1 (Rhodes Piano)",
+        "5": "Electric Piano 2 (Chorused Piano)",
+        "6": "Harpsichord",
+        "7": "Clavinet",
+        "8": "Celesta",
+        "9": "Glockenspiel",
+        "10": "Music Box",
+        "11": "Vibraphone",
+        "12": "Marimba",
+        "13": "Xylophone",
+        "14": "Tubular Bells",
+        "15": "Dulcimer (Santur)",
+        "16": "Drawbar Organ (Hammond)",
+        "17": "Percussive Organ",
+        "18": "Rock Organ",
+        "19": "Church Organ",
+        "20": "Reed Organ",
+        "21": "Accordion (French)",
+        "22": "Harmonica",
+        "23": "Tango Accordion (Band neon)",
+        "24": "Acoustic Guitar (nylon)",
+        "25": "Acoustic Guitar (steel)",
+        "26": "Electric Guitar (jazz)",
+        "27": "Electric Guitar (clean)",
+        "28": "Electric Guitar (muted)",
+        "29": "Overdriven Guitar",
+        "30": "Distortion Guitar",
+        "31": "Guitar harmonics",
+        "32": "Acoustic Bass",
+        "33": "Electric Bass (fingered)",
+        "34": "Electric Bass (picked)",
+        "35": "Fretless Bass",
+        "36": "Slap Bass 1",
+        "37": "Slap Bass 2",
+        "38": "Synth Bass 1",
+        "39": "Synth Bass 2",
+        "40": "Violin",
+        "41": "Viola",
+        "42": "Cello",
+        "43": "Contrabass",
+        "44": "Tremolo Strings",
+        "45": "Pizzicato Strings",
+        "46": "Orchestral Harp",
+        "47": "Timpani",
+        "48": "String Ensemble 1 (strings)",
+        "49": "String Ensemble 2 (slow strings)",
+        "50": "SynthStrings 1",
+        "51": "SynthStrings 2",
+        "52": "Choir Aahs",
+        "53": "Voice Oohs",
+        "54": "Synth Voice",
+        "55": "Orchestra Hit",
+        "56": "Trumpet",
+        "57": "Trombone",
+        "58": "Tuba",
+        "59": "Muted Trumpet",
+        "60": "French Horn",
+        "61": "Brass Section",
+        "62": "SynthBrass 1",
+        "63": "SynthBrass 2",
+        "64": "Soprano Sax",
+        "65": "Alto Sax",
+        "66": "Tenor Sax",
+        "67": "Baritone Sax",
+        "68": "Oboe",
+        "69": "English Horn",
+        "70": "Bassoon",
+        "71": "Clarinet",
+        "72": "Piccolo",
+        "73": "Flute",
+        "74": "Recorder",
+        "75": "Pan Flute",
+        "76": "Blown Bottle",
+        "77": "Shakuhachi",
+        "78": "Whistle",
+        "79": "Ocarina",
+        "80": "Lead 1 (square wave)",
+        "81": "Lead 2 (sawtooth wave)",
+        "82": "Lead 3 (calliope)",
+        "83": "Lead 4 (chiffer)",
+        "84": "Lead 5 (charang)",
+        "85": "Lead 6 (voice solo)",
+        "86": "Lead 7 (fifths)",
+        "87": "Lead 8 (bass + lead)",
+        "88": "Pad 1 (new age Fantasia)",
+        "89": "Pad 2 (warm)",
+        "90": "Pad 3 (polysynth)",
+        "91": "Pad 4 (choir space voice)",
+        "92": "Pad 5 (bowed glass)",
+        "93": "Pad 6 (metallic pro)",
+        "94": "Pad 7 (halo)",
+        "95": "Pad 8 (sweep)",
+        "96": "FX 1 (rain)",
+        "97": "FX 2 (soundtrack)",
+        "98": "FX 3 (crystal)",
+        "99": "FX 4 (atmosphere)",
+        "100": "FX 5 (brightness)",
+        "101": "FX 6 (goblins)",
+        "102": "FX 7 (echoes, drops)",
+        "103": "FX 8 (sci-fi, star theme)",
+        "104": "Sitar",
+        "105": "Banjo",
+        "106": "Shamisen",
+        "107": "Koto",
+        "108": "Kalimba",
+        "109": "Bag pipe",
+        "110": "Fiddle",
+        "111": "Shanai",
+        "112": "Tinkle Bell",
+        "113": "Agogo",
+        "114": "Steel Drums",
+        "115": "Woodblock",
+        "116": "Taiko Drum",
+        "117": "Melodic Tom",
+        "118": "Synth Drum",
+        "119": "Reverse Cymbal",
+        "120": "Guitar Fret Noise",
+        "121": "Breath Noise",
+        "122": "Seashore",
+        "123": "Bird Tweet",
+        "124": "Telephone Ring",
+        "125": "Helicopter",
+        "126": "Applause",
+        "127": "Gunshot"
+    }
+}

docker-compose.yml

@@ -0,0 +1,18 @@
+services:
+  rmkv_runner:
+    image: rwkv-runner:latest
+    build: .
+    # Append "--rwkv.cpp" parameter to use rwkv.cpp
+    # command: python3.10 ./backend-python/main.py  --port 27777 --host 0.0.0.0 --webui --rwkv.cpp
+    volumes:
+      - /mnt:/mnt
+    ports:
+      - "27777:27777"
+    # Comment the following lines if use rwkv.cpp
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

finetune/get_layer_and_embd.py

@@ -52,9 +52,13 @@ for x in keys:
     if "time_maa" in x:
         version = max(6, version)
 
+params = f"--vocab_size {vocab_size} --n_layer {n_layer} --n_embd {n_embd}"
+
 if version <= expected_max_version:
+    if version == 6:
+        params += ' --my_testing "x060"'
     print(
-        f"v{int(version)}/train.py --vocab_size {vocab_size} --n_layer {n_layer} --n_embd {n_embd}",
+        f"v{int(version)}/train.py {params}",
         end="",
     )
 else:

finetune/install-wsl-dep-and-train.sh

@@ -1,7 +1,7 @@
 echo $@
 
 if [[ ${cnMirror} == 1 ]]; then
-  export PIP_INDEX_URL="https://pypi.tuna.tsinghua.edu.cn/simple"
+  export PIP_INDEX_URL="https://mirrors.aliyun.com/pypi/simple"
   if grep -q "mirrors.aliyun.com" /etc/apt/sources.list; then
     echo "apt cnMirror already set"
   else
@@ -22,6 +22,12 @@ else
   sudo apt -y install python3-pip
 fi
 
+if dpkg -s "python3-dev" >/dev/null 2>&1; then
+  echo "python3-dev installed"
+else
+  sudo apt -y install python3-dev
+fi
+
 if dpkg -s "ninja-build" >/dev/null 2>&1; then
   echo "ninja installed"
 else
@@ -47,11 +53,12 @@ else
 fi
 
 echo "loading $loadModel"
-modelInfo=$(python3 ./finetune/get_layer_and_embd.py $loadModel 5.2)
+modelInfo=$(python3 ./finetune/get_layer_and_embd.py $loadModel 6.0)
 echo $modelInfo
 if [[ $modelInfo =~ "--n_layer" ]]; then
+  sudo rm -rf /root/.cache/torch_extensions
   python3 ./finetune/lora/$modelInfo $@ --proj_dir lora-models --data_type binidx --lora \
-    --lora_parts=att,ffn,time,ln --strategy deepspeed_stage_2 --accelerator gpu
+    --lora_parts=att,ffn,time,ln --strategy deepspeed_stage_2 --accelerator gpu --ds_bucket_mb 2
 else
   echo "modelInfo is invalid"
   exit 1

finetune/lora/v6/cuda/wkv5_cuda.cu

@@ -0,0 +1,202 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _w += h*_N_;
+    _u += h*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_] = {0};
+
+    __syncthreads();
+    w[i] = _w[i];
+    u[i] = float(_u[i]);
+    __syncthreads();
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const float *__restrict__ __w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gw, F *__restrict__ const _gu)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _w += h*_N_;
+    _u += h*_N_;
+    __w += h*_N_;
+
+    __shared__ float w_[_N_], u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], gy[_N_];
+    __syncthreads();
+    w_[i] = _w[i];
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float w = w_[i];
+    const float ww = __w[i];
+    const float u = u_[i];
+
+    float state[_N_] = {0}, saaaa[_N_] = {0}, sbbbb[_N_] = {0}, scccc[_N_] = {0}, sdddd[_N_] = {0};
+
+    float gw = 0, gu = 0;
+    const int t000 = b*T*C + h*_N_ + i;
+    const int t111 = (b+1)*T*C + h*_N_ + i;
+    const int t222 = t111 - 2*C;
+
+    for (int t = t000; t < t111; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+    
+    for (int t = t000; t < t222; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t + 2*C]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        float gw_ = 0;
+        
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            float& s2 = sbbbb[j];
+            float x = k * v[j];
+            
+            float tmp = w * (x + s);
+            s = tmp;
+            s2 = tmp + w * s2;
+            gw_ += s2 * gy[j];
+        }
+        gw += float(_r[t + 2*C]) * gw_;
+    }    
+    _gw[b*C + h*_N_ + i] = F(ww * gw);
+
+    for (int t = t111 - C; t >= t000; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t111 - C; t >= t000; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, float *ww, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, ww, u, gy, gr, gk, gv, gw, gu);
+}

finetune/lora/v6/cuda/wkv5_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, float *ww, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &ww, torch::Tensor &u, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), ww.data_ptr<float>(), u.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv5 forward");
+    m.def("backward", &backward, "wkv5 backward");
+}
+
+TORCH_LIBRARY(wkv5, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/cuda/wkv6_cuda.cu

@@ -0,0 +1,242 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_] = {0};
+
+    __syncthreads();
+    u[i] = float(_u[i]);
+    __syncthreads();
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        w[i] = exp(_w[t]);
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward_111(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gu)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+
+    __shared__ float u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], w_[_N_], gy[_N_];
+    __syncthreads();
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float u = u_[i];
+
+    float state[_N_] = {0}, scccc[_N_] = {0}, sdddd[_N_] = {0};
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_T_1 = t_0 + (T-1)*C;
+    const int t_T = t_0 + T*C;
+
+    float gu = 0;
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        const float w = exp(_w[t]);
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        const float w = exp(_w[t]);
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        w_[i] = exp(_w[t]);
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+}
+
+template <typename F>
+__global__ void kernel_backward_222(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ const _gy,
+    F *__restrict__ const _gw)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+
+    __shared__ float v[_N_], gy[_N_];
+    float saaaa[_N_] = {0}, sbbbb[_T_-2] = {0}, scccc[_N_] = {0};
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_1 = t_0 + C;
+    const int t_2 = t_0 + 2*C;
+    const int t_T_1 = t_0 + (T-1)*C;
+
+    for (int t = t_T_1; t > t_1; t -= C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float r = float(_r[t]);
+        const float w = exp(_w[t-C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            float x = r * gy[j];
+            s = (s + x) * w;
+            sum += s * v[j];
+        }
+        sbbbb[(t-t_2)/C] = sum * float(_k[t-2*C]);
+    }
+
+    float sss = sbbbb[0];
+    _gw[t_0] = 0;
+    _gw[t_1] = F(sss * _w[t_1]);
+
+    for (int t = t_2; t < t_T_1; t += C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float w = exp(_w[t-C]);
+        const float k = float(_k[t-2*C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = k * v[j];
+            s = (s + x) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[(t-t_1)/C] - (sum * float(_r[t]));
+        _gw[t] = F(sss * _w[t]);
+    }
+    _gw[t_T_1] = 0;
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward_111<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, gy, gr, gk, gv, gu);
+    kernel_backward_222<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, gy, gw);
+}

finetune/lora/v6/cuda/wkv6_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<float>(), u.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv6 forward");
+    m.def("backward", &backward, "wkv6 backward");
+}
+
+TORCH_LIBRARY(wkv6, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/cuda/wkv6infctx_cuda.cu

@@ -0,0 +1,311 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u, F *__restrict__ _s,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+    _s += h*_N_*_N_ + i*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_];
+
+    __syncthreads();
+    u[i] = float(_u[i]);
+    __syncthreads();
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j]);
+    }
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        w[i] = __expf(-__expf(float(_w[t])));
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        _s[j] = F(state[j]);
+}
+
+template <typename F>
+__global__ void kernel_backward_111(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ _s, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gu, F *__restrict__ const _gs)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+    _s += h*_N_*_N_ + i;
+
+    __shared__ float u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], w_[_N_], gy[_N_];
+    __syncthreads();
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float u = u_[i];
+
+    float state[_N_], scccc[_N_] = {0}, sdddd[_N_] = {0}, sssss[_N_] = {0}, swwww[_N_];
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j*_N_]);
+        swwww[j] = 1.0;
+    }
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_T_1 = t_0 + (T-1)*C;
+    const int t_T = t_0 + T*C;
+
+    float gu = 0;
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        const float w = __expf(-__expf(float(_w[t])));
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        const float w = __expf(-__expf(float(_w[t])));
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        w_[i] = __expf(-__expf(float(_w[t])));
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        w_[i] = __expf(-__expf(float(_w[t])));
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& w = swwww[j];
+            sssss[j] += gyy * w * r[j];
+            w *= w_[j];
+        }
+    }
+    for (int j = 0; j < _N_; j++)
+        _gs[b*H*_N_*_N_ + h*_N_*_N_ + i*_N_ + j] = F(sssss[j]);
+}
+
+template <typename F>
+__global__ void kernel_backward_222(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ _s, const F *__restrict__ const _gy,
+    F *__restrict__ const _gw)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _s += h*_N_*_N_ + i;
+
+    __shared__ float v[_N_], gy[_N_];
+    float state[_N_], saaaa[_N_] = {0}, sbbbb[_T_-1] = {0}, scccc[_N_] = {0};
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j*_N_]);
+    }
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_1 = t_0 + C;
+    const int t_2 = t_0 + 2*C;
+    const int t_T_1 = t_0 + (T-1)*C;
+
+    for (int t = t_T_1; t > t_1; t -= C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float r = float(_r[t]);
+        const float w = __expf(-__expf(float(_w[t-C])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            s = (s + r * gy[j]) * w;
+            sum += s * v[j];
+        }
+        sbbbb[(t-t_1)/C] = sum * float(_k[t-2*C]);
+    }
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t_1]);
+        __syncthreads();
+
+        const float r = float(_r[t_1]);
+        const float w = __expf(-__expf(float(_w[t_0])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            s = (s + r * gy[j]) * w;
+            sum += s * state[j];
+        }
+        sbbbb[0] = sum;
+    }
+
+    float sss = sbbbb[0];
+    _gw[t_0] = F(sss * -__expf(float(_w[t_0])));
+
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t_1]);
+        __syncthreads();
+
+        const float w = __expf(-__expf(float(_w[t_0])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            s = (s + state[j]) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[1] - (sum * float(_r[t_1]));
+        _gw[t_1] = F(sss * -__expf(float(_w[t_1])));
+    }
+    for (int t = t_2; t < t_T_1; t += C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float w = __expf(-__expf(float(_w[t-C])));
+        const float k = float(_k[t-2*C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            s = (s + k * v[j]) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[(t-t_0)/C] - (sum * float(_r[t]));
+        _gw[t] = F(sss * -__expf(float(_w[t])));
+    }
+    _gw[t_T_1] = 0;
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *z, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *z, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu, bf16 *gs)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward_111<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, gy, gr, gk, gv, gu, gs);
+    kernel_backward_222<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, gy, gw);
+}

finetune/lora/v6/cuda/wkv6infctx_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *s, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *s, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu, bf16 *gs);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &s, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<bf16>(), u.data_ptr<bf16>(), s.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &s, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gs) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<bf16>(), u.data_ptr<bf16>(), s.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>(), gs.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv6state forward");
+    m.def("backward", &backward, "wkv6state backward");
+}
+
+TORCH_LIBRARY(wkv6state, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/cuda/wkv6state_cuda.cu

@@ -0,0 +1,311 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u,const F *__restrict__ _s,
+                               F *__restrict__ const _y)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+    _s += h*_N_*_N_ + i*_N_;
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    float state[_N_];
+
+    __syncthreads();
+    u[i] = float(_u[i]);
+    __syncthreads();
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j]);
+    }
+
+    for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C)
+    {
+        __syncthreads();
+        w[i] = __expf(-__expf(float(_w[t])));
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        __syncthreads();
+
+        const float v = float(_v[t]);
+        float y = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j+=4)
+        {
+            const float4& r_ = (float4&)(r[j]);
+            const float4& k_ = (float4&)(k[j]);
+            const float4& w_ = (float4&)(w[j]);
+            const float4& u_ = (float4&)(u[j]);
+            float4& s = (float4&)(state[j]);
+            float4 x;
+
+            x.x = k_.x * v;
+            x.y = k_.y * v;
+            x.z = k_.z * v;
+            x.w = k_.w * v;
+
+            y += r_.x * (u_.x * x.x + s.x);
+            y += r_.y * (u_.y * x.y + s.y);
+            y += r_.z * (u_.z * x.z + s.z);
+            y += r_.w * (u_.w * x.w + s.w);
+
+            s.x = s.x * w_.x + x.x;
+            s.y = s.y * w_.y + x.y;
+            s.z = s.z * w_.z + x.z;
+            s.w = s.w * w_.w + x.w;
+        }
+        _y[t] = F(y);
+    }
+    // #pragma unroll
+    // for (int j = 0; j < _N_; j++)
+    //     _s[j] = F(state[j]);
+}
+
+template <typename F>
+__global__ void kernel_backward_111(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ _s, const F *__restrict__ const _gy,
+    F *__restrict__ const _gr, F *__restrict__ const _gk, F *__restrict__ const _gv, F *__restrict__ const _gu, F *__restrict__ const _gs)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _u += h*_N_;
+    _s += h*_N_*_N_ + i;
+
+    __shared__ float u_[_N_];
+    __shared__ float r[_N_], k[_N_], v[_N_], w_[_N_], gy[_N_];
+    __syncthreads();
+    u_[i] = float(_u[i]);
+    __syncthreads();
+
+    const float u = u_[i];
+
+    float state[_N_], scccc[_N_] = {0}, sdddd[_N_] = {0}, sssss[_N_] = {0}, swwww[_N_];
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j*_N_]);
+        swwww[j] = 1.0;
+    }
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_T_1 = t_0 + (T-1)*C;
+    const int t_T = t_0 + T*C;
+
+    float gu = 0;
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float k = float(_k[t]);
+        const float w = __expf(-__expf(float(_w[t])));
+        float gr = 0, gu_ = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            float x = k * v[j];
+
+            gr += (u * x + s) * gy[j];
+            gu_ += x * gy[j];
+            s = s * w + x;
+        }
+        _gr[t] = F(gr);
+        gu += float(_r[t]) * gu_;
+    }
+    _gu[b*C + h*_N_ + i] = F(gu);
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        v[i] = float(_v[t]);
+        gy[i] = float(_gy[t]);
+        __syncthreads();
+
+        const float rr = float(_r[t]);
+        const float w = __expf(-__expf(float(_w[t])));
+        float gk = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            float x = rr * gy[j];
+            
+            gk += (u * x + s) * v[j];
+            s = x + s * w;
+        }
+        _gk[t] = F(gk);
+    }
+
+    for (int t = t_T_1; t >= t_0; t -= C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        k[i] = float(_k[t]);
+        w_[i] = __expf(-__expf(float(_w[t])));
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+        float gv = 0;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = sdddd[j];
+            float x = gyy * r[j];
+            
+            gv += (u_[j] * x + s) * k[j];
+            s = x + s * w_[j];
+        }
+        _gv[t] = F(gv);
+    }
+
+    for (int t = t_0; t < t_T; t += C)
+    {
+        __syncthreads();
+        r[i] = float(_r[t]);
+        w_[i] = __expf(-__expf(float(_w[t])));
+        __syncthreads();
+
+        const float gyy = float(_gy[t]);
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& w = swwww[j];
+            sssss[j] += gyy * w * r[j];
+            w *= w_[j];
+        }
+    }
+    for (int j = 0; j < _N_; j++)
+        _gs[b*H*_N_*_N_ + h*_N_*_N_ + i*_N_ + j] = F(sssss[j]);
+}
+
+template <typename F>
+__global__ void kernel_backward_222(const int B, const int T, const int C, const int H,
+    const F *__restrict__ const _r, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ _w, const F *__restrict__ _u, const F *__restrict__ _s, const F *__restrict__ const _gy,
+    F *__restrict__ const _gw)
+{
+    const int b = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _s += h*_N_*_N_ + i;
+
+    __shared__ float v[_N_], gy[_N_];
+    float state[_N_], saaaa[_N_] = {0}, sbbbb[_T_-1] = {0}, scccc[_N_] = {0};
+    for (int j = 0; j < _N_; j++) {
+        state[j] = float(_s[j*_N_]);
+    }
+
+    const int t_0 = b*T*C + h*_N_ + i;
+    const int t_1 = t_0 + C;
+    const int t_2 = t_0 + 2*C;
+    const int t_T_1 = t_0 + (T-1)*C;
+
+    for (int t = t_T_1; t > t_1; t -= C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float r = float(_r[t]);
+        const float w = __expf(-__expf(float(_w[t-C])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            s = (s + r * gy[j]) * w;
+            sum += s * v[j];
+        }
+        sbbbb[(t-t_1)/C] = sum * float(_k[t-2*C]);
+    }
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t_1]);
+        __syncthreads();
+
+        const float r = float(_r[t_1]);
+        const float w = __expf(-__expf(float(_w[t_0])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = saaaa[j];
+            s = (s + r * gy[j]) * w;
+            sum += s * state[j];
+        }
+        sbbbb[0] = sum;
+    }
+
+    float sss = sbbbb[0];
+    _gw[t_0] = F(sss * -__expf(float(_w[t_0])));
+
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t_1]);
+        __syncthreads();
+
+        const float w = __expf(-__expf(float(_w[t_0])));
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            s = (s + state[j]) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[1] - (sum * float(_r[t_1]));
+        _gw[t_1] = F(sss * -__expf(float(_w[t_1])));
+    }
+    for (int t = t_2; t < t_T_1; t += C)
+    {
+        __syncthreads();
+        gy[i] = float(_gy[t]);
+        v[i] = float(_v[t-2*C]);
+        __syncthreads();
+
+        const float w = __expf(-__expf(float(_w[t-C])));
+        const float k = float(_k[t-2*C]);
+        float sum = 0.0f;
+
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = scccc[j];
+            s = (s + k * v[j]) * w;
+            sum += s * gy[j];
+        }
+        sss += sbbbb[(t-t_0)/C] - (sum * float(_r[t]));
+        _gw[t] = F(sss * -__expf(float(_w[t])));
+    }
+    _gw[t_T_1] = 0;
+}
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *z, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, y);
+}
+
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *z, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu, bf16 *gs)
+{
+    assert(H*_N_ == C);
+    assert(_N_%4 == 0);
+    kernel_backward_111<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, gy, gr, gk, gv, gu, gs);
+    kernel_backward_222<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, r, k, v, w, u, z, gy, gw);
+}

finetune/lora/v6/cuda/wkv6state_op.cpp

@@ -0,0 +1,22 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+
+void cuda_forward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *s, bf16 *y);
+void cuda_backward(int B, int T, int C, int H, bf16 *r, bf16 *k, bf16 *v, bf16 *w, bf16 *u, bf16 *s, bf16 *gy, bf16 *gr, bf16 *gk, bf16 *gv, bf16 *gw, bf16 *gu, bf16 *gs);
+
+void forward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &s, torch::Tensor &y) {
+    cuda_forward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<bf16>(), u.data_ptr<bf16>(), s.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void backward(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &s, torch::Tensor &gy, torch::Tensor &gr, torch::Tensor &gk, torch::Tensor &gv, torch::Tensor &gw, torch::Tensor &gu, torch::Tensor &gs) {
+    cuda_backward(B, T, C, H, r.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), w.data_ptr<bf16>(), u.data_ptr<bf16>(), s.data_ptr<bf16>(), gy.data_ptr<bf16>(), gr.data_ptr<bf16>(), gk.data_ptr<bf16>(), gv.data_ptr<bf16>(), gw.data_ptr<bf16>(), gu.data_ptr<bf16>(), gs.data_ptr<bf16>());
+}
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward, "wkv6state forward");
+    m.def("backward", &backward, "wkv6state backward");
+}
+
+TORCH_LIBRARY(wkv6state, m) {
+    m.def("forward", forward);
+    m.def("backward", backward);
+}

finetune/lora/v6/demo/demo-lora-merge.sh

@@ -0,0 +1,16 @@
+
+base_model='/home/rwkv/JL/model/rwkv-x060-7b-world-v2.1-36%trained-20240413-ctx4k.pth'
+lora_init='/home/rwkv/JL/out_model/nf4/init_lora.pth'
+lora_checkpoint='/home/rwkv/JL/out_model/nf4/rwkv-0.pth'
+output='/home/rwkv/JL/model/nf4-world.pth'
+QUANT='nf4' #follow train
+TYPE='lora'
+Lora_alpha=128
+
+python merge/merge.py --base_model $base_model \
+--lora_init $lora_init \
+--lora_checkpoint $lora_checkpoint \
+--output $output \
+--quant $QUANT \
+--type $TYPE \
+--lora_alpha $Lora_alpha

finetune/lora/v6/demo/demo-lora.sh

@@ -0,0 +1,27 @@
+load_model='/home/rwkv/JL/model/rwkv-x060-7b-world-v2.1-36%trained-20240413-ctx4k.pth'
+proj_dir='/home/rwkv/JL/out_model/nf4'
+data_file='/home/rwkv/JL/data/roleplay'
+
+QUANT='nf4'  #4bit nf4 fp4 none
+
+lora_r=64
+lora_alpha=128
+
+n_layer=32
+n_embd=4096
+
+micro_bsz=8
+epoch_save=1
+epoch_steps=1000
+ctx_len=1024
+
+python train.py --load_model $load_model \
+--proj_dir $proj_dir --data_file $data_file \
+--data_type binidx --vocab_size 65536 \
+--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 20 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+--n_layer $n_layer --n_embd $n_embd \
+--pre_ffn 0 --head_qk 0 --lr_init 5e-5 --lr_final 5e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
+--my_testing "x060" \
+--lora_load rwkv-0 --lora --lora_r $lora_r --lora_alpha $lora_alpha --lora_dropout 0.01 --lora_parts=att,ffn,time,ln \
+--quant $QUANT

finetune/lora/v6/demo/demo-pissa-merge.sh

@@ -0,0 +1,15 @@
+
+
+base_model='/home/rwkv/JL/model/RWKV-x060-World-1B6-v2-20240208-ctx4096.pth'
+lora_init='/home/rwkv/JL/out_model/nf4/init_lora.pth'
+lora_checkpoint='/home/rwkv/JL/out_model/nf4/rwkv-0.pth'
+output='/home/rwkv/JL/model/end-world.pth'
+QUANT='nf4' #follow train
+TYPE='pissa'
+
+python merge/merge.py --base_model $base_model \
+--lora_init $lora_init \
+--lora_checkpoint $lora_checkpoint \
+--output $output \
+--quant $QUANT \
+--type $TYPE 

finetune/lora/v6/demo/demo-pissa.sh

@@ -0,0 +1,40 @@
+
+load_model='/home/rwkv/JL/model/RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth'
+proj_dir='/home/rwkv/JL/out_model/nf4'
+data_file='/home/rwkv/JL/data/end_text_document'
+
+QUANT='nf4'  #4bit nf4 fp4 none
+svd_niter=4  
+lora_r=64
+
+n_layer=24
+n_embd=2048
+
+micro_bsz=8
+epoch_save=1
+epoch_steps=1000
+ctx_len=1024
+
+python train.py --load_model $load_model \
+--proj_dir $proj_dir --data_file $data_file \
+--data_type binidx --vocab_size 65536 \
+--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 1 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+--n_layer $n_layer --n_embd $n_embd \
+--pre_ffn 0 --head_qk 0 --lr_init 5e-5 --lr_final 5e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
+--my_testing "x060" \
+--lora_load rwkv-0 --lora --lora_r $lora_r --lora_alpha 128 --lora_dropout 0.01 --lora_parts=att,ffn,time,ln \
+--PISSA --svd_niter $svd_niter \
+--dataload pad
+
+###remove load_model
+# python train.py --proj_dir $proj_dir --data_file $data_file \
+# --data_type binidx --vocab_size 65536 \
+# --ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 20 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+# --n_layer $n_layer --n_embd $n_embd \
+# --pre_ffn 0 --head_qk 0 --lr_init 5e-5 --lr_final 5e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+# --accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
+# --my_testing "x060" \
+# --lora_load rwkv-0 --lora --lora_r $lora_r --lora_alpha 128 --lora_dropout 0.01 --lora_parts=att,ffn,time,ln \
+# --PISSA --svd_niter $svd_niter \
+# --quant $QUANT

finetune/lora/v6/demo/demo-qpissa-pt.sh

@@ -0,0 +1,27 @@
+load_model='/home/rwkv/JL/model/rwkv-x060-7b-world-v2.1-36%trained-20240413-ctx4k.pth'
+proj_dir='/home/rwkv/JL/out_model/nf4'
+data_file='/home/rwkv/JL/data/roleplay'
+
+QUANT='nf4'  #4bit nf4 fp4 none
+svd_niter=4  
+lora_r=64
+
+n_layer=32
+n_embd=4096
+
+micro_bsz=4
+epoch_save=1
+epoch_steps=1000
+ctx_len=1024
+
+
+python train.py --proj_dir $proj_dir --data_file $data_file \
+--data_type binidx --vocab_size 65536 \
+--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 20 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+--n_layer $n_layer --n_embd $n_embd \
+--pre_ffn 0 --head_qk 0 --lr_init 5e-5 --lr_final 5e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
+--my_testing "x060" \
+--lora_load rwkv-0 --lora --lora_r $lora_r --lora_alpha 128 --lora_dropout 0.01 --lora_parts=att,ffn,time,ln \
+--PISSA --svd_niter $svd_niter \
+--quant $QUANT

finetune/lora/v6/demo/demo-state-merge.sh

@@ -0,0 +1,8 @@
+base_model='/home/rwkv/JL/model/RWKV-x060-World-3B-v2.1-20240417-ctx4096.pth'
+state_checkpoint='/home/rwkv/JL/out_model/state/rwkv-9.pth'
+output='/home/rwkv/JL/model/state-0.pth'
+
+
+python merge/merge_state.py --base_model $base_model \
+--state_checkpoint $state_checkpoint \
+--output $output

finetune/lora/v6/demo/demo-state-tuning.sh

@@ -0,0 +1,22 @@
+load_model='/home/rwkv/JL/model/RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth'
+proj_dir='/home/rwkv/JL/out_model/state'
+data_file='/home/rwkv/JL/data/end_text_document'
+
+
+n_layer=24
+n_embd=2048
+
+micro_bsz=1
+epoch_save=1
+epoch_steps=1000
+ctx_len=1024
+
+python train.py --load_model $load_model \
+--proj_dir $proj_dir --data_file $data_file \
+--data_type binidx --vocab_size 65536 \
+--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 1 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+--n_layer $n_layer --n_embd $n_embd \
+--pre_ffn 0 --head_qk 0 --lr_init 1 --lr_final 1e-1 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 0 \
+--my_testing "x060" \
+--train_type "state"  --dataload pad --wandb fla --fla

finetune/lora/v6/demo/demo-training-prepare.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+# Create data directory
+
+mkdir -p data
+
+# Download minipile (1498226207 tokens, around 3GB)
+
+wget --continue -O data/minipile.idx https://huggingface.co/datasets/BlinkDL/minipile-tokenized/resolve/main/rwkv_vocab_v20230424/minipile.idx
+wget --continue -O data/minipile.bin https://huggingface.co/datasets/BlinkDL/minipile-tokenized/resolve/main/rwkv_vocab_v20230424/minipile.bin
+
+# Generate initial model (L12-D768 = 169M)
+
+BASE_NAME="model/0.1-1"
+N_LAYER="12"
+N_EMBD="768"
+
+# magic_prime = the largest 3n+2 prime smaller than datalen/ctxlen-1 (= 1498226207/512-1 = 2926222.06 in this case)
+# use https://www.dcode.fr/prime-numbers-search
+
+python train.py --wandb "" --proj_dir $BASE_NAME \
+ --data_file "data/minipile" --data_type "binidx" --vocab_size 65536 \
+ --ctx_len 512 --my_pile_stage 1 --epoch_count 1 --epoch_begin 0 \
+ --epoch_save 1 --weight_decay 0 --head_size_a 64 \
+ --num_nodes 1 --micro_bsz 1 --n_layer $N_LAYER --n_embd $N_EMBD --pre_ffn 0 --head_qk 0 --my_exit_tokens 1498226207 --magic_prime 2926181 \
+ --lr_init 1e-5 --lr_final 1e-5 --warmup_steps 10 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 --my_pile_edecay 0 \
+ --accelerator cpu --devices 1 --precision bf16 --strategy deepspeed_stage_2 --grad_cp 0 --enable_progress_bar False --ds_bucket_mb 200

finetune/lora/v6/demo/demo-training-run.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+BASE_NAME="model/0.1-1"
+N_LAYER="12"
+N_EMBD="768"
+M_BSZ="16" # takes 16G VRAM (reduce this to save VRAM)
+LR_INIT="6e-4"
+LR_FINAL="6e-5"
+GRAD_CP=0 # set to 1 to save VRAM (will be slower)
+EPOCH_SAVE=10
+
+# magic_prime = the largest 3n+2 prime smaller than datalen/ctxlen-1 (= 1498226207/512-1 = 2926222.06 in this case)
+# use https://www.dcode.fr/prime-numbers-search
+
+python train.py --load_model "0" --wandb "RWKV-5-Test" --proj_dir $BASE_NAME \
+ --ctx_len 512 --my_pile_stage 3 --epoch_count 999999 --epoch_begin 0 \
+ --data_file "data/minipile" --my_exit_tokens 1498226207 --magic_prime 2926181 \
+ --num_nodes 1 --micro_bsz $M_BSZ --n_layer $N_LAYER --n_embd $N_EMBD --pre_ffn 0 --head_qk 0 \
+ --lr_init $LR_INIT --lr_final $LR_FINAL --warmup_steps 10 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 --my_pile_edecay 0 --data_type "binidx" --vocab_size 65536 \
+ --weight_decay 0.001 --epoch_save $EPOCH_SAVE --head_size_a 64 \
+ --accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_2 --grad_cp $GRAD_CP --enable_progress_bar True --ds_bucket_mb 200

finetune/lora/v6/demo/infctx.sh

@@ -0,0 +1,25 @@
+load_model='/home/rwkv/JL/model/RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth'
+proj_dir='/home/rwkv/JL/out_model/infctx'
+data_file='/home/rwkv/JL/data/roleplay'
+
+
+n_layer=24
+n_embd=2048
+
+micro_bsz=8
+epoch_save=5
+epoch_steps=1000
+ctx_len=16384
+chunk_ctx=2048
+
+
+python train.py --load_model $load_model \
+--proj_dir $proj_dir --data_file $data_file \
+--data_type binidx --vocab_size 65536 \
+--ctx_len $ctx_len --epoch_steps $epoch_steps --epoch_count 1 --epoch_begin 0 --epoch_save $epoch_save --micro_bsz $micro_bsz \
+--n_layer $n_layer --n_embd $n_embd \
+--pre_ffn 0 --head_qk 0 --lr_init 1e-4 --lr_final 1e-4 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
+--accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 1 \
+--lora_load rwkv-0 --lora --lora_r 64 --lora_alpha 128 --lora_dropout 0.01 --lora_parts=att,ffn,time,ln \
+--my_testing "x060"  --dataload pad \
+--train_type infctx --chunk_ctx $chunk_ctx --fla --wandb infctx

finetune/lora/v6/fla/__init__.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+from fla.layers import (ABCAttention, BasedLinearAttention, DeltaNet,
+                        GatedLinearAttention, HGRN2Attention, LinearAttention,
+                        MultiScaleRetention, ReBasedLinearAttention)
+from fla.models import (ABCForCausalLM, ABCModel, DeltaNetForCausalLM,
+                        DeltaNetModel, GLAForCausalLM, GLAModel,
+                        HGRN2ForCausalLM, HGRN2Model, HGRNForCausalLM,
+                        HGRNModel, LinearAttentionForCausalLM,
+                        LinearAttentionModel, RetNetForCausalLM, RetNetModel,
+                        RWKV6ForCausalLM, RWKV6Model, TransformerForCausalLM,
+                        TransformerModel)
+from fla.ops import (chunk_gla, chunk_retention, fused_chunk_based,
+                     fused_chunk_gla, fused_chunk_retention)
+
+__all__ = [
+    'ABCAttention',
+    'BasedLinearAttention',
+    'DeltaNet',
+    'HGRN2Attention',
+    'GatedLinearAttention',
+    'LinearAttention',
+    'MultiScaleRetention',
+    'ReBasedLinearAttention',
+    'ABCForCausalLM',
+    'ABCModel',
+    'DeltaNetForCausalLM',
+    'DeltaNetModel',
+    'HGRNForCausalLM',
+    'HGRNModel',
+    'HGRN2ForCausalLM',
+    'HGRN2Model',
+    'GLAForCausalLM',
+    'GLAModel',
+    'LinearAttentionForCausalLM',
+    'LinearAttentionModel',
+    'RetNetForCausalLM',
+    'RetNetModel',
+    'RWKV6ForCausalLM',
+    'RWKV6Model',
+    'TransformerForCausalLM',
+    'TransformerModel',
+    'chunk_gla',
+    'chunk_retention',
+    'fused_chunk_based',
+    'fused_chunk_gla',
+    'fused_chunk_retention'
+]
+
+__version__ = '0.1'

finetune/lora/v6/fla/layers/__init__.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+from .abc import ABCAttention
+from .based import BasedLinearAttention
+from .delta_net import DeltaNet
+from .gla import GatedLinearAttention
+from .hgrn import HGRNAttention
+from .hgrn2 import HGRN2Attention
+from .linear_attn import LinearAttention
+from .multiscale_retention import MultiScaleRetention
+from .rebased import ReBasedLinearAttention
+from .rwkv6 import RWKV6Attention
+
+__all__ = [
+    'ABCAttention',
+    'BasedLinearAttention',
+    'DeltaNet',
+    'GatedLinearAttention',
+    'HGRNAttention',
+    'HGRN2Attention',
+    'LinearAttention',
+    'MultiScaleRetention',
+    'ReBasedLinearAttention',
+    'RWKV6Attention'
+]

finetune/lora/v6/fla/layers/abc.py

@@ -0,0 +1,195 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.cache_utils import Cache
+
+from fla.modules import (FusedRMSNormSwishGate, RMSNorm, RotaryEmbedding,
+                         ShortConvolution)
+from fla.modules.activations import swiglu, swish
+from fla.modules.convolution import proj_then_conv1d
+from fla.ops.abc.chunk import chunk_abc
+
+
+class ABCAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        num_slots: Optional[int] = None,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        gate_low_rank_dim: int = 16,
+        gate_logit_normalizer: int = 16,
+        use_input_gate: bool = False,
+        use_output_gate: bool = True,
+        use_norm: bool = True,
+        clamp_min: Optional[float] = -32,
+        clamp_max: Optional[float] = 32,
+        layer_idx: Optional[int] = None,
+        **kwargs
+    ) -> ABCAttention:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.key_dim = int(self.hidden_size * self.expand_k)
+        self.value_dim = int(self.hidden_size * self.expand_v)
+        self.head_k_dim = self.key_dim // self.num_heads
+        self.head_v_dim = self.value_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.use_input_gate = use_input_gate
+        self.use_output_gate = use_output_gate
+        self.use_norm = use_norm
+
+        if num_slots is None:
+            num_slots = self.head_k_dim
+        self.num_slots = num_slots
+
+        self.norm_eps = norm_eps
+
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.layer_idx = layer_idx
+
+        if layer_idx is None:
+            warnings.warn(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+        self.s_proj = nn.Linear(self.hidden_size, self.num_heads * self.num_slots, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+                self.k_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+                self.v_conv1d = ShortConvolution(self.value_dim, conv_size, activation='silu')
+
+        if self.use_norm:
+            if self.use_output_gate:
+                self.g_norm = FusedRMSNormSwishGate(self.head_v_dim, elementwise_affine, norm_eps)
+            else:
+                self.g_norm = RMSNorm(self.head_v_dim, elementwise_affine, norm_eps)
+
+        if self.use_rope:
+            self.rotary = RotaryEmbedding(self.head_k_dim)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                hidden_states = self.h_conv1d(hidden_states)
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+            else:
+                q = proj_then_conv1d(hidden_states, self.q_proj.weight, self.q_conv1d.weight, self.q_conv1d.bias)
+                k = proj_then_conv1d(hidden_states, self.k_proj.weight, self.k_conv1d.weight, self.k_conv1d.bias)
+                v = proj_then_conv1d(hidden_states, self.v_proj.weight, self.v_conv1d.weight, self.v_conv1d.bias)
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        if self.use_input_gate:
+            q, k, v = map(lambda x: swish(x), (q, k, v))
+
+        if self.use_rope:
+            q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads)
+            k = rearrange(k, '... (h d) -> ... h d', h=self.num_heads)
+            seqlen_offset = 0
+            if past_key_values is not None:
+                seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            q, k = self.rotary(q, k, seqlen_offset)
+            q = rearrange(q, 'b n h d -> b h n d', h=self.num_heads)
+            k = rearrange(k, 'b n h d -> b h n d', h=self.num_heads)
+        else:
+            q = rearrange(q, 'b n (h d) -> b h n d', h=self.num_heads)
+            k = rearrange(k, 'b n (h d) -> b h n d', h=self.num_heads)
+        v = rearrange(v, 'b n (h d) -> b h n d', h=self.num_heads)
+
+        # [batch_size, n_heads, seq_len, num_slots]
+        s = rearrange(self.s_proj(hidden_states), 'b t (h m) -> b h t m', h=self.num_heads)
+        s = s.clamp_(self.clamp_min, self.clamp_max)
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        o, last_state = chunk_abc(q, k, v, s, initial_state=last_state, output_final_state=use_cache)
+        if past_key_values is not None and last_state is not None:
+            past_key_values.update(last_state, self.layer_idx, q.shape[2])
+
+        o = rearrange(o, 'b h t d -> b t h d')
+        if self.use_norm and not self.use_output_gate:
+            o = self.g_norm(o)
+        elif self.use_output_gate:
+            g = rearrange(self.g_proj(hidden_states), 'b t (h d) -> b t h d', h=self.num_heads)
+            o = self.g_norm(o, g) if self.use_norm else swiglu(g, o)
+        o = rearrange(o, 'b t h d -> b t (h d)')
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_k_dim, self.num_slots),
+                  param.new_zeros(batch_size, self.num_heads, self.num_slots, self.head_v_dim))
+        return state
+
+    def state_size(self, sequence_length: int = 2048):
+        return self.num_heads * self.key_dim * self.head_v_dim

finetune/lora/v6/fla/layers/based.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+
+"""
+Linear attention in Based.
+https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py
+"""
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules.feature_map import TaylorFeatureMap
+from fla.ops.based import parallel_based
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn
+
+
+class BasedLinearAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        l_max: int = 2048,
+        feature_dim: int = 16,
+        num_key_value_heads: int = 12,
+        num_heads: int = 12,
+        feature_name: str = "taylor_exp",
+        eps: float = 1e-12,
+        causal: bool = True,
+        mode: str = "parallel",
+    ):
+        super().__init__()
+        self.hidden_size
+        self.l_max = l_max
+        self.mode = mode
+        assert self.mode in ["fused_chunk", "parallel", 'chunk']
+
+        # linear attention
+        self.feature_name = feature_name
+        self.feature_dim = feature_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_key_value_heads
+        self.causal = causal
+
+        self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.dropout = nn.Identity()
+        self.feature_map = TaylorFeatureMap(feature_dim)
+        self.eps = eps
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        mode = self.mode
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        q, k, v = map(lambda x: rearrange(x, "b l (h d) -> b h l d", h=self.num_heads), [q, k, v])
+        if mode == "fused_chunk":
+            q, k = self.feature_map(q), self.feature_map(k)
+            o = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'chunk':
+            q, k = self.feature_map(q), self.feature_map(k)
+            o = chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'parallel':
+            assert q.shape[-1] <= 128
+            o = parallel_based(q, k, v, True, True)
+        o = rearrange(o, "b h l d -> b l (h d)")
+        o = self.o_proj(o)
+        o = self.dropout(o)
+        return o
+
+    # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119
+
+    def forward_reference(self, hidden_states: torch.Tensor, filters: torch.Tensor = None, *args, **kwargs):
+        """
+        x (torch.Tensor): tensor of shape (b, d, l)
+        y (torch.Tensor): tensor of shape (b, d, l)
+        """
+        # hidden_states = hidden_states.transpose(1, 2)
+        b, l, _ = hidden_states.size()
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = q.view(b, l, self.num_heads, self.feature_dim).transpose(1, 2)
+        k = k.view(b, l, self.num_key_value_heads, self.feature_dim).transpose(1, 2)
+        v = v.view(b, l, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # Linear attention
+        q, k = self.feature_map(q), self.feature_map(k)
+        q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)
+
+        # Compute attention
+        if self.causal:
+            y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps))
+        else:
+            y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps))
+        y = rearrange(y, 'b h l d -> b l (h d)')
+        y = self.o_proj(y.to(hidden_states.dtype))
+        y = self.dropout(y)
+        return y.to(hidden_states.dtype)
+
+
+if __name__ == '__main__':
+    batch = 4
+    seq_len = 1024
+    hidden_size = 1024
+    dtype = torch.float32
+    x = torch.randn(batch, seq_len, hidden_size).to(dtype).cuda().requires_grad_(True)
+    dy = torch.randn(batch, seq_len, hidden_size).to(dtype).cuda()
+    model = BasedLinearAttention(hidden_size, mode='chunk').to(dtype).cuda()
+    y = model(x)
+    y.backward(dy, retain_graph=True)
+    x_grad, x.grad = x.grad, None
+    y2 = model.forward_reference(x)
+    y2.backward(dy)
+    assert y.allclose(y2, 0, 1e-4), breakpoint()
+    assert x_grad.allclose(x.grad, 0, 1e-4), breakpoint()
+    print("Pass")

finetune/lora/v6/fla/layers/delta_net.py

@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+
+# Sect4.2 of Linear Transformers Are Secretly Fast Weight Programmers https://arxiv.org/abs/2102.11174
+
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.cache_utils import Cache
+
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution, LayerNorm
+from fla.modules.rotary import RotaryEmbedding
+from fla.ops.delta_rule import (fused_chunk_delta_rule,
+                                fused_recurrent_linear_attn_delta_rule,
+                                chunk_delta_rule)
+from torch.nn import functional as F
+
+
+def simple_norm(x):
+    return (F.normalize(x, dim=-1) * x.shape[-1] ** 0.5).to(x)
+
+
+# @torch.jit.script
+def elu_p1(x):
+    return (F.elu(x, 1., False) + 1.).to(x)
+
+
+# @torch.jit.script
+def sum_norm(x):
+    return (x / x.sum(-1, keepdim=True)).to(x)
+
+
+# @torch.jit.script
+def elu_norm(x):
+    dtype = x.dtype
+    x = F.elu(x, 1., False) + 1.
+    return (x / x.sum(-1, keepdim=True)).to(dtype)
+
+
+
+
+# https://github.com/IDSIA/recurrent-fwp/blob/master/algorithmic/layers.py#L86C1-L146C1
+class DeltaNet(nn.Module):
+    def __init__(
+        self,
+        d_model: int = None,
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        mode: str = 'fused_chunk',
+        chunk_size: int = 16,
+        use_beta: bool = True,
+        use_gate: bool = True,
+        use_rope: bool = False,
+        use_output_norm: bool = True,
+        use_elu: bool = False,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = False,
+        layer_idx: int = None,
+        qk_activation: str = 'silu',
+        qk_norm: str = None,
+        save_memory: str = False,
+        **kwargs
+    ) -> DeltaNet:
+        super().__init__()
+        self.mode = mode
+        self.qk_activation = qk_activation
+        self.qk_norm = qk_norm
+        assert self.qk_activation in ['silu', 'relu', 'elu', 'identity']
+        assert self.qk_norm in ['l2', 'sum']
+        if d_model is not None:
+            hidden_size = d_model
+        self.hidden_size = hidden_size    
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.chunk_size = chunk_size
+        self.use_gate = use_gate
+        self.use_output_norm = use_output_norm
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+        self.layer_idx = layer_idx
+
+        self.silu = torch.nn.SiLU()
+
+        assert mode in ['chunk', 'fused_chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.use_beta = use_beta
+        self.use_elu = use_elu
+        if self.use_beta:
+            self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False)
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation=None)
+            else:
+                self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu' if qk_activation == 'silu' else None)
+                self.k_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu' if qk_activation == 'silu' else None)
+                self.v_conv1d = ShortConvolution(self.value_dim, conv_size, activation='silu')
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        if self.use_gate:
+            self.norm = FusedRMSNormSwishGate(self.head_v_dim)
+        else:
+            self.norm = RMSNorm(self.head_v_dim)
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+
+        # change to inference mode.
+        mode = 'fused_recurrent' if hidden_states.shape[1] < 64 else self.mode    
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+
+        if attention_mask is not None:
+            if attention_mask.shape[-1] != hidden_states.shape[-2]:
+                attention_mask = attention_mask[:, -1:]
+        
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+            else:
+                conv_state_q = last_state[0] if use_cache else None
+                conv_state_k = last_state[1] if use_cache else None
+                conv_state_v = last_state[2] if use_cache else None
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+                q = self.q_proj(hidden_states)
+                q = self.q_conv1d(q, attention_mask, conv_state_q)
+                k = self.k_conv1d(k, attention_mask, conv_state_k)
+                v = self.v_conv1d(v, attention_mask, conv_state_v)
+        else:
+            q = (self.q_proj(hidden_states))
+            k = (self.k_proj(hidden_states))
+            v = self.silu(self.v_proj(hidden_states))
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask.unsqueeze(-1))
+        
+        q, k, v = map(lambda x: rearrange(x, 'b l (h d) -> b h l d', h=self.num_heads), (q, k, v))
+
+        if self.qk_activation != 'silu':
+            if self.qk_activation == 'relu':
+                q, k = q.relu(), k.relu()
+            elif self.qk_activation == 'elu':
+                q, k = elu_p1(q), elu_p1(k)
+            elif self.qk_activation == 'identity':
+                pass
+            else:
+                raise NotImplementedError
+
+        if self.qk_norm is not None:
+            if self.qk_norm == 'l2':
+                k = torch.nn.functional.normalize(k, dim=-1, p=2).to(v) #auto mixed precision type transfer is annoying.
+                q = torch.nn.functional.normalize(q, dim=-1, p=2).to(v)
+            elif self.qk_norm == 'sum':
+                q = sum_norm(q).to(v)
+                k = sum_norm(k).to(v)      
+            
+        if self.use_beta:
+            beta = rearrange(self.b_proj(hidden_states), 'b l h -> b h l').sigmoid()
+        else:
+            beta = q.new_ones(q.shape[0], q.shape[1], q.shape[2])
+        state = past_key_values[self.layer_idx][-1] if use_cache else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_linear_attn_delta_rule(q, k, v, beta, state, output_final_state=use_cache)
+        elif mode == 'fused_chunk':
+            assert self.chunk_size in [16, 32, 64]
+            o, recurrent_state = fused_chunk_delta_rule(q, k, v, beta, self.chunk_size, state, output_final_state=use_cache)
+        elif mode == 'chunk':
+            assert self.chunk_size in [16, 32, 64]
+            o, recurrent_state = chunk_delta_rule(q, k, v, beta, self.chunk_size, state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    state = (conv_state, recurrent_state)
+                else:
+                    state = (conv_state_q, conv_state_k, conv_state_v, recurrent_state)
+            else:
+                state = (recurrent_state,)
+            past_key_values.update(state, self.layer_idx)
+
+        o = rearrange(o, 'b h l d -> b l h d')
+        if self.use_gate:
+            g = rearrange(self.g_proj(hidden_states), 'b l (h d) -> b l h d', h=self.num_heads)
+            o = self.norm(o, g)
+        else:
+            o = self.norm(o)
+        o = rearrange(o, 'b l h d -> b l (h d)')
+        o = self.o_proj(o)
+            
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                # for q/k/v each
+                state += (param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.value_dim, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_qk_dim, self.head_v_dim),)
+        return state

finetune/lora/v6/fla/layers/gated_abc.py

@@ -0,0 +1,234 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from transformers.cache_utils import Cache
+
+from fla.modules import (FusedRMSNormSwishGateLinear, RMSNormLinear,
+                         RotaryEmbedding, ShortConvolution)
+from fla.modules.activations import ACT2FN, swiglu_linear, swish
+from fla.ops.abc.chunk_gate import chunk_gated_abc
+
+
+class GatedABCAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        expand_k: float = 1.,
+        expand_v: float = 1.,
+        num_heads: int = 4,
+        num_kv_heads: Optional[int] = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        num_slots: Optional[int] = None,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        gate_low_rank_dim: Optional[int] = None,
+        gate_logit_normalizer: int = 16,
+        feature_map: str = 'swish',
+        use_rope: bool = False,
+        use_output_gate: bool = False,
+        use_norm: bool = True,
+        layer_idx: Optional[int] = None,
+        **kwargs
+    ) -> GatedABCAttention:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.key_dim_per_group = self.key_dim // self.num_kv_groups
+        self.value_dim_per_group = self.value_dim // self.num_kv_groups
+        self.head_k_dim = self.key_dim // self.num_heads
+        self.head_v_dim = self.value_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+
+        if gate_low_rank_dim is None:
+            gate_low_rank_dim = self.hidden_size // 16
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.feature_map = feature_map
+        self.use_rope = use_rope
+        self.use_output_gate = use_output_gate
+        self.use_norm = use_norm
+
+        if num_slots is None:
+            num_slots = self.head_k_dim
+        self.num_slots = num_slots
+
+        self.layer_idx = layer_idx
+
+        if layer_idx is None:
+            warnings.warn(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.key_dim_per_group, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.value_dim_per_group, bias=False)
+        self.f_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.num_slots, bias=False)
+
+        if use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+                self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu')
+                self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu')
+
+        if self.use_norm:
+            if self.use_output_gate:
+                self.g_norm = FusedRMSNormSwishGateLinear(self.hidden_size, elementwise_affine, norm_eps)
+            else:
+                self.g_norm = RMSNormLinear(self.hidden_size, elementwise_affine, norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+        if self.use_rope:
+            self.rotary = RotaryEmbedding(self.head_k_dim)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        lower_bound: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+            else:
+                conv_state_q = last_state[0] if use_cache else None
+                conv_state_k = last_state[1] if use_cache else None
+                conv_state_v = last_state[2] if use_cache else None
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+                q = self.q_conv1d(q, attention_mask, conv_state_q)
+                k = self.k_conv1d(k, attention_mask, conv_state_k)
+                v = self.v_conv1d(v, attention_mask, conv_state_v)
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+        f = self.f_proj(hidden_states)
+
+        if self.use_rope:
+            q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads)
+            k = rearrange(k, '... (h d) -> ... h d', h=self.num_kv_heads)
+            seqlen_offset = 0
+            if past_key_values is not None:
+                seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            q, k = self.rotary(q, k, seqlen_offset)
+            q = rearrange(q, 'b n h d -> b h n d', h=self.num_heads)
+            k = rearrange(k, 'b n h d -> b h n d', h=self.num_kv_heads)
+        else:
+            q = rearrange(q, 'b n (h d) -> b h n d', h=self.num_heads)
+            if self.num_kv_groups > 1:
+                k = repeat(k, 'b n (h d) -> b (h g) n d', h=self.num_kv_heads, g=self.num_kv_groups)
+            else:
+                k = rearrange(k, 'b n (h d) -> b h n d', h=self.num_kv_heads)
+        if self.num_kv_groups > 1:
+            v = repeat(v, 'b n (h d) -> b (h g) n d', h=self.num_kv_heads, g=self.num_kv_groups)
+            f = repeat(f, 'b n (h m) -> b (h g) n m', h=self.num_kv_heads, g=self.num_kv_groups)
+        else:
+            v = rearrange(v, 'b n (h d) -> b h n d', h=self.num_kv_heads)
+            f = rearrange(f, 'b n (h m) -> b h n m', h=self.num_kv_heads)
+
+        if self.feature_map is not None:
+            q, k, v = map(lambda x: ACT2FN[self.feature_map](x), (q, k, v))
+        f = F.logsigmoid(f) / self.gate_logit_normalizer
+        s = (1 - f.exp()).to(f.dtype)
+        # dealing with left-padding
+        if attention_mask is not None:
+            s = s.mul_(attention_mask.view(attention_mask.shape[0], 1, -1, 1))
+            v = v.mul_(attention_mask.view(attention_mask.shape[0], 1, -1, 1))
+
+        recurrent_state = last_state[-2:] if use_cache else None
+        o, recurrent_state = chunk_gated_abc(q, k, v, s, f,
+                                             initial_state=recurrent_state,
+                                             output_final_state=use_cache)
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    last_state = (conv_state,) + recurrent_state
+                else:
+                    last_state = (conv_state_q, conv_state_k, conv_state_v) + recurrent_state
+            else:
+                last_state = recurrent_state
+            past_key_values.update(last_state, self.layer_idx, q.shape[2])
+
+        o = rearrange(o, 'b h t d -> b t (h d)')
+        if self.use_norm and not self.use_output_gate:
+            o = swish(o)
+            o = self.g_norm(o, self.o_proj.weight, self.o_proj.bias)
+        elif self.use_output_gate and not self.use_norm:
+            o = swiglu_linear(self.g_proj(hidden_states), o, self.o_proj.weight, self.o_proj.bias)
+        elif self.use_output_gate and self.use_norm:
+            o = self.g_norm(o, self.g_proj(hidden_states), self.o_proj.weight, self.o_proj.bias)
+        else:
+            o = self.o_proj(o)
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                state += (param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.value_dim, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_k_dim, self.num_slots),
+                  param.new_zeros(batch_size, self.num_heads, self.num_slots, self.head_v_dim))
+        return state
+
+    def state_size(self, sequence_length: int = 2048):
+        return self.num_heads * self.key_dim * self.head_v_dim

finetune/lora/v6/fla/layers/gla.py

@@ -0,0 +1,268 @@
+# -*- coding: utf-8 -*-
+
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from transformers.cache_utils import Cache
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
+from fla.modules.activations import ACT2FN
+from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+
+
+class GatedLinearAttention(nn.Module):
+    r"""
+    The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635).  # noqa
+
+    Args:
+        mode (str, Optional):
+            Which GLA kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`.
+            Default: `chunk`.
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 0.5.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 1.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        num_kv_heads (int, Optional):
+            The number of key/value heads, used for MQA. Default: None.
+        feature_map (str, Optional):
+            Feature map function applied to queries/keys. Default: None.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `False`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        share_conv_kernel (bool, Optional):
+            Whether to apply convolutions berfore q/k/v mapping, only taking effects when `use_short_conv`. Default: `True`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        gate_logit_normalizer (int, Optional):
+            The normalizer for the gate logits, appied after `logsigmoid`. Default: 16.
+        gate_low_rank_dim (int, Optional):
+            The low rank dim for the gate projection. Default: 16.
+        clamp_min (float, Optional):
+            The minimum value for the gate logits. Default: None.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        num_kv_heads: Optional[int] = None,
+        feature_map: Optional[str] = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        use_output_gate: bool = True,
+        gate_fn: str = 'swish',
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        gate_logit_normalizer: int = 16,
+        gate_low_rank_dim: int = 16,
+        clamp_min: Optional[float] = None,
+        fuse_norm: bool = True,
+        layer_idx: int = None,
+    ) -> GatedLinearAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+        self.use_output_gate = use_output_gate
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.key_dim_per_group = self.key_dim // self.num_kv_groups
+        self.value_dim_per_group = self.value_dim // self.num_kv_groups
+        self.clamp_min = clamp_min
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False)
+        if self.use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+                self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu')
+                self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu')
+
+        self.gk_proj = nn.Sequential(nn.Linear(hidden_size, gate_low_rank_dim, bias=False),
+                                     nn.Linear(gate_low_rank_dim, self.key_dim_per_group, bias=True))
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm and use_output_gate:
+            self.g_norm_swish_gate = FusedRMSNormSwishGate(self.head_v_dim, elementwise_affine, norm_eps)
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(self.head_v_dim, elementwise_affine, norm_eps)
+            self.gate_fn = ACT2FN[gate_fn]
+
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+            else:
+                conv_state_q = last_state[0] if use_cache else None
+                conv_state_k = last_state[1] if use_cache else None
+                conv_state_v = last_state[2] if use_cache else None
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+                q = self.q_conv1d(q, attention_mask, conv_state_q)
+                k = self.k_conv1d(k, attention_mask, conv_state_k)
+                v = self.v_conv1d(v, attention_mask, conv_state_v)
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+        gk = self.gk_proj(hidden_states)
+
+        if self.feature_map_fn is not None:
+            q, k = map(self.feature_map_fn, (q, k))
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask.unsqueeze(-1))
+        q = rearrange(q, 'b l (h d) -> b h l d', h=self.num_heads)
+        if self.num_kv_groups > 1:
+            k, v, gk = (repeat(x, 'b l (h d) -> b (h g) l d', h=self.num_kv_heads, g=self.num_kv_groups) for x in (k, v, gk))
+        else:
+            k, v, gk = (rearrange(x, 'b l (h d) -> b h l d', h=self.num_kv_heads) for x in (k, v, gk))
+        gk = F.logsigmoid(gk) / self.gate_logit_normalizer
+
+        if self.clamp_min is not None:
+            gk = torch.clamp_min(gk, self.clamp_min)
+
+        recurrent_state = last_state[-1] if use_cache else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(q, k, v, gk, initial_state=recurrent_state, output_final_state=use_cache)
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_gla(q, k, v, gk, initial_state=recurrent_state, output_final_state=use_cache)
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gla(q, k, v, gk, initial_state=recurrent_state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    last_state = (conv_state, recurrent_state)
+                else:
+                    last_state = (conv_state_q, conv_state_k, conv_state_v, recurrent_state)
+            else:
+                last_state = (recurrent_state,)
+            past_key_values.update(last_state, self.layer_idx, q.shape[2])
+
+        o = rearrange(o, 'b h l d -> b l h d')
+        if self.use_output_gate:
+            g = self.g_proj(hidden_states)
+            if self.fuse_norm_and_gate:
+                g = rearrange(g, 'b l (h d) -> b l h d', h=self.num_heads)
+                o = self.g_norm_swish_gate(o, g)
+                o = rearrange(o, 'b l h d -> b l (h d)')
+            else:
+                o = rearrange(self.g_norm(o), 'b l h d -> b l (h d)')
+                o = o * self.gate_fn(g)
+        else:
+            o = rearrange(self.g_norm(o), 'b l h d -> b l (h d)')
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                state += (param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.value_dim, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_qk_dim, self.head_v_dim),)
+        return state
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.key_dim * self.head_v_dim
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

finetune/lora/v6/fla/layers/hgrn.py

@@ -0,0 +1,165 @@
+# -*- coding: utf-8 -*-
+
+# "Hierarchically Gated Recurrent Neural Network for Sequence Modeling" [https://arxiv.org/abs/2311.04823]
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.cache_utils import Cache
+
+from fla.modules import FusedRMSNormSwishGate, ShortConvolution
+from fla.modules.activations import swiglu
+from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn
+
+
+class HGRNAttention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        num_heads: Optional[int] = None,
+        expand_ratio: Optional[int] = 1,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None
+    ) -> HGRNAttention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.expand_ratio = expand_ratio
+        self.input_dim = int(hidden_size * expand_ratio)
+        self.head_dim = self.input_dim // self.num_heads
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.hidden_size % num_heads == 0, f"hidden size must be divisible by num_heads of {num_heads}"
+
+        self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+        self.f_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+        self.g_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.input_dim, conv_size, activation='silu')
+                self.f_conv1d = ShortConvolution(self.input_dim, conv_size, activation='silu')
+                self.i_conv1d = ShortConvolution(self.input_dim, conv_size, activation='silu')
+
+        self.g_norm = FusedRMSNormSwishGate(self.input_dim, elementwise_affine, norm_eps)
+        self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        lower_bound: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                i = self.i_proj(hidden_states)
+                f = self.f_proj(hidden_states)
+            else:
+                conv_state_i = last_state[2] if use_cache else None
+                conv_state_f = last_state[1] if use_cache else None
+                i = self.i_conv1d(self.i_proj(hidden_states), attention_mask, conv_state_i)
+                f = self.f_conv1d(self.f_proj(hidden_states), attention_mask, conv_state_f)
+        else:
+            i = self.i_proj(hidden_states)
+            f = self.f_proj(hidden_states)
+
+        # the lower bound for the first layer is zero
+        if lower_bound is None or self.layer_idx == 0:
+            i, f = swiglu(i, 1 - f.sigmoid()), F.logsigmoid(f)
+        else:
+            g = lower_bound + (1 - lower_bound) * f.sigmoid()
+            i, f = swiglu(i, 1 - g), g.log()
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            i = i.mul_(attention_mask.unsqueeze(-1))
+        i, f = map(lambda x: rearrange(x, 'b l (h d) -> b h l d', h=self.num_heads), (i, f))
+
+        recurrent_state = last_state[-1] if use_cache else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_hgrn(i, f, initial_state=recurrent_state, output_final_state=use_cache)
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_hgrn(i, f, initial_state=recurrent_state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    last_state = (conv_state, recurrent_state)
+                else:
+                    last_state = (conv_state_i, conv_state_f, recurrent_state)
+            else:
+                last_state = (recurrent_state,)
+            past_key_values.update(last_state, self.layer_idx, i.shape[2])
+
+        o = self.g_norm(self.g_proj(hidden_states), rearrange(o, 'b h l d -> b l (h d)'))
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),
+                          param.new_zeros(batch_size, self.hidden_size, self.conv_size),
+                          param.new_zeros(batch_size, self.hidden_size, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_dim),)
+        return state
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.hidden_size
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

finetune/lora/v6/fla/layers/hgrn2.py

@@ -0,0 +1,186 @@
+# -*- coding: utf-8 -*-
+
+# "HGRN2: Gated Linear RNNs with State Expansion"[https://arxiv.org/abs/2404.07904]
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.cache_utils import Cache
+
+from fla.modules import RMSNorm, ShortConvolution
+from fla.modules.activations import swish
+from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+
+
+class HGRN2Attention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        num_heads: Optional[int] = None,
+        expand_ratio: Optional[int] = 128,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None
+    ) -> HGRN2Attention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+
+        if expand_ratio is None and num_heads is not None:
+            expand_ratio = hidden_size // num_heads
+        elif expand_ratio is not None and num_heads is None:
+            num_heads = hidden_size // expand_ratio
+        else:
+            raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.")
+        self.num_heads = num_heads
+        self.expand_ratio = expand_ratio
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+
+        self.forget_dim = int(self.num_heads * self.expand_ratio)
+        self.input_dim = hidden_size
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not suppoerted mode `{mode}`."
+        assert self.forget_dim % num_heads == 0, f"forget dim must be divisible by num_heads of {num_heads}"
+        assert self.input_dim % num_heads == 0, f"input dim must be divisible by num_heads of {num_heads}"
+
+        self.head_f_dim = self.expand_ratio
+        self.head_i_dim = self.hidden_size // num_heads
+
+        self.q_proj = nn.Linear(hidden_size, self.forget_dim, bias=False)
+        self.f_proj = nn.Linear(hidden_size, self.forget_dim, bias=False)
+        self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.forget_dim, conv_size, activation='silu')
+                self.f_conv1d = ShortConvolution(self.forget_dim, conv_size, activation='silu')
+                self.i_conv1d = ShortConvolution(self.input_dim, conv_size, activation='silu')
+
+        self.g_norm = RMSNorm(self.hidden_size, elementwise_affine, norm_eps)
+        self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        lower_bound: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                q = self.q_proj(hidden_states)
+                f = self.f_proj(hidden_states)
+                i = self.i_proj(hidden_states)
+            else:
+                conv_state_q = last_state[0] if use_cache else None
+                conv_state_f = last_state[1] if use_cache else None
+                conv_state_i = last_state[2] if use_cache else None
+                q = self.q_proj(hidden_states)
+                f = self.f_proj(hidden_states)
+                i = self.i_proj(hidden_states)
+                q = self.q_conv1d(q, attention_mask, conv_state_q)
+                f = self.f_conv1d(f, attention_mask, conv_state_f)
+                i = self.i_conv1d(i, attention_mask, conv_state_i)
+        else:
+            q = self.q_proj(hidden_states)
+            f = self.f_proj(hidden_states)
+            i = self.i_proj(hidden_states)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            i = i.mul_(attention_mask.unsqueeze(-1))
+
+        q = swish(q)
+        # the lower bound for the first layer is zero
+        if lower_bound is None or self.layer_idx == 0:
+            k, g = 1 - f.sigmoid(), F.logsigmoid(f)
+        else:
+            g = lower_bound + (1 - lower_bound) * f.sigmoid()
+            k, g = 1 - g, g.log()
+        q, k, i, g = map(lambda x: rearrange(x, 'b l (h d) -> b h l d', h=self.num_heads), (q, k, i, g))
+
+        recurrent_state = last_state[-1] if use_cache else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_gla(q, k, i, g, initial_state=recurrent_state, output_final_state=use_cache)
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_gla(q, k, i, g, initial_state=recurrent_state, output_final_state=use_cache)
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_gla(q, k, i, g, initial_state=recurrent_state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    last_state = (conv_state, recurrent_state)
+                else:
+                    last_state = (conv_state_q, conv_state_f, conv_state_i, recurrent_state)
+            else:
+                last_state = (recurrent_state,)
+            past_key_values.update(last_state, self.layer_idx, q.shape[2])
+
+        o = self.g_norm(rearrange(o, 'b h l d -> b l (h d)'))
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                state += (param.new_zeros(batch_size, self.forget_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.forget_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.input_dim, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_f_dim, self.head_i_dim),)
+        return state
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.forget_dim * self.head_i_dim
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

finetune/lora/v6/fla/layers/linear_attn.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fla.modules import RMSNorm
+from fla.modules.feature_map import (DPFPFeatureMap, HadamardFeatureMap,
+                                     HedgehogFeatureMap, T2RFeatureMap)
+from fla.ops.linear_attn import (chunk_linear_attn, fused_chunk_linear_attn,
+                                 fused_recurrent_linear_attn)
+
+
+class LinearAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: str = 1024,
+        expand_k: int = 1.0,
+        expand_v: int = 1.0,
+        num_heads: int = 8,
+        mode: str = 'chunk',
+        feature_map: str = 'elementwise_product',
+        tie_feature_map_qk: bool = False,
+        output_norm: str = 'rmsnorm',
+        norm_q: bool = False,
+        norm_k: bool = False,
+        # standard linear attention normalization
+        do_feature_map_norm: bool = False,
+        elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        **kwargs,
+    ):
+        super().__init__()
+        assert feature_map in ['elu', 'relu', 'hedgehog', 't2r', 'dpfp',
+                               'identity', 'elementwise_product'], f"Not supported feature map `{feature_map}`."
+
+        assert output_norm in ['rmsnorm', 'identity'], f"Not supported output norm `{output_norm}`."
+
+        self.hidden_size
+        self.mode = mode
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.num_heads = num_heads
+
+        assert mode in ['chunk', 'fused_chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+
+        if feature_map == 'hedgehog':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_qk_dim)
+            else:
+                self.feature_map_q = HedgehogFeatureMap(head_dim=self.head_qk_dim)
+                self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_qk_dim)
+
+        elif feature_map == 't2r':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = T2RFeatureMap(head_dim=self.head_qk_dim)
+            else:
+                self.feature_map_q = T2RFeatureMap(head_dim=self.head_qk_dim)
+                self.feature_map_k = T2RFeatureMap(head_dim=self.head_qk_dim)
+
+        elif feature_map == 'elementwise_product':
+            if tie_feature_map_qk:
+                self.feature_map_q = self.feature_map_k = HadamardFeatureMap(head_dim=self.head_qk_dim)
+            else:
+                self.feature_map_q = HadamardFeatureMap(head_dim=self.head_qk_dim)
+                self.feature_map_k = HadamardFeatureMap(head_dim=self.head_qk_dim)
+
+        elif feature_map == 'dpfp':
+            self.feature_map_q = DPFPFeatureMap(head_dim=self.head_qk_dim)
+            self.feature_map_k = DPFPFeatureMap(head_dim=self.head_qk_dim)
+
+        elif feature_map == 'elu':
+            def elu(x):
+                return F.elu(x) + 1
+            self.feature_map_q = elu
+            self.feature_map_k = elu
+
+        elif feature_map == 'relu':
+            self.feature_map_q = nn.ReLU()
+            self.feature_map_k = nn.ReLU()
+
+        elif feature_map == 'identity':
+            self.feature_map_q = nn.Identity()
+            self.feature_map_k = nn.Identity()
+        else:
+            raise NotImplementedError
+
+        self.do_feature_map_norm = do_feature_map_norm
+        if output_norm == 'rmsnorm':
+            self.norm = RMSNorm(self.head_v_dim, elementwise_affine, norm_eps)
+        elif output_norm == 'identity':
+            self.norm = nn.Identity()
+        else:
+            raise NotImplementedError
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        self.norm_q = norm_q
+        self.norm_k = norm_k
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(self, x):
+        mode = self.mode
+        q = rearrange(self.q_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        k = rearrange(self.k_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        v = rearrange(self.v_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        q = self.feature_map_q(q)
+        k = self.feature_map_k(k)
+        if self.norm_q:
+            q = q / (q.sum(-1, keepdim=True) + 1e-4)
+        if self.norm_k:
+            k = k / (k.sum(-1, keepdim=True) + 1e-4)
+
+        if mode == 'chunk':
+            o = chunk_linear_attn(q, k, v, normalize=self.do_feature_map_norm)
+        elif mode == 'fused_chunk':
+            o = fused_chunk_linear_attn(q, k, v, normalize=self.do_feature_map_norm)
+        elif mode == 'fused_recurrent':
+            o = fused_recurrent_linear_attn(q, k, v, normalize=self.do_feature_map_norm)
+        else:
+            raise NotImplementedError
+        o = self.norm(o)
+        o = rearrange(o, 'b h n d -> b n (h d)')
+        o = self.o_proj(o)
+        return o
+
+
+if __name__ == '__main__':
+    import torch
+    batch = 4
+    seq_len = 1024
+    hidden_size = 1024
+    x = torch.randn(batch, seq_len, hidden_size).to(torch.bfloat16).cuda().requires_grad_(True)
+    model = LinearAttention(hidden_size, feature_map='dplp').to(torch.bfloat16).cuda()
+    y = model(x)
+    print(y.shape)
+    y.sum().backward()
+    print(x.grad.shape)

finetune/lora/v6/fla/layers/multiscale_retention.py

@@ -0,0 +1,271 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
+from fla.modules.rotary import RotaryEmbedding
+from fla.ops.retention import (chunk_retention, fused_chunk_retention,
+                               fused_recurrent_retention, parallel_retention)
+
+
+class MultiScaleRetention(nn.Module):
+    r"""
+    The layer implementaion for [Retentive Network: A Successor to Transformer for Large Language Models](https://arxiv.org/pdf/2307.08621.pdf).  # noqa
+
+    Args:
+        mode (str, Optional):
+            Which Retention kernel to use.
+            Currently available: `chunk`, `fused_recurrent`, `parallel`, and `fused_chunk`.
+            Default: `fused_chunk`.
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 1.0.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 2.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 8.
+        num_kv_heads (int, Optional):
+            The number of key/value heads, used for MQA. Default: None.
+        feature_map (str, Optional):
+            Feature map function applied to queries/keys. Default: None.
+        use_short_conv (bool, Optional):
+            Whether to use short convolutions. Default: `False`.
+        conv_size (int, Optional):
+            The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4.
+        conv_bias (bool, Optional):
+            Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`.
+        share_conv_kernel (bool, Optional):
+            Whether to apply convolutions berfore q/k/v mapping, only taking effects when `use_short_conv`. Default: `True`.
+        use_output_gate (bool, Optional):
+            Whether to use output gate. Default: `True`.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'fused_chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 2.0,
+        num_heads: int = 8,
+        num_kv_heads: Optional[int] = None,
+        feature_map: Optional[str] = None,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        conv_bias: bool = False,
+        share_conv_kernel: bool = True,
+        use_output_gate: bool = True,
+        gate_fn: str = 'swish',
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        fuse_norm: bool = True,
+        layer_idx: int = None,
+        **kwargs
+    ) -> MultiScaleRetention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.share_conv_kernel = share_conv_kernel
+        self.use_output_gate = use_output_gate
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.key_dim_per_group = self.key_dim // self.num_kv_groups
+        self.value_dim_per_group = self.value_dim // self.num_kv_groups
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_chunk', 'parallel', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False)
+        if self.use_output_gate:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if use_short_conv:
+            self.conv_size = conv_size
+            if share_conv_kernel:
+                self.h_conv1d = ShortConvolution(hidden_size, conv_size, activation='silu')
+            else:
+                self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu')
+                self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu')
+                self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu')
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm and use_output_gate:
+            self.g_norm_swish_gate = FusedRMSNormSwishGate(self.head_v_dim, elementwise_affine, norm_eps)
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(self.head_v_dim, elementwise_affine, norm_eps)
+            self.gate_fn = ACT2FN[gate_fn]
+
+        # TODO: fix this issue
+        # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py#L180
+        # Ideally, we would want to support arbitrary d_head_qk
+        assert self.head_qk_dim <= 256, "head_qk_dim must be less than or equal to 256"
+        self.rotary = RotaryEmbedding(dim=self.head_qk_dim)
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        if self.use_short_conv:
+            conv_state = last_state[0] if use_cache else None
+            if self.share_conv_kernel:
+                # conv state is updated inplace
+                hidden_states = self.h_conv1d(hidden_states, attention_mask, conv_state)
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+            else:
+                conv_state_q = last_state[0] if use_cache else None
+                conv_state_k = last_state[1] if use_cache else None
+                conv_state_v = last_state[2] if use_cache else None
+                q = self.q_proj(hidden_states)
+                k = self.k_proj(hidden_states)
+                v = self.v_proj(hidden_states)
+                q = self.q_conv1d(q, attention_mask, conv_state_q)
+                k = self.k_conv1d(k, attention_mask, conv_state_k)
+                v = self.v_conv1d(v, attention_mask, conv_state_v)
+        else:
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask.unsqueeze(-1))
+        q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads)
+        k = rearrange(k, '... (h d) -> ... h d', h=self.num_kv_heads)
+        if self.feature_map_fn is not None:
+            q, k = map(self.feature_map_fn, (q, k))
+
+        seqlen_offset, max_seqlen = 0, None
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            max_seqlen = q.shape[1] + seqlen_offset
+        if attention_mask is not None:
+            # to deliminate the offsets of padding tokens
+            seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]
+            max_seqlen = q.shape[1] + max(seqlen_offset)
+        q, k = self.rotary(q, k, seqlen_offset, max_seqlen)
+        q = q.transpose(1, 2)
+        if self.num_kv_groups > 1:
+            k = repeat(k, 'b t h d -> b (h g) t d', h=self.num_kv_heads, g=self.num_kv_groups)
+            v = repeat(v, 'b t (h d) -> b (h g) t d', h=self.num_kv_heads, g=self.num_kv_groups)
+        else:
+            k, v = rearrange(k, 'b t h d -> b h t d'), rearrange(v, 'b t (h d) -> b h t d', h=self.num_kv_heads)
+
+        state = last_state[-1] if use_cache else None
+        if mode == 'chunk':
+            o, recurrent_state = chunk_retention(q, k, v, initial_state=state, output_final_state=use_cache)
+        elif mode == 'fused_chunk':
+            o, recurrent_state = fused_chunk_retention(q, k, v, initial_state=state, output_final_state=use_cache)
+        elif mode == 'parallel':
+            o, recurrent_state = parallel_retention(q, k, v, initial_state=state, output_final_state=use_cache)
+        elif mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_retention(q, k, v, initial_state=state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            if self.use_short_conv:
+                if self.share_conv_kernel:
+                    last_state = (conv_state, recurrent_state)
+                else:
+                    last_state = (conv_state_q, conv_state_k, conv_state_v, recurrent_state)
+            else:
+                last_state = (recurrent_state,)
+            past_key_values.update(last_state, self.layer_idx, q.shape[2])
+
+        o = rearrange(o, 'b h l d -> b l h d')
+        if self.use_output_gate:
+            g = self.g_proj(hidden_states)
+            if self.fuse_norm_and_gate:
+                g = rearrange(g, 'b l (h d) -> b l h d', h=self.num_heads)
+                o = self.g_norm_swish_gate(o, g)
+                o = rearrange(o, 'b l h d -> b l (h d)')
+            else:
+                o = rearrange(self.g_norm(o), 'b l h d -> b l (h d)')
+                o = o * self.gate_fn(g)
+        else:
+            o = rearrange(self.g_norm(o), 'b l h d -> b l (h d)')
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = tuple()
+        if self.use_short_conv:
+            if self.share_conv_kernel:
+                state += (param.new_zeros(batch_size, self.hidden_size, self.conv_size),)
+            else:
+                state += (param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.key_dim, self.conv_size),
+                          param.new_zeros(batch_size, self.value_dim, self.conv_size))
+        state += (param.new_zeros(batch_size, self.num_heads, self.head_qk_dim, self.head_v_dim),)
+        return state
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.key_dim * self.head_v_dim
+        for module in self.children():
+            if isinstance(module, ShortConvolution):
+                state_size += module.state_size
+        return state_size

finetune/lora/v6/fla/layers/rebased.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+
+"""
+https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/layers/rebased_fast.py
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fla.modules.feature_map import RebasedFeatureMap
+from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn
+from fla.ops.rebased import parallel_rebased
+
+
+class ReBasedLinearAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        l_max: int = 2048,
+        feature_dim: int = 16,
+        num_key_value_heads: int = 16,
+        num_heads: int = 16,
+        use_gamma: Optional[bool] = True,
+        use_beta: Optional[bool] = True,
+        normalize: Optional[bool] = True,
+        causal: bool = True,
+        eps: float = 1e-5,
+        mode: str = "parallel",
+        layer_idx: Optional[int] = None,
+        **kwargs
+    ) -> ReBasedLinearAttention:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.l_max = l_max
+        self.mode = mode
+        assert self.mode in ["fused_chunk", "parallel", 'chunk']
+
+        # linear attention
+        self.feature_dim = feature_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_heads = num_heads
+        self.head_dim = self.hidden_size // self.num_key_value_heads
+        self.use_gamma = use_gamma
+        self.use_beta = use_beta
+        self.normalize = normalize
+        self.causal = causal
+
+        self.feature_map = RebasedFeatureMap(self.feature_dim, use_gamma, use_beta, normalize)
+        self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.dropout = nn.Identity()
+        self.eps = eps
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs):
+        mode = self.mode
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+        q, k, v = map(lambda x: rearrange(x, "b l (h d) -> b h l d", h=self.num_heads), [q, k, v])
+        q, k = self.feature_map(q, flatten=(mode != 'parallel')), self.feature_map(k, flatten=(mode != 'parallel'))
+        if mode == "fused_chunk":
+            o = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'chunk':
+            o = chunk_linear_attn(q, k, v, normalize=True, scale=1)
+        elif mode == 'parallel':
+            assert q.shape[-1] <= 128
+            o = parallel_rebased(q, k, v, self.eps, True, True)
+        o = rearrange(o, "b h l d -> b l (h d)")
+        o = self.o_proj(o)
+        o = self.dropout(o)
+        return o
+
+    # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119
+    def forward_reference(self, hidden_states: torch.Tensor, filters: torch.Tensor = None, *args, **kwargs):
+        """
+        x (torch.Tensor): tensor of shape (b, d, l)
+        y (torch.Tensor): tensor of shape (b, d, l)
+        """
+        # hidden_states = hidden_states.transpose(1, 2)
+        b, l, _ = hidden_states.size()
+        q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+
+        q = q.view(b, l, self.num_heads, self.feature_dim).transpose(1, 2)
+        k = k.view(b, l, self.num_key_value_heads, self.feature_dim).transpose(1, 2)
+        v = v.view(b, l, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # Linear attention
+        q, k = self.feature_map(q), self.feature_map(k)
+        q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1)
+
+        # Compute attention
+        if self.causal:
+            y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps))
+        else:
+            y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps))
+        y = rearrange(y, 'b h l d -> b l (h d)')
+        y = self.o_proj(y.to(hidden_states.dtype))
+        y = self.dropout(y)
+        return y.to(hidden_states.dtype)
+
+
+if __name__ == '__main__':
+    batch = 4
+    seq_len = 1024
+    hidden_size = 1024
+    dtype = torch.float32
+    x = torch.randn(batch, seq_len, hidden_size).to(dtype).cuda().requires_grad_(True)
+    dy = torch.randn(batch, seq_len, hidden_size).to(dtype).cuda()
+    model = ReBasedLinearAttention(hidden_size=hidden_size, mode='parallel').to(dtype).cuda()
+
+    y = model(x)
+    y.backward(dy, retain_graph=True)
+    x_grad, x.grad = x.grad, None
+    print(model.mode)
+    model.mode = 'fused_chunk'
+    y2 = model(x)
+    print(model.mode)
+    y2.backward(dy)
+    # assert y.allclose(y2, 0, 1e-4), breakpoint()
+    # assert x_grad.allclose(x.grad, 0, 1e-4), breakpoint()
+    print("Pass")

finetune/lora/v6/fla/layers/rwkv6.py

@@ -0,0 +1,264 @@
+# -*- coding: utf-8 -*-
+
+# "Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence"[https://arxiv.org/abs/2404.05892]
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+
+from fla.modules import FusedLayerNormSwishGate, LayerNorm
+from fla.ops.rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6
+
+
+class RWKV6Attention(nn.Module):
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 0.5,
+        expand_v: float = 1.0,
+        num_heads: int = 4,
+        gate_fn: str = 'swish',
+        proj_low_rank_dim: int = 32,
+        gate_low_rank_dim: int = 64,
+        fuse_norm: bool = True,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        layer_idx: int = None,
+        **kwargs
+    ) -> RWKV6Attention:
+        super().__init__()
+
+        self.mode = mode
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.num_heads = num_heads
+        self.proj_low_rank_dim = proj_low_rank_dim
+        self.gate_low_rank_dim = gate_low_rank_dim
+
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        self.layer_idx = layer_idx
+
+        assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        self.x_proj = nn.Sequential(
+            LerpLinear(hidden_size, proj_low_rank_dim * 5),
+            nn.Tanh(),
+            nn.Linear(proj_low_rank_dim * 5, hidden_size, bias=True)
+        )
+        self.r_proj = DDLerpLinear(hidden_size, self.key_dim)
+        self.w_proj = DDLerpLinear(hidden_size, self.key_dim, low_rank_dim=gate_low_rank_dim)
+        self.k_proj = DDLerpLinear(hidden_size, self.key_dim)
+        self.v_proj = DDLerpLinear(hidden_size, self.value_dim)
+        self.g_proj = DDLerpLinear(hidden_size, self.value_dim)
+        self.bonus = nn.Parameter(torch.zeros(num_heads, self.head_qk_dim))
+
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm:
+            self.g_norm_swish_gate = FusedLayerNormSwishGate(self.head_v_dim, elementwise_affine, norm_eps)
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = LayerNorm(self.head_v_dim, elementwise_affine, norm_eps)
+            self.gate_fn = ACT2FN[gate_fn]
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        if isinstance(module, nn.Parameter):
+            nn.init.xavier_uniform_(module, gain=2 ** -2.5)
+        module._is_hf_initialized = True
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+        batch_size, seq_len, hidden_size = hidden_states.size()
+        # launching the triton kernel for just one token will actually be slower
+        mode = 'fused_recurrent' if hidden_states.shape[1] == 1 else self.mode
+
+        delta = self.time_shift(hidden_states) - hidden_states
+        x = self.x_proj[0](hidden_states, delta).view(batch_size, seq_len, -1, self.proj_low_rank_dim)
+        r, w, k, v, g = torch.einsum('b l n r, n r d-> b l n d',
+                                     self.x_proj[1](x),
+                                     self.x_proj[2].weight.view(5, -1, hidden_size)).unbind(-2)
+        r = self.r_proj(hidden_states, r, delta)
+        w = self.w_proj(hidden_states, w, delta)
+        k = self.k_proj(hidden_states, k, delta)
+        v = self.v_proj(hidden_states, v, delta)
+        g = self.g_proj(hidden_states, g, delta)
+
+        # dealing with left-padding
+        if attention_mask is not None:
+            v = v.mul_(attention_mask.unsqueeze(-1))
+        r, w, k, v = map(lambda x: rearrange(x, 'b l (h d) -> b h l d', h=self.num_heads), (r, w, k, v))
+        w = -torch.exp(w)
+        u = self.bonus
+
+        last_state = past_key_values[self.layer_idx] if use_cache else None
+        state = last_state[-1] if use_cache else None
+        if mode == 'fused_recurrent':
+            o, recurrent_state = fused_recurrent_rwkv6(r, k, v, w, u, initial_state=state, output_final_state=use_cache)
+        elif mode == 'chunk':
+            o, recurrent_state = chunk_rwkv6(r, k, v, w, u, initial_state=state, output_final_state=use_cache)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        if past_key_values is not None:
+            past_key_values.update((recurrent_state,), self.layer_idx, r.shape[2])
+
+        o = rearrange(o, 'b h l d -> b l h d')
+        if self.fuse_norm_and_gate:
+            g = rearrange(g, 'b l (h d) -> b l h d', h=self.num_heads)
+            o = self.g_norm_swish_gate(o, g)
+            o = rearrange(o, 'b l h d -> b l (h d)')
+        else:
+            o = self.g_norm(o)
+            o = rearrange(o, 'b l h d -> b l (h d)')
+            o = o * self.gate_fn(g)
+        o = self.o_proj(o)
+
+        return o, None, past_key_values
+
+    def init_state(self, batch_size: int) -> Tuple[torch.Tensor]:
+        param = next(self.parameters())
+        state = (param.new_zeros(batch_size, self.num_heads, self.head_qk_dim, self.head_v_dim),)
+        return state
+
+    def state_size(self, **kwargs) -> int:
+        state_size = self.key_dim * self.head_v_dim
+        return state_size
+
+
+class LoRA(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: int,
+        bias: Optional[bool] = True
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+        self.bias = bias
+
+        self.lora = nn.Sequential(
+            nn.Linear(input_dim, low_rank_dim, bias=False),
+            nn.Tanh(),
+            nn.Linear(low_rank_dim, output_dim, bias=bias)
+        )
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}("
+        s += f"input_dim={self.input_dim}, low_rank_dim={self.low_rank_dim}, output_dim={self.output_dim}"
+        if not self.bias:
+            s += f", bias={self.bias}"
+        s += ")"
+        return s
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lora(x)
+
+
+class LerpLinear(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: Optional[int] = None
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        if low_rank_dim is None:
+            self.linear = nn.Linear(input_dim, output_dim, bias=False)
+        else:
+            self.linear = LoRA(input_dim, output_dim, low_rank_dim)
+        self.mu = nn.Parameter(torch.zeros(input_dim))
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}"
+        if self.low_rank_dim is not None:
+            s += f", low_rank_dim={self.low_rank_dim}"
+        s += ")"
+        return s
+
+    def forward(self, x: torch.Tensor, delta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if delta is None:
+            shifted = self.time_shift(x)
+            if len(shifted.shape) == 2:
+                shifted = shifted.unsqueeze(1)
+            delta = shifted - x
+        return self.linear(x + delta * self.mu)
+
+
+class DDLerpLinear(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        low_rank_dim: Optional[int] = None
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.low_rank_dim = low_rank_dim
+
+        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+        if low_rank_dim is None:
+            self.linear = nn.Linear(input_dim, output_dim, bias=False)
+        else:
+            self.linear = LoRA(input_dim, output_dim, low_rank_dim)
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}({self.input_dim}, {self.output_dim}"
+        if self.low_rank_dim is not None:
+            s += f", low_rank_dim={self.low_rank_dim}"
+        s += ")"
+        return s
+
+    def forward(self, x: torch.Tensor, mu: torch.Tensor, delta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if delta is None:
+            shifted = self.time_shift(x)
+            if len(shifted.shape) == 2:
+                shifted = shifted.unsqueeze(1)
+            delta = shifted - x
+        return self.linear(x + delta * mu)

finetune/lora/v6/fla/layers/simple_gla.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.activations import ACT2FN
+
+from fla.modules import FusedRMSNormSwishGate, RMSNorm
+from fla.ops.simple_gla import chunk_simple_gla
+
+
+class SimpleGatedLinearAttention(nn.Module):
+    r"""
+    The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635).  # noqa
+    This layer calls the simplified GLA kernel in which the gating is head-wise instead of elementwise.
+
+    Args:
+        mode (str, Optional):
+            Which GLA kernel to use.
+            Currently available: `chunk`.
+            Default: `chunk`.
+        hidden_size (int, Optional):
+            The hidden size of the input. Default: 1024.
+        expand_k (float, Optional):
+            The expansion ratio for the key dim. Default: 0.5.
+        expand_v (float, Optional):
+            The expansion ratio for the value dim. Default: 1.0.
+        num_heads (int, Optional):
+            The number of heads. Default: 4.
+        gate_fn (str, Optional):
+            The activation function for the output gate. Default: `swish`.
+        elementwise_affine (bool, Optional):
+            If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`.
+        norm_eps (float, Optional):
+            The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5.
+        gate_logit_normalizer (int, Optional):
+            The normalizer for the gate logits, appied after `logsigmoid`. Default: 16.
+        fuse_norm (bool, Optional):
+            Whether to fuse the norm and the output gate for better memory footprint. Default: `True`.
+        layer_idx (int, Optional):
+            The index of the layer. Default: None.
+    """
+
+    def __init__(
+        self,
+        mode: str = 'chunk',
+        hidden_size: int = 1024,
+        expand_k: float = 1.0,
+        expand_v: float = 2.0,
+        num_heads: int = 4,
+        gate_fn: str = 'swish',
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-5,
+        gate_logit_normalizer: int = 16,
+        fuse_norm: bool = True,
+        **kwargs
+    ) -> SimpleGatedLinearAttention:
+        super().__init__()
+        self.hidden_size = hidden_size
+
+        self.mode = mode
+        self.key_dim = int(hidden_size * expand_k)
+        self.value_dim = int(hidden_size * expand_v)
+        assert mode in ['chunk'], f"Not suppoerted mode `{mode}`."
+        assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}"
+        assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}"
+        self.num_heads = num_heads
+        self.head_qk_dim = self.key_dim // num_heads
+        self.head_v_dim = self.value_dim // num_heads
+        self.gate_fn = ACT2FN[gate_fn]
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+        self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        self.gk_proj = nn.Linear(hidden_size, self.num_heads)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+        if gate_fn == 'swish' and fuse_norm:
+            self.g_norm_swish_gate = FusedRMSNormSwishGate(self.head_v_dim, elementwise_affine, norm_eps)
+            self.fuse_norm_and_gate = True
+        else:
+            self.fuse_norm_and_gate = False
+            self.g_norm = RMSNorm(self.head_v_dim, elementwise_affine, norm_eps)
+
+        self.gate_logit_normalizer = gate_logit_normalizer
+
+        self.apply(self._initialize_weights)
+
+    def _initialize_weights(self, module: nn.Module):
+        if getattr(module, "_is_hf_initialized", False):
+            return
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        module._is_hf_initialized = True
+
+    def forward(self, x):
+        mode = self.mode
+        q = rearrange(self.q_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        k = rearrange(self.k_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        v = rearrange(self.v_proj(x), 'b n (h d) -> b h n d', h=self.num_heads)
+        gk = rearrange(self.gk_proj(x), 'b n h -> b h n')
+        gk = (F.logsigmoid(gk) / self.gate_logit_normalizer)
+
+        if mode == 'chunk':
+            o = chunk_simple_gla(q, k, v, gk)
+        else:
+            raise NotImplementedError(f"Not supported mode `{mode}`.")
+
+        o = rearrange(o, 'b h l d -> b l h d')
+        g = self.g_proj(x)
+
+        if self.fuse_norm_and_gate:
+            g = rearrange(g, 'b l (h d) -> b l h d', h=self.num_heads)
+            o = self.g_norm_swish_gate(o, g)
+            o = rearrange(o, 'b l h d -> b l (h d)')
+        else:
+            o = self.g_norm(o)
+            o = rearrange(o, 'b l h d -> b l (h d)')
+            o = o * self.gate_fn(g)
+        o = self.o_proj(o)
+        return o
+
+
+if __name__ == '__main__':
+    batch = 4
+    seq_len = 1024
+
+    hidden_size = 2048
+    x = torch.randn(batch, seq_len, hidden_size).to(torch.bfloat16).cuda().requires_grad_(True)
+    model = SimpleGatedLinearAttention(hidden_size=hidden_size, mode='chunk').to(torch.bfloat16).cuda()
+    y = model(x)
+    print(y.shape)
+    y.sum().backward()
+    print(x.grad.shape)

finetune/lora/v6/fla/models/__init__.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+
+from fla.models.abc import ABCConfig, ABCForCausalLM, ABCModel
+from fla.models.delta_net import (DeltaNetConfig, DeltaNetForCausalLM,
+                                  DeltaNetModel)
+from fla.models.gla import GLAConfig, GLAForCausalLM, GLAModel
+from fla.models.hgrn import HGRNConfig, HGRNForCausalLM, HGRNModel
+from fla.models.hgrn2 import HGRN2Config, HGRN2ForCausalLM, HGRN2Model
+from fla.models.linear_attn import (LinearAttentionConfig,
+                                    LinearAttentionForCausalLM,
+                                    LinearAttentionModel)
+from fla.models.mamba import MambaConfig, MambaForCausalLM, MambaModel
+from fla.models.retnet import RetNetConfig, RetNetForCausalLM, RetNetModel
+from fla.models.rwkv6 import RWKV6Config, RWKV6ForCausalLM, RWKV6Model
+from fla.models.transformer import (TransformerConfig, TransformerForCausalLM,
+                                    TransformerModel)
+
+__all__ = [
+    'ABCConfig', 'ABCForCausalLM', 'ABCModel',
+    'DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel',
+    'GLAConfig', 'GLAForCausalLM', 'GLAModel',
+    'HGRNConfig', 'HGRNForCausalLM', 'HGRNModel',
+    'HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model',
+    'LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel',
+    'MambaConfig', 'MambaForCausalLM', 'MambaModel',
+    'RetNetConfig', 'RetNetForCausalLM', 'RetNetModel',
+    'RWKV6Config', 'RWKV6ForCausalLM', 'RWKV6Model',
+    'TransformerConfig', 'TransformerForCausalLM', 'TransformerModel'
+]

finetune/lora/v6/fla/models/abc/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.abc.configuration_abc import ABCConfig
+from fla.models.abc.modeling_abc import ABCForCausalLM, ABCModel
+
+AutoConfig.register(ABCConfig.model_type, ABCConfig)
+AutoModel.register(ABCConfig, ABCModel)
+AutoModelForCausalLM.register(ABCConfig, ABCForCausalLM)
+
+
+__all__ = ['ABCConfig', 'ABCForCausalLM', 'ABCModel']

finetune/lora/v6/fla/models/abc/configuration_abc.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class ABCConfig(PretrainedConfig):
+
+    model_type = 'abc'
+    keys_to_ignore_at_inference = ['past_key_values']
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 2048,
+        gate_low_rank_dim: int = 16,
+        clamp_min: float = -32,
+        clamp_max: float = 32,
+        hidden_ratio: Optional[int] = 4,
+        intermediate_size: Optional[int] = None,
+        num_hidden_layers: int = 24,
+        num_heads: int = 4,
+        num_slots: Optional[int] = 64,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        share_conv_kernel: bool = True,
+        exapnd_k: float = 0.5,
+        exapnd_v: float = 1,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 2048,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        pad_token_id: int = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        fuse_norm: bool = True,
+        fuse_cross_entropy: bool = True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.gate_low_rank_dim = gate_low_rank_dim
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.num_slots = num_slots
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.share_conv_kernel = share_conv_kernel
+        self.expand_k = exapnd_k
+        self.expand_v = exapnd_v
+        self.hidden_act = hidden_act
+        self.elementwise_affine = elementwise_affine
+        self.norm_eps = norm_eps
+        self.use_cache = use_cache
+        self.initializer_range = initializer_range
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.fuse_norm = fuse_norm
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

finetune/lora/v6/fla/models/abc/modeling_abc.py

@@ -0,0 +1,394 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from fla.layers.abc import ABCAttention
+from fla.models.abc.configuration_abc import ABCConfig
+from fla.models.utils import RecurrentCache
+from fla.modules import FusedCrossEntropyLoss, RMSNorm
+from fla.modules.activations import swiglu_linear
+
+logger = logging.get_logger(__name__)
+
+
+class ABCMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        hidden_ratio: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        hidden_act: str = 'swish'
+    ) -> ABCMLP:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        # the final number of params is `hidden_ratio * hidden_size^2`
+        # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio`
+        if hidden_ratio is None:
+            hidden_ratio = 4
+        if intermediate_size is None:
+            intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
+            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        y = self.gate_proj(x)
+        gate, y = y.chunk(2, -1)
+        return swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
+
+
+class ABCBlock(nn.Module):
+    def __init__(self, config: ABCConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attn_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.attn = ABCAttention(
+            hidden_size=config.hidden_size,
+            expand_k=config.expand_k,
+            expand_v=config.expand_v,
+            num_heads=config.num_heads,
+            num_slots=config.num_slots,
+            use_short_conv=config.use_short_conv,
+            conv_size=config.conv_size,
+            share_conv_kernel=config.share_conv_kernel,
+            gate_fn=config.hidden_act,
+            elementwise_affine=config.elementwise_affine,
+            norm_eps=config.norm_eps,
+            clamp_min=config.clamp_min,
+            clamp_max=config.clamp_max,
+            fuse_norm=config.fuse_norm,
+            layer_idx=layer_idx
+        )
+        self.mlp_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.mlp = ABCMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions
+        )
+        hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs
+
+
+class ABCPreTrainedModel(PreTrainedModel):
+
+    config_class = ABCConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['ABCBlock']
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(
+        self,
+        module: nn.Module,
+        rescale_prenorm_residual: bool = True,
+        num_residuals_per_layer: int = 2,
+    ):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+        if rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["o_proj.weight", "down_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    with torch.no_grad():
+                        p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers)
+
+
+class ABCModel(ABCPreTrainedModel):
+
+    def __init__(self, config: ABCConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([ABCBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # noqa
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        if output_attentions:
+            warnings.warn("`ABCModel` does not `output_attentions` now, setting it to `False`.")
+            output_attentions = False
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+        hidden_states = inputs_embeds
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = [layer.attn.init_state(batch_size) for layer in self.layers]
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attentions, past_key_values = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values,
+                    use_cache,
+                    output_attentions
+                )
+            else:
+                hidden_states, attentions, past_key_values = layer(
+                    hidden_states,
+                    attention_mask,
+                    past_key_values=past_key_values,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions
+                )
+
+            if output_attentions:
+                all_attns += (attentions,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = past_key_values.to_legacy_cache()
+        if not return_dict:
+            return tuple(x for x in [hidden_states, next_cache, all_hidden_states, all_attns] if x is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_attns
+        )
+
+
+class ABCForCausalLM(ABCPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ABCModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def generate(self, *args, **kwargs):
+        try:
+            return super().generate(*args, **kwargs)
+        except AttributeError as exception:
+            if 'past_key_values' in str(exception):
+                raise AttributeError(
+                    f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, "
+                    f"which is not supported for {self.__class__.__name__}. "
+                    f"Try another generation strategy instead. "
+                    f"For the available generation strategies, check this doc: "
+                    f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies"
+                )
+            else:
+                raise exception
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs
+    ):
+        # only last token for `inputs_ids` if the `past_key_values` is passed along.
+        if past_key_values is not None:
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values, input_ids.shape[1] - 1)
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            model_inputs = {'input_ids': input_ids}
+        model_inputs['past_key_values'] = past_key_values
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            if self.config.fuse_cross_entropy:
+                loss_fct = FusedCrossEntropyLoss(inplace_backward=True)
+            else:
+                loss_fct = nn.CrossEntropyLoss()
+            # Enable model parallelism
+            labels = labels.to(logits.device)
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], loss_fct.ignore_index)), 1)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

finetune/lora/v6/fla/models/delta_net/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.delta_net.configuration_delta_net import \
+    DeltaNetConfig
+from fla.models.delta_net.modeling_delta_net import (
+    DeltaNetForCausalLM, DeltaNetModel)
+
+AutoConfig.register(DeltaNetConfig.model_type, DeltaNetConfig)
+AutoModel.register(DeltaNetConfig, DeltaNetModel)
+AutoModelForCausalLM.register(DeltaNetConfig, DeltaNetForCausalLM)
+
+__all__ = ['DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel']

finetune/lora/v6/fla/models/delta_net/configuration_delta_net.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class DeltaNetConfig(PretrainedConfig):
+
+    model_type = 'delta_net'
+    keys_to_ignore_at_inference = ['past_key_values']
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 2048,
+        expand_k: int = 1,
+        expand_v: int = 1,
+        use_gate: bool = False,
+        use_short_conv: bool = True,
+        conv_size: int = 4,
+        share_conv_kernel: bool = False,
+        use_rope: bool = False,
+        use_beta: bool = True,
+        use_output_norm: bool = True,
+        hidden_ratio: Optional[int] = 4,
+        intermediate_size: Optional[int] = None,
+        num_hidden_layers: int = 24,
+        num_heads: int = 4,
+        attn_mode: str = "chunk",
+        qk_norm: str = 'l2',
+        qk_activation: str = 'silu',
+        chunk_size: int = 64,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 2048,
+        rms_norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        pad_token_id: int = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        initializer_range: float = 0.02,
+        fuse_cross_entropy: bool = True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.attn_mode = attn_mode
+        self.hidden_act = hidden_act
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.initializer_range = initializer_range
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.use_gate = use_gate
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.share_conv_kernel = share_conv_kernel
+        self.use_rope = use_rope
+        self.use_beta = use_beta
+        self.use_output_norm = use_output_norm
+        self.qk_norm = qk_norm
+        self.qk_activation = qk_activation
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

finetune/lora/v6/fla/models/delta_net/modeling_delta_net.py

@@ -0,0 +1,405 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from fla.layers.delta_net import DeltaNet
+from fla.models.delta_net.configuration_delta_net import DeltaNetConfig
+from fla.models.utils import RecurrentCache
+from fla.modules import FusedCrossEntropyLoss, RMSNorm
+from fla.modules.activations import swiglu_linear
+
+logger = logging.get_logger(__name__)
+
+
+class DeltaNetMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        hidden_ratio: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        hidden_act: str = 'swish'
+    ) -> DeltaNetMLP:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        # the final number of params is `hidden_ratio * hidden_size^2`
+        # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio`
+        if hidden_ratio is None:
+            hidden_ratio = 4
+        if intermediate_size is None:
+            intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
+            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        y = self.gate_proj(x)
+        gate, y = y.chunk(2, -1)
+        return swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
+
+
+class DeltaNetBlock(nn.Module):
+    def __init__(self, config: DeltaNetConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attn_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = DeltaNet(
+            mode=config.attn_mode,
+            hidden_size=config.hidden_size,
+            expand_k=config.expand_k,
+            expand_v=config.expand_v,
+            num_heads=config.num_heads,
+            use_gate=config.use_gate,
+            use_rope=config.use_rope,
+            use_beta=config.use_beta,
+            use_short_conv=config.use_short_conv,
+            use_output_norm=config.use_output_norm,
+            conv_size=config.conv_size,
+            share_conv_kernel=config.share_conv_kernel,
+            layer_idx=layer_idx,
+            qk_norm=config.qk_norm,
+            qk_activation=config.qk_activation
+        )
+        self.mlp_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = DeltaNetMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions
+        )
+        hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs
+
+
+class DeltaNetPreTrainedModel(PreTrainedModel):
+
+    config_class = DeltaNetConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['DeltaNetBlock']
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(
+        self,
+        module: nn.Module,
+        rescale_prenorm_residual: bool = True,
+        num_residuals_per_layer: int = 2,
+    ):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+        if rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["o_proj.weight", "down_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    with torch.no_grad():
+                        p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers)
+
+
+class DeltaNetModel(DeltaNetPreTrainedModel):
+
+    def __init__(self, config: DeltaNetConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([DeltaNetBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # noqa
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        if output_attentions:
+            warnings.warn("`DeltaNetModel` does not `output_attentions` now, setting it to `False`.")
+            output_attentions = False
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+        hidden_states = inputs_embeds
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = [layer.attn.init_state(batch_size) for layer in self.layers]
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attentions, past_key_values = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values,
+                    use_cache,
+                    output_attentions
+                )
+            else:
+                hidden_states, attentions, past_key_values = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions
+                )
+
+            if output_attentions:
+                all_attns += (attentions,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = past_key_values
+        # if use_cache:
+            # next_cache = past_key_values.to_legacy_cache()
+        if not return_dict:
+            return tuple(x for x in [hidden_states, next_cache, all_hidden_states, all_attns] if x is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_attns
+        )
+
+
+class DeltaNetForCausalLM(DeltaNetPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = DeltaNetModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def generate(self, *args, **kwargs):
+        try:
+            return super().generate(*args, **kwargs)
+        except AttributeError as exception:
+            if 'past_key_values' in str(exception):
+                raise AttributeError(
+                    f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, "
+                    f"which is not supported for {self.__class__.__name__}. "
+                    f"Try another generation strategy instead. "
+                    f"For the available generation strategies, check this doc: "
+                    f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies"
+                )
+            else:
+                raise exception
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs
+    ):
+        # only last token for `inputs_ids` if the `past_key_values` is passed along.
+        if past_key_values is not None:
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values, input_ids.shape[1] - 1)
+            # breakpoint()
+            input_ids, attention_mask = input_ids[:, -1:], attention_mask[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard.
+            # Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {'input_ids': input_ids.contiguous()}
+        
+        model_inputs.update({
+            'past_key_values': past_key_values,
+            'use_cache': kwargs.get('use_cache'),
+            'attention_mask': attention_mask,
+        })
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            if self.config.fuse_cross_entropy:
+                loss_fct = FusedCrossEntropyLoss(inplace_backward=True)
+            else:
+                loss_fct = nn.CrossEntropyLoss()
+            # Enable model parallelism
+            labels = labels.to(logits.device)
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], loss_fct.ignore_index)), 1)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

finetune/lora/v6/fla/models/gla/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.gla.configuration_gla import GLAConfig
+from fla.models.gla.modeling_gla import GLAForCausalLM, GLAModel
+
+AutoConfig.register(GLAConfig.model_type, GLAConfig)
+AutoModel.register(GLAConfig, GLAModel)
+AutoModelForCausalLM.register(GLAConfig, GLAForCausalLM)
+
+
+__all__ = ['GLAConfig', 'GLAForCausalLM', 'GLAModel']

finetune/lora/v6/fla/models/gla/configuration_gla.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class GLAConfig(PretrainedConfig):
+
+    model_type = 'gla'
+    keys_to_ignore_at_inference = ['past_key_values']
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 2048,
+        expand_k: int = 0.5,
+        expand_v: int = 1,
+        hidden_ratio: Optional[int] = 4,
+        intermediate_size: Optional[int] = None,
+        num_hidden_layers: int = 24,
+        num_heads: int = 4,
+        num_kv_heads: Optional[int] = None,
+        feature_map: Optional[str] = None,
+        attn_mode: str = "chunk",
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        share_conv_kernel: bool = True,
+        use_output_gate: bool = True,
+        clamp_min: Optional[float] = None,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 2048,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-6,
+        use_gk: bool = True,
+        use_gv: bool = False,
+        use_cache: bool = True,
+        pad_token_id: int = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        initializer_range: float = 0.02,
+        fuse_norm: bool = True,
+        fuse_cross_entropy: bool = True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.expand_k = expand_k
+        self.expand_v = expand_v
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.feature_map = feature_map
+        self.attn_mode = attn_mode
+        self.clamp_min = clamp_min
+        self.hidden_act = hidden_act
+        self.elementwise_affine = elementwise_affine
+        self.norm_eps = norm_eps
+        self.use_gk = use_gk
+        self.use_gv = use_gv
+        self.use_cache = use_cache
+        self.initializer_range = initializer_range
+        self.fuse_norm = fuse_norm
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.share_conv_kernel = share_conv_kernel
+        self.use_output_gate = use_output_gate
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

finetune/lora/v6/fla/models/gla/modeling_gla.py

@@ -0,0 +1,403 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from fla.layers.gla import GatedLinearAttention
+from fla.models.gla.configuration_gla import GLAConfig
+from fla.models.utils import RecurrentCache
+from fla.modules import FusedCrossEntropyLoss, RMSNorm
+from fla.modules.activations import swiglu_linear
+
+logger = logging.get_logger(__name__)
+
+
+class GLAMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        hidden_ratio: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        hidden_act: str = 'swish'
+    ) -> GLAMLP:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        # the final number of params is `hidden_ratio * hidden_size^2`
+        # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio`
+        if hidden_ratio is None:
+            hidden_ratio = 4
+        if intermediate_size is None:
+            intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
+            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        y = self.gate_proj(x)
+        gate, y = y.chunk(2, -1)
+        return swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
+
+
+class GLABlock(nn.Module):
+    def __init__(self, config: GLAConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attn_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.attn = GatedLinearAttention(
+            mode=config.attn_mode,
+            hidden_size=config.hidden_size,
+            expand_k=config.expand_k,
+            expand_v=config.expand_v,
+            num_heads=config.num_heads,
+            num_kv_heads=config.num_kv_heads,
+            feature_map=config.feature_map,
+            use_short_conv=config.use_short_conv,
+            conv_size=config.conv_size,
+            share_conv_kernel=config.share_conv_kernel,
+            use_output_gate=config.use_output_gate,
+            gate_fn=config.hidden_act,
+            elementwise_affine=config.elementwise_affine,
+            norm_eps=config.norm_eps,
+            clamp_min=config.clamp_min,
+            fuse_norm=config.fuse_norm,
+            layer_idx=layer_idx
+        )
+        self.mlp_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.mlp = GLAMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions
+        )
+        hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs
+
+
+class GLAPreTrainedModel(PreTrainedModel):
+
+    config_class = GLAConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['GLABlock']
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(
+        self,
+        module: nn.Module,
+        rescale_prenorm_residual: bool = True,
+        num_residuals_per_layer: int = 2,
+    ):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+        if rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["o_proj.weight", "down_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    with torch.no_grad():
+                        p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers)
+
+
+class GLAModel(GLAPreTrainedModel):
+
+    def __init__(self, config: GLAConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([GLABlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # noqa
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        if output_attentions:
+            warnings.warn("`GLAModel` does not `output_attentions` now, setting it to `False`.")
+            output_attentions = False
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+        hidden_states = inputs_embeds
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = [layer.attn.init_state(batch_size) for layer in self.layers]
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attentions, past_key_values = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values,
+                    use_cache,
+                    output_attentions
+                )
+            else:
+                hidden_states, attentions, past_key_values = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions
+                )
+
+            if output_attentions:
+                all_attns += (attentions,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = past_key_values.to_legacy_cache()
+        if not return_dict:
+            return tuple(x for x in [hidden_states, next_cache, all_hidden_states, all_attns] if x is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_attns
+        )
+
+
+class GLAForCausalLM(GLAPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GLAModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def generate(self, *args, **kwargs):
+        try:
+            return super().generate(*args, **kwargs)
+        except AttributeError as exception:
+            if 'past_key_values' in str(exception):
+                raise AttributeError(
+                    f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, "
+                    f"which is not supported for {self.__class__.__name__}. "
+                    f"Try another generation strategy instead. "
+                    f"For the available generation strategies, check this doc: "
+                    f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies"
+                )
+            else:
+                raise exception
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        # only last token for `inputs_ids` if the `past_key_values` is passed along.
+        if past_key_values is not None:
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values, input_ids.shape[1] - 1)
+            input_ids, attention_mask = input_ids[:, -1:], attention_mask[:, -1:]
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard.
+            # Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {'input_ids': input_ids.contiguous()}
+
+        model_inputs.update({
+            'past_key_values': past_key_values,
+            'use_cache': kwargs.get('use_cache'),
+            'attention_mask': attention_mask,
+        })
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            if self.config.fuse_cross_entropy:
+                loss_fct = FusedCrossEntropyLoss(inplace_backward=True)
+            else:
+                loss_fct = nn.CrossEntropyLoss()
+            # Enable model parallelism
+            labels = labels.to(logits.device)
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], loss_fct.ignore_index)), 1)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

finetune/lora/v6/fla/models/hgrn/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+from fla.models.hgrn.configuration_hgrn import HGRNConfig
+from fla.models.hgrn.modeling_hgrn import HGRNForCausalLM, HGRNModel
+
+AutoConfig.register(HGRNConfig.model_type, HGRNConfig)
+AutoModel.register(HGRNConfig, HGRNModel)
+AutoModelForCausalLM.register(HGRNConfig, HGRNForCausalLM)
+
+
+__all__ = ['HGRNConfig', 'HGRNForCausalLM', 'HGRNModel']

finetune/lora/v6/fla/models/hgrn/configuration_hgrn.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class HGRNConfig(PretrainedConfig):
+
+    model_type = 'hgrn'
+    keys_to_ignore_at_inference = ['past_key_values']
+
+    def __init__(
+        self,
+        attn_mode: str = "chunk",
+        vocab_size: int = 32000,
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 24,
+        num_heads: Optional[int] = 1,
+        expand_ratio: Optional[int] = 1,
+        use_short_conv: bool = False,
+        conv_size: int = 4,
+        share_conv_kernel: bool = True,
+        use_lower_bound: bool = True,
+        hidden_ratio: Optional[int] = 4,
+        intermediate_size: Optional[int] = None,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 2048,
+        elementwise_affine: Optional[bool] = True,
+        norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        pad_token_id: int = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        initializer_range: float = 0.02,
+        fuse_cross_entropy: bool = True,
+        **kwargs
+    ):
+        self.attn_mode = attn_mode
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_heads = num_heads
+        self.expand_ratio = expand_ratio
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.share_conv_kernel = share_conv_kernel
+        self.use_lower_bound = use_lower_bound
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.elementwise_affine = elementwise_affine
+        self.norm_eps = norm_eps
+        self.use_cache = use_cache
+        self.initializer_range = initializer_range
+        self.fuse_cross_entropy = fuse_cross_entropy
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

finetune/lora/v6/fla/models/hgrn/modeling_hgrn.py

@@ -0,0 +1,407 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from fla.layers.hgrn import HGRNAttention
+from fla.models.hgrn.configuration_hgrn import HGRNConfig
+from fla.models.utils import RecurrentCache
+from fla.modules import FusedCrossEntropyLoss, RMSNorm
+from fla.modules.activations import swiglu_linear
+
+logger = logging.get_logger(__name__)
+
+
+class HGRNMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        hidden_ratio: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        hidden_act: str = 'swish'
+    ) -> HGRNMLP:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        # the final number of params is `hidden_ratio * hidden_size^2`
+        # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio`
+        if hidden_ratio is None:
+            hidden_ratio = 4
+        if intermediate_size is None:
+            intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
+            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        y = self.gate_proj(x)
+        gate, y = y.chunk(2, -1)
+        return swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
+
+
+class HGRNBlock(nn.Module):
+    def __init__(self, config: HGRNConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attn_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.attn = HGRNAttention(
+            mode=config.attn_mode,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_heads,
+            expand_ratio=config.expand_ratio,
+            use_short_conv=config.use_short_conv,
+            conv_size=config.conv_size,
+            share_conv_kernel=config.share_conv_kernel,
+            elementwise_affine=config.elementwise_affine,
+            norm_eps=config.norm_eps,
+            layer_idx=layer_idx
+        )
+        self.mlp_norm = RMSNorm(hidden_size=config.hidden_size, eps=config.norm_eps)
+        self.mlp = HGRNMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+        lower_bound: Optional[torch.Tensor] = False,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            lower_bound=lower_bound
+        )
+        hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attentions, past_key_values)
+
+        return outputs
+
+
+class HGRNPreTrainedModel(PreTrainedModel):
+
+    config_class = HGRNConfig
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['HGRNBlock']
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(
+        self,
+        module: nn.Module,
+        rescale_prenorm_residual: bool = True,
+        num_residuals_per_layer: int = 2,
+    ):
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+        if rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["o_proj.weight", "down_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    with torch.no_grad():
+                        p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers)
+
+
+class HGRNModel(HGRNPreTrainedModel):
+
+    def __init__(self, config: HGRNConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        if config.use_lower_bound:
+            self.lower_bounds = nn.Parameter(torch.zeros(config.num_hidden_layers, config.hidden_size))
+        self.layers = nn.ModuleList([HGRNBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # noqa
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        if output_attentions:
+            warnings.warn("`HGRNModel` does not `output_attentions` now, setting it to `False`.")
+            output_attentions = False
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+        hidden_states = inputs_embeds
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = [layer.attn.init_state(batch_size) for layer in self.layers]
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values)
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attns = () if output_attentions else None
+
+        if self.config.use_lower_bound:
+            lower_bounds = self.lower_bounds.softmax(0)
+            lower_bounds = lower_bounds.cumsum(0) - lower_bounds[0]
+        for i, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            lower_bound = lower_bounds[i] if self.config.use_lower_bound else None
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attentions, past_key_values = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values,
+                    use_cache,
+                    output_attentions,
+                    lower_bound
+                )
+            else:
+                hidden_states, attentions, past_key_values = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    lower_bound=lower_bound
+                )
+
+            if output_attentions:
+                all_attns += (attentions,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = past_key_values.to_legacy_cache()
+        if not return_dict:
+            return tuple(x for x in [hidden_states, next_cache, all_hidden_states, all_attns] if x is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_attns
+        )
+
+
+class HGRNForCausalLM(HGRNPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = HGRNModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def generate(self, *args, **kwargs):
+        try:
+            return super().generate(*args, **kwargs)
+        except AttributeError as exception:
+            if 'past_key_values' in str(exception):
+                raise AttributeError(
+                    f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, "
+                    f"which is not supported for {self.__class__.__name__}. "
+                    f"Try another generation strategy instead. "
+                    f"For the available generation strategies, check this doc: "
+                    f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies"
+                )
+            else:
+                raise exception
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        # only last token for `inputs_ids` if the `past_key_values` is passed along.
+        if past_key_values is not None:
+            if not isinstance(past_key_values, RecurrentCache):
+                past_key_values = RecurrentCache.from_legacy_cache(past_key_values, input_ids.shape[1] - 1)
+            input_ids, attention_mask = input_ids[:, -1:], attention_mask[:, -1:]
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard.
+            # Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {'input_ids': input_ids.contiguous()}
+
+        model_inputs.update({
+            'past_key_values': past_key_values,
+            'use_cache': kwargs.get('use_cache'),
+            'attention_mask': attention_mask,
+        })
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Tuple[List[torch.Tensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            if self.config.fuse_cross_entropy:
+                loss_fct = FusedCrossEntropyLoss(inplace_backward=True)
+            else:
+                loss_fct = nn.CrossEntropyLoss()
+            # Enable model parallelism
+            labels = labels.to(logits.device)
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], loss_fct.ignore_index)), 1)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )