Slow Training Speed

[adrianjoshua-strutt]

I'm experiencing extremely slow training times while training BiFormer Small on a custom dataset with an image size of 224x244. The training takes up to a minute per batch with a batch size of 32 on an RTX 3070 GPU. I'm unsure if this is an expected training speed or if there might be an issue with my training setup implementation.

I have verified that the GPU is being utilized during training. Other models have not shown similar slowdowns on the same hardware and dataset.

Any guidance or suggestions to improve the training speed would be greatly appreciated!

Thank you!

[LiJichen0114]

Hello, could you please let me know how to train this model on a single GPU without using Slurm? What parameters do I need to modify?

[gtebbutt]

For anyone else running into performance issues: the main bottleneck seems to be the backward pass on torch.gather. If you're seeing extremely slow training, try setting torch.use_deterministic_algorithms(False) (this should be the default, but it's always possible another library or a setting in your environment has changed it) - the standard torch.gather implementation is nondeterministic, and the deterministic version triggers a CUDA sync event that seems to slow down training by a factor of 5x to 10x.

That said, even the standard nondeterministic version seems to be somewhat slow in training, and obviously if you need determinism in your experiments that's no use to you anyway, so the ideal is going to be a proper hardware accelerated kernel specifically for bi-level routing attention.