NVIDIA-NeMo / NVIDIA-NeMo/RL

SLURM_JOB_ID support for checkpoint output directory

Open
#1,010 4 comments 0 reactions 0 assignees View on GitHub
external x-inflection
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**

If the same script is launched twice (i.e. with different hyperparams), the checkpoint output dir is not separated unless user manually changed the `checkpointing.checkpoint_dir`.

**Describe the solution you'd like**

I'd suggest an option to append `SLURM_JOB_ID` to `checkpointing.checkpoint_dir` same way as it's done for logs at ray.sub here https://github.com/NVIDIA-NeMo/RL/blob/main/ray.sub#L98

**Describe alternatives you've considered**

One way is to have an env var like `CKPT_DIR_SUFFIX` which causes ray.sub to append smth like "10001-run/" to `checkpointing.checkpoint_dir`. It is quite hardcoded though.

Another way would be to support interpolations with key words in nemo-rl configs, i.e. if user config has `checkpointing.checkpoint_dir: output/ckpts/${slurm_job_id}-run`, it will be replaced with `SLURM_JOB_ID` at ray.sub stage. That would take more time to implement.

The downside is that ray.sub will mess with user defined args which is not the case at the moment

**Additional context**

torchtune does interpolations -> https://docs.pytorch.org/torchtune/stable/deep_dives/configs.html#referencing-other-config-fields-with-interpolations but without key words

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.