help with IterableDataModule

[zEpoch]

class MyIterableDataset(IterableDataset):
    def __init__(self, parquet_dir, batch_size):
        self.parquet_dir = parquet_dir
        self.parquet_files = [
            os.path.join(parquet_dir, f) for f in os.listdir(parquet_dir) if f.endswith(".parquet")
        ]
        self.batch_size = batch_size

    def __iter__(self):
        for parquet_file in self.parquet_files:
            parquet_file_obj = pq.ParquetFile(parquet_file)
            for batch in parquet_file_obj.iter_batches(self.batch_size):
                # df = batch.to_pandas()
                # for _, row in df.iterrows():
                #     yield row.to_dict()
                yield batch['sequence']

class MyIterableDataModule(L.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.train_batch_size = config.train_batch_size
        self.val_batch_size = config.val_batch_size
        self.test_batch_size = config.test_batch_size
        self.train_parquet_dir = config.train_parquet_dir
        self.val_parquet_dir = config.val_parquet_dir
        self.test_parquet_dir = config.test_parquet_dir

    def setup(self, stage=None):
        if stage in (None, "fit"):
            self.train_dataset = MyIterableDataset(self.train_parquet_dir, batch_size=self.train_batch_size)
        
        if stage in (None, "validate") and self.val_parquet_dir:
            self.val_dataset = MyIterableDataset(self.val_parquet_dir, batch_size=self.val_batch_size)
        
        if stage in (None, "test") and self.test_parquet_dir:
            self.test_dataset = MyIterableDataset(self.test_parquet_dir, batch_size=self.test_batch_size)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.train_batch_size)

    def val_dataloader(self):
        if hasattr(self, "val_dataset"):
            return DataLoader(self.val_dataset, batch_size=self.val_batch_size)
        return None

    def test_dataloader(self):
        if hasattr(self, "test_dataset"):
            return DataLoader(self.test_dataset, batch_size=self.test_batch_size)
        return None
dm = MyIterableDataModule(config)
trainer = Trainer(
        max_epochs=5,
        accelerator="gpu",
        devices='auto',
        strategy="deepspeed",
        num_nodes=config.n_nodes, 
        val_check_interval=1,
        logger=logger if config.use_wandb else None, #type: ignore
        accumulate_grad_batches=4, 
        precision=16, 
        log_every_n_steps=1, 
    )
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params}")
    trainer.fit(model, datamodule=dm)

but it get wrong with:

Traceback (most recent call last):
  File "/home/share/huadjyin/home/zhoutao3/MoBADNA/train.py", line 87, in <module>
    main(config)
  File "/home/share/huadjyin/home/zhoutao3/MoBADNA/train.py", line 60, in main
    trainer.fit(model, datamodule=dm)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 47, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
Traceback (most recent call last):
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py", line 386, in _check_dataloader_iterable
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1054, in _run_stage
    self._run_sanity_check()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1083, in _run_sanity_check
    val_loop.run()
Traceback (most recent call last):
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py", line 386, in _check_dataloader_iterable
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py", line 179, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 120, in run
    self.setup_data()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 189, in setup_data
    _check_dataloader_iterable(dl, source, trainer_fn)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py", line 402, in _check_dataloader_iterable
    raise TypeError(
TypeError: An invalid dataloader was returned from `MyIterableDataModule.val_dataloader()`. Found None.
    iter(dataloader)  # type: ignore[call-overload]    
TypeErroriter(dataloader)  # type: ignore[call-overload]: 'NoneType' object is not iterable

[zEpoch]

class MyIterableDataset(IterableDataset):
    def __init__(self, parquet_dir, batch_size):
        self.parquet_dir = parquet_dir
        self.parquet_files = [
            os.path.join(parquet_dir, f) for f in os.listdir(parquet_dir) if f.endswith(".parquet")
        ]
        self.batch_size = batch_size

    def __iter__(self):
        for parquet_file in self.parquet_files:
            parquet_file_obj = pq.ParquetFile(parquet_file)
            for batch in parquet_file_obj.iter_batches(self.batch_size):
                # df = batch.to_pandas()
                # for _, row in df.iterrows():
                #     yield row.to_dict()
                yield batch['sequence']

class MyIterableDataModule(L.LightningDataModule):
    def __init__(self, config):
        super().__init__()
        self.train_batch_size = config.train_batch_size
        self.val_batch_size = config.val_batch_size
        self.test_batch_size = config.test_batch_size
        self.train_parquet_dir = config.train_parquet_dir
        self.val_parquet_dir = config.val_parquet_dir
        self.test_parquet_dir = config.test_parquet_dir

    def setup(self, stage=None):
        # if stage in (None, "fit"):
        self.train_dataset = MyIterableDataset(self.train_parquet_dir, batch_size=self.train_batch_size)
        
        # if stage in (None, "validate") and self.val_parquet_dir:
        self.val_dataset = MyIterableDataset(self.val_parquet_dir, batch_size=self.val_batch_size)
        
        # if stage in (None, "test") and self.test_parquet_dir:
        self.test_dataset = MyIterableDataset(self.test_parquet_dir, batch_size=self.test_batch_size)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.train_batch_size)

    def val_dataloader(self):
        if hasattr(self, "val_dataset"):
            return DataLoader(self.val_dataset, batch_size=self.val_batch_size)
        return None

    def test_dataloader(self):
        if hasattr(self, "test_dataset"):
            return DataLoader(self.test_dataset, batch_size=self.test_batch_size)
        return None

and then:

Traceback (most recent call last):
  File "/home/share/huadjyin/home/zhoutao3/MoBADNA/train.py", line 87, in <module>
    main(config)
  File "/home/share/huadjyin/home/zhoutao3/MoBADNA/train.py", line 61, in main
    trainer.fit(model, datamodule=dm)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 47, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1054, in _run_stage
    self._run_sanity_check()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1083, in _run_sanity_check
    val_loop.run()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py", line 179, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 138, in run
    batch, batch_idx, dataloader_idx = next(data_fetcher)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/fetchers.py", line 134, in __next__
    batch = super().__next__()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/loops/fetchers.py", line 61, in __next__
    batch = next(self.iterator)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/utilities/combined_loader.py", line 341, in __next__
    out = next(self._iterator)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/pytorch_lightning/utilities/combined_loader.py", line 142, in __next__
    out = next(self.iterators[0])
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 42, in fetch
    return self.collate_fn(data)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 277, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
  File "/home/share/huadjyin/home/zhoutao3/.conda/envs/MoBA/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 152, in collate
    raise TypeError(default_collate_err_msg_format.format(elem_type))
TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class 'pyarrow.lib.StringArray'>