[PoC] Training Restart - V2 #8337

tchaton · 2021-07-08T15:53:22Z

What does this PR do?

Different Loop API compared to #8131

TODOS
[] Resolve batch_idx not properly reset in loops
[] Add ProgressBar proper indices
[] Resolve _map_dl_idx_sampler_states should be updated after the batch is processed
[] Save extras for on epoch_end
[] Add ResultCollection to state dict and reload automatically on Base Loop
[] Enable non-fault tolerant reload.

Does your PR introduce any breaking changes ? If yes, please list them.

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)
Did you list all the breaking changes introduced by this pull request?

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

…TorchLightning/pytorch-lightning into refactor/logger-connector-poc

…s_everywhere_train

rename files and classes

for more information, see https://pre-commit.ci

…s_everywhere_train

awaelchli · 2021-07-09T12:18:29Z

pytorch_lightning/loops/base.py

        if self.restarting:
-            self.restore()
-            self.restarting = False
+            if not is_overridden("restore", self, Loop):
+                warning_cache.warn(f"{self.__class__.__name__} Loop doesn't override the restore function.")
+            is_restore_finished = self.restore()
+            if is_restore_finished is None or is_restore_finished:
+                self.restarting = False


instead of

if restarting self.restore() else: self.reset()

would it make sense to do:

self.reset() if restarting: self.restore()

?

Then we don't need the self._initialize() in the subclassed loops and we can assume after reset() the state is properly initialized and we can restore if appropritae.

awaelchli

review part 1/n

this PR is integrating progress tracking
do you see a way to break it out or is it rather not so easy?

awaelchli · 2021-07-09T15:36:03Z

pytorch_lightning/loops/base.py

+        Restore the internal state of the loop the beginning of run if restarting is ``True``.
+
+        Returns:
+            is_restore_finished: Whether the restoration actually finished.


in which case will a restoration not finish.
is this this about a success/fail return signal?

This is just a mechanism to enable a loop to delay the restarting variable to be set to False.

In TrainingBatchLoop, this property is required to know if we should fast-forward the dataloaders, therefore by return False, I indicate to the loop I will manually switch off the restarting boolean later on in the loop.

Example:
https://github.com/PyTorchLightning/pytorch-lightning/pull/8337/files#diff-eba04421e2e60c9d7b55f56ca57c9a319a4e3de38e3eb96c931324ecacbb73e0R175

awaelchli · 2021-07-09T15:38:49Z

pytorch_lightning/loops/base.py

+    def load_state_dict(self, state_dict: Dict) -> None:
+        """Reload Loop state"""
+
+    def get_state_dict(self, destination: Optional[OrderedDict] = None, prefix: Optional[str] = '') -> OrderedDict:


as discussed, could we consider merging several of these state_dict related APIs into one.
IMO we should aim at keeping the base loop class as simple as possible.

awaelchli · 2021-07-09T15:40:30Z

pytorch_lightning/loops/batch/training_batch_loop.py

+    def restore(self) -> bool:
+        """Restore the loop state"""
+        self._initialize()
+
+        # restarting isn't finished yet.
+        return False
+
+    def reset(self) -> None:
+        """Resets the loop state"""
+        self._initialize()
+
+        self.optim_progress.optimizer.reset_on_epoch()


these _initialize() calls could be simplified as proposed in the earlier comment

awaelchli · 2021-07-09T15:41:06Z

pytorch_lightning/loops/batch/training_batch_loop.py

+        self.progress.increment_completed()
+        return super().on_advance_end()


we should probably increment after the super call?

awaelchli · 2021-07-09T15:41:59Z

pytorch_lightning/loops/batch/training_batch_loop.py

+                # handle optimization restart
+                if self.restarting:
+                    if len(active_optimizers) > 1 and opt_idx < self.progress.current.completed:
+                        continue
+                    self.restarting = False


okay I see.
this is something a optimizer loop could handle for restarting.

Yes, exactly !

awaelchli · 2021-07-09T15:45:37Z

pytorch_lightning/loops/epoch/training_epoch_loop.py

    @property
    def done(self) -> bool:
        """Returns whether the training should be stopped.
        The criteria are that the number of steps reached the max steps,
        the last batch is reached or the trainer signals to stop (e.g. by early stopping).
        """
-        max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps
+        max_steps_reached = self.max_steps is not None and (self.total_optimizer_step) >= self.max_steps


max steps in the trainer does not represent optimizer steps (today). It represents how many times training_step() was called (roughly). If we want to change that I think it should be done in a separate PR with a proper evaluation of what the effects are.

awaelchli · 2021-07-09T15:49:36Z

pytorch_lightning/plugins/training_type/ddp.py

+        # share ddp pids to all processes
+        self.share_pids()
+
+    def share_pids(self):


leftover from a merge :)

awaelchli · 2021-07-09T15:50:30Z

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -141,6 +144,13 @@ def restore_model(self) -> None:
        # restore model state_dict
        self.trainer.training_type_plugin.load_model_state_dict(self._loaded_checkpoint)

+        gradients = self._loaded_checkpoint.get("gradients", None)


what would we call this feature? gradient checkpointing 🤣

If the users had gradient accumulation activated.

awaelchli · 2021-07-09T15:54:30Z

pytorch_lightning/trainer/connectors/loop_connector.py

+        self.trainer.fit_loop = self.on_fit_loop_init(*args, **kwargs)
+        self.trainer.validate_loop = self.on_validate_loop_init(*args, **kwargs)
+        self.trainer.test_loop = self.on_test_loop_init(*args, **kwargs)
+        self.trainer.predict_loop = self.on_predict_loop_init(*args, **kwargs)


a loop connector could be useful in general, but we said we wanted to avoid Trainer state management outside the trainer. If Trainer owns loops, imo it would be safer to let the Trainer manage its state.

cc @ananthsub @justusschock

awaelchli · 2021-07-09T15:58:31Z

pytorch_lightning/trainer/data_loading.py

@@ -111,6 +115,41 @@ def auto_add_worker_init_fn(self, dataloader: DataLoader) -> None:
        if int(os.environ.get("PL_SEED_WORKERS", 0)) and dataloader.worker_init_fn is None:
            dataloader.worker_init_fn = partial(pl_worker_init_function, rank=self.global_rank)

+    def add_samplers_to_iterable_dataset(self, dataloader: DataLoader):


here we have a little bit of code duplication from the replace_sampler method.
We could consider sharing some code for the DataLoader attribute discovery and patching.

awaelchli · 2021-07-09T16:01:26Z

pytorch_lightning/trainer/data_loading.py

+        # detect iterable dataset
+        contains_iterable_dataset = False
+
+        def detect_iterable_dataset(dataloader: DataLoader):
+            nonlocal contains_iterable_dataset
+            if isinstance(dataloader.dataset, IterableDataset):
+                contains_iterable_dataset = True
+
+        apply_to_collection(loaders, DataLoader, detect_iterable_dataset)


could potentially be factored out into a helper function.

awaelchli · 2021-07-09T16:13:09Z

pytorch_lightning/trainer/progress.py

        obj = cls()
        obj.load_state_dict(state_dict)
        return obj


+class ProgressDict(Dict):


isn't this what a regular dict can do?

old_dict = dict(...) new_dict = dict(old_dict) new_dict.update(...)

awaelchli · 2021-07-09T16:16:38Z

pytorch_lightning/trainer/supporters.py

+                partial(cycle_to_next_worker, state_dict=self.loaders_iter_state_dict, count=count)
+            )
+
+            count = 0


these local vars can be removed, right?

awaelchli · 2021-07-09T16:18:54Z

pytorch_lightning/trainer/supporters.py

+    if trainer:
+        trainer.current_iterator = it


can you explain what this workaround is about?
I can see that you are using it in the checkpoint connector but it is not clear to me what problem it is trying to solve.

This can be removed ! This was used to not lose the Iterator reference as it contains the DataLoader Iterators.

awaelchli · 2021-07-09T16:21:24Z

pytorch_lightning/trainer/trainer.py

@@ -995,12 +995,41 @@ def _run_train(self) -> None:
            if distributed_available() and self.world_size > 1:
                # try syncing remaing processes, kill otherwise
                self.training_type_plugin.reconciliate_processes(traceback.format_exc())
+            # save a checkpoint for fault tolerant training
+            self.fit_loop._check_checkpoint_callback(True)


@carmocca that's what I meant the other day when we talked about why we have the check_checkpoint_callback in the first place.

this usage here was the old way Lightning would save checkpoints on interrupt.

We shouldnt manually call the ModelCheckpoint callback like this as it will be moved to a fully hook-based callback next week.

If training is stopped and we need to save, we should do it by calling the on_train_end hook

Adding to this:

If we need to manually save a checkpoint, we can always call trainer.save_checkpoint() and this doesn't go through the ModelCheckpoint :)

I would prefer for the ModelCheckpoint to save the checkpoint as the last_model_path should be the one saved on failure.

awaelchli · 2021-07-09T16:24:15Z

pytorch_lightning/utilities/auto_restart.py

+                    state_dict.append(iterable_dataset_state_dict)
+                    return data["data"]
+                out.append((k, CaptureIterableDataset._sanetize_batch_from_sampler_state(v, state_dict)))
+            return elem_type(OrderedDict(out))


is there a specific reason for using OrderedDict vs. just dict?

Following apply_to_collection implementation. I believe there were a reason why OrderedDict was selected over it, but not sure why.

stale · 2021-07-25T12:06:08Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. If you need further help see our docs: https://pytorch-lightning.readthedocs.io/en/latest/generated/CONTRIBUTING.html#pull-request or ask the assistance of a core contributor here or on Slack. Thank you for your contributions.

stale · 2021-07-30T13:29:38Z

This pull request is going to be closed. Please feel free to reopen it create a new from the actual master.

carmocca and others added 30 commits June 9, 2021 00:10

Update logger connector

37a0b9d

Update result

aaea387

Update imports

e2f69ce

Update after rename

6037833

Merge branch 'refactor/logger-connector-poc' of https://github.com/Py…

3804963

…TorchLightning/pytorch-lightning into refactor/logger-connector-poc

Refactor reduce_fx and op

499da76

Fix test after rename

6eb448a

mypy

f871cbd

manual merge poc changes

5631b53

integrate latest changes from logger connector poc

d10d5c7

Fix test

7b6803a

Refactor test

9bfedc9

Deprecate self.log(sync_dist_op) in favor of self.log(reduce_fx)

c9c7829

Undo field

e3dde0b

Merge branch 'refactor/logger-connector-poc' into refactor/loops/loop…

bae2139

…s_everywhere_train

rename

2c167cc

rename files and classes

[pre-commit.ci] auto fixes from pre-commit.com hooks

99db497

for more information, see https://pre-commit.ci

rename

832dfb9

imports

f92e01d

loop hygiene

b15fc34

yapf on loops

7175a50

protected new loop trigger

59d6227

Merge branch 'master' into refactor/logger-connector-poc

e1d4fd2

Replace code

a7c3555

Fix names and imports

501224d

Remove metric_attribute

dee7e5f

Merge branch 'refactor/logger-connector-poc' into refactor/loops/loop…

4eb9757

…s_everywhere_train

integrate latest logger connector changes

d4bb357

resolve todo dataloading reset

c9b4e9e

re-add notebooks

a3ef0aa

tchaton added 5 commits July 8, 2021 20:43

update

2dace4c

cleanup

b05350f

resolve some tests

29b0c0e

update

c70a024

remove cache state

c65a87c

awaelchli reviewed Jul 9, 2021

View reviewed changes

tchaton added 16 commits July 9, 2021 19:21

resolve tests

76785cd

cleanup dataloader checkpoiting

664829a

resolve comments

7f8000f

update

4153b81

update

d6280e0

update changelog

c499c24

Merge branch 'loop_improvement' into training_restart_2

cd3d95b

update

3cb6df2

Merge branch 'loop_improvement' into training_restart_2

a8cbd78

update

2ee2217

resolve bugs

1fa94ea

update

e8c12e9

remove space

df4b1ba

update

ee8d9b8

Merge branch 'loop_improvement' into training_restart_2

f203b1e

resolve bugs

af4ee87

tchaton mentioned this pull request Jul 10, 2021

[Feat] Add Loops Restart #8131

Closed

11 tasks

stale bot added the won't fix This will not be worked on label Jul 25, 2021

stale bot closed this Jul 30, 2021

Borda deleted the training_restart_2 branch November 4, 2021 23:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PoC] Training Restart - V2 #8337

[PoC] Training Restart - V2 #8337

tchaton commented Jul 8, 2021 •

edited

Loading

awaelchli Jul 9, 2021

awaelchli left a comment

awaelchli Jul 9, 2021

tchaton Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

tchaton Jul 9, 2021

awaelchli Jul 9, 2021

tchaton Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

tchaton Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

awaelchli Jul 9, 2021

tchaton Jul 9, 2021

awaelchli Jul 9, 2021

carmocca Jul 9, 2021

awaelchli Jul 9, 2021

tchaton Jul 11, 2021

awaelchli Jul 9, 2021

tchaton Jul 9, 2021 •

edited

Loading

stale bot commented Jul 25, 2021

stale bot commented Jul 30, 2021

		self.progress.increment_completed()
		return super().on_advance_end()

[PoC] Training Restart - V2 #8337

[PoC] Training Restart - V2 #8337

Conversation

tchaton commented Jul 8, 2021 • edited Loading

What does this PR do?

Does your PR introduce any breaking changes ? If yes, please list them.

Before submitting

PR review

Did you have fun?

Choose a reason for hiding this comment

awaelchli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tchaton Jul 9, 2021 • edited Loading

Choose a reason for hiding this comment

stale bot commented Jul 25, 2021

stale bot commented Jul 30, 2021

tchaton commented Jul 8, 2021 •

edited

Loading

tchaton Jul 9, 2021 •

edited

Loading