run_rm.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import pprint

from omegaconf import OmegaConf

from nemo_rl.algorithms.rm import MasterConfig, rm_train, setup
from nemo_rl.algorithms.utils import get_tokenizer
from nemo_rl.data.utils import setup_preference_data
from nemo_rl.distributed.virtual_cluster import init_ray
from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry
from nemo_rl.utils.config import load_config, parse_hydra_overrides
from nemo_rl.utils.logger import get_next_experiment_dir


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Run RM training with configuration")
    parser.add_argument(
        "--config", type=str, default=None, help="Path to YAML config file"
    )

    # Parse known args for the script
    args, overrides = parser.parse_known_args()

    return args, overrides


def main():
    """Main entry point."""
    # Parse arguments
    args, overrides = parse_args()

    if not args.config:
        args.config = os.path.join(os.path.dirname(__file__), "configs", "rm.yaml")

    config = load_config(args.config)
    print(f"Loaded configuration from: {args.config}")

    if overrides:
        print(f"Overrides: {overrides}")
        config = parse_hydra_overrides(config, overrides)

    config = OmegaConf.to_container(config, resolve=True)
    config = MasterConfig(**config)
    print("Applied CLI overrides")

    # Print config
    print("Final config:")
    pprint.pprint(config)

    assert config.policy["reward_model_cfg"]["enabled"]

    config.logger.log_dir = get_next_experiment_dir(config.logger.log_dir)
    print(f"📊 Using log directory: {config.logger.log_dir}")
    if config.checkpointing["enabled"]:
        print(
            f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}"
        )

    # Initialise telemetry on the driver BEFORE init_ray() so the resolved
    # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited
    # by every worker. No-op unless telemetry is on.
    init_telemetry_driver(config, algorithm="rm")

    try:
        init_ray()

        # setup tokenizer
        tokenizer = get_tokenizer(config.policy["tokenizer"])

        # setup data
        dataset, val_dataset = setup_preference_data(tokenizer, config.data)

        (
            policy,
            cluster,
            train_dataloader,
            val_dataloader,
            loss_fn,
            logger,
            checkpointer,
            rm_save_state,
            master_config,
        ) = setup(config, tokenizer, dataset, val_dataset)

        # The checkpointer owns background async-checkpoint finalization threads;
        # the context manager guarantees they are flushed (rename + delete) on exit.
        with checkpointer:
            rm_train(
                policy,
                train_dataloader,
                val_dataloader,
                tokenizer,
                loss_fn,
                master_config,
                logger,
                checkpointer,
                rm_save_state,
            )
    finally:
        # Flush on the failure paths too, and before cluster teardown: the OTel
        # SDK's own atexit hook is registered ahead of Ray's and so runs after
        # it. No-op when telemetry is inactive.
        shutdown_telemetry()


if __name__ == "__main__":
    main()