Slow PyTorch Training

[Lundez]

Hi!

I'm trying to use daft to train a PyTorch model. But it's really slow compared to native PyTorch code. Sometimes it crashes on memory too.

I thought daft was supposed to lazy load items, but it seems not? Or am I doing something wrong?

Dependencies:

[dependencies]
python = ">=3.12.7,<3.13"
dataset = ">=1.6.2,<1.7"
ray-data = ">=2.37.0,<2.38"
datasets = ">=3.0.2,<3.1"
pillow = ">=11.0.0,<11.1"
torchvision = ">=0.19.1,<0.20"
lightning = ">=2.4.0,<2.5"
torchmetrics = ">=1.5.1,<1.6"
getdaft = "0.3.9"

import datasets
import ray
import torch
import daft
import utils  # Assuming this contains your PREPROCESS_TRANSFORMS and AUGMENTATIONS
from torchvision import transforms
import torch.nn.functional as F
import lightning as L


def load_imagenette_datasets_daft(dataset_path="./imagenette_full_size"):
    # Load the Arrow dataset with Daft
    # ds = load_from_disk(dataset_path)
    ds = datasets.load_from_disk(dataset_path)

    ds_train = daft.from_arrow(ds["train"].data.table).select(
        "label", daft.col("image").struct.get("bytes").alias("image")
    )

    ds_valid = daft.from_arrow(ds["validation"].data.table).select(
        "label", daft.col("image").struct.get("bytes").alias("image")
    )

    def transform(ds: daft.DataFrame):
        ds = ds.with_column(
            "image",
            daft.col("image").image.decode(mode="RGB").image.resize(224, 224),
        )

        return ds

    ds_train = transform(ds_train)

    def to_f32_tensor(ds: daft.DataFrame):
        return ds.with_column(
            "image",
            daft.col("image").apply(
                lambda x: (x / 255.0).transpose(2, 0, 1),
                return_dtype=daft.DataType.tensor(daft.DataType.float32()),
            ),
        )

    ds_valid = transform(ds_valid)
    ds_train = to_f32_tensor(ds_train)
    ds_valid = to_f32_tensor(ds_valid)

    ds_train = ds_train.to_torch_iter_dataset()
    ds_valid = ds_valid.to_torch_iter_dataset()

    dls_train = torch.utils.data.DataLoader(ds_train, batch_size=32)
    dls_valid = torch.utils.data.DataLoader(ds_valid, batch_size=32)
    return dls_train, dls_valid


def train_model():
    dls_train, dls_valid = load_imagenette_datasets_daft()

    # Initialize trainer and model
    trainer = L.Trainer(max_epochs=5)
    model = utils.SimpleModel()

    # Train the model
    trainer.fit(model, dls_train)


if __name__ == "__main__":
    train_model()

When it runs correctly it takes ~14m to train, but usually it crashes. PyTorch DataLoader takes ~3m30s in comparison (running transforms on a wrapped Dataset around the HF Dataset). Running on M1 macbook.

[Lundez]

Adding daft.set_execution_config(enable_native_executor=True, default_morsel_size=256) seems to make it run more persistently and I see the following runtime: 3m24s.

That's really good to see! But it's not clear that this is available from docs, found some PR's regarding this setting that I wasn't aware of (using the new Swordfish engine I suppose`?)

EDIT: Adding only daft.set_execution_config(default_morsel_size=256) seems to help running consistently but runtime ends up at 13m7s. The native_executor really seem to step things up!

[jaychia]

Yep you've found our new execution engine 😝

The current engine relies on partitioning for parallelism. The way your dataset is loaded (as a single PyArrow table) means that there is only 1 partition to parallelize/pipeline over.

The new engine is streaming/morsel-based and so you just have to configure a morsel size. This helps with memory stability and also is just a better user experience overall.

I think there are still some limitations for the new execution engine that @colin-ho is working on, especially wrt reading of local data. It will become the default engine soon in about 1-2 weeks.

[Lundez]

Awesome to hear!

Some further findings, when I ran with full imagesize dataset the new streaming engine consistently turned out faster than PyTorch DataLoader, phenomenal too see!

Love that its fast locally on single node while able to scale big 🙌

[Lundez]

Found the native executed that makes it fast. Closing for now.