Repository: oumi-ai/oumi
File: src/oumi/core/trainers/verl_grpo_trainer.py (Python)

Code before:
        config.actor_rollout_ref.rollout.gpu_memory_utilization = (
            grpo_params.vllm_gpu_memory_utilization
        )

        # Normally, training steps is determined by the number of epochs.
        # If max_steps is set, it will override this.
        config.trainer.total_epochs = training_params.num_train_epochs
        if training_params.max_steps != -1:
            config.trainer.total_training_steps = training_params.max_steps

        if training_params.eval_strategy == "steps":
            config.trainer.test_freq = training_params.eval_steps
        if not training_params.save_epoch:
            config.trainer.save_freq = training_params.save_steps

        config.trainer.logger = []
        if training_params.logging_strategy != "no":
            config.trainer.logger.append("console")
        if training_params.enable_wandb:
            config.trainer.logger.append("wandb")
        config.trainer.project_name = os.environ.get("WANDB_PROJECT", "oumi_verl")
        config.trainer.experiment_name = training_params.run_name
        config.trainer.default_local_dir = training_params.output_dir

        config.critic.model.path = model_name
        config.reward_model.model.path = model_name

        # 3. Apply user overrides
        overrides_config = OmegaConf.create(training_params.verl_config_overrides)
        config = cast(DictConfig, OmegaConf.merge(config, overrides_config))

        # 4. Finalize and validate config.

        # Resolves the value of all interpolation fields in the config.
        # ex. `prompt_length: ${data.max_prompt_length}`
        # https://omegaconf.readthedocs.io/en/2.3_branch/usage.html#omegaconf-resolve
        OmegaConf.resolve(config)

        if (
            config.actor_rollout_ref.actor.strategy == "fsdp"
            and config.actor_rollout_ref.actor.strategy != config.critic.strategy
        ):
            raise ValueError(
                "Actor and critic must use the same strategy when using FSDP."
            )
        return config

    def _setup_verl_trainer(self):
        """Sets up verl's RayPPOTrainer."""
        if ray is None:
            raise RuntimeError(

Review comment (suggestion):
Ditto, the reward model is a different thing as well from the actor model (the model we're training). AFAIK, GRPO generally doesn't use a reward model, and instead uses reward fns. If we want to expose this field in the Oumi configs natively, we'd need to add a new field.

Code after:
        eval_dataset.to_parquet(val_file)
        self._val_filepath = str(val_file)

    def _create_config(self) -> DictConfig:
        """Creates a verl config."""
        model_params = self._oumi_config.model
        model_name = model_params.model_name

        # 1. Read verl default dict config from YAML.
        yaml_path = Path(__file__).parent / "verl_trainer_config.yaml"
        config = OmegaConf.load(yaml_path)
        config = cast(DictConfig, config)

        # 2. Set config values, ex. from Oumi config values
        config.algorithm.adv_estimator = "grpo"
        config.data.train_files = self._train_filepath
        config.data.val_files = self._val_filepath

        grpo_params = self._oumi_config.training.grpo
        training_params = self._oumi_config.training

        config.data.max_response_length = grpo_params.max_completion_length
        config.actor_rollout_ref.model.path = model_name
        config.actor_rollout_ref.actor.optim.lr = training_params.learning_rate
        config.actor_rollout_ref.model.enable_gradient_checkpointing = (
            training_params.enable_gradient_checkpointing
        )
        if grpo_params.use_vllm:
            config.actor_rollout_ref.rollout.name = "vllm"
        else:
            config.actor_rollout_ref.rollout.name = "hf"
        config.actor_rollout_ref.rollout.temperature = grpo_params.temperature
        config.actor_rollout_ref.rollout.gpu_memory_utilization = (
            grpo_params.vllm_gpu_memory_utilization
        )

        # Normally, training steps is determined by the number of epochs.
        # If max_steps is set, it will override this.
        config.trainer.total_epochs = training_params.num_train_epochs
        if training_params.max_steps != -1:
            config.trainer.total_training_steps = training_params.max_steps

        if training_params.eval_strategy == "steps":
            config.trainer.test_freq = training_params.eval_steps
        if not training_params.save_epoch:
            config.trainer.save_freq = training_params.save_steps

        config.trainer.logger = []
        if training_params.logging_strategy != "no":
            config.trainer.logger.append("console")
        if training_params.enable_wandb: