[Train] Unifying Lightning and AIR CheckpointConfig #36368

woshiyyya · 2023-06-13T08:42:02Z

Why are these changes needed?

Currently the users have to specify checkpoint configs twice with same parameters, one in LightningConfigBuilder.checkpointing(), one in AIR CheckpointConfig.

If the users only provides configures in LightningConfigBuilder, by default AIR will save all checkpoints, which is not desired. This PR is intended to automatically generate an AIR CheckpointConfig for the user if they forget to provide one.

(Another minor fix, raised a RuntimeError when users forget to provide a datasets_iter_config with datasets.)

Related issue number

Closes #35920
Closes #36509

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

…ning_configs

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

justinvyu

Left some comments.

One other thing: can we document what should be passed in through LightningTrainer(datasets)? I can't find any mentions of the special "val" key that needs to be passed in to enable evaluation.

Also, we have this EVALUATION_DATASET_KEY that TransformersTrainer uses. Does it make sense to use that instead of "val"?

python/ray/train/tests/test_lightning_trainer.py

python/ray/train/lightning/lightning_trainer.py

justinvyu · 2023-06-23T18:39:59Z

python/ray/train/lightning/lightning_trainer.py

@@ -427,6 +447,17 @@ def _check_checkpoint_configs(
                "through `LightningConfigBuilder.checkpointing()`."
            )

+        # Auto-fill the AIR CheckpointConfig if the user didn't specify it.
+        if air_ckpt_config == CheckpointConfig():
+            save_top_k = ptl_ckpt_config.get("save_top_k", 1)


We now default AIR to save 1 checkpoint when using it with lightning trainer -- is this intentional? PTL only saves 1 checkpoint by default?

Yeah, in PTL, the default value of save_top_k is 1. Also, by default monitor is None which saves a checkpoint only for the last epoch. We are applying default behavior of PTL to AIR now.

python/ray/train/lightning/lightning_trainer.py

matthewdeng · 2023-06-23T21:01:08Z

python/ray/train/lightning/lightning_trainer.py

@@ -427,6 +447,17 @@ def _check_checkpoint_configs(
                "through `LightningConfigBuilder.checkpointing()`."
            )

+        # Auto-fill the AIR CheckpointConfig if the user didn't specify it.
+        if air_ckpt_config == CheckpointConfig():


One thing to be aware of is that if the user does explicitly pass in a default CheckpointConfig, then this will also get overridden. Is that expected?

Yes, it will be overridden, and this is painful because there's no way to determine whether the users explicitly passed a CheckpointConfig or not. But I think few people will provide an empty CheckpointConfig?
But yeah, I'll think of a workaround for this corner case.

Co-authored-by: matthewdeng <matthew.j.deng@gmail.com> Signed-off-by: Yunxuan Xiao <xiaoyunxuan1998@gmail.com>

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

…ning_configs

justinvyu

Just one nit. Want to avoid deepcopy since run config can take in a bunch of user objects like callbacks and stoppers.

Also, any thoughts on the eval dataset thing? We can leave that for a separate PR.

python/ray/train/lightning/lightning_trainer.py

woshiyyya · 2023-06-27T17:11:48Z

@justinvyu We introduced the val key in the LightningTrainer's docstring. I agree to change the eval dataset key from val to evaluation, but we might decide whether to keep both as valid keys, or raise an error for val.

In my opinion, we should keep both for now, and throw an api change warning if they are using val. Created an issue here, #36873, let's do it in a separate PR.

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

…ning_configs

justinvyu

Thank you!

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com> Signed-off-by: Yunxuan Xiao <xiaoyunxuan1998@gmail.com> Co-authored-by: matthewdeng <matthew.j.deng@gmail.com> Signed-off-by: e428265 <arvind.chandramouli@lmco.com>

woshiyyya added 9 commits June 13, 2023 01:41

auto complete an AIR CheckpointConfig

da18ef3

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

raise error when datasets_iter_config is empty

53dda7e

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

reorder

9224827

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

unify ckpt configs

df422f8

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

fix wording

52eecc6

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

wip

bb1688e

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

w

700d2e5

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

Merge remote-tracking branch 'upstream/master' into train/unify_light…

e9443c3

…ning_configs

add unittests

e732915

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

woshiyyya marked this pull request as ready for review June 23, 2023 17:36

woshiyyya requested review from justinvyu and matthewdeng June 23, 2023 17:36

justinvyu reviewed Jun 23, 2023

View reviewed changes

matthewdeng reviewed Jun 23, 2023

View reviewed changes

woshiyyya and others added 3 commits June 23, 2023 15:17

Apply suggestions from code review

6c7f94e

Co-authored-by: matthewdeng <matthew.j.deng@gmail.com> Signed-off-by: Yunxuan Xiao <xiaoyunxuan1998@gmail.com>

address comments

24a967e

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

Merge remote-tracking branch 'upstream/master' into train/unify_light…

6fca042

…ning_configs

woshiyyya requested review from justinvyu and matthewdeng June 26, 2023 16:57

justinvyu reviewed Jun 26, 2023

View reviewed changes

python/ray/train/lightning/lightning_trainer.py Outdated Show resolved Hide resolved

woshiyyya added 2 commits June 27, 2023 10:29

change deepcopy to copy

991b7ff

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

Merge remote-tracking branch 'upstream/master' into train/unify_light…

8442a45

…ning_configs

woshiyyya assigned matthewdeng Jun 27, 2023

justinvyu approved these changes Jun 27, 2023

View reviewed changes

add comments

7b32870

Signed-off-by: woshiyyya <xiaoyunxuan1998@gmail.com>

matthewdeng merged commit 9e716a4 into ray-project:master Jun 28, 2023

akshay-anyscale mentioned this pull request Jul 21, 2023

Add service deployment instructions to stable diffusion template #37645

Closed

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Train] Unifying Lightning and AIR CheckpointConfig #36368

[Train] Unifying Lightning and AIR CheckpointConfig #36368

woshiyyya commented Jun 13, 2023 •

edited

Loading

justinvyu left a comment

justinvyu Jun 23, 2023

woshiyyya Jun 23, 2023 •

edited

Loading

matthewdeng Jun 23, 2023

woshiyyya Jun 23, 2023

justinvyu left a comment

woshiyyya commented Jun 27, 2023 •

edited

Loading

justinvyu left a comment

[Train] Unifying Lightning and AIR CheckpointConfig #36368

[Train] Unifying Lightning and AIR CheckpointConfig #36368

Conversation

woshiyyya commented Jun 13, 2023 • edited Loading

Why are these changes needed?

Related issue number

Checks

justinvyu left a comment

Choose a reason for hiding this comment

justinvyu Jun 23, 2023

Choose a reason for hiding this comment

woshiyyya Jun 23, 2023 • edited Loading

Choose a reason for hiding this comment

matthewdeng Jun 23, 2023

Choose a reason for hiding this comment

woshiyyya Jun 23, 2023

Choose a reason for hiding this comment

justinvyu left a comment

Choose a reason for hiding this comment

woshiyyya commented Jun 27, 2023 • edited Loading

justinvyu left a comment

Choose a reason for hiding this comment

woshiyyya commented Jun 13, 2023 •

edited

Loading

woshiyyya Jun 23, 2023 •

edited

Loading

woshiyyya commented Jun 27, 2023 •

edited

Loading