[algo, perf] feat: Vectorize RLOO Advantage Estimator - 20x Speedup

[EduardDurech]

Vectorize RLOO advantage estimator
130ms -> 6ms
Similar method can be done for other advantage estimators, I just don't have time

Implements

$$r_i - \frac{\sum_{j\ne i} r_j}{G-1} = \frac{(G-1)r_i - \sum_{j\ne i} r_j}{G-1} = \frac{G r_i - \sum_{j\in g} r_j}{G-1}$$

[gemini-code-assist]

Code Review

This pull request vectorizes the RLOO advantage estimator, leading to a significant performance improvement. The previous implementation using Python loops and dictionaries has been replaced with an efficient, vectorized approach using torch.bincount. This not only boosts performance but also enhances correctness by correctly handling groups with a single response, for which the advantage is now properly set to zero. The new code is more concise and idiomatic for tensor computations. Overall, this is an excellent and well-executed optimization.

[vermouth1992]

Could you help add a CI to ensure that they are consistent?

[vermouth1992]

We may want to introduce a new adv estimator called rloo_vector instead of directly over-write the original one

[EduardDurech]

@vermouth1992 I don't know your guys' standard workflow YAML but there is a PyTest if you want to add to ci

$ pytest tests/trainer/ppo/test_core_algos_on_cpu.py::test_rloo_and_vectorized_equivalence -q -s
> [RLOO] seed=0 groups=5 shape=torch.Size([64, 128]) mask_tokens=4147 adv_max_diff=1.907e-06 ret_max_diff=1.907e-06
> [RLOO] seed=1 groups=8 shape=torch.Size([128, 256]) mask_tokens=16364 adv_max_diff=1.907e-06 ret_max_diff=1.907e-06
> [RLOO] seed=2 groups=10 shape=torch.Size([512, 512]) mask_tokens=130968 adv_max_diff=3.815e-06 ret_max_diff=3.815e-06