jgrusewski b6fb720acd perf(ml-alpha): eliminate all dtoh from training hot path ...

GPU-resident grad-norm + clip-scale; mapped-pinned loss readback.
Replaces 9× memcpy_dtoh per Mamba2 AdamW step (grad-norm host roundtrip)
+ 1× per-step download() (loss). Saves ~10 stream-sync barriers/step.

New kernels (cuda/grad_norm.cu):
  - grad_norm_sq_phase1: per-block tree-reduce of x[i]^2 (no atomicAdd)
  - grad_norm_sq_phase2: cross-tensor accumulator (sequential stream-ordered)
  - grad_clip_scale: writes min(1, max_norm/sqrt(norm_sq)) to device ptr
  - mamba2_alpha_adamw_step_devscale: reads grad_scale from device pointer
    instead of host scalar, allowing AdamW kernels to launch async without
    waiting for a CPU-side norm computation.

Trainer changes (perception.rs):
  - loss_d kept device-side; mapped-pinned MappedF32Buffer shadow.
  - Single stream.synchronize() at end of step (was 2: post-bwd + download).
  - DtoD copy loss_d → loss_host_d queued, then sync flushes both kernels
    + copy in one barrier. Loss read via host_ptr (no dtoh).

Mamba2 AdamW (mamba2_block.rs):
  - step_from_buffers_gpu_clip(): all grad-norm tensors processed via
    phase1+phase2 chain, scale computed on-device, AdamW launches with
    devscale variant. Zero host roundtrips.
  - Pre-allocated block_partials_d, grad_norm_sq_d, grad_scale_d.

Optimizer (optim.rs): removed redundant stream.synchronize() per AdamW step.
Each per-tensor AdamW kernel is stream-ordered; sync only needed before
host reads, which the trainer handles centrally.

Synthetic overfit smoke: initial=0.30 → final=0.0006 (matches pre-refactor
trajectory). Full ml-alpha test suite passes (45 tests across lib +
integration).

Honors:
  - feedback_no_htod_htoh_only_mapped_pinned.md (MappedF32Buffer only)
  - feedback_no_atomicadd.md (block tree-reduce only)
  - feedback_no_legacy_aliases.md (step_from_buffers replaced, not aliased)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

2026-05-17 13:57:02 +02:00

3.1 KiB

Raw Blame History

View Raw