HEALPix Curriculum - #2795
Conversation
github-actions
Bot
added
infra
…d on curriculum restart
Here are the results I got from my test run. It works pretty good! I had some problems with the configs and auto-restarting after a transition, but was able to make everything work with some help from AI. I will do a self-review of my code on Friday. |
Took me some time to get back to it, but it looks good. @kctezcan do you have time to review it? |
clessig
left a comment
clessig
left a comment
There was a problem hiding this comment.
Thanks for the implementation. Some high-level comments for the moment.
| dist.all_reduce(l_seed, op=torch.distributed.ReduceOp.SUM) | ||
| cf.data_loader_rng_seed = l_seed.item() | ||
|
|
||
| if dist.is_initialized(): |
There was a problem hiding this comment.
Why is this change in the PR? It seems unrelated to the functionality that is targeted.
| assert isinstance(c, Config) | ||
| c = _sanitize_time_keys(c) | ||
|
|
||
| if c.get("healpix_curriculum"): |
There was a problem hiding this comment.
Can we avoid this in config.py? And very specific functionality there turned out to be problematic.
| @@ -0,0 +1,11 @@ | |||
| healpix_curriculum: | |||
There was a problem hiding this comment.
Could you explain what this config means / describes
|
|
||
| try: | ||
| trainer.run(cf, devices) | ||
| from_run_id_iter = None |
There was a problem hiding this comment.
This block does not belong to run_train.py. What is done here?
| if self.cf.general.istep > 0 and is_root(): | ||
| logger.info(f"Continuing run with learning rate: {self.lr_scheduler.get_lr()}") | ||
|
|
||
| if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum and is_root(): |
There was a problem hiding this comment.
This should be performed in a separate function
|
|
||
| self.cf.general.istep += 1 | ||
|
|
||
| if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum: |
There was a problem hiding this comment.
Let's try to avoid this complexity in the main training loop. Can you make a suggestion?
2 participants
Description
This adds a HEALPix curriculum to our training pipeline. During my bachelor thesis, I found that starting at a lower HEALPix Level (HL) and stepping it up during training cuts early compute time by almost 95%, all without hurting the final model performance.
How it works
healpix_curriculumparameter to set how manyistepseach HL should run for.streams_directory: "${curriculum_streams.${healpix_level}}") so the datasets swap automatically as the HL goes up.run_train.pynow handles the transitions. When a stage ends, it stops the mini-epoch, saves a checkpoint and immediately starts the next stage within the same Slurm job.FYI
config_curriculum.ymljust to make reviewing the OmegaConf logic easier. Before merging it can be deleted so we don't clutter the repo with new configs.Issue Number
Closes #2794
@kctezcan
Is this PR a draft? Mark it as draft.
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60FastEvaluation