Repository: Lightning-AI/pytorch-lightning
Pull request: FCCV Docs
## What does this PR do?

Adds docs for custom data iterables
Fixes #11538 

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


cc @borda @justusschock @awaelchli

Problem statement:
enable ffcv with PL
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:
```python
# 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

Hints from issue discussion:
@ethanwharris This could be built-in within Flash.
https://github.com/libffcv/ffcv/issues/72
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.


Any update on this?
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 `LightningDataModule`s 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](https://github.com/libffcv/ffcv-imagenet/blob/main/train_imagenet.py#L236)) 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!
Ok - awesome! Thanks so much!
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?
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 `Loader`s:
(https://docs.ffcv.io/making_dataloaders.html#dataset-ordering) inside `*_dataloader` methods

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?
@leejiahe You don't, this was just to showcase the customization. But it shouldn't be necessary

Commit message: nit

Patch to review:
diff --git a/docs/source-pytorch/data/custom_data_iterables.rst b/docs/source-pytorch/data/custom_data_iterables.rst
new file mode 100644
index 0000000000000..2b793b6b79f8f
--- /dev/null
+++ b/docs/source-pytorch/data/custom_data_iterables.rst
@@ -0,0 +1,109 @@
+.. _dataiters:
+
+##################################
+Injecting 3rd Party Data Iterables
+##################################
+
+When training a model on a specific task, data loading and preprocessing might become a bottleneck.
+Lightning does not enforce a specific data loading approach nor does it try to control it.
+The only assumption Lightning makes is that the data is returned as an iterable of batches.
+
+For PyTorch-based programs, these iterables are typically instances of :class:`~torch.utils.data.DataLoader`.
+
+However, Lightning also supports other data types such as plain list of batches, generators or other custom iterables.
+
+.. code-block:: python
+
+    # random list of batches
+    data = [(torch.rand(32, 3, 32, 32), torch.randint(0, 10, (32,))) for _ in range(100)]
+    model = LitClassifier()
+    trainer = Trainer()
+    trainer.fit(model, data)
+
+Examples for custom iterables include `NVIDIA DALI <https://github.com/NVIDIA/DALI>`__ or `FFCV <https://github.com/libffcv/ffcv>`__ for computer vision.
+Both libraries offer support for custom data loading and preprocessing (also hardware accelerated) and can be used with Lightning.
+
+
+For example, taking the example from FFCV's readme, we can use it with Lightning by just removing the hardcoded ``ToDevice(0)``
+as Lightning takes care of GPU placement. In case you want to use some data transformations on GPUs, change the
+``ToDevice(0)`` to ``ToDevice(self.trainer.local_rank)`` to correctly map to the desired GPU in your pipeline.
+
+.. code-block:: python
+
+    from ffcv.loader import Loader, OrderOption
+    from ffcv.transforms import ToTensor, ToDevice, ToTorchImage, Cutout
+    from ffcv.fields.decoders import IntDecoder, RandomResizedCropRGBImageDecoder
+
+
+    class CustomClf(LitClassifier):
+        def train_dataloader(self):
+
+            # Random resized crop
+            decoder = RandomResizedCropRGBImageDecoder((224, 224))
+
+            # Data decoding and augmentation
+            image_pipeline = [decoder, Cutout(), ToTensor(), ToTorchImage()]
+            label_pipeline = [IntDecoder(), ToTensor()]
+
+            # Pipeline for each data field
+            pipelines = {"image": image_pipeline, "label": label_pipeline}
+
+            # Replaces PyTorch data loader (`torch.utils.data.Dataloader`)
+            loader = Loader(
+                write_path, batch_size=bs, num_workers=num_workers, order=OrderOption.RANDOM, pipelines=pipelines
+            )
+
+            return loader
+
+When moving data to a specific device, you can always refer to ``self.trainer.local_rank`` to get the accelerator
+used by the current process.
+
+By just changing ``device_id=0`` to ``device_id=self.trainer.local_rank`` we can also leverage DALI's GPU decoding:
+
+.. code-block:: python
+
+        from nvidia.dali.pipeline import pipeline_def
+        import nvidia.dali.types as types
+        import nvidia.dali.fn as fn
+        from nvidia.dali.plugin.pytorch import DALIGenericIterator
+        import os
+
+
+        class CustomLitClassifier(LitClassifier):
+            def train_dataloader(self):
+
+                # To run with different data, see documentation of nvidia.dali.fn.readers.file
+                # points to https://github.com/NVIDIA/DALI_extra
+                data_root_dir = os.environ["DALI_EXTRA_PATH"]
+                images_dir = os.path.join(data_root_dir, "db", "single", "jpeg")
+
+                @pipeline_def(num_threads=4, device_id=self.trainer.local_rank)
+                def get_dali_pipeline():
+                    images, labels = fn.readers.file(file_root=images_dir, random_shuffle=True, name="Reader")
+                    # decode data on the GPU
+                    images = fn.decoders.image_random_crop(images, device="mixed", output_type=types.RGB)
+                    # the rest of processing happens on the GPU as well
+                    images = fn.resize(images, resize_x=256, resize_y=256)
+                    images = fn.crop_mirror_normalize(
+                        images,
+                        crop_h=224,
+                        crop_w=224,
+                        mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
+                        std=[0.229 * 255, 0.224 * 255, 0.225 * 255],
+                        mirror=fn.random.coin_flip(),
+                    )
+                    return images, labels
+
+                train_data = DALIGenericIterator(
+                    [get_dali_pipeline(batch_size=16)],
+                    ["data", "label"],
+                    reader_name="Reader",
+                )
+
+                return train_data
+
+
+Lightning works seamlessly with all kinds of custom data iterables,
+but unfortunately it cannot support the entire featureset with arbitrary iterables as some are specific to dataloaders.
+These features are mainly automatic replacement of the sampler and fully fault-tolerant training as these dataloaders
+typically don't expose sampling APIs to fast-forward or save and load states.
diff --git a/docs/source-pytorch/index.rst b/docs/source-pytorch/index.rst
index 77149ffbd6004..50b02c483997f 100644
--- a/docs/source-pytorch/index.rst
+++ b/docs/source-pytorch/index.rst
@@ -206,6 +206,7 @@ Current Lightning Users
    Train on single or multiple TPUs <accelerators/tpu>
    Train on MPS <accelerators/mps>
    Use a pretrained model <advanced/pretrained>
+   Inject Custom Data Iterables <data/custom_data_iterables>
    model/own_your_loop
 
 .. toctree::


Code review comments:
diff_hunk: @@ -0,0 +1,109 @@
+.. _dataiters:
+
+##################################
+Injecting 3rd Party Data Iterables
+##################################
+
+When training a model on a specific task, data loading and preprocessing might become a bottleneck.
+Lightning does not enforce a specific data loading approach nor does it try to control it.
+The only assumption Lightning makes is that the data is returned as an iterable of batches.
+
+For PyTorch-based programs, these iterables are typically instances of :class:`~torch.utils.data.DataLoader`.
+
+However, Lightning also supports other data types such as plain list of batches, generators or other custom iterables.
+
+.. code-block:: python
+
+    # random list of batches
+    data = [(torch.rand(32, 3, 32, 32), torch.randint(0, 10, (32,))) for _ in range(100)]
+    model = LitClassifier()
+    trainer = Trainer()
+    trainer.fit(model, data)
+
+Examples for custom iterables include `NVIDIA DALI <https://github.com/NVIDIA/DALI>`__ or `FFCV <https://github.com/libffcv/ffcv>`__ for computer vision.
+Both libraries offer support for custom data loading and preprocessing (also hardware accelerated) and can be used with Lightning.
+
+
+For example, taking the example from FFCV's readme, we can use it with Lightning by just removing the hardcoded ``ToDevice(0)``
+as Lightning takes care of GPU placement. In case you want to use some data transformations on GPUs, change the
+``ToDevice(0)`` to ``ToDevice(self.trainer.local_rank)`` to correctly map to the desired GPU in your pipeline.
+
+.. code-block:: python
+
+    from ffcv.loader import Loader, OrderOption
+    from ffcv.transforms import ToTensor, ToDevice, ToTorchImage, Cutout
+    from ffcv.fields.decoders import IntDecoder, RandomResizedCropRGBImageDecoder
+
+
+    class CustomClf(LitClassifier):
path: docs/source-pytorch/data/custom_data_iterables.rst
text: @user1:
```suggestion
    class CustomClassifier(LitClassifier):
```