petals环境搭建

首先挂上colab网址:这里.
果不其然刚开始就出现了问题.

inputs = tokenizer('A cat in French is "', return_tensors="pt")["input_ids"].cuda()
outputs = model.generate(inputs, max_new_tokens=3)
print(tokenizer.decode(outputs[0]))

报错是:RuntimeError: at::cuda::blas::gemm: not implemented for N3c108BFloat16E.简单来说,就是torch.bfloat16类型的矩阵乘在cublas中还没实现...
不过后来经过我的仔细搜索,在CUDA官方文档的这里发现可能是硬件的问题.如图:

即只对8.x及以上CC的设备提供支持.在后面的具体说明中也看到了这一点.

嗯...看来只能换设备了.

另外提供一个检查cublas版本的方法:cat /usr/local/cuda/include/cublas_api.h | grep CUBLAS_VER.相信聪明的你肯定知道输出的内容是什么意思!