search

google-deepmind / synthid-text

Apache License 2.0
372 stars 27 forks source link

The model device in the default settings may not be as expected. #14

Closed LOVEANDNINENINE closed 2 weeks ago

LOVEANDNINENINE commented 2 weeks ago

This is just a very minor issue. I personally think your code is very well-structured. However, it seems that multi-GPU machines were not considered (even though I’m only using one GPU). In Colab, the model is loaded with the 'auto' mode for device_map, but when the global DEVICE is set, it is configured as 'cuda:0'. It would be more user-friendly to assign the device_map to the global variable DEVICE for better accessibility, especially for beginners.

RyanMullins commented 2 weeks ago

Thanks for the feedback, @LOVEANDNINENINE.

You're correct that supporting multi-accelerator configurations is not a design goal for this codebase, which is focused on reproducibility of the associated research paper. Instead, we assume that users will be limited to whatever is supported by Colab, which today is a single GPU per runtime.

Given this assumption, setting device_map="auto" will infer the correct device map and fill the available accelerator (i.e., cuda:0). We find to be sufficiently user-friendly for our goals, and definitely more maintainable than defining custom device_maps using the torch.device identifier from DEVICE for each model/accelerator configuration. If users find that their model is not fitting on the accelerators available to them, they will need to consider purchasing a subscription.

We strongly encourage that users interested in multi-accelerator configurations first consider using the production-grade SynthID Text implementation available in Hugging Face Transformers.

Researchers looking to directly reproduce this work on multi-accelerator setups (e.g., VMs or enterprise resources) will need to fork this repo and update the code accoridngly.