Running the benchmark suite on very big models

Hi, thank you very much for the OLMES work, it is great!

I am having some troubles in running the run_eval.py script on a 340B Megatron model. This model typically requires 16 A100 GPUs for the forward pass. in particular, I keep getting an out of memory error while I load the checkpoint files. I am running the script on a SLURM cluster, and it seems that simply increasing the number of assigned compute nodes does not solve the issue. Any tips on how to make the script work on very big models?