search

lucidrains / ring-attention-pytorch

Implementation of 💍 Ring Attention, from Liu et al. at Berkeley AI, in Pytorch
MIT License
474 stars 27 forks source link

RingAttention exception from Triton #17

Closed GindaChen closed 2 months ago

GindaChen commented 2 months ago

I was trying to use the assert_attn.py script on CUDA. I turn on the use_cuda flag, and got the following exception:

python: /project/include/triton/Conversion/TritonGPUToLLVM/Utility.h:379: Value mlir::LLVM::getSharedMemoryBase(Location, RewriterBase &, Operation *): Assertion `op->hasAttr("allocation.offset")' failed.
  #0 0x00007f9efbd15088 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/azureuser/ring-attention-pytorch/env2/lib/python3.10/site-packages/triton/_C/libtriton.so+0x6b33088)
  #1 0x00007f9ef8251880 triton_stacktrace_signal_handler(void*) /project/python/src/llvm.cc:424:3
  #2 0x00007f9efbd12bac llvm::sys::RunSignalHandlers() (/home/azureuser/ring-attention-pytorch/env2/lib/python3.10/site-packages/triton/_C/libtriton.so+0x6b30bac)
  #3 0x00007f9efbd1573d SignalHandler(int) (/home/azureuser/ring-attention-pytorch/env2/lib/python3.10/site-packages/triton/_C/libtriton.so+0x6b3373d)
...
#9 0x00007f9ef80826a5 mlir::Operation::hasAttr(llvm::StringRef) /root/.triton/llvm/llvm-657ec732-almalinux-x64/include/mlir/IR/Operation.h:567:12
...

Environment setup. Here's how I setup my environment:

mamba create -p ./env python=3.10
mamba install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia
pip install -e .
pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly

Full error log. error.txt.

It looks like the error is from triton, so maybe related to the way I install it in my environment. But still got no clue at this point what had failed.

lucidrains commented 2 months ago

@GindaChen could you try on the triton-nightly?

$ pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly

just ran it for triton-nightly-3.0.0.post20240716052845 and it works

edit: oh strange, you are doing that already, i'm not sure

GindaChen commented 2 months ago

Thanks! I just tried reverted to triton 3.0.0 and everything seems to work again. 

pip install 'triton==3.0.0'

What is the reason we need to run on triton-nightly (since nightly changes a little more frequently)?

lucidrains commented 2 months ago

@GindaChen at some point I couldn't get things working without it being nightly haha, but this was a while back

lucidrains commented 2 months ago

it is strange as I had just installed nightly, what type of Nvidia gpu are you using?

GindaChen commented 2 months ago

I'm using H100. In retrospection, it may just be a very temporary issue in different "nightly" build...

Mark as resolved for now. Thanks for getting back to me @lucidrains !