state-spaces / mamba

Mamba SSM architecture
Apache License 2.0
12.66k stars 1.06k forks source link

Error when using FP16 or Mixed precision #474

Open Vaccummer opened 2 months ago

Vaccummer commented 2 months ago

When I train the model of dtype fp16 or mixed precision on V100 sxm2 16G, loss.backward() return error:

python: /project/lib/Analysis/Allocation.cpp:40: std::pair<llvm::SmallVector, llvm::SmallVector > mlir::triton::getCvtOrder(mlir::Attribute, mlir::Attribute): Assertion `!(srcMmaLayout && dstMmaLayout) && "Unexpected mma -> mma layout conversion"' failed.

tridao commented 2 months ago

Triton doens't support V100 very well

Vaccummer commented 2 months ago

Triton doens't support V100 very well

Thanks for reply, and is there any solution or alternative of Triton?

tridao commented 2 months ago

There's a reference implementation in pytorch but would probably be quite a bit slower