feat: ZTENSOR_DETERMINISTIC mode -- deterministic GPU training runs (T4.1)

compute/gpu_fused_encoder.go

@@ -5,6 +5,7 @@ import (
 	"unsafe"
 
 	"github.com/zerfoo/ztensor/internal/cublas"
+	"github.com/zerfoo/ztensor/internal/cuda"
 	"github.com/zerfoo/ztensor/internal/gpuapi"
 )
 
@@ -46,6 +47,12 @@ func (e *GPUEngine[T]) FusedEncoderForward(
 }
 
 // FusedEncoderBackward computes all gradients for one fused encoder layer.
+//
+// dScale/dBias accumulate via atomicAdd across row blocks in
+// fused_encoder_bwd.cu (order-dependent, no deterministic variant exists).
+// Under ZTENSOR_DETERMINISTIC=1 this refuses to run rather than silently
+// return order-dependent gradients -- see docs/design.md "ZTENSOR_DETERMINISTIC
+// scope" (plan-gpu-training-hardening.md T4.1) and zerfoo docs/lore.md.
 func (e *GPUEngine[T]) FusedEncoderBackward(
 	weights *[16]unsafe.Pointer,
 	weightT *[6]unsafe.Pointer,
@@ -55,6 +62,12 @@ func (e *GPUEngine[T]) FusedEncoderBackward(
 	dOutput, dInput, input unsafe.Pointer,
 	totalRows, dModel, nHeads, headDim, ffnDim, bsC, numPatches int,
 ) error {
+	if cuda.DeterministicEnabled() {
+		return fmt.Errorf("FusedEncoderBackward: dScale/dBias gradient accumulation uses " +
+			"order-dependent atomicAdd across row blocks (fused_encoder_bwd.cu) with no " +
+			"deterministic variant; refusing under ZTENSOR_DETERMINISTIC=1 (unset it, or use " +
+			"the unfused per-op backward path)")
+	}
 	h := blasHandlePtr(e.blas)
 	if h == nil {
 		return fmt.Errorf("FusedEncoderBackward: cuBLAS handle not available")

compute/gpu_fused_encoder_test.go

@@ -0,0 +1,45 @@
+package compute
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/zerfoo/ztensor/internal/cuda"
+)
+
+// TestFusedEncoderBackward_RefusesUnderDeterministicMode proves the
+// ZTENSOR_DETERMINISTIC=1 honesty guard (T4.1): fused_encoder_bwd.cu's
+// dScale/dBias gradient accumulation uses atomicAdd across row blocks with no
+// deterministic variant, so FusedEncoderBackward must refuse to run under the
+// flag instead of silently returning order-dependent gradients. The guard
+// fires before any CUDA/cuBLAS handle is touched, so this runs without a GPU.
+func TestFusedEncoderBackward_RefusesUnderDeterministicMode(t *testing.T) {
+	restore := cuda.SetDeterministicEnabledForTesting(true)
+	defer restore()
+
+	e := &GPUEngine[float32]{}
+	err := e.FusedEncoderBackward(nil, nil, nil, nil, nil, nil, nil, nil, 0, 0, 0, 0, 0, 0, 0)
+	if err == nil {
+		t.Fatal("FusedEncoderBackward: expected an error under ZTENSOR_DETERMINISTIC=1, got nil")
+	}
+	if !strings.Contains(err.Error(), "ZTENSOR_DETERMINISTIC") {
+		t.Fatalf("FusedEncoderBackward error does not name ZTENSOR_DETERMINISTIC as the cause: %v", err)
+	}
+}
+
+// TestFusedEncoderBackward_AllowedWithoutDeterministicMode confirms the guard
+// is inert (falls through to the ordinary cuBLAS-handle-missing error, not
+// the determinism error) when the flag is unset.
+func TestFusedEncoderBackward_AllowedWithoutDeterministicMode(t *testing.T) {
+	restore := cuda.SetDeterministicEnabledForTesting(false)
+	defer restore()
+
+	e := &GPUEngine[float32]{}
+	err := e.FusedEncoderBackward(nil, nil, nil, nil, nil, nil, nil, nil, 0, 0, 0, 0, 0, 0, 0)
+	if err == nil {
+		t.Fatal("FusedEncoderBackward: expected an error (no cuBLAS handle on a bare engine), got nil")
+	}
+	if strings.Contains(err.Error(), "ZTENSOR_DETERMINISTIC") {
+		t.Fatalf("FusedEncoderBackward: determinism guard fired with the flag unset: %v", err)
+	}
+}

docs/design.md

@@ -121,6 +121,43 @@ All methods accept a `context.Context` and an optional variadic `dst` tensor for
 
 Accumulation dtype is preserved per site (float32 data reduces in float32, float64 in float64). Widening reduced-precision element types (float16/bfloat16) to a float32 accumulator, and a fixed-order rewrite of the GPU reduction kernels and the arm64 NEON RMSNorm SIMD path (both currently deterministic-but-SIMD-strided), are deferred to the deterministic-reductions mode (`ZTENSOR_DETERMINISTIC`, plan-gpu-training-hardening T4.1).
 
+#### `ZTENSOR_DETERMINISTIC` scope (T4.1)
+
+`ZTENSOR_DETERMINISTIC=1` is an env-gated debug/verification mode for
+bit-reproducible GPU training runs: the same seed, model, and data must
+produce bitwise-identical per-epoch losses across repeated runs. It is read
+once at process init (`internal/cuda.DeterministicEnabled`), off by default,
+and is a scope statement, not a rewrite of the numerics -- the T135.2
+nondeterminism inventory (below) found that most of the training path was
+already deterministic run-to-run; the flag closes the two remaining gaps and
+gives the rest an honest, checked-in disposition instead of an implicit
+assumption.
+
+**Nondeterminism inventory and disposition:**
+
+| Source | Status under `ZTENSOR_DETERMINISTIC=1` | Notes |
+|---|---|---|
+| CPU fp32 reductions (`numeric.Sum`, `compute.CPUEngine` Sum/ReduceSum/ReduceMean/Softmax, xblas RMSNorm sum-of-squares, arm64 NEON SIMD path) | **Deterministic already, unconditionally** (T135.2) | Fixed-order pairwise/tree accumulation; a pure function of length, not `GOMAXPROCS` or goroutine scheduling. The flag changes nothing here. |
+| GPU custom kernels: softmax, RMSNorm, GEMV family (`sgemv_m1`, `gemv_warp`, `gemv_q4k*`), flash attention/decode, `fused_adamw` | **Deterministic already** | Warp-shuffle/tree reductions inside a launch configuration that is a pure function of input shape; audited across the full kernel inventory (`internal/cuda/kernels/*.cu`) and found free of cross-block atomics in the reduction path. The flag changes nothing here. |
+| cuBLAS `Sgemm`/`SgemmStridedBatched`/`GemmEx` (the dense f32/bf16/fp16 GEMM path) | **Routed to a deterministic configuration** | `internal/cublas.CreateHandle` sets `CUBLAS_WORKSPACE_CONFIG` (if unset) and calls `cublasSetMathMode(CUBLAS_PEDANTIC_MATH)` on the handle -- the two documented NVIDIA/PyTorch levers that disable TF32 tensor-core downcasting and any workspace-dependent split-K/atomics reduction algorithm cuBLAS's heuristic might otherwise select. Best-effort: a workspace env var set after an earlier cuBLAS handle already exists in the process, or a cuBLAS build missing `cublasSetMathMode`, only warns (stderr), never fails handle creation. |
+| `fused_encoder_bwd.cu` `dScale`/`dBias` (LayerNorm/RMSNorm-style bias-gradient accumulation across rows) | **NOT covered -- refuses to run** | Uses `atomicAdd` across row blocks; the addition order depends on block-completion order, which is not fixed. No deterministic variant exists. `compute.GPUEngine.FusedEncoderBackward` checks the flag first and returns an error instead of silently producing order-dependent gradients. As of T135.4, this path is not reachable from zerfoo's PatchTST training (the Go-side wiring is a stub that always falls back to the unfused per-op path -- #522), so this exclusion does not block the GB10 bitwise-identity proof on that model; it would block any future consumer that wires the fused path in. |
+| `counter.cu`'s `atomicAdd` (GPU-resident decode-position counter) | **Not a determinism concern** | Single-thread kernel (`<<<1,1>>>`); atomicity is for correctness under concurrent access, not a multi-thread reduction order. |
+| Arena reuse / stream ordering (host-access sync, reset-epoch frees) | **Out of scope for T4.1** | These are correctness invariants (L-0002/L-0003/L-0004 in zerfoo's lore), not reduction-order nondeterminism; already enforced unconditionally, independent of this flag. |
+
+**What this mode does NOT guarantee:** bitwise identity across different GPUs,
+driver versions, or cuBLAS versions (cuBLAS's chosen algorithm is a function
+of all three, even under `CUBLAS_PEDANTIC_MATH`); bitwise identity for any
+code path that reaches `FusedEncoderBackward` (it errors instead); bitwise
+identity against the CPU engine (CPU pairwise order and GPU warp-shuffle
+order are different, equally valid parenthesizations of the same sum -- see
+`docs/kernel-tolerances.md` for the accuracy gap, which is a separate
+concern from run-to-run determinism).
+
+**Cost:** `CUBLAS_PEDANTIC_MATH` forgoes TF32 tensor-core paths for f32
+GEMMs, so expect a GEMM slowdown when the flag is set; it is a debug/
+verification tool, off by default, not intended for production training
+throughput. See devlog for the measured GB10 double-run proof.
+
 ### Optional Engine Capabilities
 
 Beyond the core interface, engines may implement optional capability interfaces discovered via type assertion:

internal/cublas/cublas_purego.go

@@ -2,6 +2,7 @@ package cublas
 
 import (
 	"fmt"
+	"os"
 	"sync"
 	"unsafe"
 
@@ -32,6 +33,20 @@ const (
 	CublasCompute32F CublasComputeType = 68 // CUBLAS_COMPUTE_32F
 )
 
+// cuBLAS math-mode constants (cublasMath_t). Only the modes this package
+// uses are declared; see NVIDIA cuBLAS docs for the full enum.
+const (
+	// CublasDefaultMath is cuBLAS's default: on Ampere+ GPUs (including the
+	// GB10) this permits TF32 tensor-core downcasting for float32 GEMMs.
+	CublasDefaultMath = 0 // CUBLAS_DEFAULT_MATH
+	// CublasPedanticMath disables TF32 downcasting and other
+	// reduced-precision/algorithm-selection paths; NVIDIA's cuBLAS
+	// documentation describes it as producing reproducible results for a
+	// fixed problem size and GPU. Used by ZTENSOR_DETERMINISTIC=1
+	// (internal/cuda.DeterministicEnabled).
+	CublasPedanticMath = 2 // CUBLAS_PEDANTIC_MATH
+)
+
 // cuBLAS status codes.
 const cublasStatusSuccess = 0
 
@@ -43,6 +58,7 @@ type cublasLib struct {
 	sgemm               uintptr // cublasSgemm_v2
 	gemmEx              uintptr // cublasGemmEx
 	sgemmStridedBatched uintptr // cublasSgemmStridedBatched
+	setMathMode         uintptr // cublasSetMathMode (optional, best-effort)
 }
 
 var (
@@ -92,6 +108,16 @@ func loadCublas() (*cublasLib, error) {
 		}
 		*s.ptr = addr
 	}
+
+	// cublasSetMathMode is resolved best-effort: it backs the debug-only
+	// ZTENSOR_DETERMINISTIC mode (T4.1), so a cuBLAS build stripped of it
+	// must not break BLAS loading for everyone else. lib.setMathMode stays
+	// zero if the symbol is unavailable; SetMathMode reports that as an
+	// error at call time.
+	if addr, err := cuda.Dlsym(handle, "cublasSetMathMode"); err == nil {
+		lib.setMathMode = addr
+	}
+
 	return lib, nil
 }
 
@@ -108,21 +134,100 @@ type Handle struct {
 }
 
 // Ptr returns the raw cuBLAS handle pointer for passing to C functions
-// (e.g., the fused encoder kernel orchestrator).
-func (h *Handle) Ptr() unsafe.Pointer { return unsafe.Pointer(h.ptr) }
+// (e.g., the fused encoder kernel orchestrator). Same pre-existing
+// uintptr<->unsafe.Pointer wrapping pattern used throughout the purego
+// bindings (internal/cuda/runtime_purego.go and siblings); not specific to
+// T4.1, just newly visible once package-scoped linting was fixed to load
+// full packages instead of individual staged files.
+func (h *Handle) Ptr() unsafe.Pointer { return unsafe.Pointer(h.ptr) } //nolint:govet
 
 // CreateHandle creates a new cuBLAS context handle.
+//
+// Under ZTENSOR_DETERMINISTIC=1 (internal/cuda.DeterministicEnabled), this
+// also sets CUBLAS_WORKSPACE_CONFIG (if the process has not already set one)
+// before creating the handle, and CUBLAS_PEDANTIC_MATH on the handle after
+// creation -- the two documented levers for bit-reproducible cuBLAS GEMMs.
+// Both are best-effort: a workspace-config env var set too late (cuBLAS
+// reads it once, at the FIRST handle creation in the process) or a cuBLAS
+// build missing cublasSetMathMode only degrades determinism guarantees and
+// is reported via a stderr warning, never a hard failure -- this handle is
+// still needed for ordinary (non-deterministic-mode) training and inference.
 func CreateHandle() (*Handle, error) {
 	lib, err := getCublasLib()
 	if err != nil {
 		return nil, err
 	}
+	if cuda.DeterministicEnabled() {
+		ensureDeterministicWorkspaceConfig()
+	}
 	var h uintptr
 	status := cuda.Ccall(lib.create, uintptr(unsafe.Pointer(&h)))
 	if status != cublasStatusSuccess {
 		return nil, fmt.Errorf("cublasCreate failed with status %d", status)
 	}
-	return &Handle{ptr: h}, nil
+	handle := &Handle{ptr: h}
+	if cuda.DeterministicEnabled() {
+		if merr := SetMathMode(handle, CublasPedanticMath); merr != nil {
+			determinismWarnFn("cublasSetMathMode(CUBLAS_PEDANTIC_MATH) unavailable: %v -- "+
+				"this handle's GEMM determinism relies on CUBLAS_WORKSPACE_CONFIG alone", merr)
+		}
+	}
+	return handle, nil
+}
+
+// determinismWarnFn sinks ZTENSOR_DETERMINISTIC setup warnings (workspace
+// config set late, math mode unavailable). Tests swap it to capture output;
+// mirrors internal/cuda's arenaPoisonWarnFn pattern.
+var determinismWarnFn = func(format string, args ...any) {
+	fmt.Fprintf(os.Stderr, "ztensor cublas determinism: "+format+"\n", args...)
+}
+
+var workspaceConfigOnce sync.Once
+
+// ensureDeterministicWorkspaceConfig sets CUBLAS_WORKSPACE_CONFIG to a fixed,
+// bounded value if the process has not already set one. cuBLAS reads this
+// variable once, the first time a cuBLAS context is created in the process,
+// to restrict the workspace it may allocate for a GEMM -- which forecloses
+// certain heuristically-chosen, workspace-dependent split-K/atomics
+// reduction algorithms for large-K problems. This is the same variable
+// PyTorch's determinism docs require for cuBLAS. Setting it here (at first
+// CreateHandle, rather than at process start) is a best-effort convenience:
+// it is NOT guaranteed to take effect if any other code path created a
+// cuBLAS handle earlier in the process. Prefer setting
+// CUBLAS_WORKSPACE_CONFIG in the environment before the process starts for
+// a guaranteed effect.
+func ensureDeterministicWorkspaceConfig() {
+	workspaceConfigOnce.Do(func() {
+		if os.Getenv("CUBLAS_WORKSPACE_CONFIG") == "" {
+			_ = os.Setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+			determinismWarnFn("ZTENSOR_DETERMINISTIC=1: CUBLAS_WORKSPACE_CONFIG was unset; " +
+				"set :4096:8 for this process. For a guaranteed effect, set it in the " +
+				"environment before the process starts instead.")
+		}
+	})
+}
+
+// SetMathMode sets the cuBLAS math mode for handle h (cublasSetMathMode).
+// ZTENSOR_DETERMINISTIC=1 uses this to request CUBLAS_PEDANTIC_MATH. Returns
+// an error if the loaded cuBLAS build does not export cublasSetMathMode
+// (very old builds) or the call itself fails; callers treat that as
+// best-effort and warn rather than fail handle creation.
+func SetMathMode(h *Handle, mode int) error {
+	if h == nil {
+		return fmt.Errorf("cublasSetMathMode: nil handle")
+	}
+	lib, err := getCublasLib()
+	if err != nil {
+		return err
+	}
+	if lib.setMathMode == 0 {
+		return fmt.Errorf("cublasSetMathMode: symbol not available in loaded cuBLAS")
+	}
+	status := cuda.Ccall(lib.setMathMode, h.ptr, uintptr(mode))
+	if status != cublasStatusSuccess {
+		return fmt.Errorf("cublasSetMathMode failed with status %d", status)
+	}
+	return nil
 }
 
 // Destroy releases the cuBLAS handle resources.

internal/cublas/deterministic_test.go

@@ -0,0 +1,90 @@
+package cublas
+
+import (
+	"os"
+	"sync"
+	"testing"
+
+	"github.com/zerfoo/ztensor/internal/cuda"
+)
+
+// TestSetMathMode_NilHandle exercises the argument-validation path of
+// SetMathMode, which does not require a live cuBLAS context, so it runs even
+// on hosts without CUDA (e.g. darwin dev machines).
+func TestSetMathMode_NilHandle(t *testing.T) {
+	if err := SetMathMode(nil, CublasPedanticMath); err == nil {
+		t.Fatal("SetMathMode(nil, ...): expected an error, got nil")
+	}
+}
+
+// TestEnsureDeterministicWorkspaceConfig_SetsWhenUnset proves the
+// ZTENSOR_DETERMINISTIC=1 workspace-config lever (T4.1): CUBLAS_WORKSPACE_CONFIG
+// is set to a fixed value when the process has not already set one. Does not
+// require CUDA -- this only touches the env var, not the cuBLAS library.
+func TestEnsureDeterministicWorkspaceConfig_SetsWhenUnset(t *testing.T) {
+	origWorkspace, hadWorkspace := os.LookupEnv("CUBLAS_WORKSPACE_CONFIG")
+	_ = os.Unsetenv("CUBLAS_WORKSPACE_CONFIG")
+	t.Cleanup(func() {
+		if hadWorkspace {
+			_ = os.Setenv("CUBLAS_WORKSPACE_CONFIG", origWorkspace)
+		} else {
+			_ = os.Unsetenv("CUBLAS_WORKSPACE_CONFIG")
+		}
+	})
+
+	// workspaceConfigOnce is process-lifetime sync.Once; reset it so this
+	// test's Do body actually runs regardless of test execution order.
+	workspaceConfigOnce = sync.Once{}
+
+	ensureDeterministicWorkspaceConfig()
+
+	if got := os.Getenv("CUBLAS_WORKSPACE_CONFIG"); got == "" {
+		t.Fatal("ensureDeterministicWorkspaceConfig: CUBLAS_WORKSPACE_CONFIG still unset")
+	}
+}
+
+// TestEnsureDeterministicWorkspaceConfig_LeavesExistingValue confirms an
+// operator-provided CUBLAS_WORKSPACE_CONFIG is never overwritten -- the
+// process-level env var, set before the process starts, is the authoritative
+// source per the documented cuBLAS/PyTorch determinism contract.
+func TestEnsureDeterministicWorkspaceConfig_LeavesExistingValue(t *testing.T) {
+	origWorkspace, hadWorkspace := os.LookupEnv("CUBLAS_WORKSPACE_CONFIG")
+	const want = ":16:8"
+	_ = os.Setenv("CUBLAS_WORKSPACE_CONFIG", want)
+	t.Cleanup(func() {
+		if hadWorkspace {
+			_ = os.Setenv("CUBLAS_WORKSPACE_CONFIG", origWorkspace)
+		} else {
+			_ = os.Unsetenv("CUBLAS_WORKSPACE_CONFIG")
+		}
+	})
+
+	workspaceConfigOnce = sync.Once{}
+	ensureDeterministicWorkspaceConfig()
+
+	if got := os.Getenv("CUBLAS_WORKSPACE_CONFIG"); got != want {
+		t.Fatalf("ensureDeterministicWorkspaceConfig overwrote an existing value: got %q, want %q", got, want)
+	}
+}
+
+// TestCreateHandle_DeterministicMode is the GPU proof that CreateHandle does
+// not fail under ZTENSOR_DETERMINISTIC=1 (best-effort setup: a warning, never
+// a hard error, when cublasSetMathMode or the workspace env var cannot be
+// honored). Requires cuBLAS; skips on hosts without a GPU.
+func TestCreateHandle_DeterministicMode(t *testing.T) {
+	if !Available() {
+		t.Skip("cuBLAS not available")
+	}
+	restore := cuda.SetDeterministicEnabledForTesting(true)
+	defer restore()
+
+	h, err := CreateHandle()
+	if err != nil {
+		t.Fatalf("CreateHandle under ZTENSOR_DETERMINISTIC=1: %v", err)
+	}
+	defer func() { _ = h.Destroy() }()
+
+	if err := SetMathMode(h, CublasPedanticMath); err != nil {
+		t.Fatalf("SetMathMode(CUBLAS_PEDANTIC_MATH) on a live handle: %v", err)
+	}
+}

internal/cuda/arena_testing.go

@@ -38,3 +38,14 @@ func HostPoisonFillForTesting(ptr unsafe.Pointer, byteLen int) error {
 	fillHostPoison(unsafe.Slice((*byte)(ptr), byteLen))
 	return nil
 }
+
+// SetDeterministicEnabledForTesting flips ZTENSOR_DETERMINISTIC mode, which
+// is normally read once from the env var at process init, and returns a func
+// that restores the previous value. Lets other packages (e.g. compute's
+// FusedEncoderBackward guard test) exercise the flag without a real
+// ZTENSOR_DETERMINISTIC=1 process environment.
+func SetDeterministicEnabledForTesting(enabled bool) (restore func()) {
+	orig := deterministicEnabled
+	deterministicEnabled = enabled
+	return func() { deterministicEnabled = orig }
+}