Repository: volcengine/verl
File: verl/workers/actor/dp_actor.py (Python)

Code before:
                    micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)

                self.actor_optimizer.zero_grad()

                global_snis_denom = None
                try:
                    if (
                        "rollout_is_weights" in mini_batch.batch.keys()
                        and "rollout_is_self_norm_flag" in mini_batch.batch.keys()
                    ):
                        # Compute a global SNIS denominator
                        with torch.no_grad():
                            weights_full = mini_batch.batch["rollout_is_weights"]
                            response_mask_full = mini_batch.batch["response_mask"]
                            mask_float = response_mask_full.to(dtype=weights_full.dtype)

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

                            if loss_mode_batch == "geo_mean":
                                global_snis_denom = None
                            else:
                                lam = "seq-mean-token-mean" if loss_mode_batch == "gspo" else loss_agg_mode_batch
                                if lam == "token-mean":
                                    if torch.distributed.is_available() and torch.distributed.is_initialized():
                                        global_snis_denom = verl_F.distributed_masked_mean(weights_full, mask_float)
                                    else:
                                        global_snis_denom = verl_F.masked_mean(weights_full, mask_float)
                                elif lam == "seq-mean-token-sum":
                                    global_snis_denom = None
                                elif lam == "seq-mean-token-mean":
                                    global_snis_denom = None
                                elif lam == "seq-mean-token-sum-norm":
                                    denom_sum = (weights_full * mask_float).sum()
                                    if torch.distributed.is_available() and torch.distributed.is_initialized():
                                        torch.distributed.all_reduce(denom_sum, op=torch.distributed.ReduceOp.SUM)
                                    global_snis_denom = denom_sum / weights_full.shape[-1]
                                else:
                                    raise ValueError(f"loss_agg_mode: {lam} not supported for SNIS")
                except Exception as e:
                    logger.warning(f"SNIS denominator could not be calculated. Error: {e}", exc_info=True)
                    global_snis_denom = None

                for micro_batch in micro_batches:
                    micro_batch = micro_batch.to(get_device_id())
                    micro_batch_metrics = {}
                    model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}
                    response_mask = model_inputs["response_mask"]
                    old_log_prob = model_inputs["old_log_probs"]
                    advantages = model_inputs["advantages"]


Review comment (refactor):
Can we move `global_snis_denom` calculation to `verl/trainer/ppo/mismatch_helper.py`? Better to keep `update_policy` function concise.

Code after:
        mini_batches = data.split(self.config.ppo_mini_batch_size)

        on_policy = len(mini_batches) == 1 and self.config.ppo_epochs == 1

        metrics = {}
        for _ in range(self.config.ppo_epochs):
            for batch_idx, mini_batch in enumerate(mini_batches):
                if self.config.use_dynamic_bsz:
                    max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size
                    micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len)
                else:
                    self.gradient_accumulation = (
                        self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu
                    )
                    micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)

                self.actor_optimizer.zero_grad()

                for micro_batch in micro_batches:
                    micro_batch = micro_batch.to(get_device_id())
                    micro_batch_metrics = {}
                    model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}
                    response_mask = model_inputs["response_mask"]
                    old_log_prob = model_inputs["old_log_probs"]
                    advantages = model_inputs["advantages"]

                    entropy_coeff = self.config.entropy_coeff
                    loss_agg_mode = self.config.loss_agg_mode

                    if self.config.use_dynamic_bsz:
                        loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size
                    else:
                        loss_scale_factor = 1 / self.gradient_accumulation

                    # all return: (bsz, response_length)
                    calculate_entropy = False
                    if entropy_coeff != 0:
                        calculate_entropy = True
                    entropy, log_prob = self._forward_micro_batch(
                        model_inputs, temperature=temperature, calculate_entropy=calculate_entropy
                    )

                    # for fully_async_policy recipe
                    if hasattr(self.config, "use_rollout_log_probs") and self.config.use_rollout_log_probs:
                        old_log_prob = model_inputs["old_log_probs"]
                    else:
                        if on_policy:
                            old_log_prob = log_prob.detach()
                        else:
                            old_log_prob = model_inputs["old_log_probs"]