Machine Learning

Short notes on ML papers, blogs, lectures etc.

The Hitchhiker's Guide to Frontier Reinforcement Learning (Rishabh Agarwal)

Source: YouTube

The first half of the talk is a general overview of applying RL to scientific problems with a lab in the loop. The second half is more technical and focusses on stabilizing/scaling RL training. The two techniques that are mentioned -

  1. Reducing mismatches in inference and training.

    1. Correcting for the logprobs differences between the rollout backend and the training backend. On paper, the RL gradient computation looks something like Ea∼π(θ)[R(a)∇θlog⁡π(a,θ)]E_{a \sim \pi(\theta)} [R(a) \nabla_{\theta} \log \pi(a, \theta)] where the symbols have their usual meaning. In practice, the rollouts are generated by one backend (say, vLLM) and the policy/training gradient update is done on another (say, Megatron). The two backends would invariably use different kernels/numerics under the hood and would lead to different logits, which means we end up with Ea∼πrollout(θ)[R(a)∇θlog⁡πtrain(a,θ)]E_{a \sim \pi_{rollout}(\theta)} [R(a) \nabla_{\theta} \log \pi_{train}(a, \theta)]. To correct for this mismatch, we use an importance sampling correction of πtrainπrollout\frac{\pi_{train}}{\pi_{rollout}} which gives us

      Ea∼πrollout(θ)[πtrain(a)πrollout(a)R(a)∇θlog⁡πtrain(a,θ)]=Ea∼πtrain(θ)[R(a)∇θlog⁡πtrain(a,θ)] E_{a \sim \pi_{rollout}(\theta)} \left[\frac{\pi_{train}(a)}{\pi_{rollout}(a)} R(a) \nabla_{\theta} \log \pi_{train}(a, \theta)\right] = E_{a \sim \pi_{train}(\theta)} [R(a) \nabla_{\theta} \log \pi_{train}(a, \theta)]

      To get a handle on the variance of the correction factor, a truncated importance sampling (TIS) ratio min⁡(πtrainπrollout,γ)\min\left(\frac{\pi_{train}}{\pi_{rollout}}, \gamma\right) is used, where γ\gamma is a clipping threshold. See this excellent post for details and experiments with using TIS.

    2. Rollout Router Replay (R3) for MoE architectures. MoE architectures have a higher discrepancy between πrollout(a)\pi_{rollout}(a) and πtrain(a)\pi_{train}(a) compared to dense architectures. This is because small perturbations in the inputs/outputs of the router selection layer can lead to the selection of different experts between rollouts and training. The R3 paper proposes to reuse the experts selected during the rollout phase for training (but they do use the logits from the training backend for calculating the gating weights for the selected experts). This router replay reduces the MoE discrepancy down to the levels of dense architectures.

  2. Scaling RL training with Asynchronous updates. The idea here is to keep the inference and training engines running in parallel (on separate sets of GPUs) and only pause them (as opposed to draining the unfinished rollouts) when pushing weight updates from the train backend -> inference. See the PipelineRL paper for details. Note that when resuming the rollouts for the interrupted sequences after a weight update, the KVCache would be stale and you could choose to either recompute it with the new weights or retain the old KVCache. PipelineRL’s ablation in Sec 5.1 shows (surprisingly) that retaining the stale KVCache isn’t too bad (at least for their experiments with the 7B Qwen2.5 base model).

Follow-up reading:

  • The Art of Scaling Reinforcement Learning Compute for LLMs [paper]
  • IcePop [paper]

Chinchilla Scaling Laws (Hoffman et al. 2022)

Source: arXiv

Core question the paper asks: given a fixed training budget (aka compute aka FLOPs), how do you allocate it between model size vs dataset size?

Ans: For a dense transformer architecture, we need to scale model size and dataset size equally i.e. the power law coefficient for both is ~0.5. The earlier scaling laws work (Kaplan et al. 2020) had estimated the power law co-efficient to be ~0.75 (more compute allocated to model scaling and less to data scaling) and the models at the time (early 2022) were under-trained.

How:

  • They train models with different FLOPs budget spread across different model and dataset sizes and look at the optimal training loss across different FLOPs allocations. They then fit a power law L=P+Q/Nα+R/DβL = P + Q/N^\alpha + R/D^\beta where LL is loss, NN is model size, and DD is dataset size. They estimate the power law coefficients (they actually do this estimation 3 different ways and the coefficients roughly agree across all 3 methods) and end up with α≈β≈0.5\alpha \approx \beta \approx 0.5. In contrast, Kaplan et al. had estimated α≈0.75\alpha \approx 0.75 and β≈0.25\beta \approx 0.25 (more compute allocated to model scaling and less to data scaling).
  • And then to make their point that the models at the time (~2021-2022) were under-trained, they take the compute budget of the Gopher 280B model that was trained on ~300B tokens and use it to train Chinchilla, a 70B model (4x smaller) and train it on 1.4 T tokens (4x higher) and show that Chinchilla has a much better performance on downstream tasks than Gopher (despite using the same FLOPs budget).

Question: How did the Kaplan paper end up with such a different estimate of the power law coefficients?

Ans: Because of the following differences in methodology (see Pearce and Song 2024 for details):

  1. They trained smaller models in Kaplan et al. The biggest model they trained for deriving the scaling laws was 1.5B compared to 16B for the Chinchilla paper.
  2. Kaplan et al. did not include the embedding parameters in the total parameter count.

These two differences ended in different estimates of the power law coefficients.

Follow-up reading:

  • MoE scaling laws. See the Joint MoE Scaling Laws paper.
  • Including inference costs in the calculation. Need to read Beyond Chinchilla Optimal. Since small models are cheaper to run inference on and you only train once, choosing a smaller model (compared to what you would if you were following Chinchilla scaling laws) and training it for longer (since you are no longer “Chinchilla optimal”) is not a bad idea.