Support for newer torch versions

[atonderski]

Hi!

Thanks for this very cool work :)

I noticed that the environment requires torch > 2.4, and when i try running with newer torch it predictably fails. Is there a fundamental issue or can i work around it with some effort? I'm on a 5090, which requires cuda 12.8 for which the oldest torch version is 2.7.

Appreciate any advice before I dive into resolving this. Happy to upstream if I get it to work.

Cheers,

[gene-chou]

Hey it shouldn't be too hard but the two packages you need to look out for are Mamba2 (specifically, its compatibility with torch.compile) and flash-attn. I think Mamba2 might have gotten an update to work with torch 2.7. Not sure about flash-attn.

[ndethe]

@atonderski Hi I have the resolution for this. I modified the setup_env.sh file.

python -m pip install -U pip setuptools wheel packaging build
pip install torch==2.7 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
pip install -U xformers --index-url https://download.pytorch.org/whl/cu126

pip install ipdb tqdm wandb
pip install matplotlib einops scipy h5py OpenEXR
pip install hydra-core
pip install opencv-python pillow

pip install flash-attn --no-build-isolation


cd mamba
export MAMBA_FORCE_BUILD=TRUE
python -m pip install --no-build-isolation .
cd ..

sudo apt install -y ffmpeg

this was tested on an nvidia A6000 and A100

[ndethe]

xformers is the trouble maker. removing xformers from the setup or combining it with the cuda version you have leads to no errors in building mamba2 or flash-attn.