post-train: M4 — clipped_pg_loss + scale_rows (GRPO policy-gradient op)
The GRPO (M4) token-level loss op + the one primitive it needs: - scale_rows(x[r,c], s[r]): per-row scale (new ~5-line CUDA kernel). The clipped-PG backward scales each completion token's row of (probs − onehot) by its own per-token coefficient, which cross_entropy_backward's single scalar scale can't express. - clipped_pg_loss(logits, target, logp_old, logp_ref, A, eps, beta): per-token ρ_t = exp(logπθ_t − logp_old_t), L = −mean min(ρA, clip(ρ,1±ε)A) + β·mean KL (k3 estimator), masked to completion tokens. Backward reuses the CE machinery (probs − onehot) + scale_rows. Gates: grad-check the active PG path + the A=0 (KL-only) path; degenerate value checks ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -152,6 +152,15 @@ unsafe extern "C" {
|
||||
pos0: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Per-row scale: y[r,c] = x[r,c] * s[r] (GRPO policy-gradient backward).
|
||||
pub fn launch_scale_rows_f32(
|
||||
x: *const f32,
|
||||
s: *const f32,
|
||||
y: *mut f32,
|
||||
rows: i32,
|
||||
cols: i32,
|
||||
stream: CudaStream,
|
||||
);
|
||||
pub fn launch_rope_dx_f32(
|
||||
dy: *const f32,
|
||||
dx: *mut f32,
|
||||
|
||||
Reference in New Issue
Block a user