[torchbench] Training benchmarks failing with: OOM

[ysiraichi]

This post has two lists of training benchmarks failing with OOM in a NVIDIA A100 40GB GPU:

• Eager-mode
• Dynamo+openxla

These lists were put together by running the upstreamed benchmarking scripts. More specifically, the following command:

python xla/benchmarks/experiment_runner.py \
       --suite-name torchbench \
       --accelerator cuda \
       --xla PJRT --xla None \
       --dynamo openxla --dynamo None \
       --test train \
       --repeat 30 --iterations-per-run 5 \
       --print-subprocess \
       --no-resume

Eager-mode

• demucs
• densenet121
• hf_GPT2_large
• hf_T5_base
• llama_v2_7b_16h (skipped -- torchbench.yaml)
• stable_diffusion_unet
• timm_nfnet
• timm_vision_transformer_large

Dynamo+openxla

• demucs
• densenet121
• llama_v2_7b_16h (skipped -- torchbench.yaml)
• stable_diffusion_unet
• timm_vision_transformer
• timm_vision_transformer_large