For ModelWrappers that use Fabric, etc., figure out how to properly handle device

[k-chrispens]

The RF3 wrapper uses Lightning Fabric, which will set up the model on its own device. I haven't yet figured out how to select the device manually, but this is something that needs to be fixed at some point.

[k-chrispens]

@marcuscollins do you have any ideas on this? Should we close this as stale?

[marcuscollins]

@k-chrispens this might be a good one for Michael or Moody to look into. I guess it isn't urgent, but it might become a problem if we scale up another 10x.

[marcuscollins]

The problem here is that we don't have a way to control which GPU RosettaFold3 chooses, and generally it chooses 0. So, we can't really run parallel jobs, for instance if I run

pixi run -e rf3 python run_grid_search.py       --proteins proteins.csv       --models rf3       --scalers pure_guidance       --ensemble-sizes "8"   
    --gradient-weights "0.0 0.005 0.01 0.02 0.035 0.05 0.1"       --output-dir /data/sampleworks-exp/occ_sweep/grid_search_results       --augmentatio
n --align-to-input --model-checkpoint checkpoints/rf3_foundry_01_24_latest.ckpt | tee ../karson-rf3-run.log

with 8 available GPUs, this will result in 8 jobs, but they all run on GPU 0, which slows us down considerably.