Training script

[nviolante25]

From #42

[kvuong2711]

Hi @nviolante25,

Thanks for the PR. I'm trying to test your PR out, but kept running into issues with float16/float32, therefore the attention function crashed out. I'm running python main.py --base configs/example_training/seva-clipl_dl3dv.yaml. I'm also happy to contribute to the PR if you have some clues on where to fix potentially fix this.

[nviolante25]

Hi @kvuong2711
I'm having the same issue, the only solution so far was to disable with sdpa_kernel(SDPBackend.FLASH_ATTENTION) . Not sure that is causing it to fail in the first place. I thought it could be the torch version but unfortunately changing it didn't work.

[jensenz-sai]

@kvuong2711 @nviolante25,

Yup, it's a bit tricky to set up flash attention for mixed-precision training. You could replace it with xformer's attention here implemented in generative-models codebase.