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

Code before:
            Role.ActorRollout: global_pool_id,
            Role.Critic: global_pool_id,
            Role.RefPolicy: global_pool_id,
        }

        resource_pool_manager = ResourcePoolManager(
            resource_pool_spec=resource_pool_spec, mapping=mapping
        )

        # Create reward function manager
        compute_score = self._reward_funcs[0] if len(self._reward_funcs) > 0 else None
        reward_fn = NaiveRewardManager(
            tokenizer=tokenizer, num_examine=0, compute_score=compute_score
        )
        # num_examine=1 means to print 1 example per batch for analysis.
        val_reward_fn = NaiveRewardManager(
            tokenizer=tokenizer, num_examine=1, compute_score=compute_score
        )

        self._verl_trainer = RayPPOTrainer(
            config=self._verl_config,
            tokenizer=tokenizer,
            processor=(
                # verl trainer uses private methods and properties of `transformers`
                # processor, so we need to pass the raw processor here.
                self._processor.raw_processor if self._processor is not None else None
            ),
            role_worker_mapping=role_worker_mapping,
            resource_pool_manager=resource_pool_manager,
            reward_fn=reward_fn,
            val_reward_fn=val_reward_fn,
        )

    def train(self, resume_from_checkpoint: Optional[str] = None) -> None:
        """Trains the model using verl's RayPPOTrainer.

        Args:
            resume_from_checkpoint: Optional path to a checkpoint to resume from.
        """
        if resume_from_checkpoint:
            raise NotImplementedError("Resuming from checkpoint is not implemented.")

        logger.info("Initializing verl trainer workers...")
        self._verl_trainer.init_workers()
        logger.info("Starting verl training...")
        self._verl_trainer.fit()

    # TODO: OPE-1192 - Implement saving model/trainer state. verl training should
    # already handle saving models, including the final checkpoint.

    def save_state(self) -> None:

Review comment (suggestion):
nit: Consider moving this logic into the init fn, and have the self._processor you save be the raw processor. This keeps this initialization lighter, and makes it clear early on which processor you'll be using throughout this file.

Code after:
                "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(
                "ray is not installed. "
                "Please install it with 'pip install `oumi[gpu]`'."
            )
        self._verl_config = self._create_config()
        logger.info(f"verl config: {pformat(self._verl_config)}")

        tokenizer = self._processing_class

        role_worker_mapping = {
            Role.ActorRollout: ray.remote(ActorRolloutRefWorker),
            Role.Critic: ray.remote(CriticWorker),
            Role.RefPolicy: ray.remote(ActorRolloutRefWorker),
        }

        # Create resource pool manager
        global_pool_id = "global_pool"
        resource_pool_spec = {
            global_pool_id: [self._verl_config.trainer.n_gpus_per_node]
            * self._verl_config.trainer.nnodes,
        }
        mapping = {
            Role.ActorRollout: global_pool_id,
            Role.Critic: global_pool_id,
            Role.RefPolicy: global_pool_id,
        }

        resource_pool_manager = ResourcePoolManager(
            resource_pool_spec=resource_pool_spec, mapping=mapping
        )

        # Create reward function manager
        compute_score = self._reward_funcs[0] if len(self._reward_funcs) > 0 else None
        reward_fn = NaiveRewardManager(
            tokenizer=tokenizer, num_examine=0, compute_score=compute_score
        )
        # num_examine=1 means to print 1 example per batch for analysis.
        val_reward_fn = NaiveRewardManager(
            tokenizer=tokenizer, num_examine=1, compute_score=compute_score
        )

        self._verl_trainer = RayPPOTrainer(
            config=self._verl_config,
            tokenizer=tokenizer,