Please refer to Rosetta Imagen, NVIDIA's project that enables seamless training of LLMs, CV models and multimodal models in JAX, for information about running Imagen models and experiments on GPUs.