Create an eval-only script for existing ckpts #736
Conversation
olmo/train.py
Outdated
| if wandb.run is not None: | ||
| wandb.finish(exit_code=exit_code, quiet=True) | ||
| # if wandb.run is not None: | ||
| # wandb.finish(exit_code=exit_code, quiet=True) |
scripts/eval.py
Outdated
| # train_loader = build_train_dataloader(cfg) | ||
| train_loader = None |
There was a problem hiding this comment.
Is this always going to be None? If so, we don't need it.
scripts/eval.py
Outdated
| if 'step' in cfg.load_path.split('/')[-1]: | ||
| load_paths = [cfg.load_path] | ||
| else: | ||
| # This globbing does not work with remote paths. |
There was a problem hiding this comment.
How is that problem handled then?
There was a problem hiding this comment.
Not handled. I will assume the checkpoints are on WEKA.
There was a problem hiding this comment.
At least throw an exception then.
scripts/eval.py
Outdated
| log.info(f"Number of non-embedding parameters: {olmo_model.num_params(include_embedding=False):,d}") | ||
| log.info(f"Peak GPU Memory (MB) before {cfg.distributed_strategy}: {int(peak_gpu_memory() or 0)}") | ||
|
|
||
| olmo_model.set_activation_checkpointing(cfg.activation_checkpointing) |
There was a problem hiding this comment.
If we only ever eval, we don't need this.
| optim = build_optimizer(cfg, dist_model) | ||
| scheduler = build_scheduler(cfg) |
There was a problem hiding this comment.
We don't need optimizers and schedulers if we're just evaluating.
There was a problem hiding this comment.
So you're creating these only so that you can produce a Trainer object?
How hard is it to pull the stuff you need out of the Trainer object, so we don't have to do so many things we don't need? It makes me particularly uncomfortable that you're creating a trainer with a None data loader, which isn't supposed to work. It just happens to work.
There was a problem hiding this comment.
The Trainer class has too many precious helper functions and it's kinda dumb to unroll them all. I do wanna keep at least a dummy Trainer object. Let me see if I can create it w/o the optim/scheduler/etc stuff.
There was a problem hiding this comment.
I think you might find that you don't need most of that stuff when you're doing inference only.
scripts/eval.py
Outdated
| if 'step' in cfg.load_path.split('/')[-1]: | ||
| load_paths = [cfg.load_path] | ||
| else: | ||
| # This globbing does not work with remote paths. |
There was a problem hiding this comment.
At least throw an exception then.
|
Let me know when this is ready for another review? |
| if cfg.load_path is None: | ||
| raise OLMoConfigurationError("To run eval you must provide a load_path") | ||
| elif "://" in cfg.load_path: | ||
| raise OLMoConfigurationError("Eval does not support remote paths. Please specify a local path or WEKA mounted path.") |
There was a problem hiding this comment.
throwing exception for remote paths here
This PR adds
scripts/eval.py, which evaluates one or more existing ckpts while bypassing the training steps.It seems impossible to backfill evals back to the original wandb run, because "step" must always increase. Rewinding the run will truncate the log, which we don't want. Therefore, this script logs things to a new wandb run.
Starting from a training setup:
XXX.shfile intoXXX-eval.sh, point toscripts/eval.sh, add a flag--wandb.group=XXXto ensure it logs to the same group, and specify--load_pathto be either a single ckpt or all ckpts under a directory.XXX-launch.shfile intoXXX-eval-launch.sh, change--task-nametoXXX-eval, and change the command so it runsXXX-eval.sh.See an example in
peteish1-eval.shandpeteish1-eval-launch.sh.