How to get correct estimated_stepping_batches in IterableDataset

[ShaneTian]

When I use IterableDataset without __len__，the estimated_stepping_batches is always -1.
So how to get correct estimated_stepping_batches?

[rohitgr7]

it's not possible to determine that since no one can really tell how many iterations will happen in such a case.

[ShaneTian]

OK. So I think it is necessary to write it in the doc.

[rohitgr7]

may I know how are you running it with the iterable dataset? like what is your trainer configuration?
asking because, it doesn't return -1, but value for max_steps and since in your case, I am assuming this is -1, that means max_steps is not set so I am curious to see the trainer configuration :)

[ShaneTian]

This is the pseudocode:

class MyDataset(IterableDataset):
    def __init__(self):
        # ...

    def __iter__(self):
        # ...


class MyDataModule(LightningDataModule):
    def __init__(self):
        # ...

    def train_dataloader(self):
        # ...

    def val_dataloader(self):
        # ...


trainer = pl.Trainer(
    gpus=-1,
    accelerator="gpu",
    strategy="deepspeed_stage_2",
    logger=tb_logger,
    callbacks=callbacks,
    # use `max_epochs` instead of `max_steps`
    max_epochs=10,
    accumulate_grad_batches=1,
    gradient_clip_val=1,
    log_every_n_steps=100,
    val_check_interval=100,
    weights_summary="top"
)
trainer.fit(model, datamodule=mydatamodule)

[rohitgr7]

well in this case it's not possible to determine the total length of the dataset since it's running based on max_epochs and there must be some count that it iterates over each epoch. If you know about it, you can set it inside your IterableDataset and Lightning will fetch it.

class IterableDS(IterableDataset):
    def __len__(self):
        return some_count

[ShaneTian]

Once I set the dataset length, two problems will arise:

• A warning about multi-process data loading
• Same issue as Can't use estimated_stepping_batches in configure_optimizers with DDP #12317

Loading `train_dataloader` to estimate number of stepping batches.
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/data.py:107: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-process data loading (when num_workers > 1), `__len__` could be inaccurate if each worker is not configured independently to avoid having duplicate data.
  "Your `IterableDataset` has `__len__` defined."
Traceback (most recent call last):
  File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/root/workspace/env_run/shane-transformers/train.py", line 161, in <module>
    train(args)
  File "/root/workspace/env_run/shane-transformers/train.py", line 141, in train
    trainer.fit(model, datamodule=model.datamodule)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 772, in fit
    self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 724, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 812, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1218, in _run
    self.strategy.setup(self)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/deepspeed.py", line 362, in setup
    self.init_deepspeed()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/deepspeed.py", line 461, in init_deepspeed
    self._initialize_deepspeed_train(model)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/deepspeed.py", line 490, in _initialize_deepspeed_train
    optimizer, lr_scheduler, _ = self._init_optimizers()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/deepspeed.py", line 466, in _init_optimizers
    optimizers, lr_schedulers, optimizer_frequencies = _init_optimizers_and_lr_schedulers(self.lightning_module)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/optimizer.py", line 180, in _init_optimizers_and_lr_schedulers
    optim_conf = model.trainer._call_lightning_module_hook("configure_optimizers", pl_module=model)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1596, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/root/workspace/env_run/shane-transformers/models/t5.py", line 187, in configure_optimizers
    print("self.trainer.estimated_stepping_batches", self.trainer.estimated_stepping_batches)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 2816, in estimated_stepping_batches
    self.reset_train_dataloader()
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py", line 1900, in reset_train_dataloader
    if has_len_all_ranks(self.train_dataloader, self.strategy, module)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/data.py", line 124, in has_len_all_ranks
    total_length = training_type.reduce(torch.tensor(local_length).to(model.device), reduce_op="sum")
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/ddp.py", line 339, in reduce
    tensor = sync_ddp_if_available(tensor, group, reduce_op=reduce_op)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/distributed.py", line 94, in sync_ddp_if_available
    return sync_ddp(result, group=group, reduce_op=reduce_op)
  File "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/distributed.py", line 136, in sync_ddp
    torch.distributed.all_reduce(result, op=op, group=group, async_op=False)
  File "/usr/local/lib/python3.7/dist-packages/torch/distributed/distributed_c10d.py", line 1287, in all_reduce
Traceback (most recent call last):
  File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    work = group.allreduce([tensor], opts)
RuntimeError: Tensors must be CUDA and dense

It will happen under both DDP and DeepSpeed.