Calculating Two Gradient for two different loss in same functions

[erfanzar]

i want to Calculating Two Gradient for two different loss in same functions

i have a ppo function im trying to implement RLHF for jax/flax models in my library EasyDel for the learn function in my trainer usually in pytorch you just calculate the loss for two model and do .step' and .zero_grad` but when i calculate the loss and policy loss i have no idea how to do that i search for that for a long time but all of the current implementation was not good enough (they were calculating or computing some algorithms twine) so i implemented my method like this

this is my trainStateclass

class TrainStateActorAndCritic(struct.PyTreeNode):
    step: int
    apply_fn_critic: Callable = struct.field(pytree_node=False)
    apply_fn_actor: Callable = struct.field(pytree_node=False)

    actor_params: core.FrozenDict[str, Any] = struct.field(pytree_node=True)
    critic_params: core.FrozenDict[str, Any] = struct.field(pytree_node=True)

    actor_optim: optax.GradientTransformation = struct.field(pytree_node=False)
    critic_optim: optax.GradientTransformation = struct.field(pytree_node=False)

    actor_opt_state: optax.OptState = struct.field(pytree_node=True)
    critic_opt_state: optax.OptState = struct.field(pytree_node=True)

    def apply_gradients(self,
                        *,
                        grads_critic,
                        grad_actor,
                        **kwargs
                        ):
        updates_critic, new_state_critic = self.critic_optim.update(
            grads_critic, self.critic_opt_state, self.critic_params)
        critic_params = optax.apply_updates(self.critic_params, updates_critic)

        updates_actor, new_state_actor = self.actor_optim.update(
            grad_actor, self.actor_opt_state, self.actor_params
        )
        actor_params = optax.apply_updates(self.actor_params, updates_actor)

        return self.replace(
            critic_opt_state=new_state_critic,
            actor_opt_state=new_state_actor,

            critic_params=critic_params,
            actor_params=actor_params,

            step=self.step + 1,
            **kwargs
        )

    @classmethod
    def create(cls,
               *,
               apply_fn_critic,
               apply_fn_actor,
               actor_params,
               critic_params,
               actor_optim: optax.GradientTransformation,
               critic_optim: optax.GradientTransformation,
               **kwargs
               ):
        actor_opt_state = actor_optim.init(actor_params)
        critic_opt_state = critic_optim.init(critic_params)
        return cls(
            step=0,

            apply_fn_actor=apply_fn_actor,
            apply_fn_critic=apply_fn_critic,

            actor_params=actor_params,
            critic_params=critic_params,

            actor_optim=actor_optim,
            critic_optim=critic_optim,

            actor_opt_state=actor_opt_state,
            critic_opt_state=critic_opt_state,
            **kwargs,
        )

and this is my forward or step function

        def forward(
                train_state: TrainStateActorAndCritic,
                input_ids,
                pm,
                rewards,
                old_values,
                attention_mask,
                old_action_probs,
                old_log_probs,
        ):
            def calculate_loss(params):
                global rewards
                global old_values

                action_masks = ~pm & attention_mask
                action_logits, values = train_state.apply_fn_actor(
                    params["params"],
                    input_ids=input_ids,
                    attention_mask=attention_mask
                )
                action_logits = shift(action_logits, shift=1, axis=-2)
                action_len = old_log_probs.shape[-1]
                action_probs = jax.nn.softmax(action_logits, axis=-1)
                action_log_probs = log_prob(action_probs, input_ids)
                action_log_probs = action_log_probs[:, -action_len:]
                entropies = masked_entropy(action_probs, attention_mask=action_masks)
                kl_penalty = masked_mean(
                    jnp.sum((old_action_probs * (jnp.log(old_action_probs) - jnp.log(action_probs))), axis=-1),
                    attention_mask=attention_mask
                ) * self.config.kl_div_loss_weight

                rewards = rewards - kl_penalty
                normalize_kwargs = dict()

                if old_values.ndim == 2:
                    old_values, values = map(lambda t: shift(t, shift=1, axis=-2), (old_values, values))

                    old_values = old_values[:, -action_len:]
                    values = values[:, -action_len:]
                    rewards = einops.rearrange(rewards, 'b -> b 1')
                    normalize_kwargs = dict(axis=-1, attention_mask=action_masks[:, -action_len:])

                if values.ndim < rewards.ndim:
                    values = einops.rearrange(values, '... -> ... 1')

                ratios = (action_log_probs - old_log_probs).exp()
                advantages = masked_normalize(rewards - old_values, **normalize_kwargs)

                if advantages.ndim == 1:
                    advantages = einops.rearrange(advantages, 'b -> b 1')

                surr1 = ratios * advantages
                surr2 = jnp.clip(ratios, 1 - self.config.eps_clip, 1 + self.config.eps_clip)
                policy_loss = -jnp.minimum(surr1, surr2) - self.config.beta_s * entropies  # Policy Loss
                loss = jnp.mean(policy_loss)  # Loss

                value_loss = jnp.mean(
                    clipped_value_loss(values, rewards, old_values, self.config.value_clip))  # VLoss
                return loss, value_loss  # I need gradient for these two losses

            grad, (loss_, value_loss_) = jax.value_and_grad(calculate_loss)(train_state.actor_params)
            train_state = train_state.apply_gradients(
                grad_actor=grad,  # Based on Loss
                grads_critic=...  # Based on Value Loss
            )
            return train_state, (loss_, value_loss_)

you can also check the code here

[cgarciae]

Hey @erfanzar, todo do this in JAX the easiest thing is just to create two different loss functions and calculate the gradients you need. While it may look wasteful, you will get good performance because XLA (jit) is able to optimize the computation graph, it will hopefully reuse the same computation paths to avoid duplicate work.