Out of vram and reboot

[tdzz1102]

Machine info

• Hardware
  • 64GB RAM
  • 32 vcpus
  • 2*3090 GPUs
• Software
  • ubuntu server 20.04
  • python 3.11.2
  • nvidia-driver 545
  • cuda 12.3
  • jax 0.4.19

When I set up the environment and called FlaxWhisperPipline('openai/whisper-xxx')method to load the model, the server rebooted without any error. Only for 'openai/whisper-tiny' can it work correctly, and it will crush loading the 'openai/whisper-small' model and larger. I've tried XLA_PYTHON_CLIENT_PREALLOCATE=false mentioned in issue 7 but it didn't work.

The image below shows the vRAM usage of my machine. Missing data means machine rebooted.

Is there any way to prevent linux to reboot automatically when vRAM usage is high?

[sanchit-gandhi]

Hey @tdzz1102 and sorry for the late reply! Could you try with XLA_PYTHON_CLIENT_MEM_FRACTION=.XX (where .XX is the percentage memory allocation you want to assign, e.g. .50 for 50% memory allocation) to reduce the memory fraction allocation? The docs say this should help combat OOMs that occur when the programme starts: https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

You might need to play with your value of .XX, e.g. incrementally reducing from .75 to .00

[tdzz1102]

@sanchit-gandhi I have solved this question by downgrading the nividia-driver and cuda version(but forgot what exactly the version was 😢). Now the server has been expired and I can't try this solution anylonger. The faster whisper has helped me a lot, and thank you anyway!