请问huggingface transformers的库实现的RWKV是否有些问题？我在backward的时候总是出现问题。

[Youngluc]

感谢您的工作，我借鉴您的RWKV结构在尝试实现一种多模态的VLM。 这里我令RWKV充当了类似于Qformer一样的结构，也就是ViT->RWKV->LLM这样的结构，然后按照pretrain&SFT两个阶段做训练，使用了DeepSpeed。

但是在反向传播的过程中出现了许多问题，主要有两个。

1. 基于custom cuda kernel进行forward和backward的时候，pretrain阶段没有问题（冻结ViT和LLM，只训练位于projector layer的RWKV），但是在SFT阶段解冻LLM后，会一直报错CUDA ERROR: an illegal memory access was encountered，出错的位置不固定，但是可以确定这个错误是在反向传播过程中出现的，如果我注释掉rwkv模块则不会有问题（只保留线性层转换维度进行forward）；
2. 不使用custom cuda kernel，使用rwkv_linear_attention_cpu函数进行forward（虽然这个函数是为CPU执行实现的，但是我的理解这个函数实际上实现了rwkv的核心机制的运算，而且只要key的device是cuda，那么其实这些运算还是在GPU上进行的）。但利用这个函数的问题在于，在batch的样本forward完毕后，backward过程会无限等待timeout（多卡情形下才会卡死，我怀疑是不是多卡的梯度聚合有问题，单卡的话利用这个函数是可以正常训练的）。

pengbo大佬可以给一些反馈和潜在的分析吗？LLM本身只有1.8B，batchsize也比较小，我也监测过，A10080G显卡，显存并没有超限。

P.S.: 对于基于cutsom cuda kernel的forward和backward，还有小概率出现这种错误： File python3.10/site-packages/torch/autograd/init.py", line 200, in backward Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16BF, lda, b, CUDA_R_16BF, ldb, &fbeta, c, CUDA_R_16BF, ldc, CUDA_R_32F, 大部分情况都是非法显存访问 此外利用CPU训练是正常的...

[BlinkDL]

可以先试不需要cuda的简化版rwkv5： https://github.com/BlinkDL/nanoRWKV

[Youngluc]

可以先试不需要cuda的简化版rwkv5： https://github.com/BlinkDL/nanoRWKV

got it