Parallel eval#772
Conversation
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d7bf27bd7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if main_process: | ||
| run_post_train_checkpoint_sweep(self.cfg, self.ckpt_paths) |
There was a problem hiding this comment.
When post_train_eval is enabled after CUDA/DDP training, this call starts the sweep while the Trainer object still owns the rank-0 model, optimizer, EMA state, and dataloader tensors on GPU 0. Destroying the process group releases the other ranks but does not free rank-0 CUDA allocations, and _resolve_worker_count() assigns an eval worker to GPU 0, so large runs can OOM when that worker loads a second model/rollout on the same occupied GPU; delete/move the CUDA state and empty the cache, or run the sweep after the training process exits.
Useful? React with 👍 / 👎.
| prepared_groundtruth = template_cfg.prepare_groundtruth( | ||
| LocalLocation(path=Path.cwd()) | ||
| ) |
There was a problem hiding this comment.
When the sweep is launched from training and the viz config leaves data_root unset with a relative groundtruth_location (the shipped viz configs use OM4.zarr this way), eval gets train_cfg.experiment.data_root but viz prepares ground truth relative to Path.cwd() instead. That makes post-train viz fail by looking under the repo/launch directory, or worse use a different local dataset; pass the training data root into the viz config/default root before preparing ground truth.
Useful? React with 👍 / 👎.
| raise ValueError( | ||
| f"last_n_checkpoints must be >= 1, got {last_n_checkpoints}" | ||
| ) | ||
| targets = targets[-last_n_checkpoints:] |
There was a problem hiding this comment.
When last_n_checkpoints is set, the final EMA entry has already been appended before this slice runs, so last_n_checkpoints: 10 evaluates only 9 periodic checkpoints plus final_ema (and 1 evaluates only EMA). The config/commit describes the EMA checkpoint as always added in addition to the selected checkpoints, so users silently get fewer saved epochs than requested; apply the limit to periodic targets before appending EMA.
Useful? React with 👍 / 👎.
jder
left a comment
jder
left a comment
There was a problem hiding this comment.
Very excited about this. I didn't read everything in detail but a few high-level comments first.
I also have one question about the goal: why are we causing the existing per-GPU torch processes to exit and then re-spawning them? I think will cause us to drop down to a single host for running evals/viz when doing multi-host training. Can we use the existing per-GPU processes instead (and run post_train_eval via slurm/torchrun as we do training when we want it to be standalone?). Basically just have the existing non-main processes go into a worker loop and distribute work via torch.distributed?
| eval_config_path = ( | ||
| Path(train_cfg.post_train_eval.eval_config_path).expanduser().resolve() | ||
| ) |
There was a problem hiding this comment.
I don't understand why we need expanduser()/resolve() here given it's in Config.from_yaml_and_cli
| logger.warning("No checkpoints selected for post-train eval sweep") | ||
| return [] | ||
|
|
||
| eval_cfg = EvalConfig.from_yaml_and_cli(list(eval_config_args)) |
There was a problem hiding this comment.
Probably nicer to have a path here instead and expose a Config.from_yaml()
| cfg = VizConfig.model_validate(updated) | ||
|
|
||
| start = time.perf_counter() | ||
| run_with_prepared_groundtruth(cfg, prepared_groundtruth) |
There was a problem hiding this comment.
I think this config surgery is a bit of a smell and somewhat fragile. How about:
- Create a new top-level config VizTemplateConfig with base_output_dir, dataset_name, variables, data_root, etc. But not
runsorname. It builds a VizTemplate which is basically the same content asPreparedVizGroundtruthtoday. - In
run_checkpoint_sweepwe load a VizTemplateConfig and call its build method to get a VizTemplate. - In this code we call
viz_template.instantiate(output_path, runs)or something to get a Viz. - VizConfig now extends VizTemplateConfig and adds runs + name. (VizConfig.build now calls super.build and then calls instantiate on that)
There was a problem hiding this comment.
(Alternatively, if you want to just land this without the viz stuff that works for me too and we can deal with this as a follow-up.)
There was a problem hiding this comment.
Implemented the template approach, now VizConfig extends VizTemplateConfig which builds reusable VizTemplate state and prepares ground truth once. The sweep now instantiates each checkpoint’s Viz directly from that template.
| cfg.prepare_output_dirs() | ||
| cfg.save_yaml(cfg.experiment.output_dir / "config.yaml") | ||
|
|
||
| self.cfg = cfg |
There was a problem hiding this comment.
The pattern we're trying to move towards (though I see we have failed to actually write this down anywhere) is to avoid passing around cfg values like this (and below), and instead turn the cfg types into "inflated" or "ready to run" types with build(...) as early as possible, passing in needed dependencies. Trainer is the main exception to this pattern since it predates the config system and is a beast to refactor in that way. Would it be possible to avoid saving the cfg here and to apply this pattern to the PostTrainCheckpointSweepConfig type, passing in the needed extra data from TrainConfig to the build method (e.g. the experiment output directory) here? (Perhaps producing a CheckpointSweep instance which has all the needed information to run such a thing?)
build the checkpoint sweep up front instead of passing TrainConfig around free training GPU state before evaluation workers start replace per-checkpoint viz config surgery with a reusable viz template fix checkpoint selection and resolve viz data against the training data root
jder
left a comment
jder
left a comment
There was a problem hiding this comment.
Hey @amogh-gulati thanks again for doing this! Very much looking forward to being able to use this. I left a bunch of comments but they're mostly pretty tiny things your agent should be able to do :) LMK if you want to chat through any of them.
|
|
||
|
|
||
| class PostTrainCheckpointSweepConfig(BaseConfig): | ||
| enabled: bool = False |
There was a problem hiding this comment.
Might be nicer to have post_train_eval be PostTrainEvalConfig | None above and remove the enabled bool? Then you can make eval_config_path and others always non-None
| eval_config_path: str | None = None | ||
| viz_config_path: str | None = None |
There was a problem hiding this comment.
I went on a journey here trying to make this EvalConfig rather than a path. And also splitting up EvalConfig into things we actually want here vs not (e.g. not ckpt_path). This lets us validate + build() earlier so things don't explode at the end of a run. But I think in the end it is probably more complicated than it's worth. If you're interested, you can see the it here: https://github.com/m2lines/Samudra/compare/parallel_eval...codex/refactor-eval-dependencies?expand=1
But anyway, I suggest these be Paths, not str, in the current state
| last_n_checkpoints: int | None = Field(default=None, ge=1) | ||
| checkpoints: list[int] | None = Field( | ||
| default=None, | ||
| description="Explicit list of checkpoint epochs (matching ckpt_<epoch>.pt) " | ||
| "to evaluate; the final EMA checkpoint is always added. Mutually " | ||
| "exclusive with last_n_checkpoints.", | ||
| ) |
There was a problem hiding this comment.
🐑 how about "epochs" instead of "checkpoints"?
| "to evaluate; the final EMA checkpoint is always added. Mutually " | ||
| "exclusive with last_n_checkpoints.", | ||
| ) | ||
| eval_dirname: str | None = None |
There was a problem hiding this comment.
Should this default to (or just be hard-coded below as) "evals"?
|
|
||
| @dataclass(frozen=True) | ||
| class CheckpointSweep: | ||
| """Ready-to-run checkpoint sweep built from configuration.""" |
There was a problem hiding this comment.
I don't think this is "ready-to-run" or "built from configuration" given it has a list of config paths to still to load?
| if process.is_alive(): | ||
| process.terminate() | ||
| process.join() |
There was a problem hiding this comment.
Do we need to do this? If so, why don't we need to do it in the timeout case just above?
| process.start() | ||
| processes.append(process) | ||
|
|
||
| while len(results) < len(targets): |
There was a problem hiding this comment.
Do we want a real timeout? e.g. 6 hours?
| data_root: ResolvedLocation | ||
| variables: list[str] |
There was a problem hiding this comment.
Might be somewhat more clear to have instantiate take the needed arguments to produce a VizRun (I think label + location?) rather than having data_root and variables used externally while dataset_name and prepared_groundtruth are used internally. Alternatively maybe variables should not be on VizRun at all… I think weird things will probably happen if different runs have different variables compared to each other or the baseline? (And for post-training runs I think we can guarantee they will be the same anyway?)
| basins: xr.Dataset, | ||
| groundtruth_rollout: xr.Dataset, | ||
| time_range: slice, | ||
| ) -> PreparedVizGroundtruth: |
There was a problem hiding this comment.
This is great, so glad we're not doing this over and over.
| if main_process and self.post_train_sweep is not None: | ||
| self.post_train_sweep.run() | ||
|
|
||
| def _release_train_state(self) -> None: |
There was a problem hiding this comment.
Might be less fragile and more reliable to drop the whole Trainer object + run gc instead?
2 participants
Post-train checkpoint eval/viz sweep
Adds an optional post-training sweep that evaluates selected checkpoints after training completes, always including the final EMA checkpoint, with optional visualization generation.
post_train_eval.pyto discover checkpoints insaved_nets/, shard eval jobs across available GPUs, run standalone inference per checkpoint, and writesummary.json.last_n_checkpoints: Nor an explicitcheckpoints: [...]epoch list, with validation for mutually exclusive settings and missing checkpoints.Trainer.finish()to launch the sweep on the main process via a newpost_train_evalconfig block.destroy_process_group()before the sweep so all training ranks finish cleanly and release GPUs. (still seems a little flacky)python -m samudra.post_train_evalfor models already trained.Notes:
eval.save_zarr=true.checkpoints: [50, 55, 60, 65, 70]evaluates those epochs plus the final EMA checkpoint.#Config
Add a post_train_eval block to the train config to enable the post-training checkpoint sweep:
post_train_eval:
enabled: true
eval_config_path: configs//eval.yaml # required when enabled
viz_config_path: configs//viz.yaml # optional; omit to skip viz
Pick exactly one checkpoint selection mode:
checkpoints: [50, 55, 60, 65, 70] # explicit checkpoint epochs
last_n_checkpoints: 10 # or evaluate the last N checkpoints
The final EMA checkpoint is always included automatically.
wandb is disabled for this, will a follow up PR I will add all the evals in one run (ideally with timeline scrub) which is not tested yet!