Repository: Lightning-AI/pytorch-lightning

Issue:
Step-wise processing, better support for `IterableDataset`, and others
I have been using PTL for a month. It is nice and saves a lot of time, and I intend to use it in future projects. That said, I have a list of feature requests and improvements that would be very helpful to have to support a wider set of use cases. I am not sure what the best format for this list, so I will just write them here. 

1) Better support for `IterableDataset`
- In addition to `val_check_interval`, we also need `num_val_steps` and `num_train_steps`. `num_val_steps` is needed because the validation set is also using an IterableDataset. `num_train_steps` is needed because you usually need to carefully pick number of gradient updates which has some interaction with the learning rate scheduler (`num_train_steps=inf` is not sufficient)
- For validation, keep the same `DataLoader` object instead of instantiating a new one on each validation cycle, because it is costly to construct new workers each time. 
- Some of the debugging features that run on a small percentage of the training/validation don't work because they are assuming a `Dataset` not `IterableDataset`

2) Step-wise processing
Thinking of the "gradient update" as the unit of training instead of (or in addition to) an epoch. A typical use case is pretraining a language model, where you want to control for number of gradient updates, not epochs (e.g. check the RoBERTa/BERT papers).
- Add an option to do `scheduler.step()` after every gradient update
- Have `self.trainer.num_train_steps` be available for the LR scheduler. The scheduler is usually a function of number of steps
- Checkpointing the current step, and resume from that step. Again, this is important to get the right scheduler, and also important for the tensorboard logging. It will be nice to resume from the same training example, but this is less important. 

3) Misc. These are smaller points, but nice to have.
- Having the default tensorboard logging include LR, time per steps, `allgradnorm` (check fairseq)
- `Trainer(gpus=2)` ignores `CUDA_VISIBLE_DEVICES` and always picks the first two gpus. 
- with ddp, sync validation stats across processes. This is a common mistake, and it will be nice to guard users against it. It is something like having the following line at the end of `validation_end`: 
```
val_loss = torch.distributed.all_reduce(val_loss, op=torch.distributed.ReduceOp.SUM)/self.trainer.world_size
```
- various logs refer to "batches", and it is not clear if it a "batch" or a "step". They are usually the same except with gradient accumulation. Personally, I prefer the word step because it eliminates that confusion. 

Thanks for the helpful library and sorry for the long list. 


Discussion:
> Trainer(gpus=2) ignores CUDA_VISIBLE_DEVICES and always picks the first two gpus.

Trainer cannot ignore CUDA_VISIBLE_DEVICES, as it is handled by the underlaying libraries. 

The device ids in pytorch/tensorflow applications are not consistent with the ids showed in `nvidia-smi `.  E.g, the gpu:0 in pytorch may be the gpu:2 in nvidia-smi. I guess that's why you failed to set the device id to use card you want.

You may set this env variable `export CUDA_DEVICE_ORDER=PCI_BUS_ID`, then the order of device ids will be consistent.

@matthew-z , PTL does override my `CUDA_DEVICE_ORDER` to 0,1,.., making it difficult to run multiple jobs on the same machine. Check the code here: 
https://github.com/PyTorchLightning/pytorch-lightning/blob/06242c200a318a37d1f882c786e60354ec04533f/pytorch_lightning/trainer/distrib_data_parallel.py#L251
and similarity here: https://github.com/PyTorchLightning/pytorch-lightning/blob/06242c200a318a37d1f882c786e60354ec04533f/pytorch_lightning/trainer/distrib_parts.py#L516

Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -83,6 +83,8 @@ def __init__(
             min_nb_epochs=None,  # backward compatible, todo: remove in v0.8.0
             max_epochs=1000,
             min_epochs=1,
+            max_steps=None,
+            min_steps=None,
             train_percent_check=1.0,
             val_percent_check=1.0,
             test_percent_check=1.0,
@@ -345,6 +347,20 @@ def __init__(
                 .. warning:: .. deprecated:: 0.5.0
                     Use `min_nb_epochs` instead. Will remove 0.8.0.
 
+            max_steps (int): Stop training after this number of steps. Disabled by default (None).
+                Training will stop if max_steps or max_epochs have reached (earliest).
+                Example::
+
+                    # Stop after 100 steps
+                    trainer = Trainer(max_steps=100)
+
+            min_steps(int): Force training for at least these number of steps. Disabled by default (None).
+                Trainer will train model for at least min_steps or min_epochs (latest).
+                Example::
+
+                    # Run at least for 100 steps (disable min_epochs)
+                    trainer = Trainer(min_steps=100, min_epochs=0)
+
             train_percent_check (int): How much of training dataset to check.
                 Useful when debugging or testing something that happens at the end of an epoch.
                 Example::
@@ -610,6 +626,9 @@ def __init__(
                 min_epochs = min_nb_epochs
         self.min_epochs = min_epochs
 
+        self.max_steps = max_steps
+        self.min_steps = min_steps
+
         # Backward compatibility
         if nb_sanity_val_steps is not None:
             warnings.warn("`nb_sanity_val_steps` has renamed to `num_sanity_val_steps` since v0.5.0"
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -372,10 +372,17 @@ def train(self):
                     raise MisconfigurationException(m)
                 self.reduce_lr_on_plateau_scheduler.step(val_loss, epoch=self.current_epoch)
 
+            if self.max_steps and self.max_steps == self.global_step:
+                self.main_progress_bar.close()
+                model.on_train_end()
+                return
+
             # early stopping
             met_min_epochs = epoch >= self.min_epochs - 1
+            met_min_steps = self.global_step >= self.min_steps if self.min_steps else True
+
             if (self.enable_early_stop and not self.disable_validation and is_val_epoch and
-                    (met_min_epochs or self.fast_dev_run)):
+                    ((met_min_epochs and met_min_steps) or self.fast_dev_run)):
                 should_stop = self.early_stop_callback.on_epoch_end()
                 # stop training
                 stop = should_stop and met_min_epochs
@@ -463,6 +470,10 @@ def run_training_epoch(self):
                 self.global_step += 1
             self.total_batch_idx += 1
 
+            # max steps reached, end training
+            if self.max_steps is not None and self.max_steps == self.global_step:
+                break
+
             # end epoch early
             # stop when the flag is changed or we've gone past the amount
             # requested in the batches


Test patch:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -1,3 +1,4 @@
+import math
 import os
 
 import pytest
@@ -6,6 +7,7 @@
 import tests.models.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import (
+    EarlyStopping,
     ModelCheckpoint,
 )
 from tests.models import (
@@ -447,5 +449,89 @@ class CurrentTestModel(
     trainer.test()
 
 
+def _init_steps_model():
+    """private method for initializing a model with 5% train epochs"""
+    tutils.reset_seed()
+    model, _ = tutils.get_model()
+
+    # define train epoch to 5% of data
+    train_percent = 0.05
+    # get number of samples in 1 epoch
+    num_train_samples = math.floor(len(model.train_dataloader()) * train_percent)
+
+    trainer_options = dict(
+        train_percent_check=train_percent,
+    )
+    return model, trainer_options, num_train_samples
+
+
+def test_trainer_max_steps_and_epochs(tmpdir):
+    """Verify model trains according to specified max steps"""
+    model, trainer_options, num_train_samples = _init_steps_model()
+
+    # define less train steps than epochs
+    trainer_options.update(dict(
+        max_epochs=5,
+        max_steps=num_train_samples + 10
+    ))
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training did not complete"
+
+    # check training stopped at max_steps
+    assert trainer.global_step == trainer.max_steps, "Model did not stop at max_steps"
+
+    # define less train epochs than steps
+    trainer_options['max_epochs'] = 2
+    trainer_options['max_steps'] = trainer_options['max_epochs'] * 2 * num_train_samples
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training did not complete"
+
+    # check training stopped at max_epochs
+    assert trainer.global_step == num_train_samples * trainer.max_nb_epochs \
+        and trainer.current_epoch == trainer.max_nb_epochs - 1, "Model did not stop at max_epochs"
+
+
+def test_trainer_min_steps_and_epochs(tmpdir):
+    """Verify model trains according to specified min steps"""
+    model, trainer_options, num_train_samples = _init_steps_model()
+
+    # define callback for stopping the model and default epochs
+    trainer_options.update({
+        'early_stop_callback': EarlyStopping(monitor='val_loss', min_delta=1.0),
+        'val_check_interval': 20,
+        'min_epochs': 1,
+        'max_epochs': 10
+    })
+
+    # define less min steps than 1 epoch
+    trainer_options['min_steps'] = math.floor(num_train_samples / 2)
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training did not complete"
+
+    # check model ran for at least min_epochs
+    assert trainer.global_step >= num_train_samples and \
+        trainer.current_epoch > 0, "Model did not train for at least min_epochs"
+
+    # define less epochs than min_steps
+    trainer_options['min_steps'] = math.floor(num_train_samples * 1.5)
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training did not complete"
+
+    # check model ran for at least num_train_samples*1.5
+    assert trainer.global_step >= math.floor(num_train_samples * 1.5) and \
+        trainer.current_epoch > 0, "Model did not train for at least min_steps"
+
 # if __name__ == '__main__':
 #     pytest.main([__file__])