Update world RWKV CPU

infer/rwkv/cuda/gemm_fp16_cublas.cpp

@@ -0,0 +1,75 @@
+#include <cublas_v2.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <ATen/cuda/CUDAContext.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__));
+
+/*
+  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(torch::Tensor a, torch::Tensor b, torch::Tensor c) {
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  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 = at::cuda::getCurrentCUDABlasHandle();
+
+#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));
+  }
+}

infer/rwkv/cuda/operators.cu

@@ -0,0 +1,246 @@
+#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);
+}

infer/rwkv/cuda/rwkv5.cu

@@ -0,0 +1,88 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+typedef at::Half fp16;
+typedef float fp32;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H, float *__restrict__ _state,
+                               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_;
+    _state += h*_N_*_N_ + i*_N_; // wrong if B > 1 !!!
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    
+    float state[_N_];
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        state[j] = _state[j];
+    
+    __syncthreads();
+    u[i] = float(_u[i]);
+    w[i] = _w[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);
+    }
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        _state[j] = state[j];
+}
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, fp32 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}

infer/rwkv/cuda/rwkv5_op.cpp

@@ -0,0 +1,34 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+#include <c10/cuda/CUDAGuard.h>
+typedef at::BFloat16 bf16;
+typedef at::Half fp16;
+typedef float fp32;
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y);
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, fp32 *y);
+
+void forward_bf16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_bf16(B, T, C, H, state.data_ptr<float>(), 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 forward_fp16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_fp16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp16>(), k.data_ptr<fp16>(), v.data_ptr<fp16>(), w.data_ptr<float>(), u.data_ptr<fp16>(), y.data_ptr<fp16>());
+}
+void forward_fp32(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_fp32(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp32>(), k.data_ptr<fp32>(), v.data_ptr<fp32>(), w.data_ptr<float>(), u.data_ptr<fp32>(), y.data_ptr<fp32>());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward_bf16", &forward_bf16, "rwkv5 forward_bf16");
+    m.def("forward_fp16", &forward_fp16, "rwkv5 forward_fp16");
+    m.def("forward_fp32", &forward_fp32, "rwkv5 forward_fp32");
+}
+TORCH_LIBRARY(rwkv5, m) {
+    m.def("forward_bf16", forward_bf16);
+    m.def("forward_fp16", forward_fp16);
+    m.def("forward_fp32", forward_fp32);
+}

infer/rwkv/cuda/rwkv6.cu

@@ -0,0 +1,87 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+typedef at::BFloat16 bf16;
+typedef at::Half fp16;
+typedef float fp32;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H, float *__restrict__ _state,
+                               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_;
+    _state += h*_N_*_N_ + i*_N_; // wrong if B > 1 !!!
+
+    __shared__ float r[_N_], k[_N_], u[_N_], w[_N_];
+    
+    float state[_N_];
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        state[j] = _state[j];
+
+    __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] = _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++)
+        _state[j] = state[j];
+}
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, fp32 *y)
+{
+    assert(H*_N_ == C);
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, k, v, w, u, y);
+}

infer/rwkv/cuda/rwkv6_op.cpp

@@ -0,0 +1,34 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+#include <c10/cuda/CUDAGuard.h>
+typedef at::BFloat16 bf16;
+typedef at::Half fp16;
+typedef float fp32;
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y);
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y);
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, fp32 *y);
+
+void forward_bf16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_bf16(B, T, C, H, state.data_ptr<float>(), 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 forward_fp16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_fp16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp16>(), k.data_ptr<fp16>(), v.data_ptr<fp16>(), w.data_ptr<float>(), u.data_ptr<fp16>(), y.data_ptr<fp16>());
+}
+void forward_fp32(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(state));
+    cuda_forward_fp32(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp32>(), k.data_ptr<fp32>(), v.data_ptr<fp32>(), w.data_ptr<float>(), u.data_ptr<fp32>(), y.data_ptr<fp32>());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward_bf16", &forward_bf16, "rwkv6 forward_bf16");
+    m.def("forward_fp16", &forward_fp16, "rwkv6 forward_fp16");
+    m.def("forward_fp32", &forward_fp32, "rwkv6 forward_fp32");
+}
+TORCH_LIBRARY(rwkv6, m) {
+    m.def("forward_bf16", forward_bf16);
+    m.def("forward_fp16", forward_fp16);
+    m.def("forward_fp32", forward_fp32);
+}

infer/rwkv/cuda/rwkv7.cu

@@ -0,0 +1,77 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+
+typedef at::Half fp16;
+typedef at::BFloat16 bf16;
+typedef float fp32;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               float *__restrict__ _state, const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b,
+                               F *__restrict__ const _y)
+{
+    const int e = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _state += h*_N_*_N_ + i*_N_; // wrong if B > 1 !!!
+
+    float state[_N_];
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        state[j] = _state[j];
+
+    __shared__ float r[_N_], k[_N_], w[_N_], a[_N_], b[_N_];
+
+    for (int _t = 0; _t < T; _t++)
+    {
+        const int t = e*T*C + h*_N_ + i + _t * C;
+        __syncthreads();
+        r[i] = float(_r[t]);
+        w[i] = __expf(-__expf(float(_w[t])));
+        k[i] = float(_k[t]);
+        a[i] = float(_a[t]);
+        b[i] = float(_b[t]);
+        __syncthreads();
+
+        float sa = 0;
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            sa += a[j] * state[j];
+        }
+
+        float vv = float(_v[t]);
+        float y = 0;
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            s = s * w[j] + k[j] * vv + sa * b[j];
+            y += s * r[j];
+        }
+        _y[t] = F(y);
+    }
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        _state[j] = state[j];    
+}
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16* w, fp16 *k, fp16 *v, fp16 *a, fp16 *b, fp16 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32* w, fp32 *k, fp32 *v, fp32 *a, fp32 *b, fp32 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}

infer/rwkv/cuda/rwkv7_op.cpp

@@ -0,0 +1,26 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+
+typedef at::Half fp16;
+typedef at::BFloat16 bf16;
+typedef float fp32;
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y);
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *w, fp16 *k, fp16 *v, fp16 *a, fp16 *b, fp16 *y);
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *w, fp32 *k, fp32 *v, fp32 *a, fp32 *b, fp32 *y);
+
+void forward_bf16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_bf16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<bf16>(), w.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), a.data_ptr<bf16>(), b.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void forward_fp16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_fp16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp16>(), w.data_ptr<fp16>(), k.data_ptr<fp16>(), v.data_ptr<fp16>(), a.data_ptr<fp16>(), b.data_ptr<fp16>(), y.data_ptr<fp16>());
+}
+void forward_fp32(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_fp32(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp32>(), w.data_ptr<fp32>(), k.data_ptr<fp32>(), v.data_ptr<fp32>(), a.data_ptr<fp32>(), b.data_ptr<fp32>(), y.data_ptr<fp32>());
+}
+
+TORCH_LIBRARY(wkv7s, m) {
+    m.def("forward_bf16", forward_bf16);
+    m.def("forward_fp16", forward_fp16);
+    m.def("forward_fp32", forward_fp32);
+}

infer/rwkv/cuda/wrapper.cpp

@@ -0,0 +1,141 @@
+#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 a, Tensor b, Tensor c);
+#endif
+
+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");
+#ifndef DISABLE_CUBLAS_GEMM
+    m.def("gemm_fp16_cublas", &gemm_fp16_cublas, "gemv fp16 cublas");
+#endif
+}
+
+TORCH_LIBRARY(rwkv, m) {
+    m.def("wkv_forward", wkv_forward);
+    m.def("mm8_seq", mm8_seq);
+    m.def("mm8_one", mm8_one);
+#ifndef DISABLE_CUBLAS_GEMM
+    m.def("gemm_fp16_cublas", gemm_fp16_cublas);
+#endif
+}

infer/rwkv/model.py

@@ -0,0 +1,469 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+from typing import Optional
+import types, gc, os, time, re, math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+current_path = os.path.dirname(os.path.abspath(__file__))
+
+########################################################################################################
+
+if os.environ.get('RWKV_JIT_ON') != '0':
+    os.environ["RWKV_JIT_ON"] = '1'
+    MyModule = torch.jit.ScriptModule
+    MyFunction = torch.jit.script_method
+    MyStatic = torch.jit.script
+else:
+    MyModule = torch.nn.Module
+    def __nop(ob):
+        return ob
+    MyFunction = __nop
+    MyStatic = __nop
+
+if os.environ.get('RWKV_CUDA_ON') == '1':
+    from torch.utils.cpp_extension import load
+    try:
+        load(
+            name=f"wkv_cuda",
+            sources=[f"{current_path}/cuda/wrapper.cpp", f"{current_path}/cuda/operators.cu", f"{current_path}/cuda/gemm_fp16_cublas.cpp"],
+            verbose=True,
+            extra_ldflags=["cublas.lib" if os.name == "nt" else ""],
+            extra_cuda_cflags=["--use_fast_math", "-O3", "--extra-device-vectorization"],
+            is_python_module=False)
+        DISABLE_CUBLAS_GEMM = False
+    except:
+        print("Failed to build cuBLAS matmul, falling back to torch.matmul. Small model with fp16 will overflow.")
+        load(
+            name=f"wkv_cuda",
+            sources=[f"{current_path}/cuda/wrapper.cpp", f"{current_path}/cuda/operators.cu"],
+            verbose=True,
+            extra_cuda_cflags=["--use_fast_math", "-O3", "--extra-device-vectorization"],
+            extra_cflags=["-DDISABLE_CUBLAS_GEMM"],
+            is_python_module=False)
+        DISABLE_CUBLAS_GEMM = True
+
+    @MyStatic
+    def cuda_wkv(T: int, C: int, w, u, k, v, aa, bb, pp):
+        assert 1 * C % min(C, 32) == 0
+        assert k.dtype == v.dtype == torch.float16 or k.dtype == v.dtype == torch.float32
+        assert w.dtype == u.dtype == aa.dtype == bb.dtype == pp.dtype == torch.float32
+        w = w.contiguous()
+        u = u.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
+        y = torch.empty((T, C), device=w.device, memory_format=torch.contiguous_format, dtype=k.dtype)
+        torch.ops.rwkv.wkv_forward(1, T, C, w, u, k, v, y, aa, bb, pp)
+        return y, aa, bb, pp
+    @MyStatic
+    def cuda_mm8_seq(B: int, N: int, M: int, x, w, mx, rx, my, ry):
+        assert x.dtype == mx.dtype == rx.dtype == my.dtype == ry.dtype
+        assert x.dtype == torch.float32 or x.dtype == torch.float16
+        assert w.dtype == torch.uint8
+        assert x.shape == (B, N)
+        assert w.shape == (N, M)
+        assert rx.shape == mx.shape == (M,)
+        assert ry.shape == my.shape == (N, 1)
+        y = torch.empty((B, M), device=w.device, dtype=x.dtype)
+        torch.ops.rwkv.mm8_seq(B, N, M, x, w, mx, rx, my, ry, y)
+        return y
+    @MyStatic
+    def cuda_mm8_one(N: int, M: int, x, w, mx, rx, my, ry):
+        assert x.dtype == mx.dtype == rx.dtype == my.dtype == ry.dtype
+        assert x.dtype == torch.float32 or x.dtype == torch.float16
+        assert w.dtype == torch.uint8
+        assert x.shape == (N,)
+        assert w.shape == (N, M)
+        assert rx.shape == mx.shape == (M,)
+        assert ry.shape == my.shape == (N, 1)
+        y = torch.zeros((M,), device=w.device, dtype=torch.float32)
+        torch.ops.rwkv.mm8_one(N, M, x, w, mx, rx, my, ry, y)
+        return y.to(dtype=x.dtype)
+else:
+    os.environ["RWKV_CUDA_ON"] = '0'
+
+
+@MyStatic
+def torch_mm8_seq(x, w, mx, rx, my, ry):
+    return x @ ((w.to(dtype=x.dtype) + 0.5) * ry * rx + my + mx)
+
+@MyStatic
+def torch_mm8_one(x, w, mx, rx, my, ry):
+    return x @ ((w.to(dtype=x.dtype) + 0.5) * ry * rx + my + mx)
+
+if os.environ.get('RWKV_CUDA_ON') == '1':
+    @MyStatic
+    def mm8_seq(x, w, mx, rx, my, ry):
+        if w.device.type == 'cuda' and x.dtype == torch.float16:
+            B, N, M = x.shape[0], w.shape[0], w.shape[1]
+            return cuda_mm8_seq(B, N, M, x, w, mx, rx, my, ry)
+        else:
+            return torch_mm8_seq(x, w, mx, rx, my, ry)
+    @MyStatic
+    def mm8_one(x, w, mx, rx, my, ry):
+        if w.device.type == 'cuda':
+            N, M = w.shape[0], w.shape[1]
+            return cuda_mm8_one(N, M, x, w, mx, rx, my, ry)
+        else:
+            return torch_mm8_one(x, w, mx, rx, my, ry)
+else:
+    @MyStatic
+    def mm8_seq(x, w, mx, rx, my, ry):
+        return torch_mm8_seq(x, w, mx, rx, my, ry)
+    @MyStatic
+    def mm8_one(x, w, mx, rx, my, ry):
+        return torch_mm8_one(x, w, mx, rx, my, ry)
+
+def mm8(x: torch.Tensor, w: torch.Tensor, mx: torch.Tensor, rx: torch.Tensor, my: torch.Tensor, ry: torch.Tensor):
+    if len(x.shape) == 1:
+        return mm8_one(x, w, mx, rx, my, ry)
+    return mm8_seq(x, w, mx, rx, my, ry)
+
+def matmul(a, b, mx: Optional[torch.Tensor]=None, rx: Optional[torch.Tensor]=None, my: Optional[torch.Tensor]=None, ry: Optional[torch.Tensor]=None, output_dtype: Optional[torch.dtype]=None) -> torch.Tensor:
+    if output_dtype is None:
+        output_dtype = a.dtype
+    if b.dtype in [torch.float16, torch.bfloat16, torch.float32]:
+        assert a.dtype == b.dtype
+        return matmul_float(a, b, output_dtype=output_dtype)
+    elif b.dtype == torch.uint8:
+        assert mx is not None
+        assert rx is not None
+        assert my is not None
+        assert ry is not None
+        return mm8(a, b, mx, rx, my, ry).to(output_dtype)
+    else:
+        raise ValueError("Unsupported dtype")
+
+
+if os.environ.get('RWKV_CUDA_ON') == '1' and not DISABLE_CUBLAS_GEMM:
+    def matmul_float(a, b, output_dtype: Optional[torch.dtype]=None):
+        if output_dtype is None:
+            output_dtype = a.dtype
+        if a.dtype == b.dtype == torch.float16 and a.device.type == 'cuda':
+            if len(a.shape) == 1:
+                assert len(b.shape) == 2
+                c = torch.empty((b.shape[-1],), dtype=output_dtype, device=a.device)
+                a = a.unsqueeze(0)
+            else:
+                assert len(a.shape) == len(b.shape)
+                assert len(a.shape) == 2 or len(a.shape) == 3
+                # torch.empty((*a.shape[:-1], b.shape[-1])) doesn't work with jit
+                if len(a.shape) == 2:
+                    c = torch.empty((a.shape[0], b.shape[-1]), dtype=output_dtype, device=a.device)
+                else:
+                    c = torch.empty((a.shape[0], a.shape[1], b.shape[-1]), dtype=output_dtype, device=a.device)
+            torch.ops.rwkv.gemm_fp16_cublas(a, b, c)
+            return c
+        else:
+            return (a @ b).to(output_dtype)
+
+else:
+    def matmul_float(a, b, output_dtype: Optional[torch.dtype]=None):
+        return (a @ b).to(output_dtype)
+
+
+if os.environ.get('RWKV_DML_ON') == '1':
+    import torch_directml
+    print("PyTorch with DirectML Enabled")
+
+
+print(f'\n### RWKV-7 "Goose" enabled ###\n')
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+# torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
+# torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
+torch._C._jit_set_autocast_mode(False)
+
+MyModule = torch.jit.ScriptModule
+MyFunction = torch.jit.script_method
+MyStatic = torch.jit.script
+from typing import List
+
+DTYPE = None
+DEVICE = None
+HEAD_SIZE = 64
+
+if os.environ.get('RWKV_CUDA_ON') == '1':
+    from torch.utils.cpp_extension import load
+    load(name="wkv7s", sources=[f"{current_path}/cuda/rwkv7_op.cpp", f"{current_path}/cuda/rwkv7.cu"], is_python_module=False,
+                        verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}"])
+    class WKV_7(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, state, r, w, k, v, a, b):
+            with torch.no_grad():
+                T, C = r.size()
+                H = C // HEAD_SIZE
+                N = HEAD_SIZE
+                assert HEAD_SIZE == C // H
+                assert all(x.dtype == DTYPE for x in [r,w,k,v,a,b])
+                assert all(x.is_contiguous() for x in [r,w,k,v,a,b])
+                y = torch.empty((T, C), device=DEVICE, dtype=r.dtype, requires_grad=False, memory_format=torch.contiguous_format)
+
+                if DTYPE == torch.float16:
+                    torch.ops.wkv7s.forward_fp16(1, T, C, H, state, r, w, k, v, a, b, y)
+                elif DTYPE == torch.bfloat16:
+                    torch.ops.wkv7s.forward_bf16(1, T, C, H, state, r, w, k, v, a, b, y)
+                elif DTYPE == torch.float32:
+                    torch.ops.wkv7s.forward_fp32(1, T, C, H, state, r, w, k, v, a, b, y)
+
+                return y
+    def RWKV7_OP(state, r, w, k, v, a, b):
+        return WKV_7.apply(state, r, w, k, v, a, b)
+
+########################################################################################################
+
+class RWKV(MyModule):
+    def __init__(self, model, strategy):
+        global DTYPE, DEVICE
+        super().__init__()
+        self.eval()
+        args = types.SimpleNamespace()
+        self.args = args
+        # args.MODEL_NAME = model
+
+        # print(f'Loading {model} ({strategy})\n')
+
+        ss = strategy.split(' ')
+        DEVICE = ss[0]
+        if ss[1] == 'fp16':
+            DTYPE = torch.half
+        elif ss[1] == 'fp32':
+            DTYPE = torch.float32
+        elif ss[1] == 'bf16':
+            DTYPE = torch.bfloat16
+        else:
+            assert False, "currently rwkv7 strategy must be: cuda/cpu fp16/fp32/bf16"
+        
+        # self.z = torch.load(args.MODEL_NAME + '.pth', map_location=DEVICE)
+        self.z = model
+
+        # for k,v in self.z.items():
+        #     print(k, v.shape)
+        z = self.z
+
+        self.n_head, self.head_size = z['blocks.0.att.r_k'].shape
+        args.head_size = self.head_size
+        args.vocab_size, args.n_embd = z['emb.weight'].shape
+
+        args.n_layer = 0
+        keys = list(z.keys())
+        for k in keys:
+            layer_id = int(k.split('.')[1]) if ('blocks.' in k) else 0
+            args.n_layer = max(args.n_layer, layer_id+1)
+            if 'key.weight' in k or 'value.weight' in k or 'receptance.weight' in k or 'output.weight' in k or 'head.weight' in k:
+                z[k] = z[k].t()
+            z[k] = z[k].squeeze().to(dtype=DTYPE)
+            if k.endswith('att.r_k'): z[k] = z[k].flatten()
+
+        self.n_embd = args.n_embd
+        self.n_layer = args.n_layer
+
+        # z['emb.weight'] = F.layer_norm(z['emb.weight'], (args.n_embd,), weight=z['blocks.0.ln0.weight'], bias=z['blocks.0.ln0.bias'])
+        z['blocks.0.att.v0'] = z['blocks.0.att.a0'] # actually ignored
+        z['blocks.0.att.v1'] = z['blocks.0.att.a1'] # actually ignored
+        z['blocks.0.att.v2'] = z['blocks.0.att.a2'] # actually ignored
+
+    def forward(self, idx, state, full_output=False, sign=None):
+        if state == None:
+            state = [None for _ in range(self.args.n_layer * 3)]
+            for i in range(self.args.n_layer): # state: 0=att_x_prev 1=att_kv 2=ffn_x_prev
+                state[i*3+0] = torch.zeros(self.args.n_embd, dtype=DTYPE, requires_grad=False, device=DEVICE)
+                state[i*3+1] = torch.zeros((self.args.n_embd // self.args.head_size, self.args.head_size, self.args.head_size), dtype=torch.float, requires_grad=False, device=DEVICE)
+                state[i*3+2] = torch.zeros(self.args.n_embd, dtype=DTYPE, requires_grad=False, device=DEVICE)
+
+        x = self.z['emb.weight'][idx]
+        if isinstance(sign, torch.Tensor):
+            sign = sign.squeeze(0)
+            # sign = F.layer_norm(sign, (self.args.n_embd,), weight=self.z['blocks.0.ln0.weight'], bias=self.z['blocks.0.ln0.bias'])
+
+            x = torch.cat((sign,x.to(DEVICE)), dim=0)
+
+        x = F.layer_norm(x, (self.args.n_embd,), weight=self.z['blocks.0.ln0.weight'], bias=self.z['blocks.0.ln0.bias'])
+        # if isinstance(sign, torch.Tensor):
+        #     print(x)
+
+        if type(idx) is list:
+            if len(idx) > 1:
+                return self.forward_seq(x, state, full_output)
+            else:
+                return self.forward_one(x[0], state)
+        else:
+            return self.forward_one(x, state)
+
+    @MyFunction
+    def forward_one(self, x, state:List[torch.Tensor]):
+        with torch.no_grad(): 
+            z = self.z
+            #x = z['emb.weight'][idx]
+
+            v_first = torch.empty_like(x)
+            for i in range(self.n_layer):
+                bbb = f'blocks.{i}.'
+                att = f'blocks.{i}.att.'
+                ffn = f'blocks.{i}.ffn.'
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln1.weight'], bias=z[bbb+'ln1.bias'])
+
+                xx, state[i*3+0], state[i*3+1], v_first = RWKV_x070_TMix_one(i, self.n_head, self.head_size, xx, state[i*3+0], v_first, state[i*3+1],
+                    z[att+'x_r'], z[att+'x_w'], z[att+'x_k'], z[att+'x_v'], z[att+'x_a'], z[att+'x_g'],
+                    z[att+'w0'], z[att+'w1'], z[att+'w2'], z[att+'a0'], z[att+'a1'], z[att+'a2'], z[att+'v0'], z[att+'v1'], z[att+'v2'],
+                    z[att+'g1'], z[att+'g2'], z[att+'k_k'], z[att+'k_a'], z[att+'r_k'],
+                    z[att+'receptance.weight'], z[att+'key.weight'], z[att+'value.weight'], z[att+'output.weight'],
+                    z[att+'ln_x.weight'], z[att+'ln_x.bias'])
+                x = x + xx
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln2.weight'], bias=z[bbb+'ln2.bias'])
+
+                xx, state[i*3+2] = RWKV_x070_CMix_one(xx, state[i*3+2], z[ffn+'x_k'], z[ffn+'key.weight'], z[ffn+'value.weight'])
+                x = x + xx
+            
+                # if math.isnan(torch.min(x).item()): print(idx, i)
+
+            x = F.layer_norm(x, (self.n_embd,), weight=z['ln_out.weight'], bias=z['ln_out.bias'])
+            x = x @ z['head.weight']
+            return x, state
+        
+    @MyFunction
+    def forward_seq(self, x, state:List[torch.Tensor], full_output:bool=False):
+        with torch.no_grad(): 
+            z = self.z
+            #x = z['emb.weight'][idx]
+
+            v_first = torch.empty_like(x)
+            for i in range(self.n_layer):
+                bbb = f'blocks.{i}.'
+                att = f'blocks.{i}.att.'
+                ffn = f'blocks.{i}.ffn.'
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln1.weight'], bias=z[bbb+'ln1.bias'])
+
+                xx, state[i*3+0], state[i*3+1], v_first = RWKV_x070_TMix_seq(i, self.n_head, self.head_size, xx, state[i*3+0], v_first, state[i*3+1],
+                    z[att+'x_r'], z[att+'x_w'], z[att+'x_k'], z[att+'x_v'], z[att+'x_a'], z[att+'x_g'],
+                    z[att+'w0'], z[att+'w1'], z[att+'w2'], z[att+'a0'], z[att+'a1'], z[att+'a2'], z[att+'v0'], z[att+'v1'], z[att+'v2'],
+                    z[att+'g1'], z[att+'g2'], z[att+'k_k'], z[att+'k_a'], z[att+'r_k'],
+                    z[att+'receptance.weight'], z[att+'key.weight'], z[att+'value.weight'], z[att+'output.weight'],
+                    z[att+'ln_x.weight'], z[att+'ln_x.bias'])
+                x = x + xx
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln2.weight'], bias=z[bbb+'ln2.bias'])
+
+                xx, state[i*3+2] = RWKV_x070_CMix_seq(xx, state[i*3+2], z[ffn+'x_k'], z[ffn+'key.weight'], z[ffn+'value.weight'])
+                x = x + xx
+            
+            if not full_output: x = x[-1,:]
+            x = F.layer_norm(x, (self.n_embd,), weight=z['ln_out.weight'], bias=z['ln_out.bias'])
+            x = x @ z['head.weight']
+            return x, state
+
+########################################################################################################
+
+@MyStatic
+def RWKV_x070_TMix_one(layer_id: int, H:int, N:int, x, x_prev, v_first, state, x_r, x_w, x_k, x_v, x_a, x_g, w0, w1, w2, a0, a1, a2, v0, v1, v2, g1, g2, k_k, k_a, r_k, R_, K_, V_, O_, ln_w, ln_b):
+    xx = x_prev - x
+    xr, xw, xk, xv, xa, xg = x+xx*x_r, x+xx*x_w, x+xx*x_k, x+xx*x_v, x+xx*x_a, x+xx*x_g
+
+    r = xr @ R_
+    w = torch.tanh(xw @ w1) @ w2
+    k = xk @ K_
+    v = xv @ V_
+    a = torch.sigmoid(a0 + (xa @ a1) @ a2)
+    g = torch.sigmoid(xg @ g1) @ g2
+
+    kk = torch.nn.functional.normalize((k * k_k).view(H,N), dim=-1, p=2.0).view(H*N)
+    k = k * (1 + (a-1) * k_a)
+    if layer_id == 0: v_first = v
+    else: v = v + (v_first - v) * torch.sigmoid(v0 + (xv @ v1) @ v2)
+    w = torch.exp(-0.606531 * torch.sigmoid((w0 + w).float())) # 0.606531 = exp(-0.5)
+
+    vk = v.view(H,N,1) @ k.view(H,1,N)
+    ab = (-kk).view(H,N,1) @ (kk*a).view(H,1,N)
+    state = state * w.view(H,1,N) + state @ ab.float() + vk.float()
+    xx = (state.to(dtype=x.dtype) @ r.view(H,N,1))
+
+    xx = torch.nn.functional.group_norm(xx.view(1,H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(H*N)    
+    xx = xx + ((r * k * r_k).view(H,N).sum(dim=-1, keepdim=True) * v.view(H,N)).view(H*N)
+    return (xx * g) @ O_, x, state, v_first
+
+if os.environ.get('RWKV_CUDA_ON') == '1':
+    @MyStatic
+    def RWKV_x070_TMix_seq(layer_id: int, H:int, N:int, x, x_prev, v_first, state, x_r, x_w, x_k, x_v, x_a, x_g, w0, w1, w2, a0, a1, a2, v0, v1, v2, g1, g2, k_k, k_a, r_k, R_, K_, V_, O_, ln_w, ln_b):
+        T = x.shape[0]
+        xx = torch.cat((x_prev.unsqueeze(0), x[:-1,:])) - x
+        xr, xw, xk, xv, xa, xg = x+xx*x_r, x+xx*x_w, x+xx*x_k, x+xx*x_v, x+xx*x_a, x+xx*x_g
+
+        r = xr @ R_
+        w = torch.tanh(xw @ w1) @ w2
+        k = xk @ K_
+        v = xv @ V_
+        a = torch.sigmoid(a0 + (xa @ a1) @ a2)
+        g = torch.sigmoid(xg @ g1) @ g2
+
+        kk = torch.nn.functional.normalize((k * k_k).view(T,H,N), dim=-1, p=2.0).view(T,H*N)
+        k = k * (1 + (a-1) * k_a)
+        if layer_id == 0: v_first = v
+        else: v = v + (v_first - v) * torch.sigmoid(v0 + (xv @ v1) @ v2)
+
+        w = -torch.nn.functional.softplus(-(w0 + w)) - 0.5
+        xx = RWKV7_OP(state, r, w, k, v, -kk, kk*a)
+
+        xx = torch.nn.functional.group_norm(xx.view(T,H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(T,H*N)
+        xx = xx + ((r * k * r_k).view(T,H,N).sum(dim=-1, keepdim=True) * v.view(T,H,N)).view(T,H*N)
+        return (xx * g) @ O_, x[-1,:], state, v_first
+else:
+    @MyStatic
+    def RWKV_x070_TMix_seq(layer_id: int, H:int, N:int, x, x_prev, v_first, state, x_r, x_w, x_k, x_v, x_a, x_g, w0, w1, w2, a0, a1, a2, v0, v1, v2, g1, g2, k_k, k_a, r_k, R_, K_, V_, O_, ln_w, ln_b):
+        T = x.shape[0]
+        xx = torch.cat((x_prev.unsqueeze(0), x[:-1,:])) - x
+        xr, xw, xk, xv, xa, xg = x+xx*x_r, x+xx*x_w, x+xx*x_k, x+xx*x_v, x+xx*x_a, x+xx*x_g
+
+        r = xr @ R_
+        w = torch.tanh(xw @ w1) @ w2
+        k = xk @ K_
+        v = xv @ V_
+        a = torch.sigmoid(a0 + (xa @ a1) @ a2)
+        g = torch.sigmoid(xg @ g1) @ g2
+
+        kk = torch.nn.functional.normalize((k * k_k).view(T,H,N), dim=-1, p=2.0).view(T,H*N)
+        k = k * (1 + (a-1) * k_a)
+        if layer_id == 0: v_first = v
+        else: v = v + (v_first - v) * torch.sigmoid(v0 + (xv @ v1) @ v2)
+
+        w = torch.exp(-0.606531 * torch.sigmoid((w0 + w).float())) # 0.606531 = exp(-0.5)
+        for t in range(T):
+            r_, w_, k_, v_, kk_, a_ = r[t], w[t], k[t], v[t], kk[t], a[t]
+            vk = v_.view(H,N,1) @ k_.view(H,1,N)
+            ab = (-kk_).view(H,N,1) @ (kk_*a_).view(H,1,N)
+            state = state * w_.view(H,1,N) + state @ ab.float() + vk.float()
+            xx[t] = (state.to(dtype=x.dtype) @ r_.view(H,N,1)).view(H*N)
+
+        xx = torch.nn.functional.group_norm(xx.view(T,H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(T,H*N)
+        xx = xx + ((r * k * r_k).view(T,H,N).sum(dim=-1, keepdim=True) * v.view(T,H,N)).view(T,H*N)
+        return (xx * g) @ O_, x[-1,:], state, v_first
+
+########################################################################################################
+
+@MyStatic
+def RWKV_x070_CMix_one(x, x_prev, x_k, K_, V_):
+    xx = x_prev - x
+    k = x + xx * x_k
+    k = torch.relu(k @ K_) ** 2
+    return k @ V_, x
+
+@MyStatic
+def RWKV_x070_CMix_seq(x, x_prev, x_k, K_, V_):
+    xx = torch.cat((x_prev.unsqueeze(0), x[:-1,:])) - x
+    k = x + xx * x_k
+    k = torch.relu(k @ K_) ** 2
+    return k @ V_, x[-1,:]
+
+########################################################################################################
+
+
+

infer/rwkv/rwkv_tokenizer.py

@@ -0,0 +1,103 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+class TRIE:
+    __slots__ = tuple("ch,to,values,front".split(","))
+    to:list
+    values:set
+    def __init__(self, front=None, ch=None):
+        self.ch = ch
+        self.to = [None for ch in range(256)]
+        self.values = set()
+        self.front = front
+
+    def __repr__(self):
+        fr = self
+        ret = []
+        while(fr!=None):
+            if(fr.ch!=None):
+                ret.append(fr.ch)
+            fr = fr.front
+        return "<TRIE %s %s>"%(ret[::-1], self.values)
+    
+    def add(self, key:bytes, idx:int=0, val=None):
+        if(idx == len(key)):
+            if(val is None):
+                val = key
+            self.values.add(val)
+            return self
+        ch = key[idx]
+        if(self.to[ch] is None):
+            self.to[ch] = TRIE(front=self, ch=ch)
+        return self.to[ch].add(key, idx=idx+1, val=val)
+    
+    def find_longest(self, key:bytes, idx:int=0):
+        u:TRIE = self
+        ch:int = key[idx]
+        
+        while(u.to[ch] is not None):
+            u = u.to[ch]
+            idx += 1
+            if(u.values):
+                ret = idx, u, u.values
+            if(idx==len(key)):
+                break
+            ch = key[idx]
+        return ret
+
+class TRIE_TOKENIZER():
+    def __init__(self, file_name):
+        self.idx2token = {}
+        sorted = [] # must be already sorted
+        with open(file_name, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        for l in lines:
+            idx = int(l[:l.index(' ')])
+            x = eval(l[l.index(' '):l.rindex(' ')])
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(l[l.rindex(' '):])
+            sorted += [x]
+            self.idx2token[idx] = x
+
+        self.token2idx = {}
+        for k,v in self.idx2token.items():
+            self.token2idx[v] = int(k)
+
+        self.root = TRIE()
+        for t, i in self.token2idx.items():
+            _ = self.root.add(t, val=(t, i))
+
+    def encodeBytes(self, src:bytes):
+        idx:int = 0
+        tokens = []
+        while (idx < len(src)):
+            _idx:int = idx
+            idx, _, values = self.root.find_longest(src, idx)
+            assert(idx != _idx)
+            _, token = next(iter(values))            
+            tokens.append(token)
+        return tokens
+
+    def decodeBytes(self, tokens):
+        return b''.join(map(lambda i: self.idx2token[i], tokens))
+
+    def encode(self, src):
+        return self.encodeBytes(src.encode("utf-8"))
+
+    def decode(self, tokens):
+        try:
+            return self.decodeBytes(tokens).decode('utf-8')
+        except:
+            return '\ufffd' # bad utf-8
+
+    def printTokens(self, tokens):
+        for i in tokens:
+            s = self.idx2token[i]
+            try:
+                s = s.decode('utf-8')
+            except:
+                pass
+            print(f'{repr(s)}{i}', end=' ')
+        print()

infer/rwkv/utils.py

@@ -0,0 +1,148 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+import os, sys
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+class PIPELINE_ARGS():
+    def __init__(self, temperature=1.0, top_p=0.85, top_k=0, alpha_frequency=0.2, alpha_presence=0.2, alpha_decay=0.996, token_ban=[], token_stop=[], chunk_len=256):
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+        self.alpha_frequency = alpha_frequency # Frequency Penalty (as in GPT-3)
+        self.alpha_presence = alpha_presence # Presence Penalty (as in GPT-3)
+        self.alpha_decay = alpha_decay # gradually decay the penalty
+        self.token_ban = token_ban # ban the generation of some tokens
+        self.token_stop = token_stop # stop generation whenever you see any token here
+        self.chunk_len = chunk_len # split input into chunks to save VRAM (shorter -> slower)
+
+class PIPELINE():
+    def __init__(self, model, WORD_NAME):
+        self.model = model
+        if WORD_NAME == 'cl100k_base':
+            import tiktoken
+            self.tokenizer = tiktoken.get_encoding(WORD_NAME)
+        elif WORD_NAME == 'rwkv_vocab_v20230424':
+            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+            from rwkv_tokenizer import TRIE_TOKENIZER
+            self.tokenizer = TRIE_TOKENIZER(os.path.dirname(os.path.abspath(__file__)) + '/rwkv_vocab_v20230424.txt')        
+        else:
+            from tokenizers import Tokenizer
+            self.tokenizer = Tokenizer.from_file(WORD_NAME)
+
+    def refine_context(self, context):
+        context = context.strip().split('\n')
+        for c in range(len(context)):
+            context[c] = context[c].strip().strip('\u3000').strip('\r')
+        context = list(filter(lambda c: c != '', context))
+        context = '\n' + ('\n'.join(context)).strip()
+        if context == '':
+            context = '\n'
+        return context
+
+    def encode(self, x):
+        if 'Tokenizer' in str(type(self.tokenizer)):
+            return self.tokenizer.encode(x).ids
+        else:
+            return self.tokenizer.encode(x)
+    
+    def decode(self, x):
+        return self.tokenizer.decode(x)
+
+    def sample_logits(self, logits, temperature=1.0, top_p=0.85, top_k=0):
+        if temperature == 0:
+            temperature = 1.0
+            top_p = 0
+        probs = F.softmax(logits.float(), dim=-1)
+        top_k = int(top_k)
+        # 'privateuseone' is the type of custom devices like `torch_directml.device()`
+        if probs.device.type in ['cpu', 'privateuseone']:
+            probs = probs.cpu().numpy()
+            sorted_ids = np.argsort(probs)
+            sorted_probs = probs[sorted_ids][::-1]
+            cumulative_probs = np.cumsum(sorted_probs)
+            cutoff = float(sorted_probs[np.argmax(cumulative_probs >= top_p)])
+            probs[probs < cutoff] = 0
+            if top_k < len(probs) and top_k > 0:
+                probs[sorted_ids[:-top_k]] = 0
+            if temperature != 1.0:
+                probs = probs ** (1.0 / temperature)
+            probs = probs / np.sum(probs)
+            out = np.random.choice(a=len(probs), p=probs)
+            return int(out)
+        else:
+            sorted_ids = torch.argsort(probs)
+            sorted_probs = probs[sorted_ids]
+            sorted_probs = torch.flip(sorted_probs, dims=(0,))
+            cumulative_probs = torch.cumsum(sorted_probs, dim=-1).cpu().numpy()
+            cutoff = float(sorted_probs[np.argmax(cumulative_probs >= top_p)])
+            probs[probs < cutoff] = 0
+            if top_k < len(probs) and top_k > 0:
+                probs[sorted_ids[:-top_k]] = 0
+            if temperature != 1.0:
+                probs = probs ** (1.0 / temperature)
+            out = torch.multinomial(probs, num_samples=1)[0]
+            return int(out)
+    
+    def generate(self, ctx, token_count=100, args=PIPELINE_ARGS(), callback=None, state=None, sign=None):
+        all_tokens = []
+        out_last = 0
+        out_str = ''
+        occurrence = {}
+        for i in range(token_count):
+
+            # forward & adjust prob.
+            tokens = self.encode(ctx) if i == 0 else [token]
+            while len(tokens) > 0:
+                out, state = self.model.forward(tokens[:args.chunk_len], state, sign=sign)
+                tokens = tokens[args.chunk_len:]
+            
+            for n in args.token_ban:
+                out[n] = -float('inf')
+            for n in occurrence:
+                out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
+            
+            # sampler
+            token = self.sample_logits(out, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k)
+            
+            if token in args.token_stop:
+                break
+            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] = www
+            else:
+                occurrence[token] += www
+            # print(occurrence) # debug
+            
+            # output
+            tmp = self.decode(all_tokens[out_last:])
+            if '\ufffd' not in tmp: # is valid utf-8 string?
+                if callback:
+                    callback(tmp)
+                out_str += tmp
+                out_last = i + 1
+            sign =None
+        return out_str, state
+    
+    def prefill(self, ctx, token_count=100, args=PIPELINE_ARGS(), callback=None, state=None, sign=None):
+        tokens = self.encode(ctx)
+        out, state = self.model.forward(tokens[:args.chunk_len], state, sign=sign,full_output=True)
+        max_indices = torch.argmax(out, dim=-1)
+        token = self.sample_logits(out[19,:], temperature=args.temperature, top_p=args.top_p, top_k=args.top_k)
+        print(token, self.decode([token]),out[19,:])
+        print(state[0].view(-1))
+        print(self.decode(max_indices.tolist()))
+        return self.decode(max_indices.tolist())

infer/worldmodel.py

@@ -0,0 +1,79 @@
+
+import numpy as np
+
+import os, sys, torch, time
+import numpy as np
+np.set_printoptions(precision=4, suppress=True, linewidth=200)
+import torch
+print(torch.__version__)
+print(torch.version.cuda)
+
+# set these before import RWKV
+# os.environ['RWKV_JIT_ON'] = '1'
+# os.environ["RWKV_CUDA_ON"] = '1' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries
+from infer.rwkv.model import RWKV # pip install rwkv
+from infer.rwkv.utils import PIPELINE, PIPELINE_ARGS
+
+
+from world.world_encoder import WorldEncoder
+
+class Worldinfer():
+    def __init__(self, model_path, encoder_type, encoder_path, strategy='cpu bf16', args=None):
+
+        ss = strategy.split(' ')
+        DEVICE = ss[0]
+        if ss[1] == 'fp16':
+            self.DTYPE = torch.half
+        elif ss[1] == 'fp32':
+            self.DTYPE = torch.float32
+        elif ss[1] == 'bf16':
+            self.DTYPE = torch.bfloat16
+        else:
+            assert False, "currently rwkv7 strategy must be: cuda/cpu fp16/fp32/bf16"
+        
+        self.model_weight = torch.load(model_path + '.pth', map_location=DEVICE)
+        modality_dict = {}
+        for key, value in self.model_weight.items():
+            if 'emb.weight' in key:
+                _, n_embd = value.shape
+            if 'modality' in key:
+                k = key.replace('modality.world_encoder.', '')
+                modality_dict[k] = value 
+        model = RWKV(model=self.model_weight, strategy=strategy)
+        self.pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
+
+        if args==None:
+            self.args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.0, top_k=0, # top_k = 0 then ignore
+                                alpha_frequency = 0.0,
+                                alpha_presence = 0.0,
+                                token_ban = [0], # ban the generation of some tokens
+                                token_stop = [24], # stop generation whenever you see any token here
+                                chunk_len = 256) # split input into chunks to save VRAM (shorter -> slower)
+        else:
+            self.args=args
+        print('RWKV finish!!!')
+
+        config = {
+            'encoder_type': encoder_type,
+            'encoder_path': encoder_path,
+            'project_dim' : n_embd
+        }
+        self.modality = WorldEncoder(**config).to(DEVICE, torch.bfloat16)        
+        self.modality.load_checkpoint(modality_dict)
+
+
+    def generate(self, text, modality='none', state=None):
+        if isinstance(modality, str):
+            y=None
+        else:
+            y = self.modality(modality).to(self.DTYPE)
+        result, state = self.pipeline.generate(text, token_count=500, args=self.args, callback=None, state=state, sign=y)
+        return result, state
+
+    # def prefill(self, text, modality='none', state=None):
+    #     if isinstance(modality, str):
+    #         y=None
+    #     else:
+    #         y = self.modality(modality).to(self.DTYPE)
+    #     result, state = self.pipeline.forward(text, token_count=500, args=self.args, callback=None, state=state, sign=y)
+    #     return result, state

world/block.py

@@ -0,0 +1,34 @@
+import os
+import torch
+from src.infctx_module import *
+import torch.nn as nn
+from torch.nn import functional as F
+
+from src.rwkv7.Channel_mix import RWKV_CMix_x070
+from src.rwkv7.Time_mix import RWKV_Tmix_x070
+
+class Block(nn.Module):
+        def __init__(self, args, layer_id):
+            super().__init__()
+            self.args = args
+            self.layer_id = layer_id
+
+            self.ln1 = nn.LayerNorm(args.n_embd)
+            self.ln2 = nn.LayerNorm(args.n_embd)
+
+            if self.layer_id == 0:
+                self.ln0 = nn.LayerNorm(args.n_embd)
+
+            self.att = RWKV_Tmix_x070(args, layer_id)  
+            self.ffn = RWKV_CMix_x070(args, layer_id)
+
+
+        def forward(self, x, v_first):
+            if self.layer_id == 0:
+                x = self.ln0(x)
+
+            x_attn, v_first = self.att(self.ln1(x), v_first)
+            x = x + x_attn
+
+            x = x + self.ffn(self.ln2(x))
+            return x, v_first

world/cat.py

@@ -0,0 +1,173 @@
+
+import torch
+from torch.nn import functional as F
+def pad_mod(self, tensor_list, signal_list):
+    """
+    对一个包含不同长度张量的列表进行填充，使所有张量的长度相同且为16的整数倍，并生成掩码。
+    参数:
+        tensor_list (list of torch.Tensor): 输入的张量列表，每个张量形状为 [seq_len]。
+        pad_value (int, optional): 填充值，默认值为 0。
+    返回:
+        padded_tensor (torch.Tensor): 填充后的张量，形状为 [batch_size, target_len]。
+        mask (torch.Tensor): 填充掩码，1 表示有效数据，0 表示填充部分。
+    """
+
+    modality_list = []
+    #max_len = max((token.size(0) + signal.size(1)) for token, signal in zip(tensor_list, modality_list))
+    max_len = 0
+    for token, signal in zip(tensor_list, signal_list):
+
+        modality = self.modality(signal)
+        if modality is False:
+            modality_list.append(False)
+            continue
+        modality_list.append(modality)
+        max_len = max(token.size(0) + modality.size(1), max_len)
+
+    # 计算目标长度（向上取整到16的整数倍）
+    target_len = ((max_len + 15) // 16 * 16)+1
+
+    if self.args.ctx_len is not None:
+        target_len = min(target_len, self.args.ctx_len+1)
+
+    masks = torch.zeros((len(tensor_list), target_len-1), dtype=torch.int32)
+    x = []
+    y = []
+    s = []
+    m = []
+    for token, signal, mask in zip(tensor_list, modality_list, masks):
+        if signal is False:
+            continue
+        pad_len = target_len-(token.size(0) + signal.size(1))
+        
+        padded_token = F.pad(token, (0, pad_len), value=0)
+
+        x_token = padded_token[:-1]
+        y_token = F.pad(padded_token, (signal.size(1)-1, 0), value=0)
+
+        mask[signal.size(1) : -pad_len] = 1 
+        
+        s.append(signal)
+        x.append(x_token)
+        y.append(y_token)
+        m.append(mask)
+
+    y = torch.stack(y, dim=0)
+    m = torch.stack(m, dim=0).cuda()
+    
+    return s, x, y, m
+
+
+
+def mod_pad_text(self, signal_list, text_inputs, text_labels):
+    """
+    对一个包含不同长度张量的列表进行填充，使所有张量的长度相同且为16的整数倍，并生成掩码。
+    参数:
+        tensor_list (list of torch.Tensor): 输入的张量列表，每个张量形状为 [seq_len]。
+        pad_value (int, optional): 填充值，默认值为 0。
+    返回:
+        padded_tensor (torch.Tensor): 填充后的张量，形状为 [batch_size, target_len]。
+        mask (torch.Tensor): 填充掩码，1 表示有效数据，0 表示填充部分。
+    """
+
+    modality_list = []
+    #max_len = max((token.size(0) + signal.size(1)) for token, signal in zip(tensor_list, modality_list))
+    max_len = 0
+    for i, (signal, token, label) in enumerate(zip(signal_list, text_inputs, text_labels)):
+
+        modality = self.modality(signal)
+        modality_list.append(modality)
+        mod_label = torch.full((modality.size(1),), -100, device='cuda')
+        text_labels[i] = torch.cat([mod_label, label])
+        max_len = max(token.size(0) + modality.size(1), max_len)
+
+    # 计算目标长度（向上取整到16的整数倍）
+    target_len = ((max_len + 15) // 16 * 16)+1
+
+    if self.args.ctx_len is not None:
+        target_len = min(target_len, self.args.ctx_len+1)
+
+ 
+    for i, (signal, token, label) in enumerate(zip(modality_list , text_inputs, text_labels)):
+        pad_len = target_len-(token.size(0) + signal.size(1))
+        
+        text_inputs[i] = F.pad(token, (0, pad_len), value=0)[:-1]
+        text_labels[i] = F.pad(label, (0, pad_len), value=-100)[1:]
+
+
+    targets = torch.stack(text_labels, dim=0).cuda()
+    
+    return modality_list, text_inputs, targets
+
+
+
+def cat_tts(self, tensor_list, signal_list):
+    """
+    对一个包含不同长度张量的列表进行填充，使所有张量的长度相同且为16的整数倍，并生成掩码。
+    参数:
+        tensor_list (list of torch.Tensor): 输入的张量列表，每个张量形状为 [seq_len]。
+        pad_value (int, optional): 填充值，默认值为 0。
+    返回:
+        padded_tensor (torch.Tensor): 填充后的张量，形状为 [batch_size, target_len]。
+        mask (torch.Tensor): 填充掩码，1 表示有效数据，0 表示填充部分。
+    """
+
+    modality_list = []
+    atokens = []
+    labels_list = [] #多模态拼接标签
+    #max_len = max((token.size(0) + signal.size(1)) for token, signal in zip(tensor_list, modality_list))
+    max_len = 0
+    for token, signal in zip(tensor_list, signal_list):
+        global_tokens, semantic_tokens = self.modality.world_encoder.encoder(signal)
+        # print(global_tokens.squeeze(0).squeeze(0), global_tokens.squeeze(0).squeeze(0)+8192)
+        global_tokens = global_tokens.squeeze(0).squeeze(0)+8194
+        # global_tokens = F.pad(global_tokens.squeeze(0).squeeze(0)+8194, (0, 1), value=8193)
+        semantic_tokens = F.pad(semantic_tokens.squeeze(0), (0, 1), value=8192)
+        audio_token = torch.cat([global_tokens,semantic_tokens]) 
+        mask_gt = torch.full_like(global_tokens, -100)
+        label = torch.cat([global_tokens-1,semantic_tokens-1]) 
+        # modality = self.modality.encoder(audio_token)
+
+        # if modality is False:
+        #     modality_list.append(False)
+        #     continue
+        
+        mask_t = torch.full_like(token, -100)
+        label = torch.cat([mask_t,label])
+        atokens.append(audio_token)
+        # modality_list.append(modality)
+        labels_list.append(label)
+        max_len = max(label.size(0), max_len)
+
+    # 计算目标长度（向上取整到16的整数倍）
+    target_len = ((max_len + 15) // 16 * 16)+1
+
+    if self.args.ctx_len is not None:
+        target_len = min(target_len, self.args.ctx_len+1)
+
+    text_token = []
+    labels = []
+    mod_token = []
+
+    for token, atoken, mask in zip(tensor_list, atokens, labels_list):
+
+        pad_len = target_len-(token.size(0) + atoken.size(0))
+        
+        padded_atoken = F.pad(atoken, (0, pad_len), value=8192)
+
+        atoken = padded_atoken[:-1]
+        mod = self.modality(atoken)
+        # padded_token = F.pad(signal, (0, 0, 0, pad_len), value=0)
+
+        # pad_mod = padded_token[:,:-1,:]
+
+        pad_mask = F.pad(mask, (0, pad_len), value=-100)[1:]
+        
+        mod_token.append(mod)
+
+        labels.append(pad_mask)
+
+    labels = torch.stack(labels, dim=0)
+
+    
+    return mod_token, tensor_list, labels

world/dataset.py

@@ -0,0 +1,354 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+import torch.nn.functional as F
+
+import json
+import math
+import random
+import os
+import sys
+import numpy as np
+import torch
+import lightning as L
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+from lightning_utilities.core.rank_zero import rank_zero_info
+from infer.rwkv.utils import PIPELINE
+pipeline = PIPELINE('rwkv', "rwkv_vocab_v20230424")
+from PIL import Image
+
+import pandas as pd
+import librosa
+import io
+import soundfile as sf
+# 读取parquet文件
+from torchvision import transforms
+
+
+
+transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor()  # 将图像转换为张量
+])
+
+def process_conversation_text(conversations):
+    conversation_text = f"\x16"
+
+    for conv in conversations:
+        role = conv.get('from', '').lower()
+        content = conv.get('value', '')
+        
+        if role == 'human':
+            conversation_text += f"User: {content}\x17"
+        elif role in ['assistant', 'gpt']:
+            conversation_text += f"Assistant: {content}\x17"
+
+    return conversation_text
+
+def process_tokens(conversations):
+    # conversation_text = f"\x16"
+    inputs = []
+    labels = []
+    for conv in conversations:
+        role = conv.get('from', '').lower()
+        content = conv.get('value', '')
+        
+        if role in ['human', 'user']:
+            question = f"\x16User: {content}\x17"
+            input = torch.tensor(pipeline.encode(question))
+            label = torch.full_like(input, -100)
+        elif role in ['assistant', 'gpt']:
+            answer = f"\x16Assistant: {content}\x17"
+            input= torch.tensor(pipeline.encode(answer))
+            label = input
+        inputs.append(input)
+        labels.append(label)
+    inputs =torch.cat(inputs)
+    labels =torch.cat(labels)
+    return inputs, labels
+
+def bytes_to_audio(audio_bytes):
+    with io.BytesIO(audio_bytes) as buf:
+        # 使用 soundfile 读取音频数据
+        audio_array, sr = sf.read(buf)
+        
+        # 确保是单声道
+        if len(audio_array.shape) > 1:
+            audio_array = audio_array.mean(axis=1)
+        
+        # 确保是 float32 类型
+        audio_array = audio_array.astype(np.float32)
+        
+        return {
+            'array': audio_array,
+            'sampling_rate': sr
+        }
+
+
+
+def get_data_by_l_version(trainer: L.Trainer, args):
+    if L.__version__[0] == '2':
+        train_data = MyDataModule(args)
+    else:
+        raise ValueError(f"Unsupported PyTorch Lightning version: {L.__version__}")
+    return train_data
+
+class GlobalIndexManager:
+    def __init__(self, rank=0, device_num=1, shuffle=True):
+        self.current_idx = 0
+        self.rank = rank
+        self.device_num = device_num
+        self.shuffle = shuffle
+        
+    def get_next_idx(self, idx_t):
+        if self.shuffle:
+            idx = idx_t
+        else:
+            idx = self.current_idx * self.device_num + self.rank 
+            self.current_idx += 1
+        return idx
+
+class MyDataModule(L.LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.train_data = None
+        
+        
+    def setup(self, stage=None):
+        self.train_data = MyDataset(self.args)
+        self.args.vocab_size = self.train_data.vocab_size
+        self.train_data.real_epoch = self.trainer.current_epoch
+        self.train_data.rank = self.trainer.global_rank
+        self.train_data.world_size = self.trainer.world_size
+        self.train_data.setup(self.trainer.global_rank, self.trainer.world_size, 
+                              int(self.args.devices), self.args.data_shuffle)
+        
+    def train_dataloader(self):
+        # must set shuffle=False, persistent_workers=False (because worker is in another thread)
+        return DataLoader(
+            self.train_data,
+            shuffle=self.args.data_shuffle,
+            pin_memory=True,
+            batch_size=self.args.micro_bsz,
+            num_workers=1,
+            persistent_workers=False,
+            drop_last=True
+        )
+
+class WorldDataset(Dataset):
+    def __init__(self, args, emb=None):
+        self.args = args
+        self.rank = 0
+        self.real_epoch = 0
+        self.world_size = 0
+        self.index_manager = None
+        self.emb = emb
+
+        if args.data_type =='wav':
+            import jsonlines
+
+            # 打开并读取 JSON 文件
+            #with open(f'{args.data_file}/answer.jsonl', 'r') as file:
+            with jsonlines.open(f'{args.data_file}/answer.jsonl') as file:
+                self.data = list(file)
+        elif args.data_type =='img':
+            import jsonlines
+
+            # 打开并读取 JSON 文件
+            #with open(f'{args.data_file}/answer.jsonl', 'r') as file:
+            with jsonlines.open(f'{args.data_file}/answer.jsonl') as file:
+                self.data = list(file)
+        elif args.data_type=='hf_img':
+            import jsonlines
+            # with open(f'{args.data_file}/chat.json', 'r', encoding='utf-8') as file:
+            #     self.data = json.load(file)          
+            with jsonlines.open(f'{args.data_file}/chat.jsonl') as file:
+                self.data = list(file)
+        elif args.data_type=='visual':
+            import jsonlines
+            # with open(f'{args.data_file}/chat.json', 'r', encoding='utf-8') as file:
+            #     self.data = json.load(file)          
+            with jsonlines.open(f'{args.data_file}/chat.jsonl') as file:
+                self.data = list(file)
+        elif args.data_type == 'visual-r1-cs':
+           
+            llava_path = os.path.join(args.data_file, 'vision_r1_llava_cot_full.json')
+            mulberry_path = os.path.join(args.data_file, 'vision_r1_mulberry_sft_full.json')
+            
+            import json
+            with open(f'{llava_path}', 'r', encoding='utf-8') as file:
+                llava_data = json.load(file)
+            with open(f'{mulberry_path}', 'r', encoding='utf-8') as file:
+                mulberry_data = json.load(file)
+            
+            # 合并数据集并添加来源标识
+            for item in llava_data:
+                item['_source'] = 'llava_cot'
+            for item in mulberry_data:
+                item['_source'] = 'mulberry'
+            
+            self.data = llava_data + mulberry_data
+
+        elif args.data_type =='hf' or args.data_type =='qa' or args.data_type =='cnqa' or args.data_type =='cnasr' or args.data_type =='tts':
+            from datasets import load_dataset, concatenate_datasets
+
+            def list_subdirectories(base_path):
+                return [
+                    name for name in os.listdir(base_path)
+                    if os.path.isdir(os.path.join(base_path, name)) and not name.startswith('.')
+                ]
+
+            datasets = []
+            files = list_subdirectories(args.data_file)
+            if not files:
+                datasets = load_dataset(args.data_file, split="train")
+            else:
+                for file in files:
+                    dataset = load_dataset(f'{args.data_file}/{file}', split="train")
+                    datasets.append(dataset)
+                datasets = concatenate_datasets(datasets)
+            self.data = datasets
+            print(len(datasets))
+            
+        elif args.data_type == "jsonl":
+            import jsonlines
+
+            with jsonlines.open(args.data_file) as file:
+                self.data = list(file)
+
+        else:
+            self.data = pd.read_parquet(args.data_file)
+
+        
+
+    def setup(self, rank, world_size, devices, shuffle):
+        self.rank = rank
+        self.world_size = world_size
+        self.index_manager = GlobalIndexManager(rank=rank, device_num=devices, shuffle=shuffle)
+    
+    def __len__(self):
+        return self.args.epoch_steps * self.args.micro_bsz
+
+
+    def __getitem__(self, idx):
+        idx = self.index_manager.get_next_idx(idx_t=idx) if self.index_manager else idx
+        args = self.args
+        if args.data_type =='wav':
+
+            mod_name = self.data[idx]['file_name']
+            data_answer = self.data[idx]['answer']
+            mod_path = f'{args.data_file}/{mod_name}'
+            audio, sample_rate = librosa.load(mod_path, sr=16000)  # sr=None 保持原采样率
+            #sign,_ = self.speech_encoder(audio)
+            sign = audio
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        elif args.data_type =='hf':
+            sample = self.data[idx]
+            audio = sample['audio']
+            data_answer = sample['text'] #####caption
+            audio = librosa.resample(audio['array'],orig_sr= audio['sampling_rate'],target_sr= 16000)  # sr=None 保持原采样率
+            sign = audio
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        elif args.data_type =='tts':
+            sample = self.data[idx]
+            audio = sample['audio']
+            data_answer = sample['text'] #####caption
+            audio = librosa.resample(audio['array'],orig_sr= audio['sampling_rate'],target_sr= 16000)  # sr=None 保持原采样率
+            sign = audio
+            token = torch.tensor(pipeline.encode(f'User: {data_answer}\x17Assistant:'))
+        elif args.data_type =='qa':
+            sample = self.data[idx]
+            # audio = sample['speech_cosy'][0]
+            # data_answer = sample['answer'] 
+
+            audio = sample['question_audio']
+            data_answer = sample['answer']
+            sign = librosa.resample(audio['array'],orig_sr= audio['sampling_rate'],target_sr= 16000)  # sr=None 保持原采样率
+
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        elif args.data_type =='cnqa':
+            sample = self.data[idx]
+            audio = sample['audio']
+            data_answer = sample['answer']
+            sign = librosa.resample(audio['array'],orig_sr= audio['sampling_rate'],target_sr= 16000)  # sr=None 保持原采样率
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        elif args.data_type =='cnasr':
+            sample = self.data[idx]
+            audio = sample['audio']
+            data_answer = sample['transcript']
+            sign = librosa.resample(audio['array'],orig_sr= audio['sampling_rate'],target_sr= 16000)  # sr=None 保持原采样率
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        elif args.data_type == "jsonl":
+            ctx_len = args.ctx_len
+            req_len = ctx_len + 1
+            ctx = self.data[idx]['text']
+            token = torch.tensor(pipeline.encode(ctx))
+            token_len = len(token)
+            pad_len = req_len - token_len
+        
+            dix = F.pad(token, (0, pad_len), value=0)
+            x = dix[:-1]
+            y = dix[1:]
+            mask = torch.zeros(req_len - 1)
+            mask[:token_len - 1] = 1
+            return x, y, mask
+        elif args.data_type == "img":
+
+            mod_name = self.data[idx]['file_name']
+            data_answer = self.data[idx]['answer']
+            mod_path = f'{args.data_file}/{mod_name}'
+            token = torch.tensor(pipeline.encode(f'\n\nAssistant: {data_answer}\x17'))
+            image = Image.open(mod_path).convert('RGB')
+            sign = transform(image)
+        elif args.data_type == 'visual':
+
+            img_name = self.data[idx]['image']
+            conversation_text = self.data[idx]['conversations']
+
+            mod_path = f'{args.data_file}/images/{img_name}' 
+            image = Image.open(mod_path).convert('RGB')
+            sign = image
+            text_tokens, text_labels = process_tokens(conversation_text)
+            return sign, text_tokens, text_labels
+        elif args.data_type== 'hf_img':
+            
+            img_name = self.data[idx]['image']
+            conversation_text = self.data[idx]['conversations']
+            conversation_text = process_conversation_text(conversation_text)
+
+            mod_path = f'{args.data_file}/images/{img_name}' 
+            token = torch.tensor(pipeline.encode(conversation_text)) 
+            image = Image.open(mod_path).convert('RGB')
+            sign = image
+        elif args.data_type == 'visual-r1-cs':
+            item = self.data[idx]
+            conversations = item['conversations']
+            
+            # 根据来源处理图像路径
+            if item['_source'] == 'llava_cot':
+                img_name = item['image']
+                
+            else:  
+                img_name = item['images']  
+                
+            
+            
+            mod_path = f'{self.args.data_file}/{img_name}'
+            image = Image.open(mod_path).convert('RGB')
+            sign = image  
+            
+            # 处理文本对话
+            text_tokens, text_labels = process_tokens(conversations)
+    
+            return sign, text_tokens, text_labels
+        else:
+            data_audio = bytes_to_audio(self.data['question_audio'][idx]['bytes'])
+            data_answer = self.data['answer'][idx]
+            audio = librosa.resample(data_audio['array'],orig_sr= 48000,target_sr= 16000)
+            #sign,_ = self.speech_encoder(audio)
+            sign = audio
+            token = torch.tensor(pipeline.encode(f'\x16Assistant: {data_answer}\x17'))
+        #print(idx, f'Assistant: {data_answer}\x17')
+        return sign, token

world/encoder/clip_encoder.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+
+
+
+class VisualAdapter(nn.Module):
+    """
+    2D Image to Patch Embedding
+    """
+    def __init__(self, encoder_dim, project_dim, hidden_dim=None):
+
+        super().__init__()
+        self.encoder_dim = encoder_dim
+        self.project_dim = project_dim
+        self.hidden_dim = hidden_dim
+
+        if self.hidden_dim==None:
+            self.hidden_dim = project_dim*2
+
+        self.pre_norm = nn.LayerNorm(self.project_dim)
+        self.proj = nn.Sequential(
+            nn.Linear(self.encoder_dim, self.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(self.hidden_dim, self.project_dim),
+        )
+        # self.proj = nn.Sequential(
+        #     nn.Linear(self.encoder_dim, self.hidden_dim),
+        #     nn.GELU(),
+        #     nn.Linear(self.hidden_dim, self.hidden_dim),
+        #     nn.GELU(),
+        #     nn.Linear(self.hidden_dim, self.project_dim),
+        # )
+
+    
+    def forward(self, x):        
+        x = self.proj(x)
+        return x + self.pre_norm(x)
+
+
+
+class ClipEncoder(nn.Module):
+    
+    def __init__(
+        self,
+        encoder_path,
+        project_dim,
+        train_mode="adapter",
+        device="cuda",) -> None:
+        super(ClipEncoder, self).__init__()
+
+        self.device = device
+        self.image_processor = CLIPImageProcessor.from_pretrained(encoder_path)
+        self.model = CLIPVisionModel.from_pretrained(encoder_path)
+        self.encoder_dim = self.model.config.hidden_size
+
+        self.adapter = VisualAdapter(self.encoder_dim, project_dim)
+    def forward(self, x):
+
+        x= torch.from_numpy(self.image_processor(x)['pixel_values'][0]).to(self.device,dtype=torch.bfloat16)
+
+        x = self.model(x.unsqueeze(0), output_hidden_states=True).last_hidden_state
+        
+        x = self.adapter(x)
+        
+        return x

world/encoder/siglip_encoder.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from transformers import AutoModel, SiglipImageProcessor
+
+class VisualAdapter(nn.Module):
+    """
+    2D Image to Patch Embedding
+    """
+    def __init__(self, encoder_dim, project_dim, hidden_dim=None):
+
+        super().__init__()
+        self.encoder_dim = encoder_dim
+        self.project_dim = project_dim
+        self.hidden_dim = hidden_dim
+
+        if self.hidden_dim==None:
+            self.hidden_dim = project_dim*2
+
+        self.pre_norm = nn.LayerNorm(self.project_dim)
+        self.proj = nn.Sequential(
+            nn.Linear(self.encoder_dim, self.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(self.hidden_dim, self.project_dim),
+        )
+        # self.conv = nn.Conv1d(
+        #         in_channels=encoder_dim,
+        #         out_channels=encoder_dim,
+        #         bias=False,
+        #         kernel_size=5,
+        #         stride=4
+        # )
+
+    
+    def forward(self, x):
+        # x = self.conv(x.permute(0,2,1)).permute(0,2,1)
+        x = self.proj(x)
+        return x + self.pre_norm(x)
+
+
+
+class SiglipEncoder(nn.Module):
+    
+    def __init__(
+        self,
+        encoder_path,
+        project_dim,
+        encoder_device='cpu',
+        train_mode="adapter"
+        ) -> None:
+        super(SiglipEncoder, self).__init__()
+
+        self.device = encoder_device
+        
+        self.model = AutoModel.from_pretrained(encoder_path).vision_model
+        self.image_processor = SiglipImageProcessor.from_pretrained(encoder_path)
+        self.encoder_dim = 768  #self.model.config.hidden_size
+
+        self.adapter = VisualAdapter(self.encoder_dim, project_dim)
+    def forward(self, x):
+
+        x= torch.from_numpy(self.image_processor(x)['pixel_values'][0]).to(self.device,dtype=torch.bfloat16)
+        x = self.model(x.unsqueeze(0), output_hidden_states=True).last_hidden_state
+        x = self.adapter(x)
+        
+        return x

world/encoder/speech_encoder.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+import numpy as np
+
+from transformers import AutoProcessor, AutoModel
+
+
+
+# class SpeechAdapter(nn.Module):
+#     def __init__(self, input_dim, output_dim):
+#         super(SpeechAdapter, self).__init__()
+#         self.conv = nn.Conv1d(in_channels=input_dim, out_channels=3072, kernel_size=3, stride=2)
+#         self.transformer = nn.TransformerEncoderLayer(d_model=3072, nhead=8, dim_feedforward=4096)
+#         self.linear = nn.Linear(3072, output_dim)
+#     def forward(self, x):
+#         # if x.size(1)<5 or x.size(1)>5000:
+#         #     return False
+#         # x shape: (batch_size, seq_len, input_dim)
+#         x = x.permute(0, 2, 1)
+#         # x shape: (batch_size, input_dim, seq_len)
+#         x = self.conv(x)
+#         # x shape after conv: (batch_size, input_dim, new_seq_len)
+#         x = x.permute(2, 0, 1)  # Transformer expects (seq_len, batch_size, input_dim)
+#         # x = self.transformer(x, src_key_padding_mask=mask.bool())
+#         x = self.transformer(x)
+#         x = x.permute(1, 0, 2)  # Back to (batch_size, seq_len, input_dim)
+#         x = self.linear(x)
+#         return x
+
+class SpeechAdapter(nn.Module):
+    def __init__(self, encoder_dim, project_dim, hidden_dim=None):
+        super(SpeechAdapter, self).__init__()
+        self.encoder_dim = encoder_dim
+        self.project_dim = project_dim
+        self.hidden_dim = hidden_dim
+
+        if self.hidden_dim==None:
+            self.hidden_dim = project_dim*2
+        self.conv = nn.Conv1d(in_channels=self.encoder_dim , out_channels=self.hidden_dim, kernel_size=3, stride=2, padding=2)
+        self.proj = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(self.hidden_dim, self.project_dim),
+        )
+    def forward(self, x):
+        # if x.size(1)<5 or x.size(1)>5000:
+        #     return False
+        
+        # x shape: (batch_size, seq_len, input_dim)
+        x = x.permute(0, 2, 1)
+        # x shape: (batch_size, input_dim, seq_len)
+        x = self.conv(x).permute(0, 2, 1)
+        # x shape after conv: (batch_size, input_dim, new_seq_len)
+        x = self.proj(x)
+        if x.size(1)>1023:
+            return False
+        return x
+
+class SpeechEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder_path,
+        project_dim,
+        train_mode="adapter",
+        device="cuda",
+    ):
+        assert train_mode in ["adapter", "full"]
+        super(SpeechEncoder, self).__init__()
+
+        self.device = device
+        
+        try:
+            self.processor = AutoProcessor.from_pretrained(encoder_path)
+        except:
+            self.processor = AutoProcessor.from_pretrained("facebook/hubert-large-ls960-ft")
+        
+        self.time_reduction_factor = int(
+            self.processor.feature_extractor.sampling_rate / 50
+        )
+        self.padding_length = 320
+        
+        self.model = AutoModel.from_pretrained(encoder_path)
+        self.model.eval()
+        self.model_output_dim = self.model.config.hidden_size
+        self.project_dim = project_dim
+            
+        self.project_dim = project_dim
+        self.adapter = SpeechAdapter(self.model_output_dim, self.project_dim).to(self.device,dtype=torch.bfloat16)
+    #     self.set_gradient(train_mode)
+        
+
+   
+
+    def forward(self, x):
+        input_dict = self.processor(
+            x, return_tensors="pt", padding=True, sampling_rate=16000
+        ).to(self.device,dtype=torch.bfloat16)
+        
+        # encoder only
+        x = self.model(**input_dict).last_hidden_state
+
+        # stf encoder
+        # x = self.model(**input_dict, output_hidden_states=True).hidden_states[-1]
+        
+        x= self.adapter(x)#x:(B,T,hidden dim)
+        # mask = torch.ones(x.shape[0],x.shape[1]).to(self.device,dtype=torch.bfloat16)
+        return x

world/encoder/visual_encoder.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from diffusers import AutoencoderKL
+
+
+class Patch(nn.Module):
+    def __init__(self, Imgsize=64, Patchsize=16) -> None:
+        super(Patch, self).__init__()
+        self.Patchsize = Patchsize
+        self.Imgsize = Imgsize
+    def encoder(self, x):
+        assert x.size(3)==self.Imgsize
+        imgsize = self.Imgsize
+        patchsize = self.Patchsize
+        x = x.unfold(2, patchsize, patchsize).unfold(3, patchsize, patchsize)
+        x = x.contiguous().reshape(x.size(0), x.size(1), int(pow(imgsize/patchsize, 2)), -1)
+        x = x.transpose(1, 2)
+        x = x.reshape(x.size(0), x.size(1), -1)
+        return x
+
+    def decoder(self, x):
+        imgsize = self.Imgsize
+        patchsize = self.Patchsize
+        x = x.reshape(x.size(0), x.size(1), x.size(2)//patchsize, patchsize)
+        x = x.transpose(1, 2)
+        x = x.unfold(2,imgsize//patchsize,imgsize//patchsize).unfold(3, patchsize, patchsize)
+        x = x.reshape(x.size(0), 4, imgsize, imgsize)
+        return x
+
+class SD_Auto():
+    def __init__(self, path="sdxl", input_dtype=torch.float32) -> None:
+        self.autoencoder = AutoencoderKL.from_pretrained(path, subfolder="vae")
+        self.autoencoder = self.autoencoder.to('cuda', input_dtype)
+
+    def encoder(self, x):
+        with torch.no_grad():
+            x = self.autoencoder.encode(x).latent_dist.sample()
+        return x
+    
+    def decoder(self, x):
+
+        with torch.no_grad():
+            x = self.autoencoder.decode(x).sample
+        #x = self.autoencoder.decode(x).sample
+        return x
+    
+def kld_loss(mu, logvar):
+    KLD = - 0.5 * torch.sum(1 + logvar - mu.pow(2) -
+                            logvar.exp()) / mu.shape[0]
+    return KLD
+
+
+class VisualAdapter(nn.Module):
+    """
+    2D Image to Patch Embedding
+    """
+    def __init__(self, img_size=512//8, patch_size=16, in_c=4, text_dim=2560, head_size=64):
+
+        super().__init__()
+        self.head_size = head_size
+        # self.img_receptance = nn.Linear((patch_size*patch_size*in_c), text_dim, bias=False)
+        # self.img_key = nn.Linear((patch_size*patch_size*in_c), text_dim, bias=False)
+        # self.img_value = nn.Linear((patch_size*patch_size*in_c), text_dim, bias=False)
+        self.linear = nn.Linear((patch_size*patch_size*in_c), text_dim, bias=False)
+        self.patch = Patch(Imgsize=img_size, Patchsize=patch_size)
+
+    
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.patch.encoder(x)
+        # r = self.img_receptance(x)
+        # k = self.img_key(x)
+        # v = self.img_value(x)
+        # r = r.view(*x.shape[:2], -1, self.head_size).transpose(1, 2)
+        # k = k.view(*x.shape[:2], -1, self.head_size).transpose(1, 2)
+        # v = v.view(*x.shape[:2], -1, self.head_size).transpose(1, 2)
+        # x_img = torch.nn.functional.scaled_dot_product_attention(
+        #     r, k, v, is_causal=True, scale=1 / self.head_size
+        # )
+        # x = x_img.transpose(1, 2).reshape(*x.shape[:2], -1)
+        return self.linear(x)
+
+
+
+class VisualEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder_path,
+        project_dim,
+        train_mode="adapter",
+        device="cuda",) -> None:
+        super(VisualEncoder, self).__init__()
+
+        self.model = AutoencoderKL.from_pretrained(path, subfolder="vae", allow_pickle=False).to('cuda', input_dtype)
+        self.adapter = VisualAdapter(text_dim=llm_dim).to('cuda',dtype=torch.bfloat16)
+
+    def forward(self, x):
+        x = self.model.encode(x.unsqueeze(0)).latent_dist.sample()
+        #print(x.view(-1))
+        x = self.adapter(x)
+        #print(x.view(-1))
+        
+        return x

world/encoder/whisper_encoder.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+
+class SpeechAdapter(nn.Module):
+    def __init__(self, encoder_dim, project_dim, hidden_dim=None):
+        super(SpeechAdapter, self).__init__()
+        self.encoder_dim = encoder_dim
+        self.project_dim = project_dim
+        self.hidden_dim = hidden_dim
+
+        if self.hidden_dim==None:
+            self.hidden_dim = project_dim*2
+        self.conv = nn.Conv1d(in_channels=self.encoder_dim , out_channels=self.hidden_dim, kernel_size=3, stride=2, padding=2)
+        self.proj = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(self.hidden_dim, self.project_dim),
+        )
+    def forward(self, x):
+        # if x.size(1)<5 or x.size(1)>5000:
+        #     return False
+        
+        # x shape: (batch_size, seq_len, input_dim)
+        x = x.permute(0, 2, 1)
+        # x shape: (batch_size, input_dim, seq_len)
+        x = self.conv(x).permute(0, 2, 1)
+        # x shape after conv: (batch_size, input_dim, new_seq_len)
+        x = self.proj(x)
+        return x
+
+
+
+class WhisperEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder_path,
+        project_dim,
+        train_mode="adapter",
+        device="cuda",
+    ):
+        assert train_mode in ["adapter", "full"]
+        super(WhisperEncoder, self).__init__()
+        self.device = device
+        self.processor = WhisperProcessor.from_pretrained(encoder_path)
+
+        self.model = WhisperForConditionalGeneration.from_pretrained(encoder_path).model.encoder
+
+        self.model_output_dim = self.model.config.d_model
+            
+        self.project_dim = project_dim
+        self.adapter = SpeechAdapter(self.model_output_dim, self.project_dim)
+
+    def forward(self, x):
+        input_dict = self.processor(
+            x, return_tensors="pt", sampling_rate=16000, return_attention_mask=True
+        ).to(self.device,dtype=torch.bfloat16)
+
+        chunk = torch.sum(input_dict['attention_mask'], dim=-1)//2+1
+        
+        x = self.model(**input_dict).last_hidden_state
+        x = x[:,:chunk,:]
+        x= self.adapter(x)#x:(B,T,hidden dim)
+        
+        return x

world/loss.py

@@ -0,0 +1,16 @@
+import torch
+class L2Wrap(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, loss, y):
+        ctx.save_for_backward(y)
+        return loss
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        y = ctx.saved_tensors[0]
+        # to encourage the logits to be close to 0
+        factor = 1e-4 / (y.shape[0] * y.shape[1])
+        maxx, ids = torch.max(y, -1, keepdim=True)
+        gy = torch.zeros_like(y)
+        gy.scatter_(-1, ids, maxx * factor)
+        return (grad_output, gy)

world/model.py

@@ -0,0 +1,317 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
+from torch.profiler import profile, record_function, ProfilerActivity
+#from adam_mini import Adam_mini
+
+import os, math, gc, importlib
+import torch
+
+import torch.nn as nn
+from torch.nn import functional as F
+import lightning as pl
+from lightning.pytorch.strategies import DeepSpeedStrategy
+if importlib.util.find_spec('deepspeed'):
+    import deepspeed
+    from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
+    
+from .block import Block
+from .loss import L2Wrap
+from .cat import mod_pad_text
+
+from rwkv.utils import PIPELINE
+pipeline = PIPELINE('rwkv6', "rwkv_vocab_v20230424")
+
+class RWKV(pl.LightningModule):
+    def __init__(self, args, modality=None):
+        super().__init__()
+        self.args = args
+        if not hasattr(args, 'dim_att'):
+            args.dim_att = args.n_embd
+        if not hasattr(args, 'dim_ffn'):
+            args.dim_ffn = args.n_embd * 4
+
+        assert args.n_embd % 32 == 0
+        assert args.dim_att % 32 == 0
+        assert args.dim_ffn % 32 == 0
+
+        self.emb = nn.Embedding(args.vocab_size, args.n_embd)
+
+        self.blocks = nn.ModuleList([Block(args, i) for i in range(args.n_layer)])
+
+        self.ln_out = nn.LayerNorm(args.n_embd)
+        self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False)
+
+        self.modality = modality
+
+
+    def pad_mod(self, tensor_list, signal_list):
+        """
+        对一个包含不同长度张量的列表进行填充，使所有张量的长度相同且为16的整数倍，并生成掩码。
+        参数:
+            tensor_list (list of torch.Tensor): 输入的张量列表，每个张量形状为 [seq_len]。
+            pad_value (int, optional): 填充值，默认值为 0。
+        返回:
+            padded_tensor (torch.Tensor): 填充后的张量，形状为 [batch_size, target_len]。
+            mask (torch.Tensor): 填充掩码，1 表示有效数据，0 表示填充部分。
+        """
+
+        modality_list = []
+        #max_len = max((token.size(0) + signal.size(1)) for token, signal in zip(tensor_list, modality_list))
+        max_len = 0
+        for token, signal in zip(tensor_list, signal_list):
+
+            modality = self.modality(signal)
+            if modality is False:
+                modality_list.append(False)
+                continue
+            modality_list.append(modality)
+            max_len = max(token.size(0) + modality.size(1), max_len)
+
+        # 计算目标长度（向上取整到16的整数倍）
+        target_len = ((max_len + 15) // 16 * 16)+1
+
+        if self.args.ctx_len is not None:
+            target_len = min(target_len, self.args.ctx_len+1)
+
+        masks = torch.zeros((len(tensor_list), target_len-1), dtype=torch.int32)
+        x = []
+        y = []
+        s = []
+        m = []
+        for token, signal, mask in zip(tensor_list, modality_list, masks):
+            if signal is False:
+                continue
+            pad_len = target_len-(token.size(0) + signal.size(1))
+            
+            padded_token = F.pad(token, (0, pad_len), value=0)
+
+            x_token = padded_token[:-1]
+            y_token = F.pad(padded_token, (signal.size(1)-1, 0), value=0)
+
+            mask[signal.size(1) : -pad_len] = 1 
+            
+            s.append(signal)
+            x.append(x_token)
+            y.append(y_token)
+            m.append(mask)
+
+        y = torch.stack(y, dim=0)
+        m = torch.stack(m, dim=0).cuda()
+       
+        return s, x, y, m
+
+
+    def forward(self, idx, signs= None):
+        args = self.args
+        #B, T = idx.size()
+        # assert T <= args.ctx_len, "Cannot forward, model ctx_len is exhausted."
+
+        x_list = []
+        if signs!=None:
+            for token, sign in zip(idx, signs):
+                sign_emb = sign#self.adapter(sign)
+                x_emb = self.emb(token)
+            # #print(sign_emb.shape, x.shape)
+                x_list.append(torch.cat([sign_emb.squeeze(0),x_emb], dim=0))
+            x = torch.stack(x_list, dim=0)
+        else:
+            x = self.emb(idx)
+
+        v_first = torch.empty_like(x)
+        for block in self.blocks:
+            if args.grad_cp == 1:
+                if args.state_tune or args.train_type == 'state' or args.peft !='none':
+                    x, v_first = torch_checkpoint(block, x, v_first ,use_reentrant=False)
+                else:
+                    x, v_first = deepspeed.checkpointing.checkpoint(block, x, v_first)
+            else:
+                x, v_first = block(x, v_first)
+
+        x = self.ln_out(x)
+        x = self.head(x)
+
+        return x
+
+    def training_step(self, batch, batch_idx):
+        args = self.args
+        if args.data_type == "jsonl": ########test
+            idx, targets, mask = batch
+            
+            mask = mask.view(-1)
+            sum_mask = torch.sum(mask).item()
+            logits = self(idx)
+
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction='none')
+            loss = torch.sum(loss * mask) / sum_mask
+            return loss
+        
+        if args.data_type in ["visual", "visual-r1-cs"]: ########test
+            signs, text_tokens, text_labels = batch
+            sign, idx, targets = mod_pad_text(self, signs, text_tokens, text_labels)
+
+            logits = self(idx,sign)
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+
+            return loss
+        
+        signs, tokens = batch
+        sign, idx, targets, mask = self.pad_mod(tokens, signs)
+
+        mask = mask.view(-1)
+        sum_mask = torch.sum(mask).item()
+        logits = self(idx,sign)
+
+        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction='none')
+        loss = torch.sum(loss * mask) / sum_mask
+    
+        return loss
+
+
+    def configure_optimizers(self):
+        args = self.args
+        
+        lr_decay = set()
+        lr_1x = set()
+        lr_2x = set()
+        lr_3x = set()
+        for n, p in self.named_parameters():
+            if not p.requires_grad:
+                continue
+            if (("_w1" in n) or ("_w2" in n)) and (args.layerwise_lr > 0):
+                lr_1x.add(n)
+            elif (("time_mix" in n) or ("time_maa" in n)) and (args.layerwise_lr > 0):
+                if args.my_pile_stage == 2:
+                    lr_2x.add(n)
+                else:
+                    lr_1x.add(n)
+            elif (("time_decay" in n) or ("time_daaaa" in n)) and (args.layerwise_lr > 0):
+                if args.my_pile_stage == 2:
+                    lr_3x.add(n)
+                else:
+                    lr_2x.add(n)
+            elif ("time_faaaa" in n) and (args.layerwise_lr > 0):
+                if args.my_pile_stage == 2:
+                    lr_2x.add(n)
+                else:
+                    lr_1x.add(n)
+            elif ("time_first" in n) and (args.layerwise_lr > 0):
+                lr_3x.add(n)
+            elif (len(p.squeeze().shape) >= 2) and (args.weight_decay > 0):
+                lr_decay.add(n)
+            else:
+                lr_1x.add(n)
+
+        lr_decay = sorted(list(lr_decay))
+        lr_1x = sorted(list(lr_1x))
+        lr_2x = sorted(list(lr_2x))
+        lr_3x = sorted(list(lr_3x))
+
+        param_dict = {n: p for n, p in self.named_parameters()}
+        
+        if args.layerwise_lr > 0:
+            if args.my_pile_stage == 2:
+                optim_groups = [
+                    {"params": [param_dict[n] for n in lr_1x], "weight_decay": 0.0, "my_lr_scale": 1.0},
+                    {"params": [param_dict[n] for n in lr_2x], "weight_decay": 0.0, "my_lr_scale": 5.0},# test: 2e-3 / args.lr_init},
+                    {"params": [param_dict[n] for n in lr_3x], "weight_decay": 0.0, "my_lr_scale": 5.0},# test: 3e-3 / args.lr_init},
+                ]
+            else:
+                optim_groups = [
+                    {"params": [param_dict[n] for n in lr_1x], "weight_decay": 0.0, "my_lr_scale": 1.0},
+                    {"params": [param_dict[n] for n in lr_2x], "weight_decay": 0.0, "my_lr_scale": 2.0},
+                    {"params": [param_dict[n] for n in lr_3x], "weight_decay": 0.0, "my_lr_scale": 3.0},
+                ]
+        else:
+            optim_groups = [{"params": [param_dict[n] for n in lr_1x], "weight_decay": 0.0, "my_lr_scale": 1.0}]
+
+        if args.weight_decay > 0:
+            optim_groups += [{"params": [param_dict[n] for n in lr_decay], "weight_decay": args.weight_decay, "my_lr_scale": 1.0}]
+            
+            if self.deepspeed_offload:
+                return DeepSpeedCPUAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adamw_mode=True, amsgrad=False)
+            return FusedAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adam_w_mode=True, amsgrad=False)
+        else:
+            
+            if self.deepspeed_offload:
+                return DeepSpeedCPUAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adamw_mode=False, weight_decay=0, amsgrad=False)
+            return FusedAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adam_w_mode=False, weight_decay=0, amsgrad=False)
+        # return ZeroOneAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, weight_decay=0, amsgrad=False, cuda_aware=False)
+
+    @property
+    def deepspeed_offload(self) -> bool:
+        strategy = self.trainer.strategy
+        if isinstance(strategy, DeepSpeedStrategy):
+            cfg = strategy.config["zero_optimization"]
+            return cfg.get("offload_optimizer") or cfg.get("offload_param")
+        return False
+
+    def generate_init_weight(self):
+        print(
+            f"""
+############################################################################
+#
+# Init model weight (slow for large models)...
+#
+############################################################################
+"""
+        )
+        m = {}
+        for n in self.state_dict():
+            p = self.state_dict()[n]
+            shape = p.shape
+
+            gain = 1.0
+            scale = 1.0
+            if "ln_" in n or ".ln" in n or "time_" in n or "_mask" in n or "pos_emb" in n or '.mask.' in n:
+                if 'ln_x.weight' in n:
+                    layer_scale = (1+int(n.split('.')[1])) / self.args.n_layer
+                    m[n] = (p * 0.0) + (layer_scale ** 0.7)
+                else:
+                    m[n] = p
+            else:
+                if n == "emb.weight":
+                    scale = -1 * self.args.lr_init
+                else:
+                    if shape[0] > shape[1]:
+                        gain = math.sqrt(shape[0] / shape[1])
+
+                    zero = [".att.output.", ".ffn.value.", ".ffn.receptance.", ".ffnPre.value.", ".ffnPre.receptance.", "head_q.", '.oo.', '.rr.']
+
+                    for kk in zero:
+                        if kk in n:
+                            scale = 0
+                    if n == "head.weight":
+                        scale = 0.5
+                    if "head_k." in n:
+                        scale = 0.1
+                    if "head_q." in n:
+                        scale = 0
+
+                print(f"{str(shape[0]).ljust(5)} {str(shape[1]).ljust(5)} {str(scale).ljust(4)} {n}")
+
+                if self.args.accelerator.upper() == "GPU":
+                    m[n] = torch.empty((shape[0], shape[1]), device="cuda")
+                else:
+                    m[n] = torch.empty((shape[0], shape[1]))
+
+                if scale == 0:
+                    nn.init.zeros_(m[n])
+                elif scale < 0:
+                    nn.init.uniform_(m[n], a=scale, b=-scale)
+                else:
+                    nn.init.orthogonal_(m[n], gain=gain * scale)
+
+            m[n] = m[n].cpu()
+            if os.environ["RWKV_FLOAT_MODE"] == "fp16":
+                m[n] = m[n].half()
+            elif os.environ["RWKV_FLOAT_MODE"] == "bf16":
+                m[n] = m[n].bfloat16()
+
+            # if n == "emb.weight":
+            #     print(m[n])
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        return m

world/world_encoder.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from world.encoder.speech_encoder import SpeechEncoder
+# from world.encoder.visual_encoder import VisualEncoder
+from world.encoder.whisper_encoder import WhisperEncoder
+from world.encoder.clip_encoder import ClipEncoder
+from world.encoder.siglip_encoder import SiglipEncoder
+
+
+class WorldEncoder(nn.Module):
+    def __init__(self, encoder_type: str, **kwargs):
+        super().__init__()
+        self.world_encoder = self._build_encoder(encoder_type, kwargs)
+        
+    def _build_encoder(self, encoder_type: str, config: dict):
+        encoder_map = {
+            "clip": ClipEncoder,
+            "whisper": WhisperEncoder,
+            # "visual": VisualEncoder,
+            "speech": SpeechEncoder,
+            "siglip": SiglipEncoder
+        }
+        
+        if encoder_type not in encoder_map:
+            raise ValueError(f"Unsupported encoder type: {encoder_type}")
+            
+        # 动态过滤有效参数
+        encoder_class = encoder_map[encoder_type]
+        
+        return encoder_class(**config)
+    
+    def forward(self, x):
+        return self.world_encoder(x)
+    
+    def load_checkpoint(self, state_dict):
+        self.world_encoder.load_state_dict(state_dict, strict=False)

world/world_load.py

@@ -0,0 +1,30 @@
+from world.model import RWKV
+from world.world_encoder import WorldEncoder
+def WorldLoading(args):
+    config = {
+        'encoder_type': args.encoder_type,
+        'encoder_path': args.encoder_path,
+        'project_dim' : args.n_embd
+        }
+    modality = WorldEncoder(**config)
+    
+    model = RWKV(args, modality=modality)
+    #model = RWKV(args)
+    print(model)
+
+    if 'moda' not in args.train_step:
+        for param in model.modality.world_encoder.model.parameters():
+            param.requires_grad = False
+    if 'adapter' not in args.train_step:
+        for param in model.modality.world_encoder.adapter.parameters():
+            param.requires_grad = False
+    if 'rwkv' not in args.train_step:
+        for param in model.emb.parameters():
+            param.requires_grad = False
+        for param in model.blocks.parameters():
+            param.requires_grad = False
+        for param in model.ln_out.parameters():
+            param.requires_grad = False
+        for param in model.head.parameters():
+            param.requires_grad = False
+    return model