enable ffcv with PL

[williamFalcon]

A team introduced ffcv which does optimizations in the dataloader level.
As long as it can be a drop in replacement for dataloader it should add the benefits to any PyTorch and Lightning script.

ie:

# current
data = Dataloader(...)

trainer.fit(..., data)

# 10x faster with ffcv
data = FFCVDataloader

trainer.fit(..., data)

Let's make sure when users want to get the benefits of both FFCV and PL they can do that!
cc @Borda @rohitgr7 @Felonious-Spellfire @justusschock @awaelchli @ninginthecloud @otaj @tchaton @carmocca

[tchaton]

@ethanwharris This could be built-in within Flash.

[paantya]

libffcv/ffcv#72

[justusschock]

This is related to #10696.

When we support arbitrary iterables, ffcv should be supported automatically. Maybe not all features we support with basic torch loaders (haven't looked at ffcv details so far), but definitely usable for training.

[siddk]

Any update on this?

[carmocca]

FFCV should be perfectly usable with no modifications to the PL source code after the 1.6 release.

Before, some patching to the loops was necessary to disable PL's pre-fetching as it conflicted with the pre-fetching done by FFCV. Now, PL avoids pre-fetching whenever it is possible.

From my own benchmarks, the speed up from FFCV applies perfectly when used with PL. We'll be releasing public benchmarks in the future.

We'll also explore making or internal data processing more flexible so that maybe we can provide an opinionated LightningDataModule integration that makes switching existing LightningDataModules easier.

A caveat I've noticed are that you'll want to remove the ToDevice(device) from the label pipeline if it's the last operation in your transformations (example) as PL already takes care of moving the tensors to the correct devices.

Features like fault-tolerance do not support FFCV.

Note that there's a lot of surface for bugs here and I think most of the battle-testing has been done on map-style datasets. If you encounter any memory or speed issues, please open a separate issue and you can ping me on it.

Cheers!

[siddk]

Ok - awesome! Thanks so much!

[siddk]

Second question (sorry for the delayed follow-up); is the recommended pattern to couple FFCV dataloading within a LightningDataModule, or should I just pass the FFCV loaders directly to the Trainer?

[carmocca]

It's up to you, but both will work.

The advantage of the LightningDataModule is that you can put the FFCV specific data-writing code (https://docs.ffcv.io/writing_datasets.html#writing-a-dataset-to-ffcv-format) inside the setup hook and also save in the instance state any specific arguments for your FFCV Loaders:
(https://docs.ffcv.io/making_dataloaders.html#dataset-ordering) inside *_dataloader methods

[leejiahe]

Hi,

Do we need to write a custom FitLoop to define how we want to pre-fetch the batch, as documented in the 1.6.0 release page to be compatible with FFCV dataloader?

[carmocca]

@leejiahe You don't, this was just to showcase the customization. But it shouldn't be necessary

[codestar12]

A caveat I've noticed are that you'll want to remove the ToDevice(device) from the label pipeline if it's the last operation in your transformations (example) as PL already takes care of moving the tensors to the correct devices.

do you want to remove all ToDevice(device)? I'm a bit confused because I've mostly been able to ignore manually controlling device placement since switching to PyTorch lightning.

[codestar12]

I have actually been having a lot of trouble getting this to work with DDP and multiple GPUs. What is the best way to get the GPU rank from a lighting data module? Their data loader seems pretty opinionated about these transforms and really wants to move it to the device.

[Data-drone]

How did you go in the end @codestar12?

[carmocca]

do you want to remove all ToDevice(device)?

No, only if it's the last operation in the pipeline. Otherwise, you'll want to keep it because doing more operations with the data on device could be faster.

What is the best way to get the GPU rank from a lighting data module?

You can use self.trainer.local_rank

[codestar12]

How did you go in the end @codestar12?

I basically just removed a lot of FFCV transforms from the pipeline and just kept the data preprocessing bit. It looks like @carmocca finally has the answer though

[paantya]

Please add somewhere in the Lightning or bolt an example of how to use the Lightning in conjunction with FFCV

[rmchurch]

Do I understand right from the docs though that you can't use FFCV + Lightning when using ddp? I pass the distributed arg to the ffcv loader and get errors when running with Lightning. Is this a fundamental issue preventing, or could multi-GPU training with FFCV + Lightning potentially be implemented some day?

[carmocca]

It was working last time I tried it. What errors do you see? FFCV requires quite a bit of configuration to work as you expect. This comment contains the script I used in the past: #15598 (review)

[rmchurch]

Ah, I had seen your example, but thought it was incomplete because for gpus>1 it was marked FIXME: distributed, and distributed=False input into FFCV Loader. Thinking it was old code, I was trying distributed=True, and getting the init_process_group not initialized error. But is it the case that I should put distributed=False into the FFCV Loader still, and Lightning will insert the DistributedSampler? If that is the case, perhaps the paragraph below in the docs should be rephrased, or additional detail given that ddp will work with FFCV, but you have to leave distributed=False.

In a distributed multi-GPU setting (ddp), Lightning automatically replaces the DataLoader’s sampler with its distributed counterpart. This makes sure that each GPU sees a different part of the dataset. As sampling can be implemented in arbitrary ways with custom iterables, there is no way for Lightning to know, how to replace the sampler.

[carmocca]

No, Lightning will not insert it automatically for FFCV. You might need to initialize the Loader class under one of the *_dataloader hooks so that the process groups are initialized before you create it.