|
60 | 60 | // ════════════════════════════════════════════════════════════════════════════ |
61 | 61 |
|
62 | 62 | use "stdlib/flame/tensor_lib" |
| 63 | +use "stdlib/flame/flame_math" // dt_exp — deterministic CE-bwd exp (F-OP19), matches FWD nn_ce_loss_allpos + device forge_dispatch_ce_grad |
63 | 64 | use "stdlib/flame/conv_lib" |
64 | 65 | use "stdlib/flame/gn_lib" |
65 | 66 | use "stdlib/flame/moe_lib" |
@@ -426,6 +427,45 @@ fn train_fwd(tokens: int, E: int, L: int, |
426 | 427 | // bdiv = minibatch size B so accumulating the seed over B windows yields the MEAN |
427 | 428 | // per-window gradient (grad-accumulation; bdiv=1 = the old single-window 1/T seed). |
428 | 429 | fn train_ce_grad(logits: int, targets: int, dlogits_out: int, T: int, V: int, bdiv: int) { |
| 430 | + // bdiv==1 (the canon single-window step, incl. MODE_CANON 303M smoke) → wire the |
| 431 | + // EXISTING hexa-lang device CE-grad kernel forge_dispatch_ce_grad (self/runtime.c |
| 432 | + // L15894 host-weak dt_exp + CUDA _hx_k_ce_grad), CLM_PROD_DEVRESIDENT / HEXA_FUSE_ALL |
| 433 | + // gated. This folds the trainer's host O(T·V)=T·256 single-thread softmax+grad seed |
| 434 | + // (#2598 GPU-idle root cause) onto device on a CUDA host. exp = dt_exp (F-OP19 |
| 435 | + // Taylor) on BOTH the device kernel AND the host fallback below — bit-exact |
| 436 | + // host↔device (max|Δ|=0) AND consistent with the FWD loss nn_ce_loss_allpos (which |
| 437 | + // already uses dt_exp); the prior local `exp` (libm) builtin was the lone inconsistency |
| 438 | + // (cross-platform 1-ULP drift). MODE_VERIFY $0 lossF byte-identical (measured CPU). |
| 439 | + if bdiv == 1 { |
| 440 | + if env("CLM_PROD_DEVRESIDENT") != "" || env("HEXA_FUSE_ALL") != "" { |
| 441 | + let rc = forge_dispatch_ce_grad(logits, targets, dlogits_out, T, V) |
| 442 | + if rc == 0 { return } |
| 443 | + } |
| 444 | + // host byte-eq fallback (dt_exp, 1/T) — == device kernel, == FWD loss exp. |
| 445 | + let mut t = 0 |
| 446 | + while t < T { |
| 447 | + let base = t * V |
| 448 | + let mut mx = t_get(logits, base) |
| 449 | + let mut i = 1 |
| 450 | + while i < V { let v = t_get(logits, base + i); if v > mx { mx = v }; i = i + 1 } |
| 451 | + let mut sm = 0.0 |
| 452 | + let mut i2 = 0 |
| 453 | + while i2 < V { sm = sm + dt_exp(t_get(logits, base + i2) - mx); i2 = i2 + 1 } |
| 454 | + let invT = 1.0 / to_float(T) |
| 455 | + let mut v3 = 0 |
| 456 | + while v3 < V { |
| 457 | + let p = dt_exp(t_get(logits, base + v3) - mx) / sm |
| 458 | + t_set(dlogits_out, base + v3, p * invT) |
| 459 | + v3 = v3 + 1 |
| 460 | + } |
| 461 | + let tgt = to_int(t_get(targets, t)) |
| 462 | + t_set(dlogits_out, base + tgt, t_get(dlogits_out, base + tgt) - invT) |
| 463 | + t = t + 1 |
| 464 | + } |
| 465 | + return |
| 466 | + } |
| 467 | + // bdiv>1 grad-accumulation: the device kernel seeds 1/T only, so the 1/(T·bdiv) |
| 468 | + // MEAN-over-minibatch scale keeps the host loop (unchanged baseline, libm exp). |
429 | 469 | let mut t = 0 |
430 | 470 | while t < T { |
431 | 471 | let base = t * V |
|
0 commit comments