Will more dataloaders/datasets cause GPU/CPU memory usage to increase?

[24-solar-terms]

I ran into a problem while training my model.

I have 21 datasets, each with thousand samples, and Out of CUDA Memory occurs when I return a list of dataloaders using LightningDataModule's train_dataloader (self) function.

So I tried using 1 dataset, which returned a list of 1 dataloder, it worked fine and cost 23G GPU memory.
When I used 2 datasets, which returned a list of 2 dataloaders, it cost 42G GPU memory.

I tried use CombinedLoader in pytorch lightning to deal with list of dataloaders, but the problem wasn't solved.
However, when I used ConcatDataset in pytorch to combine datasets and created 1 dataloader, CUDA memory usage was stable, but CPU memory increasing while training until OOM occured.

Here is my pytorch dataset：

def fast_load_dataset(dataset_name_or_path, split):
    """
    only load metadata.jsonl
    Args:
        dataset_name_or_path (str): dataset basedir
        split (str): （train, validation, test）

    Raises:
        Exception: _description_

    Returns:
        list: list of dict
    """
    with open(os.path.join(dataset_name_or_path, split, 'metadata.jsonl'), 'r', encoding='utf-8') as f:
        jsonl_reader = jsonlines.Reader(f)
        dataset = list(jsonl_reader)
    return dataset


class DonutDataset(Dataset):
    """
    DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
    Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
    and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string)

    Args:
        dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
        ignore_id: ignore_index for torch.nn.CrossEntropyLoss
        task_start_token: the special token to be fed to the decoder to conduct the target task
    """

    def __init__(
        self,
        dataset_name_or_path: str,
        donut_model: PreTrainedModel,
        max_length: int,
        split: str = "train",
        ignore_id: int = -100,
        task_start_token: str = "<s>",
        prompt_end_token: str = None,
        sort_json_key: bool = True,
    ):
        super().__init__()

        self.donut_model = donut_model
        self.max_length = max_length
        self.dataset_name_or_path = dataset_name_or_path
        self.split = split
        self.ignore_id = ignore_id
        self.task_start_token = task_start_token
        self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
        self.sort_json_key = sort_json_key

        self.dataset = fast_load_dataset(dataset_name_or_path, split=self.split)  # list of dict
        self.dataset_length = len(self.dataset)

        self.gt_token_sequences = []
        for sample in self.dataset:
            ground_truth = json.loads(sample["ground_truth"])
            if "gt_parses" in ground_truth:  # when multiple ground truths are available, e.g., docvqa
                assert isinstance(ground_truth["gt_parses"], list)
                gt_jsons = ground_truth["gt_parses"]
            else:
                assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
                gt_jsons = [ground_truth["gt_parse"]]

            self.gt_token_sequences.append(
                [
                    task_start_token
                    + self.donut_model.json2token(
                        gt_json,
                        update_special_tokens_for_json_key=self.split == "train",
                        sort_json_key=self.sort_json_key,
                    )
                    + self.donut_model.decoder.tokenizer.eos_token
                    for gt_json in gt_jsons  # load json from list of json
                ]
            )

        self.donut_model.decoder.add_special_tokens([self.task_start_token, self.prompt_end_token])
        self.prompt_end_token_id = self.donut_model.decoder.tokenizer.convert_tokens_to_ids(self.prompt_end_token)

    def __len__(self) -> int:
        return self.dataset_length

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Load image from image_path of given dataset_path and convert into input_tensor and labels.
        Convert gt data into input_ids (tokenized string)

        Returns:
            input_tensor : preprocessed image
            input_ids : tokenized gt_data
            labels : masked labels (model doesn't need to predict prompt and pad token)
        """
        sample = self.dataset[idx]
        # match fast_load_dataset，read img during training
        try:
            sample["image"] = Image.open(os.path.join(self.dataset_name_or_path, self.split, sample["file_name"]))
        except Exception as e:
            # use opencv if img corrupt
            tmp_img = cv2.imread(os.path.join(self.dataset_name_or_path, self.split, sample["file_name"]))
            tmp_img = cv2.cvtColor(tmp_img, cv2.COLOR_BGR2RGB)
            sample["image"] = Image.fromarray(tmp_img)
        # input_tensor
        input_tensor = self.donut_model.encoder.prepare_input(sample["image"], random_padding=self.split == "train")

        # input_ids
        processed_parse = random.choice(self.gt_token_sequences[idx])  # can be more than one, e.g., DocVQA Task 1
        input_ids = self.donut_model.decoder.tokenizer(
            processed_parse,
            add_special_tokens=False,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )["input_ids"].squeeze(0)

        if self.split == "train":
            labels = input_ids.clone()
            labels[
                labels == self.donut_model.decoder.tokenizer.pad_token_id
            ] = self.ignore_id  # model doesn't need to predict pad token
            labels[
                : torch.nonzero(labels == self.prompt_end_token_id).sum() + 1
            ] = self.ignore_id  # model doesn't need to predict prompt (for VQA)
            return input_tensor, input_ids, labels
        else:
            prompt_end_index = torch.nonzero(
                input_ids == self.prompt_end_token_id
            ).sum()  # return prompt end index instead of target output labels
            return input_tensor, input_ids, promp

Here is my LightningDataModule

class DonutDataPLModule(pl.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.train_batch_sizes = self.config.train_batch_sizes
        self.val_batch_sizes = self.config.val_batch_sizes
        self.train_datasets = []
        self.val_datasets = []
        self.g = torch.Generator()
        self.g.manual_seed(self.config.seed)

    def train_dataloader(self):
        concat_dataset = ConcatDataset(self.train_datasets)
        return DataLoader(
            concat_dataset,
            batch_size=self.train_batch_sizes[0],
            num_workers=self.config.num_workers,
            pin_memory=False,
            worker_init_fn=self.seed_worker,
            generator=self.g,
            shuffle=True,
        )
        # loaders = list()
        # for train_dataset, batch_size in zip(self.train_datasets, self.train_batch_sizes):
        #     loaders.append(
        #         DataLoader(
        #             train_dataset,
        #             batch_size=batch_size,
        #             num_workers=self.config.num_workers,
        #             pin_memory=True,
        #             worker_init_fn=self.seed_worker,
        #             generator=self.g,
        #             shuffle=True,
        #         )
        #     )
        # return loaders

    def val_dataloader(self):
        concat_dataset = ConcatDataset(self.val_datasets)
        return DataLoader(
            concat_dataset,
            batch_size=self.val_batch_sizes[0],
            pin_memory=False,
            shuffle=False,
        )
        # loaders = list()
        # for val_dataset, batch_size in zip(self.val_datasets, self.val_batch_sizes):
        #     loaders.append(
        #         DataLoader(
        #             val_dataset,
        #             batch_size=batch_size,
        #             pin_memory=True,
        #             shuffle=False,
        #         )
        #     )
        # return loaders

    @staticmethod
    def seed_worker(wordker_id):
        worker_seed = torch.initial_seed() % 2 ** 32
        np.random.seed(worker_seed)
        random.seed(worker_seed)

Here is train.py

data_module = DonutDataPLModule(config)
...
for i, dataset_name_or_path in enumerate(config.dataset_name_or_paths):
        ...
        for split in ["train", "validation"]:
            datasets[split].append(
                DonutDataset(
                    dataset_name_or_path=dataset_name_or_path,
                    donut_model=model_module.model,
                    max_length=config.max_length,
                    split=split,
                    task_start_token=config.task_start_tokens[i]
                    if config.get("task_start_tokens", None)
                    else f"<s_{task_name}>",
                    prompt_end_token="<s_synthdog>",
                    sort_json_key=config.sort_json_key,
                )
            )
            # prompt_end_token is used for ignoring a given prompt in a loss function
            # for docvqa task, i.e., {"question": {used as a prompt}, "answer": {prediction target}},
            # set prompt_end_token to "<s_answer>"
    data_module.train_datasets = datasets["train"]
    data_module.val_datasets = datasets["validation"]

Where is the potential for a memory leak？