Some commands and tips, often used when working on a server, I feel tired of asking gpt again and again.

1. Kill Jupyter Notebook Kernel

When you use Jupyter Notebook on a server,especially when in VSCode, sometimes the kernel may get stuck or become unresponsive. Unfortunately, VSCode does not provide a direct way to kill or shutdown the kernel from its interface.This make a lot of jupyter processes running on the server, consuming resources and causing confusion.

Check the running Jupyter processes with:

1

pgrep -a -u username -f jupyter

and then kill

1

kill -9 PID

And if you want to be more aggressive, the most brutal method is:

1

pkill -9 -u username -f jupyter

2. Check Processes

To check all the processes running under your username, you can use:

1

ps -u username -o pid,ppid,cmd,%mem,%cpu --sort=-rss | head -n 20

This command will show you the top 20 processes consuming the most memory, along with their PID, parent PID, command, memory usage, and CPU usage.

To check all the processes running on the server, you can use:

1

ps -a -o pid,ppid,user,cmd,%mem,%cpu --sort=-rss | head -n 20

This command will show you the top 20 processes consuming the most memory, along with their PID, parent PID, user, command, memory usage, and CPU usage.

Let me see who’s zombie processes occupying the server resources. 💢

The simplified code for the RLHF with PPO algorithm in TRL is as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

for batch in dataloader:
    # Sample
    with torch.no_grad():
        old_policy_outputs,logprobs_old_policy = policy_model(batch)
        logprobs_reference = reference_model(batch)
        scores = reward_model(query+old_policy_outputs)
        values = value_model(query+old_policy_outputs)

        kl = logprobs_old_policy - logprobs_reference
        # reward function have different implementation methods
        rewards = scores - kl * kl_coef 
        lastgaelam = 0
        advantages_reversed = []
        gen_length = responses.shape[1]
        # cal advantages
        for t in reversed(range(gen_length)):
            nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
            delta = rewards[:, t] + args.gamma * nextvalues - values[:, t]
            lastgaelam = delta + args.gamma * args.lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)
        returns = advantages + values
    # Train
    for ppo_epoch_idx in range(num_ppo_epochs):
        logprobs_new_policy, v_pred = policy_model(batch)
        # compute value_loss
        v_pred_clipped = clip(v_pred, 
                            values - cliprange_value, 
                            values + cliprange_value)
        vf_loss1 = (v_pred - returns) ** 2
        vf_loss2 = (v_pred_clipped - returns) ** 2
        vf_loss = 0.5 * torch.max(vf_loss1, vf_loss2).mean()
        # compute policy_loss
        important_ratio = logprobs_new_policy - logprobs_old_policy
        policy_loss1 = -(exp(important_ratio) * advantages)
        policy_loss2 = -(clip(exp(important_ratio), 
                            1.0 - cliprange_ratio, 
                            1.0 + cliprange_ratio) \ 
                        * advantages)
        policy_loss = torch.max(policy_loss1, policy_loss2).mean()
        loss = policy_loss + vf_loss * args.vf_coef

2. Theory

To express the PPO algorithm explicitly, I will ignore the kl-penalty term and clip operation for simplicity.

2.1. Policy Loss

2.1.1 Policy Gradient

The optimization objective of reinforcement learning is to maximize the expected return:

where $\tau$ is a trajectory sampled from the policy $\pi_\theta$, and $R(\tau)$ is the return of the trajectory.

Then the policy gradient can be derived as follows(details can be found in the [1]):

where $T_n$ is the length of the trajectory, and $s$ is the state at time step $t$.

2.1.2. Improvement of $R(\tau)$ - GAE

As we consider Eq 2, the return $R(\tau)$ is a sum of rewards on the trajectory, but the $\pi_\theta(a|s)$ is a probability distribution on single step. We may think this is not a good choice to use the sum of rewards as the return. In fact there are many ways to improve the return $R(\tau)$, but the most common one is Generalized Advantage Estimation (GAE).

where $A_\phi(s_t,a_t)$ is the advantage function, $r_t$ is the reward at time step $t$, $\gamma$ is the discount factor, and $V_\phi(s)$ is the value function.

The new policy gradient can be written as:

According to the gradient Eq 4, we can get the new optimization objective:

Because the solution is equal between the expectation of the trajectory $\tau$ and the expectation of the time step $t$, we can derive line 2 to line3 in Eq 5.

2.1.3. Importance Sampling

In the above equation, we can see that the policy gradient is computed by the advantage function and the log probability of the action taken by the policy. However, in practice, we often have a old policy $\pi_{\text{old}}(a|s)$, which is used to generate the trajectory $\tau$. We will use these trajectories to update the policy $\pi_\theta(a|s)$ for ppo_epochs times.

The Policy Gradient can be rewritten as:

At last, we can get the final optimization objective:

2.2 Critic Loss

References

1. 猛猿. 人人都能看懂的RL-PPO理论知识 .
2. TRL: PPO Trainer Implementation .

1
2

cd /Users/<username>/Zotero
ln -s /Users/<username>/Library/Mobile\ Documents/com~apple~CloudDocs/zotero_storage storage

Reference

1. Xi Chen. Using Zotero With iCloud As The Storage .
2. Northword and l0o0 Zotero 中文社区-百度翻译接口申请和配置 .

Before we start the derivation in detail, we first introduce the motivation and overview of DPO.

In RLHF, we first need to train a reward model to predict the reward of a response given a prompt. Then, we use the reward model to train the policy model under the restriction of the reference model. However, this two-step process is not very efficient. Especially, we at least need to maintain two model weights, the reward model and the reference model,and train one model or maybe two models at the same time.

To address this issue, DPO proposes a more efficient approach by directly optimizing the policy model in the processing when training the reward model.

Derivation

We start with the objective function of the reward model, the most popular choice to model human preferences is Bradley-Terry (BT) model,can be written as:

The reward model is trained to maximize the likelihood of the preference data, which can be written as:

To optimize the policy model in the same time, we need to find a way to convert the reward model into a policy model. In other words, we need to find a way to express the policy model $\pi_{\theta}(y|x)$ in terms of the reward model $r_{\phi}(x, y)$.

All RLHF algorithms can be summarized as the following optimization problem:

where $\beta$ is a hyperparameter to control the trade-off between the reward and the KL divergence, $r_{\phi}(x, y)$ is the reward model, and $\pi_{\text{ref}}(y|x)$ is the reference model.

To solve this optimization problem(Eq 3), we can get the optimal policy model $p^*(y|x)$ expressed by the reward model $r_\phi(x, y)$.

Now, we find the Eq 4 like a KL divergence term which is easy for us to optimize, but with a non-probability distribution in the denominator. So we need to convert the non-probability distribution into a probability distribution by add a normalization term $Z(x)$, which is the partition function:

so that we can rewrite the denominator as:

where $\pi^*(y|x) = \frac {\pi_{\text{ref}}(y|x) \exp\left(\frac{1}{\beta}r_{\phi}(x, y)\right)}{\sum_y \pi_{\text{ref}}(y|x) \exp\left(\frac{1}{\beta}r_{\phi}(x, y)\right)}$, means the probability distribution of the reference model after being re-weighted by the reward function.

we substitute Eq 6 into Eq 4:

Because the second term is a constant, we can ignore it, and the KL-divergence is minimized at 0 when $\pi_{\theta}(y|x) = \pi^*(y|x)$, the Explicit Solution of the optimal policy model is:

However, we can not directly use this solution to train the policy model, because computing the partition function $Z(x)$ is expensive.

According to Eq 8, We can get:

Finally, we substitute the Eq 9 into the loss function for reward model Eq 2, we can get the loss function for the policy model:

Reference

1. Rafailov, Rafael, et al. “Direct preference optimization: Your language model is secretly a reward model.” Advances in neural information processing systems 36 (2023): 53728-53741.
2. 猛猿. 人人都能看懂的DPO数学原理 .
3. SmallerFL. Direct Preference Optimization (DPO) 原理详解及公式推导 .