RWKV-LM icon indicating copy to clipboard operation
RWKV-LM copied to clipboard

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

Open Youngluc opened this issue 5 months ago • 2 comments

感谢您的工作,我借鉴您的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训练是正常的...

Youngluc avatar Jan 29 '24 09:01 Youngluc

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

BlinkDL avatar Jan 29 '24 15:01 BlinkDL

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

got it

Youngluc avatar Jan 30 '24 02:01 Youngluc