TPU support

[miguelalba96]

🚀 Feature

TPU support

Motivation

Does litdata supports TPU environments, specifically when using lighting fabric?

Additional context

I have >16M image-text pairs I am writing in mosaic-ml streaming format to train contrastive models, I am working with lighting fabric to train using DDP in GCP and I want to move to TPU training. mosaic-ml streaming dataset doesn't support TPU (afaik), all of this bring me to the questions:

• Does litdata work on TPU?
• Does it require to set up something in the code additional to what is provided in the available documentation. (ex. is it necessary to provide a distributed sampler? or set different env variables?)
• Do you have an example of how to setup TPU training with litdata?

[github-actions]

Hi! thanks for your contribution!, great first issue!

[tchaton]

Hey @miguelalba96,

I haven't tried with TPU. Maybe @carmocca would know more.

[carmocca]

litdata is meant to be used with a regular DataLoader, so there's nothing specific to do on a TPU machine. If you use Fabric or PyTorch Lightning, that will take care of enabling the DistributedSampler or do any required XLA steps, but these are common to all TPU runs, not just those using litdata

[dasoto]

I will recommend setting this env variables:
DATA_OPTIMIZER_GLOBAL_RANK
DATA_OPTIMIZER_NUM_WORKERS
DATA_OPTIMIZER_NUM_NODES

Otherwise the StreamDataloader will not be aware of the distribution.

[tchaton]

Yes, as @dasoto mentioned, I didn't add wiring for TPU env detection. Feel free to contribute support for it if you try litdata on TPUs.