rlox_core/llm/
ops.rs

1//! LLM advantage and KL ops — forwarding re-exports from `rlox-rl-ops`.
2//!
3//! The implementations live in `rlox-rl-ops` so that `rlox-sandbox` can depend
4//! on that slim crate without pulling in all of `rlox-core`.  Everything that
5//! was previously defined here is still **publicly accessible under the same
6//! path** (`rlox_core::llm::ops::*`, `rlox_core::llm::ops::f32_ops::*`, etc.)
7//! so existing callers — including `rlox-python` PyO3 bindings — compile and
8//! behave identically without any change.
9//!
10//! The only item that remains defined here rather than re-exported is
11//! [`DPOPair`], which is a training-plane data container unrelated to advantage
12//! estimation.
13
14// ---------------------------------------------------------------------------
15// Re-export everything from rlox-rl-ops::kl (the macro-generated f64 + f32
16// sub-modules and the module-level f64 re-exports).
17// ---------------------------------------------------------------------------
18
19// f64 sub-module — keep accessible as `rlox_core::llm::ops::f64_ops`.
20pub use rlox_rl_ops::kl::f64_ops;
21
22// f32 sub-module — keep accessible as `rlox_core::llm::ops::f32_ops`.
23pub use rlox_rl_ops::kl::f32_ops;
24
25// Module-level f64 convenience re-exports (same as before: `pub use f64_ops::*`).
26pub use rlox_rl_ops::kl::{
27    compute_batch_group_advantages, compute_batch_token_kl, compute_batch_token_kl_schulman,
28    compute_group_advantages, compute_token_kl, compute_token_kl_schulman,
29};
30
31// ---------------------------------------------------------------------------
32// DPOPair — stays in rlox-core (training-plane data container).
33// ---------------------------------------------------------------------------
34
35/// A DPO preference pair holding tokenized prompt, chosen, and rejected sequences.
36#[derive(Debug, Clone)]
37pub struct DPOPair {
38    pub prompt_tokens: Vec<u32>,
39    pub chosen_tokens: Vec<u32>,
40    pub rejected_tokens: Vec<u32>,
41}
42
43impl DPOPair {
44    pub fn new(
45        prompt_tokens: Vec<u32>,
46        chosen_tokens: Vec<u32>,
47        rejected_tokens: Vec<u32>,
48    ) -> Self {
49        Self {
50            prompt_tokens,
51            chosen_tokens,
52            rejected_tokens,
53        }
54    }
55
56    pub fn chosen_len(&self) -> usize {
57        self.chosen_tokens.len()
58    }
59
60    pub fn rejected_len(&self) -> usize {
61        self.rejected_tokens.len()
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_group_advantages_basic() {
71        let rewards = [1.0, 0.5, 0.8];
72        let adv = compute_group_advantages(&rewards);
73        assert_eq!(adv.len(), 3);
74        let mean: f64 = adv.iter().sum::<f64>() / adv.len() as f64;
75        assert!(mean.abs() < 1e-10);
76    }
77
78    #[test]
79    fn test_group_advantages_constant_rewards() {
80        let rewards = [5.0, 5.0, 5.0];
81        let adv = compute_group_advantages(&rewards);
82        assert!(adv.iter().all(|&v| v == 0.0));
83    }
84
85    #[test]
86    fn test_group_advantages_empty() {
87        let adv = compute_group_advantages(&[]);
88        assert!(adv.is_empty());
89    }
90
91    #[test]
92    fn test_token_kl_identical() {
93        let log_p = [-1.0, -2.0, -0.5];
94        let kl = compute_token_kl(&log_p, &log_p).unwrap();
95        assert!(kl.abs() < 1e-15);
96    }
97
98    #[test]
99    fn test_token_kl_known_value() {
100        let log_p = [-1.0];
101        let log_q = [-2.0];
102        let kl = compute_token_kl(&log_p, &log_q).unwrap();
103        assert!((kl - (-1.0_f64).exp()).abs() < 1e-10);
104    }
105
106    #[test]
107    fn test_token_kl_mismatched_lengths_returns_err() {
108        let result = compute_token_kl(&[1.0, 2.0], &[1.0]);
109        assert!(result.is_err());
110    }
111
112    #[test]
113    fn token_kl_mismatched_lengths_returns_err_not_panic() {
114        let log_p = vec![-1.0f64, -2.0];
115        let log_q = vec![-1.0f64];
116        let result = compute_token_kl(&log_p, &log_q);
117        assert!(result.is_err(), "mismatched lengths must return Err");
118    }
119
120    #[test]
121    fn token_kl_matching_lengths_returns_ok() {
122        let log_p = vec![-1.0f64, -2.0, -0.5];
123        let log_q = vec![-1.0f64, -2.0, -0.5];
124        let result = compute_token_kl(&log_p, &log_q);
125        assert!(result.is_ok());
126        assert!(result.unwrap().abs() < 1e-15);
127    }
128
129    #[test]
130    fn token_kl_empty_slices_returns_zero() {
131        let result = compute_token_kl(&[], &[]);
132        assert!(result.is_ok());
133        assert_eq!(result.unwrap(), 0.0);
134    }
135
136    #[test]
137    fn token_kl_nan_input_propagates_to_output() {
138        let log_p = vec![f64::NAN];
139        let log_q = vec![-1.0f64];
140        let result = compute_token_kl(&log_p, &log_q);
141        if let Ok(v) = result {
142            assert!(v.is_nan(), "NaN input should produce NaN output");
143        }
144    }
145
146    #[test]
147    fn token_kl_inf_input_does_not_panic() {
148        let log_p = vec![f64::INFINITY];
149        let log_q = vec![-1.0f64];
150        let _result = compute_token_kl(&log_p, &log_q);
151    }
152
153    #[test]
154    fn token_kl_known_value_still_correct_after_refactor() {
155        let log_p = vec![-1.0f64];
156        let log_q = vec![-2.0f64];
157        let kl = compute_token_kl(&log_p, &log_q).unwrap();
158        assert!((kl - (-1.0_f64).exp()).abs() < 1e-10);
159    }
160
161    #[test]
162    fn test_batch_group_advantages() {
163        let rewards = [1.0, 2.0, 3.0, 10.0, 10.0, 10.0];
164        let adv = compute_batch_group_advantages(&rewards, 3).unwrap();
165        assert_eq!(adv.len(), 6);
166        let g1_mean: f64 = adv[..3].iter().sum::<f64>() / 3.0;
167        assert!(g1_mean.abs() < 1e-10);
168        assert!(adv[3..6].iter().all(|&v| v == 0.0));
169    }
170
171    #[test]
172    fn test_batch_group_advantages_bad_size() {
173        assert!(compute_batch_group_advantages(&[1.0, 2.0, 3.0], 2).is_err());
174        assert!(compute_batch_group_advantages(&[1.0], 0).is_err());
175    }
176
177    #[test]
178    fn test_token_kl_schulman_identical() {
179        let log_p = [-1.0, -2.0, -0.5];
180        let kl = compute_token_kl_schulman(&log_p, &log_p).unwrap();
181        assert!(kl.abs() < 1e-15);
182    }
183
184    #[test]
185    fn test_token_kl_schulman_known_value() {
186        let log_p = [-1.0];
187        let log_q = [-2.0];
188        let kl = compute_token_kl_schulman(&log_p, &log_q).unwrap();
189        assert!((kl - (1.0_f64.exp() - 2.0)).abs() < 1e-10);
190    }
191
192    #[test]
193    fn test_token_kl_schulman_non_negative() {
194        let log_p = [-0.5, -1.0, -3.0, 0.0];
195        let log_q = [-1.0, -0.5, -0.1, -2.0];
196        let kl = compute_token_kl_schulman(&log_p, &log_q).unwrap();
197        assert!(kl >= 0.0, "Schulman KL should be non-negative, got {kl}");
198    }
199
200    #[test]
201    fn test_dpo_pair() {
202        let pair = DPOPair::new(vec![1, 2, 3], vec![4, 5], vec![6, 7, 8]);
203        assert_eq!(pair.chosen_len(), 2);
204        assert_eq!(pair.rejected_len(), 3);
205        assert_eq!(pair.prompt_tokens.len(), 3);
206    }
207
208    // --- Batched KL tests ---
209
210    #[test]
211    fn test_batch_token_kl_matches_unbatched() {
212        let log_p = vec![-1.0, -2.0, -0.5, -1.5, -0.3, -2.5];
213        let log_q = vec![-1.1, -1.9, -0.6, -1.4, -0.4, -2.4];
214        let batched = compute_batch_token_kl(&log_p, &log_q, 3).unwrap();
215        let kl0 = compute_token_kl(&log_p[..3], &log_q[..3]).unwrap();
216        let kl1 = compute_token_kl(&log_p[3..], &log_q[3..]).unwrap();
217        assert_eq!(batched.len(), 2);
218        assert!((batched[0] - kl0).abs() < 1e-12);
219        assert!((batched[1] - kl1).abs() < 1e-12);
220    }
221
222    #[test]
223    fn test_batch_token_kl_schulman_matches_unbatched() {
224        let log_p = vec![-1.0, -2.0, -0.5, -1.5, -0.3, -2.5];
225        let log_q = vec![-1.1, -1.9, -0.6, -1.4, -0.4, -2.4];
226        let batched = compute_batch_token_kl_schulman(&log_p, &log_q, 3).unwrap();
227        let kl0 = compute_token_kl_schulman(&log_p[..3], &log_q[..3]).unwrap();
228        let kl1 = compute_token_kl_schulman(&log_p[3..], &log_q[3..]).unwrap();
229        assert_eq!(batched.len(), 2);
230        assert!((batched[0] - kl0).abs() < 1e-12);
231        assert!((batched[1] - kl1).abs() < 1e-12);
232    }
233
234    #[test]
235    fn test_batch_token_kl_bad_seq_len() {
236        assert!(compute_batch_token_kl(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0], 0).is_err());
237        assert!(compute_batch_token_kl(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0], 2).is_err());
238    }
239
240    #[test]
241    fn test_batch_token_kl_mismatched_lengths() {
242        assert!(compute_batch_token_kl(&[1.0, 2.0], &[1.0], 1).is_err());
243    }
244
245    // --- f32 variant tests ---
246
247    #[test]
248    fn test_f32_token_kl_identical() {
249        let log_p: Vec<f32> = vec![-1.0, -2.0, -0.5];
250        let kl = f32_ops::compute_token_kl(&log_p, &log_p).unwrap();
251        assert!(kl.abs() < 1e-6);
252    }
253
254    #[test]
255    fn test_f32_batch_token_kl_schulman() {
256        let log_p: Vec<f32> = vec![-1.0, -2.0, -0.5, -1.5];
257        let log_q: Vec<f32> = vec![-1.1, -1.9, -0.6, -1.4];
258        let batched = f32_ops::compute_batch_token_kl_schulman(&log_p, &log_q, 2).unwrap();
259        assert_eq!(batched.len(), 2);
260        let kl0 = f32_ops::compute_token_kl_schulman(&log_p[..2], &log_q[..2]).unwrap();
261        assert!((batched[0] - kl0).abs() < 1e-6);
262    }
263
264    #[test]
265    fn test_f32_group_advantages() {
266        let rewards: Vec<f32> = vec![1.0, 2.0, 3.0];
267        let adv = f32_ops::compute_group_advantages(&rewards);
268        assert_eq!(adv.len(), 3);
269        let mean: f32 = adv.iter().sum::<f32>() / 3.0;
270        assert!(mean.abs() < 1e-5);
271    }
272}
rlox_core/llm/ops.rs

rlox_core/llm/
ops.rs
