Unexpected behaviour when using prepare_data

[bazingayu]

Bug description

Hi lightning team,

We were trying to using prepare_data() to move the data from a shared disk to local disk.
Even though we copied the correct data to the local disk. When we printed out the dataloader file list on different GPUs(I have checked the md5 for those), the filelist is different.

**GPU 0:** self.__files [**'ez13.npz', 'sc4.npz',** 'ez2.npz', 'mb9.npz', 'bc3.npz', 'lc14.npz', 'lc3.npz', 'ez10.npz', 'lo11.npz', 'ez3.npz', 'Read14_roto.npz', 'lo6.npz', 'sc6.npz', 'mb3.npz', 'lo10.npz', 'lo7.npz', 'lc8.npz', 'mb6.npz', 'bc7.npz', 'sc13.npz', 'se3.npz', 'ez12.npz', 'lc10.npz', 'lc7.npz', 'sc3.npz', 'bc15.npz', 'lc15.npz', 'bc11.npz', 'lc5.npz', 'se12.npz', 'se6.npz', 'bc9.npz', 'sc12.npz', 'lo12.npz', 'ez9.npz', 'mb7.npz', 'lo16.npz', 'se10.npz', 'se8.npz', 'mb10.npz', 'mb14.npz', 'bc12.npz', 'se5.npz', 'ez16.npz', 'se14.npz', 'Read15_roto.npz', 'bc4.npz', 'mb11.npz', 'lc6.npz', 'lo14.npz', 'bc14.npz', 'sc16.npz', 'se11.npz', 'lo2.npz', 'lo9.npz', 'sc8.npz', 'ez4.npz', 'se9.npz', 'sc9.npz', 'lo3.npz', 'lc2.npz', 'bc5.npz', 'se7.npz', 'mb1.npz', 'mb12.npz', 'se1.npz', 'sc2.npz', 'bc16.npz', 'ez8.npz', 'ez1.npz', 'bc10.npz', 'sc14.npz', 'se4.npz', 'lo13.npz', 'lc16.npz', 'mb2.npz', 'lo5.npz', 'lo15.npz', 'bc2.npz', 'ez6.npz', 'bc1.npz', 'lc11.npz', 'bc13.npz', 'mb8.npz', 'lo1.npz', 'sc11.npz', 'se2.npz', 'lc13.npz', 'ez7.npz', 'sc15.npz', 'lc9.npz', 'bc6.npz', 'sc7.npz', 'mb13.npz', 'lc1.npz', 'lo8.npz', 'sc5.npz', 'mb4.npz', 'mb5.npz', 'se13.npz', 'lc12.npz', 'lo4.npz', 'ez15.npz', 'ez5.npz', 'mb15.npz', 'ez11.npz', 'bc8.npz', 'lc4.npz', 'sc10.npz', 'sc1.npz', 'ez14.npz', 'mb16.npz']

**GPU1:** self.__files [**'bc11.npz', 'lo9.npz'**, 'se12.npz', 'bc9.npz', 'lc6.npz', 'se9.npz', 'sc13.npz', 'mb3.npz', 'se3.npz', 'bc16.npz', 'lc5.npz', 'bc6.npz', 'lo6.npz', 'Read14_roto.npz', 'bc7.npz', 'mb10.npz', 'bc5.npz', 'bc2.npz', 'lo2.npz', 'mb13.npz', 'lc7.npz', 'mb16.npz', 'lc16.npz', 'sc11.npz', 'mb6.npz', 'ez16.npz', 'lc8.npz', 'sc9.npz', 'bc1.npz', 'bc14.npz', 'lc3.npz', 'se2.npz', 'lc9.npz', 'lo3.npz', 'lc10.npz', 'sc3.npz', 'bc3.npz', 'lo10.npz', 'se4.npz', 'bc10.npz', 'sc14.npz', 'se11.npz', 'mb7.npz', 'ez1.npz', 'lo13.npz', 'lc14.npz', 'lo8.npz', 'bc4.npz', 'ez4.npz', 'ez8.npz', 'lc11.npz', 'ez6.npz', 'Read15_roto.npz', 'se6.npz', 'lo15.npz', 'lo12.npz', 'ez2.npz', 'sc7.npz', 'ez5.npz', 'lo5.npz', 'sc8.npz', 'se14.npz', 'bc15.npz', 'se1.npz', 'sc6.npz', 'se8.npz', 'ez15.npz', 'mb1.npz', 'lc13.npz', 'se13.npz', 'lc15.npz', 'se10.npz', 'lc2.npz', 'bc13.npz', 'lo7.npz', 'ez14.npz', 'sc5.npz', 'sc2.npz', 'ez13.npz', 'bc12.npz', 'lc1.npz', 'mb11.npz', 'ez11.npz', 'mb12.npz', 'sc15.npz', 'ez9.npz', 'mb4.npz', 'sc12.npz', 'se5.npz', 'mb15.npz', 'se7.npz', 'sc4.npz', 'lc12.npz', 'ez12.npz', 'lo4.npz', 'ez10.npz', 'ez7.npz', 'mb5.npz', 'sc16.npz', 'mb9.npz', 'mb8.npz', 'ez3.npz', 'bc8.npz', 'lo16.npz', 'sc10.npz', 'lo1.npz', 'lo14.npz', 'mb14.npz', 'lo11.npz', 'mb2.npz', 'lc4.npz', 'sc1.npz']

This issue caused problems when we used FSDP(protentially also for DDP). I printed the filenames in the dataloader and found that some data was missing, while other data was processed multiple times."

filename, counting
ez13.npz , 0
sc4.npz , 1
ez2.npz , 2
mb9.npz , 1
bc3.npz , 1
lc14.npz , 1
lc3.npz , 1
ez10.npz , 1
lo11.npz , 1
ez3.npz , 1
Read14_roto.npz , 1
lo6.npz , 2
sc6.npz , 1
mb3.npz , 1
lo10.npz , 1
lo7.npz , 1
lc8.npz , 0
mb6.npz , 2

I can tell the copying process happened in single thread, but still, there's something wrong within prepare_data function

What version are you seeing the problem on?

v2.4

Reproduced in studio

No response

How to reproduce the bug

def prepare_data(self):
        """
        Copy datasets to local disk for faster access during training.
        This is a LightningDataModule callback, called once on the CPU before
        dataloaders are initialized in the setup() callback.
        """

        for source_path in train_paths:
            target_path = os.path.join(LOCAL_DATA_ROOT, *(Path(source_path).parts[-2:]))
            if not os.path.exists(target_path):
                print(f"Copying training data from {source_path} to {target_path}...")
                shutil.copytree(source_path, target_path, dirs_exist_ok=True)

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

cc @lantiga @Borda

[bhimrazy]

Hi @bazingayu,

Thanks for the report! The details you’ve shared are helpful, but to help us debug this we might need a minimal reproducible example showing how prepare_data is interacting with your LightningDataModule and dataloaders.

It would also be great if you could include:

• Your full LightningDataModule implementation (especially setup() and train_dataloader())
• The exact Trainer/strategy settings (FSDP, DDP, number of processes, etc.)
• Whether the file listing is generated before or after DataLoader shuffling

Right now, the description highlights the symptoms but doesn’t give us enough context to pinpoint the cause.

If you’d like, you can also use Lightning Studio to share a full reproducible example: https://www.youtube.com/watch?v=YcW-2Zt_bFg

[SkafteNicki]

To add to @bhimrazy comment, if prepare_data is the culprit of this issue, please share what you are setting num_nodes to in your Trainer (e.g. are you doing multi-node training) because that argument in particular influences the behavior of prepare_data.

[bazingayu]

Hi @bhimrazy @SkafteNicki ,
Thanks for the quick support, it is very helpful. Here's some other details you ask for. For some credential problems, I can't share the whole file we're using, but it should have enough details you want

# LightningDataModule

class LitData(pl.LightningDataModule):
    def __init__(
        self,
        cfg: Config,
        generator: torch.Generator,
    ) -> None:
        super().__init__()
        self.__cfg = cfg
        self.__generator = generator

    def prepare_data(self):
        """
        Copy datasets to local disk for faster access during training.
        This is a LightningDataModule callback, called once on the CPU before
        dataloaders are initialized in the setup() callback.
        """

        for source_path in train_paths:
            target_path = os.path.join(LOCAL_DATA_ROOT, *(Path(source_path).parts[-2:]))
            if not os.path.exists(target_path):
                print(f"Copying training data from {source_path} to {target_path}...")
                shutil.copytree(source_path, target_path, dirs_exist_ok=True)

    def setup(self, stage) -> None:
    	self.__dataset = torch.utils.data.ConcatDataset([
            DatasetFunction(data_root=path,
                           transform = FormatNpzSample())
                    for path in target_path
                ])

        self.train_set, self.val_set, self.test_set = random_split(
            self.__dataset, [1, 0, 0], generator=self.__generator)

    def train_dataloader(self):
        return DataLoader(
            self.train_set,
            batch_size=1,
            generator=self.__generator,
            num_workers=10,
            pin_memory=True,
            shuffle=True,
            collate_fn=partial(collate_fn, target_interp_parameters=self.__cfg.target_interp_parameters),
            drop_last=True,
	    )

# Training strategy

strategy = FSDPStrategy(auto_wrap_policy=wrap_policy, 
                        activation_checkpointing_policy=checkpointing_policy,
                        )
trainer = pl.Trainer(
        strategy=strategy,
        devices="auto",
        accelerator="gpu",
        check_val_every_n_epoch=4,
        precision=cfg.precision,
        num_nodes=args.num_nodes,
        log_every_n_steps=200,
        max_steps=50000,
        callbacks=callbacks,
        logger=logger,
    )

# The filelist generated in DatasetFunction init function.

class DatasetFunction(Dataset):
    def __init__(self,
                 data_root: str):
        super().__init__()
        self.__data_root = data_root
        self.__files = [npz for npz in os.listdir(
            self.__data_root) if npz.endswith('.npz')]

And for @SkafteNicki 's question, we set devices as "auto", and num_nodes was 1. And we're training on 4 or 8 GPU.

Let me know if you need more details for debug. Thanks for the quick support again.

[SkafteNicki]

@bazingayu thanks for the further information. I will try to reproduce the error, but I do want to ask why you so sure that it is prepare_data that is the coulpit? When num_nodes=1 what essentially happens is that lightning will just call prepare_data as just a regular python function. It happens here:

pytorch-lightning/src/lightning/pytorch/trainer/connectors/data_connector.py

Lines 79 to 101 in 5a2b678

	def prepare_data(self) -> None:
	trainer = self.trainer
	# on multi-gpu jobs we only want to manipulate (download, etc) on node_rank=0, local_rank=0
	# or in the case where each node needs to do its own manipulation in which case just local_rank=0
	local_rank_zero = trainer.local_rank == 0
	global_rank_zero = trainer.local_rank == 0 and trainer.node_rank == 0
	datamodule = trainer.datamodule
	lightning_module = trainer.lightning_module
	# handle datamodule prepare data:
	if datamodule is not None and is_overridden("prepare_data", datamodule):
	prepare_data_per_node = datamodule.prepare_data_per_node
	with _InfiniteBarrier():
	if (prepare_data_per_node and local_rank_zero) or (not prepare_data_per_node and global_rank_zero):
	call._call_lightning_datamodule_hook(trainer, "prepare_data")
	# handle lightning module prepare data:
	if lightning_module is not None and is_overridden("prepare_data", lightning_module):
	prepare_data_per_node = lightning_module.prepare_data_per_node
	with _InfiniteBarrier():
	if (prepare_data_per_node and local_rank_zero) or (not prepare_data_per_node and global_rank_zero):
	call._call_lightning_module_hook(trainer, "prepare_data")

I know it seems like a lot of code but call._call_lightning_datamodule_hook(trainer, "prepare_data") basically just calls the function, no magic. So what I am trying to get to...have you tried taking the logic you have inside the prepare_data and just move it outside lightning as a kind of preprocessing step?

This would help us narrow down if this is truly a issue with prepare_data or maybe it is with FSDP

[bazingayu]

Hi @SkafteNicki , I have already tried taking the logic of prepare data and move it outside pytorch lightning, and the issue has gone.
And also I did the md5 checksum to make sure the data was correctly copied. Then I think it should be prepare_data's issue.

[SkafteNicki]

@bazingayu thanks for confirming, just wanted to make sure to not spend time debugging the wrong location of the codebase.

[bazingayu]

No worries at all, and thanks for looking into this. I'm happy to provide any information you need to resolve this bug.
It's a subtle one that most people wouldn't notice.
We only found it because the validation metrics were inconsistent across different machines, which prompted us to investigate further.