GPU out of memory

### Summary

Regardless of the configuration I attempt, training fail due to out of memory error.

### Uni-Mol Version

Uni-Mol Tools with model version 2

### Details

Hello everyone,

I have been trying to use **uni-mol v2** in a workstation with 4 GPU cards. Each card is a Nvidia GeForce RTX 2080 Ti with **11GB RAM** but every time I try to train the model on a dataset of 8k compounds with batch size 16 and the 84m model, I get an out of memory error. Despite the default is to use all the GPU's, the model is only trying to use GPU 0. My experiments are using the `unimol_tools` python package.

What kind of setup is it necessary to run your model?

Is there a way to actually make the the package to use all GPUs?

I would appreciate your help.

Start with the unimol_tools training entry point and its GPU or distributed-training configuration; reproduce the reported setup with four RTX 2080 Ti cards, an 84m model, batch size 16, and the 8k-compound dataset. Done means identifying why training uses only GPU 0 and documenting or fixing a supported multi-GPU setup that avoids the reported out-of-memory failure.

Written by the indexing model from the issue text.