fix(training): GPU-resident optimizer steps by default; clipping and anomaly guard on the device

src/NeuralNetworks/NeuralNetworkBase.cs

@@ -11706,7 +11706,9 @@ protected void TrainWithTapeStreaming(
                     // GPU resident-weight coherence: the streaming optimizer just mutated this
                     // weight's CPU backing in place; invalidate its resident device buffer so the
                     // next forward re-uploads current weights (parity with the eager/fused fixes).
-                    if (Engine is DirectGpuTensorEngine _sg) _sg.InvalidateResidentWeightBuffer(source);
+                    // A device-side streaming update leaves the device copy authoritative: keep it (see the eager step).
+                    if (Engine is DirectGpuTensorEngine _sg && !source.HasPendingGpuData)
+                        _sg.InvalidateResidentWeightBuffer(source);
                 });
         }
         else if (twoPass)
@@ -12676,7 +12678,17 @@ void UnscaleAndCheck(Tensor<T>? g)
                 if (Engine is DirectGpuTensorEngine residentGpuEngine)
                 {
                     foreach (var wtensor in trainableParams)
+                    {
+                        // A device-resident optimizer step (GpuOptimizer) leaves the parameter's device copy
+                        // authoritative, with a pending lazy host download. InvalidateResidentWeightBuffer's contract
+                        // is the opposite case -- the HOST copy was just written -- and it drops both the device buffer
+                        // and that pending download, so calling it here threw the update away: the next forward
+                        // re-uploaded the pre-step host weights and the loss stayed flat. Same predicate as the fused
+                        // path's invalidation below.
+                        if (wtensor.HasPendingGpuData)
+                            continue;
                         residentGpuEngine.InvalidateResidentWeightBuffer(wtensor);
+                    }
                 }
             }
 
@@ -12945,16 +12957,61 @@ void UnscaleAndCheck(Tensor<T>? g)
     /// <para>
     /// Reset at the start of every clipped step, so it always describes the latest one.
     /// </para>
+    /// <para>
+    /// On a GPU engine the global-norm clip runs on the device and the norm stays there; reading this property is
+    /// what reads it back (one scalar), as reading the tensor <c>clip_grad_norm_</c> returns does in PyTorch.
+    /// </para>
     /// </remarks>
-    public bool LastStepHadNonFiniteGradients { get; private set; }
+    public bool LastStepHadNonFiniteGradients
+    {
+        get => _lastDeviceGradientNorm is { } deviceNorm ? !deviceNorm.IsFinite : _lastStepHadNonFiniteGradients;
+        private set => _lastStepHadNonFiniteGradients = value;
+    }
+
+    private bool _lastStepHadNonFiniteGradients;
+
+    // The latest device-side clip's total norm (GPU engine, float): kept on the device until someone asks for it.
+    private AiDotNet.Tensors.Engines.Gpu.DeviceGradientNorm? _lastDeviceGradientNorm;
 
     private void ApplyGradientClipping(
         Dictionary<Tensor<T>, Tensor<T>> grads, double maxNorm,
         IReadOnlyList<Tensor<T>> iterationOrder)
     {
         LastStepHadNonFiniteGradients = false;
+        _lastDeviceGradientNorm?.Dispose();
+        _lastDeviceGradientNorm = null;
         if (UsesElementWiseGradientClipping)
         {
+            // GPU engine: the finiteness check is one device sum of squares and one scalar read, and the clamp runs on
+            // the device; the host loops below download every gradient. A non-finite sum is confirmed by the exact
+            // host scan (a float-accumulating backend can overflow on finite gradients).
+            if (Engine.SupportsGpu)
+            {
+                bool provenFinite = false;
+                if (typeof(T) == typeof(float))
+                {
+                    var deviceGrads = new List<AiDotNet.Tensors.LinearAlgebra.Tensor<float>>(iterationOrder.Count);
+                    for (int p = 0; p < iterationOrder.Count; p++)
+                        if (grads.TryGetValue(iterationOrder[p], out var dg) && dg is { Length: > 0 })
+                            deviceGrads.Add((AiDotNet.Tensors.LinearAlgebra.Tensor<float>)(object)dg);
+                    using var deviceSum = AiDotNet.Tensors.Engines.Gpu.GpuOptimizer.TryGradientSumOfSquares(deviceGrads);
+                    provenFinite = deviceSum is not null && deviceSum.IsFinite;
+                }
+                if (provenFinite)
+                {
+                    T lo = NumOps.FromDouble(-maxNorm), hi = NumOps.FromDouble(maxNorm);
+                    using var noGrad = new AiDotNet.Tensors.Engines.Autodiff.NoGradScope<T>();
+                    for (int p = 0; p < iterationOrder.Count; p++)
+                    {
+                        var key = iterationOrder[p];
+                        if (!grads.TryGetValue(key, out var g) || g is null || g.Length == 0) continue;
+                        // In place: the gradient tensor is shared with the published gradient set.
+                        Engine.TensorCopy(Engine.TensorClamp(g, lo, hi), g);
+                    }
+                    return;
+                }
+            }
+
             // Detect non-finite values before mutating any tensor. This preserves
             // the existing diagnostic contract: a poisoned step is left intact
             // so the originating layer can still be identified.
@@ -12995,6 +13052,28 @@ private void ApplyGradientClipping(
             return;
         }
 
+        // GPU engine: clip_grad_norm_ on the device (multi-tensor sum of squares, the coefficient computed on the
+        // device, one multi-tensor scale), as PyTorch does. The host loop below reads every gradient element through
+        // Data.Span, which downloads every device-resident gradient each step (measured on a 27M-parameter LM: 32
+        // transfers, 107 MB per step) and re-uploads them for the optimizer. A non-finite norm leaves the gradients
+        // unscaled there too, and LastStepHadNonFiniteGradients reads the device norm when asked.
+        if (typeof(T) == typeof(float) && Engine is DirectGpuTensorEngine)
+        {
+            var deviceGrads = new List<AiDotNet.Tensors.LinearAlgebra.Tensor<float>>(iterationOrder.Count);
+            for (int p = 0; p < iterationOrder.Count; p++)
+            {
+                if (!grads.TryGetValue(iterationOrder[p], out var g)) continue;
+                if (g is null || g.Length == 0) continue;
+                deviceGrads.Add((AiDotNet.Tensors.LinearAlgebra.Tensor<float>)(object)g);
+            }
+            var deviceNorm = AiDotNet.Tensors.Engines.Gpu.GpuOptimizer.TryClipGradNorm(deviceGrads, maxNorm);
+            if (deviceNorm is not null)
+            {
+                _lastDeviceGradientNorm = deviceNorm;
+                return;
+            }
+        }
+
         // Step 1: total L2 norm across all gradient tensors, iterating in
         // the caller-supplied deterministic order (NOT dict bucket order —
         // that's process-randomized for reference-keyed dicts).

src/Optimizers/AMSGradOptimizer.cs

@@ -394,9 +394,8 @@ protected override void StepCore(TapeStepContext<T> context)
         T biasCorrection1 = NumOps.FromDouble(FirstMomentCorrection(_tapeStep));
         T biasCorrection2 = NumOps.FromDouble(SecondMomentCorrection(_tapeStep));
 
-        // GPU-resident step (AIDOTNET_GPU_ADAM=1); gated off, CPU fallback per-param when not GPU-resident.
+        // GPU-resident step; CPU fallback per-param when not GPU-resident.
         bool gpuAdam = typeof(T) == typeof(float)
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         foreach (var param in context.Parameters)

src/Optimizers/AdaDeltaOptimizer.cs

@@ -468,7 +468,6 @@ protected override void StepCore(TapeStepContext<T> context)
         T epsilon = NumOps.FromDouble(_options.Epsilon);
 
         bool gpuAdam = typeof(T) == typeof(float)
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         foreach (var param in context.Parameters)

src/Optimizers/AdaMaxOptimizer.cs

@@ -479,9 +479,8 @@ protected override void StepCore(TapeStepContext<T> context)
         T alpha = NumOps.Divide(CurrentLearningRate, NumOps.FromDouble(1 - Math.Pow(_options.Beta1, _tapeStep)));
         T epsilon = NumOps.FromDouble(1e-8);
 
-        // GPU-resident step (AIDOTNET_GPU_ADAM=1); gated off, CPU fallback per-param when not GPU-resident.
+        // GPU-resident step; CPU fallback per-param when not GPU-resident.
         bool gpuAdam = typeof(T) == typeof(float)
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         foreach (var param in context.Parameters)

src/Optimizers/AdagradOptimizer.cs

@@ -409,9 +409,8 @@ protected override void StepCore(TapeStepContext<T> context)
 
         T epsilon = NumOps.FromDouble(_options.Epsilon);
 
-        // GPU-resident step (AIDOTNET_GPU_ADAM=1); gated off, CPU fallback per-param when not GPU-resident.
+        // GPU-resident step; CPU fallback per-param when not GPU-resident.
         bool gpuAdam = typeof(T) == typeof(float)
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         foreach (var param in context.Parameters)

src/Optimizers/AdamOptimizer.cs

@@ -930,23 +930,15 @@ protected override void StepCore(TapeStepContext<T> context)
         // That single gate respects AnomalyGuardMode and runs BEFORE
         // _tapeStep++ so a skipped step is a true no-op.
 
-        // GPU-RESIDENT ADAM (Phase 1, env AIDOTNET_GPU_ADAM=1): when on a GPU engine + float, run the Adam
-        // update on the GPU (GpuOptimizer.TryAdamStep -> backend.AdamUpdate) so gradients never download to
-        // host and weights/moments update in place — eliminating the per-step grad-download + weight-reupload
-        // that makes the step host-bound (and uncapturable as a CUDA graph). Moments are allocated GPU-resident
-        // (RentPinnedOnGpu) so TryGetGpuBuffer resolves them. Gated OFF by default; falls back to the CPU SIMD
-        // path per-parameter whenever any tensor isn't GPU-resident. AMSGrad uses the CPU path (no GPU vMax yet).
-        // cudaGraph-safety: the GPU-resident step is only host-read-free when NO
-        // host-side gradient scan runs first. The anomaly guard
-        // (ShouldRunAnomalyGuard -> AnyGradientIsAnomalous) and global-norm gradient
-        // clipping need the global gradient norm read back to the host (one scalar), so disable the GPU
-        // fast path whenever either is active — the step is not graph-capturable with that read in it.
-        bool hostGradientScanActive = ShouldRunAnomalyGuard()
-            || (GradientOptions.EnableGradientClipping
-                && GradientOptions.GradientClippingMethod == AiDotNet.Models.Options.GradientClippingMethod.ByNorm);
+        // GPU-RESIDENT ADAM: on a GPU engine + float, run the Adam update on the GPU (GpuOptimizer.TryAdamStep ->
+        // backend.AdamUpdate) so gradients never download and weights/moments update in place, instead of the
+        // per-step grad download + weight re-upload of the host path. Moments are allocated GPU-resident
+        // (RentPinnedOnGpu) so TryGetGpuBuffer resolves them. Falls back to the CPU SIMD path per parameter
+        // whenever a tensor is not GPU-resident. AMSGrad uses the CPU path (no GPU vMax yet).
+        // The anomaly guard and global-norm clipping above read one scalar back to the host. That keeps this eager
+        // tape step out of CUDA-graph capture, but it is not a reason to update on the host: graph capture is the
+        // compiled plan's path, which has its own fused optimizer, and this step is never captured.
         bool gpuAdam = isFloat && !_options.UseAMSGrad
-            && !hostGradientScanActive
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         // Per-parameter Adam updates are fully INDEPENDENT: each param's m/v/vMax

src/Optimizers/AdamWOptimizer.cs

@@ -738,11 +738,10 @@ protected override void StepCore(TapeStepContext<T> context)
         T biasCorrection2 = NumOps.FromDouble(1 - Math.Pow(_options.Beta2, _tapeStep));
         T weightDecay = NumOps.FromDouble(_options.WeightDecay);
 
-        // GPU-RESIDENT ADAMW (env AIDOTNET_GPU_ADAM=1): see AdamOptimizer.Step for the rationale — run the
+        // GPU-RESIDENT ADAMW: see AdamOptimizer.Step for the rationale — run the
         // update on the GPU (GpuOptimizer.TryAdamWStep -> backend.AdamWUpdate, decoupled weight decay) so grads
-        // never download. Moments are GPU-resident; falls back to the tensor-op path per-param otherwise. Gated off.
+        // never download. Moments are GPU-resident; falls back to the tensor-op path per-param otherwise.
         bool gpuAdam = configuredMask is null && typeof(T) == typeof(float) && !_options.UseAMSGrad
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         int maskOffset = 0;
@@ -806,7 +805,9 @@ protected override void StepCore(TapeStepContext<T> context)
                         (Tensor<float>)(object)m, (Tensor<float>)(object)v,
                         (float)NumOps.ToDouble(CurrentLearningRate), (float)_options.Beta1, (float)_options.Beta2,
                         (float)_options.Epsilon, (float)_options.WeightDecay, _tapeStep))
+                {
                     continue;
+                }
             }
 
             // Host fast path: one fused pass per parameter (m, v and the decayed parameter updated in place),

src/Optimizers/FTRLOptimizer.cs

@@ -310,11 +310,9 @@ protected override void StepCore(TapeStepContext<T> context)
         var lambda1 = NumOps.FromDouble(_options.Lambda1);
         var lambda2 = NumOps.FromDouble(_options.Lambda2);
 
-        // The AIDOTNET_GPU_ADAM env var is the shared opt-in gate for all GPU-resident
-        // optimizer fast paths (named for the first optimizer it enabled); here it
-        // selects the FTRL GPU step, so use a purpose-named local.
+        // On a float GPU engine the FTRL step runs on the device (falls back per parameter when a
+        // tensor is not device-resident).
         bool useGpuOptimizer = typeof(T) == typeof(float)
-            && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1"
             && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine;
 
         foreach (var param in context.Parameters)

src/Optimizers/GradientBasedOptimizerBase.cs

@@ -1983,6 +1983,10 @@ protected virtual IFullModel<T, TInput, TOutput> UpdateSolution(
                 {
                     var updated = updatedParams[pi];
                     if (updated is null) continue;
+                    // Skip a parameter whose DEVICE copy the step just made authoritative (GpuOptimizer): the version
+                    // bump below would mark that buffer stale and the invalidation would drop it with its pending host
+                    // download, losing the update (see NeuralNetworkBase's eager tape step).
+                    if (updated.HasPendingGpuData) continue;
                     // Bump the version so a GPU engine's version gate re-uploads the mutated weights
                     // from the host backing store on the next forward.
                     updated.IncrementVersion();
@@ -3736,6 +3740,20 @@ private IEnumerable<FieldInfo> EnumerateOptimizerFields()
     /// optimizer that opts in keeps that loss-explosion contained at the
     /// canonical PyTorch transformer default <c>MaxGradientNorm = 1.0</c>.
     /// </summary>
+    /// <summary>
+    /// The step's gradients as float tensors when the current engine is a GPU engine and T is float (the device
+    /// clip and finiteness paths); null otherwise.
+    /// </summary>
+    private static List<Tensor<float>>? DeviceTapeGradients(TapeStepContext<T> context)
+    {
+        if (typeof(T) != typeof(float) || AiDotNetEngine.Current is not AiDotNet.Tensors.Engines.DirectGpuTensorEngine)
+            return null;
+        var list = new List<Tensor<float>>(context.Gradients.Count);
+        foreach (var kvp in context.Gradients)
+            if (kvp.Value is { Length: > 0 } grad) list.Add((Tensor<float>)(object)grad);
+        return list;
+    }
+
     protected static void ApplyTapeGlobalNormGradientClipping(
         TapeStepContext<T> context, double maxNorm)
     {
@@ -3744,6 +3762,15 @@ protected static void ApplyTapeGlobalNormGradientClipping(
         // overwrites every gradient element with NaN, distorting the next
         // optimizer step instead of being a no-op.
         if (maxNorm <= 0.0 || double.IsNaN(maxNorm) || double.IsInfinity(maxNorm)) return;
+
+        // GPU engine: clip_grad_norm_ on the device with no host read (PyTorch's foreach path). The host loops below
+        // read every gradient through Data.Span, a download of each device-resident gradient per step. A non-finite
+        // norm leaves the gradients unscaled on the device too.
+        if (DeviceTapeGradients(context) is { } deviceGrads)
+        {
+            using var deviceNorm = AiDotNet.Tensors.Engines.Gpu.GpuOptimizer.TryClipGradNorm(deviceGrads, maxNorm);
+            if (deviceNorm is not null) return;
+        }
         var numOps = MathHelper.GetNumericOperations<T>();
         // Host engine: one parallel SIMD pass per tensor for the norm and for the scale (FusedOptimizer helpers);
         // the per-element loops remain for other layouts.
@@ -3795,6 +3822,22 @@ protected static void ApplyTapeValueGradientClipping(
         if (maxValue <= 0.0 || double.IsNaN(maxValue) || double.IsInfinity(maxValue)) return;
         var numOps = MathHelper.GetNumericOperations<T>();
 
+        // GPU engine: clamp on the device (clip_grad_value_), written back into each gradient tensor. The
+        // host loop below downloads every gradient and leaves the device copy to be re-uploaded.
+        if (AiDotNetEngine.Current.SupportsGpu)
+        {
+            var engine = AiDotNetEngine.Current;
+            T lo = numOps.FromDouble(-maxValue), hi = numOps.FromDouble(maxValue);
+            using var noGrad = new AiDotNet.Tensors.Engines.Autodiff.NoGradScope<T>();
+            foreach (var kvp in context.Gradients)
+            {
+                var grad = kvp.Value;
+                if (grad is null || grad.Length == 0) continue;
+                engine.TensorCopy(engine.TensorClamp(grad, lo, hi), grad);   // in place: callers hold the tensor
+            }
+            return;
+        }
+
         foreach (var kvp in context.Gradients)
         {
             var grad = kvp.Value;
@@ -3825,6 +3868,16 @@ protected static void ApplyTapeValueGradientClipping(
     /// </summary>
     protected static bool HasAnomalousTapeGradients(TapeStepContext<T> context)
     {
+        // GPU engine: one device sum of squares over every gradient and one scalar read (PyTorch's GradScaler also
+        // reads its found-inf flag once per step). A finite sum proves every element finite. A non-finite one is
+        // confirmed by the exact element scan below (a backend that accumulates in float can overflow on finite
+        // gradients); that scan used to run, downloading every device-resident gradient, on every step.
+        if (DeviceTapeGradients(context) is { } deviceGrads)
+        {
+            using var deviceSum = AiDotNet.Tensors.Engines.Gpu.GpuOptimizer.TryGradientSumOfSquares(deviceGrads);
+            if (deviceSum is not null && deviceSum.IsFinite) return false;
+        }
+
         var numOps = MathHelper.GetNumericOperations<T>();
         bool hostFast = !AiDotNetEngine.Current.SupportsGpu;
         foreach (var kvp in context.Gradients)