Skip to content

Gemma4 DPO and GRPO implementation for Cuda kernels - #626

Open
timkaye11 wants to merge 16 commits into
antflydb:mainfrom
timkaye11:gemma_afourb_cuda
Open

timkaye11 wants to merge 16 commits into
antflydb:mainfrom
timkaye11:gemma_afourb_cuda

Conversation

@timkaye11

Copy link
Copy Markdown
Contributor

This PR productionizes Gemma 4 CUDA preference training for both E2B and E4B,
plus the qualified 26B-A4B inference load path.

  • Adds real-checkpoint, optimizer-backed Gemma 4 DPO and GRPO with strict CUDA
    execution, batched preference gradients, device-resident AdamW, and LoRA
    artifact publication.
  • Aligns GRPO reward scaling and loss reduction with pinned TRL/Unsloth:
    group sample standard deviation (correction = 1), additive 1e-4
    denominator epsilon, and per-completion token-mean loss.
  • Honors Gemma 4's checkpoint-declared MLP topology: E2B doubles the shared-KV
    tail while E4B keeps uniform-width MLPs. Missing flags default to the
    upstream-compatible uniform topology.
  • Adds a fail-closed, local-only qualification and benchmark harness that
    verifies CUDA execution evidence, adapter movement, deterministic output,
    model/protocol fingerprints, and reports matched end-to-end performance
    against Unsloth.
  • Hardens the Gemma 4 26B-A4B Q4_0 load path with resident CUDA loading,
    bounded pinned-host pipelining, prepared-pack policy, startup prefetch,
    memory envelopes, and strict public configuration.

Qualification and Unsloth comparison

Matched end-to-end benchmarks on one NVIDIA L4 (SM89, 23,034 MiB), using local
Gemma 4 E2B and E4B BF16 SafeTensors checkpoints. Every sample is a fresh
process and includes startup, model load, 25 optimizer updates, and adapter
publication. Results are medians of three deterministic repetitions.

Protocol: LoRA rank 16 / alpha 32, q_proj + v_proj, dropout 0, sequence cap
128, AdamW (lr=1e-4, weight decay 0.01, betas 0.9/0.999, epsilon 1e-8),
constant LR, gradient clipping 1.0, and accumulation 1. GRPO uses group size 4,
four completion tokens, clip epsilon 0.2, KL coefficient 0.04, sample-std
reward scaling, and advantage epsilon 1e-4.

Model Objective Zig median Unsloth median Zig steps/s Unsloth steps/s Throughput ratio Peak GPU memory
E2B DPO 127.56 s 137.41 s 0.19599 0.18194 1.077x 11,370 vs 10,294 MiB (+10.5%)
E2B GRPO 165.04 s 165.67 s 0.15148 0.15091 1.004x 11,540 vs 10,098 MiB (+14.3%)
E4B DPO 194.38 s 175.29 s 0.12861 0.14262 0.902x 17,348 vs 15,742 MiB (+10.2%)
E4B GRPO 239.36 s 214.71 s 0.10445 0.11643 0.897x 17,510 vs 15,534 MiB (+12.7%)

E2B DPO completes in 7.2% less wall time than Unsloth and corrected-semantics
GRPO is at parity, completing in 0.4% less wall time. E4B is within 10.3% of
Unsloth throughput for both objectives while using 10.2-12.7% more peak GPU
memory. All three Zig repetitions for every model/objective benchmark produced
identical losses, gradient norms, and trained-adapter SHA-256 values. Every run
reported exactly 25 CUDA optimizer steps, changed every expected LoRA tensor,
and reported zero graph fallback, unsupported operations, interpreter
fallback, or host gradients.

Pinned baseline versions:

  • Torch 2.12.1+cu130
  • Transformers 5.9.0
  • Triton 3.7.1
  • TRL 0.28.0
  • Unsloth 2026.8.19 / Unsloth Zoo 2026.8.13

ReleaseFast CUDA binary SHA-256 used for the E4B benchmark and current E2B
regression qualification:
bd9d19f56eef478a72deebf71a6c80e1180bff84fc840f5e1b8ca3191d625700

The E2B three-repetition benchmark predates the parser-only E4B follow-up and
used binary SHA-256
6b8ec016e1c12bb4c981d5fc029497bed616213475a5390fe4a96e339e3abf3e.
The current binary was requalified on the real E2B checkpoint for both DPO and
GRPO; both objectives passed with every expected adapter tensor changed.

Validation

  • zig build inference-test-finetune
  • zig build quant-kernel-local-check -Dmetal=false -Dcuda=false
    (120 selected compiler/renderer tests)
  • Focused sparse preference-loss finite-difference gradient test
  • 29 Python qualification/baseline tests plus py_compile
  • 10/10 focused Gemma 4 parser/topology tests
  • ReleaseFast CUDA build
  • CUDA artifact source policy and 67 generated-source checks
  • Real-checkpoint E2B and E4B topology preflight
  • Strict one-update E2B and E4B DPO/GRPO CUDA smoke
  • Three-repetition E2B and E4B DPO/GRPO Zig benchmark
  • Three-repetition E2B and E4B DPO/GRPO pinned-Unsloth baseline
  • Changed-file Zig formatting and committed/uncommitted diff checks
  • OpenAPI comparison and idempotent regeneration

timkaye11 and others added 3 commits September 2, 2026 23:09
Add optimizer-backed DPO and GRPO with strict CUDA execution, batched preference gradients, device-resident AdamW, deterministic qualification, and LoRA artifact publication.

Qualify E2B and E4B BF16 checkpoints against pinned Unsloth, honor checkpoint-declared shared-tail MLP topology, harden A4B quantized loading, and regenerate the affected API and CUDA artifacts.
@ajroetker

Copy link
Copy Markdown
Contributor

looks like this needs a make generate?

Comment thread specs/openapi/inference/api.yaml Outdated
- auto
- streamed
- resident
A4bLoadStrategy:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we should have a4b specific config options

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need all these?

@ajroetker ajroetker added this to the v0.3.0 milestone Sep 8, 2026
@timkaye11 timkaye11 added the ci:gpu Select GPU smoke tests for the next approved PR CI run label Sep 17, 2026
@timkaye11

Copy link
Copy Markdown
Contributor Author

/ci run d8dfae0

# Conflicts:
#	go/pkg/sdk/oapi/client.gen.go
#	zig/pkg/antfly/src/storage/lsm_backend/cache.zig
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:gpu Select GPU smoke tests for the next approved PR CI run

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants