[megatron] refactor: refactor MegatronPPOActor

[vermouth1992]

What does this PR do?

• Make megatron related print only print on rank zero
• Remove unused code in megatron actor
• Modularize megatron loss computation so that it can be used for SFT as well

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

[gemini-code-assist]

Code Review

This pull request refactors MegatronPPOActor to improve modularity and clarity. The changes include restricting print statements to rank zero, removing unused code, and modularizing the PPO loss computation. The refactoring is well-done, but I've found a potential issue in the new compute_ppo_loss method regarding loss accumulation and some redundant code that should be addressed for robustness and clarity.

[gemini-code-assist]

This new method is a great step towards modularity. However, there are a few areas for improvement for robustness and clarity:

1. Fragile Loss Accumulation: The policy_loss is updated with policy_loss = pg_loss - ... and then policy_loss = policy_loss + .... While correct in the current order, this is fragile. If another loss component were added, it could be unintentionally overwritten. Using += and -= for accumulation would make the code more robust.
2. Redundant Definitions: loss_agg_mode is defined at lines 312 and 319, and entropy_coeff is defined at lines 318 and 346. These redundancies should be removed.

I've provided a suggestion to address these points.

    def compute_ppo_loss(self, model_output, data):
        log_prob = model_output["log_probs"]
        entropy = model_output.get("entropy", None)

        metrics = {}

        response_mask = data["response_mask"].to(bool)
        loss_agg_mode = self.config.loss_agg_mode
        entropy_coeff = self.config.entropy_coeff

        # compute policy loss
        old_log_prob = data["old_log_probs"]
        advantages = data["advantages"]

        loss_mode = self.config.policy_loss.get("loss_mode", "vanilla")

        policy_loss_fn = get_policy_loss_fn(loss_mode)
        pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = policy_loss_fn(
            old_log_prob=old_log_prob,
            log_prob=log_prob,
            advantages=advantages,
            response_mask=response_mask,
            loss_agg_mode=loss_agg_mode,
            config=self.config,
        )

        metrics.update(
            {
                "actor/pg_loss": pg_loss.detach().item(),
                "actor/pg_clipfrac": pg_clipfrac.detach().item(),
                "actor/ppo_kl": ppo_kl.detach().item(),
                "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
            }
        )
        policy_loss = pg_loss

        # add entropy loss
        if entropy is not None:
            entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
            policy_loss -= entropy_coeff * entropy_loss

        # add kl loss
        if self.config.use_kl_loss:
            ref_log_prob = data["ref_log_prob"]
            # compute kl loss
            kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type)
            kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=self.config.loss_agg_mode)

            policy_loss += kl_loss * self.config.kl_loss_coef
            metrics["actor/kl_loss"] = kl_loss.detach().item()
            metrics["actor/kl_coef"] = self.config.kl_loss_coef

        return policy_loss, metrics

[vermouth1992]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the MegatronPPOActor to improve code quality and modularity. Key changes include restricting print and log statements to rank zero in distributed settings, removing unused code and imports, and modularizing the PPO loss computation into a separate compute_ppo_loss method. This modularization makes the loss logic reusable, for instance, for Supervised Fine-Tuning (SFT). The internal data flow within forward_backward_batch and compute_log_prob has also been simplified, leading to cleaner and more readable code. The changes are well-executed and align with the stated goals of the pull request. I have no major concerns.