Lightning-AI / Lightning-AI/pytorch-lightning

`_update_dataloader` improperly copies state of subclassed dataloader with attribute names that differ from `__init__` parameters.

Open
#20,265 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug data data handling ver: 2.3.x ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

`_update_dataloader` in `lightning/fabric/utilities/data.py` does not copy a dataloader, but instead infers the arguments passed to the dataloader and instantiates a new one. The inference of arguments in `_get_dataloader_init_args_and_kwargs` does not take into account custom logic in `__init__` that cause attributes to have different names from the parameters and results in the new dataloader using default arguments instead of those passed to the original dataloader.

This problem arises during distributed training when a new dataloader is created.

In the code example, the dataloader takes an argument `x=2` and assigns it to the attribute `self._x`. When a new dataloader is created, `_get_dataloader_init_args_and_kwargs` infers `x` used the default argument instead of 2. The log below shows the output of print statements in `__init__` and `__iter__`. When the original dataloader is created `self._x` is 2, but when it is copied, `self._x` has the default argument `None`. The relevant print lines in the log are surrounded with `**`.

I assume there's a reason why deepcopy or some other copying mechanism isn't used?

Expected behavior
-------------------
State of dataloader is copied as is.

Actual behavior
----------------
Dataloader is reinitialized with wrong arguments.

Thank you for this awesome software!

### What version are you seeing the problem on?

v2.3, v2.4

### How to reproduce the bug

```python
import argparse
import sys
from pathlib import Path

import lightning as L
import torch
import torch.nn.functional as F

class MyDataLoader(torch.utils.data.DataLoader):

def __init__(self, *args, x=None, **kwargs):
super().__init__(*args, **kwargs)
self._x = x
print("DummyDataLoader.__init__(): self._x is", self._x)

def __iter__(self):
print("DummyDataLoader.__iter__(): self._x is", self._x)

for item in super().__iter__():
yield item

class MyLinear(torch.nn.Module):

def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(1, 1, "cuda")

def forward(self, X):
return self.linear(X)

class MyLightning(L.LightningModule):

def __init__(self):
super().__init__()
self.linear = MyLinear()

def forward(self, X):
return self.linear(X)

def training_step(self, batch, batch_idx):
X, y = batch
y_hat = self(X)
loss = F.mse_loss(y, y_hat)
return loss

def configure_optimizers(self):
return torch.optim.Adam(self.parameters())

def train_func():
X = y = torch.arange(4, device="cuda", dtype=torch.float32).view(-1, 1)
train_data = MyDataLoader(list(zip(X, y)), batch_size=4, x=2)

model = MyLightning()
model.to("cuda")
trainer = L.Trainer(
max_epochs=10,
devices=1,
num_nodes=1,
accelerator="gpu",
use_distributed_sampler=True,
strategy="ddp",
enable_checkpointing=False
)
trainer.fit(model, train_data)
return

train_func()
```

### Error messages and logs

```
**DummyDataLoader.__init__(): self._x is 2**
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/home/spencer.mitchell@insmed.local/.anaconda3/envs/lightning_error2/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` \
has been removed as a dependency of the `lightning.pytorch` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, u\
nless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for pe\
rformance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

| Name | Type | Params | Mode
--------------------------------------------
0 | linear | MyLinear | 2 | train
--------------------------------------------
2 Trainable params
0 Non-trainable params
2 Total params
0.000 Total estimated model params size (MB)
2 Modules in train mode
0 Modules in eval mode
**DummyDataLoader.__init__(): self._x is None**
/home/spencer.mitchell@insmed.local/.anaconda3/envs/lightning_error2/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'train_dataloader' does not have many workers which\
may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=31` in the `DataLoader` to improve performance.
/home/spencer.mitchell@insmed.local/.anaconda3/envs/lightning_error2/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:298: The number of training batches (1) is smaller than the logging interval Tr\
ainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
Epoch 0: 0%| | 0/1 [00:00
Current environment

* CUDA:
- GPU:
- NVIDIA GeForce RTX 3090
- available: True
- version: 11.8
* Lightning:
- lightning: 2.4.0
- lightning-utilities: 0.11.7
- pytorch-lightning: 2.4.0
- torch: 2.4.1
- torchmetrics: 1.4.0.post0
* Packages:
- autocommand: 2.2.2
- backports.tarfile: 1.2.0
- brotli: 1.0.9
- certifi: 2024.8.30
- cffi: 1.16.0
- charset-normalizer: 3.3.2
- colorama: 0.4.6
- filelock: 3.13.1
- fsspec: 2024.9.0
- h2: 4.1.0
- hpack: 4.0.0
- hyperframe: 6.0.1
- idna: 3.8
- importlib-metadata: 8.0.0
- importlib-resources: 6.4.0
- inflect: 7.3.1
- jaraco.context: 5.3.0
- jaraco.functools: 4.0.1
- jaraco.text: 3.12.1
- jinja2: 3.1.4
- lightning: 2.4.0
- lightning-utilities: 0.11.7
- markupsafe: 2.1.3
- mkl-fft: 1.3.10
- mkl-random: 1.2.7
- mkl-service: 2.4.0
- more-itertools: 10.3.0
- mpmath: 1.3.0
- networkx: 3.3
- numpy: 1.26.4
- ordered-set: 4.1.0
- packaging: 24.1
- pip: 24.2
- platformdirs: 4.2.2
- pycparser: 2.22
- pysocks: 1.7.1
- pytorch-lightning: 2.4.0
- pyyaml: 6.0.1
- requests: 2.32.3
- setuptools: 72.1.0
- sympy: 1.13.2
- tomli: 2.0.1
- torch: 2.4.1
- torchmetrics: 1.4.0.post0
- tqdm: 4.66.5
- triton: 3.0.0
- typeguard: 4.3.0
- typing-extensions: 4.11.0
- urllib3: 2.2.2
- wheel: 0.43.0
- zipp: 3.19.2
- zstandard: 0.22.0
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.12.4
- release: 4.18.0-553.16.1.el8_10.x86_64
- version: #1 SMP Thu Aug 1 04:16:12 EDT 2024

### More info

_No response_

cc @tchaton

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in lightning/fabric/utilities/data.py by reading _update_dataloader and _get_dataloader_init_args_and_kwargs, then reproduce the issue with the MyDataLoader example in this report. Done means a subclassed dataloader recreated during distributed training preserves the original x=2 state instead of using None or another default.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
data, distributed-systems
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.