diff --git a/src/NeuralNetworks/NeuralNetworkBase.cs b/src/NeuralNetworks/NeuralNetworkBase.cs index 5a86002b63..27568290c4 100644 --- a/src/NeuralNetworks/NeuralNetworkBase.cs +++ b/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? 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? g) /// /// Reset at the start of every clipped step, so it always describes the latest one. /// + /// + /// 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 clip_grad_norm_ returns does in PyTorch. + /// /// - 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> grads, double maxNorm, IReadOnlyList> 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>(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)(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(); + 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>(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)(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). diff --git a/src/Optimizers/AMSGradOptimizer.cs b/src/Optimizers/AMSGradOptimizer.cs index b44f46aef1..e033935a00 100644 --- a/src/Optimizers/AMSGradOptimizer.cs +++ b/src/Optimizers/AMSGradOptimizer.cs @@ -394,9 +394,8 @@ protected override void StepCore(TapeStepContext 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) diff --git a/src/Optimizers/AdaDeltaOptimizer.cs b/src/Optimizers/AdaDeltaOptimizer.cs index b39c84eeda..e60f2209e3 100644 --- a/src/Optimizers/AdaDeltaOptimizer.cs +++ b/src/Optimizers/AdaDeltaOptimizer.cs @@ -468,7 +468,6 @@ protected override void StepCore(TapeStepContext 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) diff --git a/src/Optimizers/AdaMaxOptimizer.cs b/src/Optimizers/AdaMaxOptimizer.cs index 7d78b95885..5f977264d4 100644 --- a/src/Optimizers/AdaMaxOptimizer.cs +++ b/src/Optimizers/AdaMaxOptimizer.cs @@ -479,9 +479,8 @@ protected override void StepCore(TapeStepContext 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) diff --git a/src/Optimizers/AdagradOptimizer.cs b/src/Optimizers/AdagradOptimizer.cs index 168ca045e6..dcfe9cae2a 100644 --- a/src/Optimizers/AdagradOptimizer.cs +++ b/src/Optimizers/AdagradOptimizer.cs @@ -409,9 +409,8 @@ protected override void StepCore(TapeStepContext 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) diff --git a/src/Optimizers/AdamOptimizer.cs b/src/Optimizers/AdamOptimizer.cs index 47ab55ebe3..b0ac75873d 100644 --- a/src/Optimizers/AdamOptimizer.cs +++ b/src/Optimizers/AdamOptimizer.cs @@ -930,23 +930,15 @@ protected override void StepCore(TapeStepContext 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 diff --git a/src/Optimizers/AdamWOptimizer.cs b/src/Optimizers/AdamWOptimizer.cs index b0130de6f9..d792c138ed 100644 --- a/src/Optimizers/AdamWOptimizer.cs +++ b/src/Optimizers/AdamWOptimizer.cs @@ -738,11 +738,10 @@ protected override void StepCore(TapeStepContext 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 context) (Tensor)(object)m, (Tensor)(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), diff --git a/src/Optimizers/FTRLOptimizer.cs b/src/Optimizers/FTRLOptimizer.cs index 459ca8f4e5..1aa6da0e9a 100644 --- a/src/Optimizers/FTRLOptimizer.cs +++ b/src/Optimizers/FTRLOptimizer.cs @@ -310,11 +310,9 @@ protected override void StepCore(TapeStepContext 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) diff --git a/src/Optimizers/GradientBasedOptimizerBase.cs b/src/Optimizers/GradientBasedOptimizerBase.cs index 29c00bb022..54b14e1d57 100644 --- a/src/Optimizers/GradientBasedOptimizerBase.cs +++ b/src/Optimizers/GradientBasedOptimizerBase.cs @@ -1983,6 +1983,10 @@ protected virtual IFullModel 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 EnumerateOptimizerFields() /// optimizer that opts in keeps that loss-explosion contained at the /// canonical PyTorch transformer default MaxGradientNorm = 1.0. /// + /// + /// 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. + /// + private static List>? DeviceTapeGradients(TapeStepContext context) + { + if (typeof(T) != typeof(float) || AiDotNetEngine.Current is not AiDotNet.Tensors.Engines.DirectGpuTensorEngine) + return null; + var list = new List>(context.Gradients.Count); + foreach (var kvp in context.Gradients) + if (kvp.Value is { Length: > 0 } grad) list.Add((Tensor)(object)grad); + return list; + } + protected static void ApplyTapeGlobalNormGradientClipping( TapeStepContext 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(); // 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(); + // 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(); + 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( /// protected static bool HasAnomalousTapeGradients(TapeStepContext 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(); bool hostFast = !AiDotNetEngine.Current.SupportsGpu; foreach (var kvp in context.Gradients) diff --git a/src/Optimizers/GradientDescentOptimizer.cs b/src/Optimizers/GradientDescentOptimizer.cs index 32ca3dc82b..501fca85f4 100644 --- a/src/Optimizers/GradientDescentOptimizer.cs +++ b/src/Optimizers/GradientDescentOptimizer.cs @@ -236,7 +236,6 @@ protected override IFullModel UpdateSolution( protected override void StepCore(TapeStepContext context) { 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) diff --git a/src/Optimizers/LionOptimizer.cs b/src/Optimizers/LionOptimizer.cs index 20e3d6706a..a1021484a8 100644 --- a/src/Optimizers/LionOptimizer.cs +++ b/src/Optimizers/LionOptimizer.cs @@ -423,10 +423,9 @@ protected override void StepCore(TapeStepContext context) var oneMinusBeta1 = NumOps.Subtract(NumOps.One, _currentBeta1); var oneMinusBeta2 = NumOps.Subtract(NumOps.One, _currentBeta2); - // GPU-resident step (AIDOTNET_GPU_ADAM=1): see AdamOptimizer.Step. Gated off; falls back to the CPU + // GPU-resident step: see AdamOptimizer.Step; falls back to the CPU // tensor-op path per-parameter 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) diff --git a/src/Optimizers/MiniBatchGradientDescentOptimizer.cs b/src/Optimizers/MiniBatchGradientDescentOptimizer.cs index cd87d256fa..ca5c24864d 100644 --- a/src/Optimizers/MiniBatchGradientDescentOptimizer.cs +++ b/src/Optimizers/MiniBatchGradientDescentOptimizer.cs @@ -389,7 +389,6 @@ public override void UpdateParametersGpu(IGpuBuffer parameters, IGpuBuffer gradi protected override void StepCore(TapeStepContext context) { 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) diff --git a/src/Optimizers/MomentumOptimizer.cs b/src/Optimizers/MomentumOptimizer.cs index 1663dc5a60..d7b0ec1d62 100644 --- a/src/Optimizers/MomentumOptimizer.cs +++ b/src/Optimizers/MomentumOptimizer.cs @@ -400,7 +400,6 @@ protected override void StepCore(TapeStepContext context) PrepareTapeState(context); 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) diff --git a/src/Optimizers/NadamOptimizer.cs b/src/Optimizers/NadamOptimizer.cs index 2352749298..5aa20a4835 100644 --- a/src/Optimizers/NadamOptimizer.cs +++ b/src/Optimizers/NadamOptimizer.cs @@ -469,9 +469,8 @@ protected override void StepCore(TapeStepContext context) T biasCorrectionV = NumOps.FromDouble(1 - Math.Pow(_options.Beta2, _tapeStep)); T nesterovFactor = NumOps.Divide(oneMinusBeta1, biasCorrectionM); - // 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) diff --git a/src/Optimizers/ProximalGradientDescentOptimizer.cs b/src/Optimizers/ProximalGradientDescentOptimizer.cs index c3bb843740..e61d1b0b51 100644 --- a/src/Optimizers/ProximalGradientDescentOptimizer.cs +++ b/src/Optimizers/ProximalGradientDescentOptimizer.cs @@ -825,7 +825,6 @@ protected override void StepCore(TapeStepContext context) // matching the CPU "subtract gradient, then L1 soft-threshold prox" path // exactly. Other regularizers (or non-GPU-resident tensors) fall back to CPU. bool gpuL1 = typeof(T) == typeof(float) - && System.Environment.GetEnvironmentVariable("AIDOTNET_GPU_ADAM") == "1" && AiDotNet.Tensors.Engines.AiDotNetEngine.Current is AiDotNet.Tensors.Engines.DirectGpuTensorEngine && _regularization is AiDotNet.Regularization.L1Regularization; float l1Strength = gpuL1 ? (float)_regularization.GetOptions().Strength : 0f; diff --git a/src/Optimizers/RootMeanSquarePropagationOptimizer.cs b/src/Optimizers/RootMeanSquarePropagationOptimizer.cs index 2c348e5970..7a9aeb8efb 100644 --- a/src/Optimizers/RootMeanSquarePropagationOptimizer.cs +++ b/src/Optimizers/RootMeanSquarePropagationOptimizer.cs @@ -405,7 +405,6 @@ protected override void StepCore(TapeStepContext context) // Existing GPU and sparse kernels implement only uncentered, momentum-free RMSProp. bool gpuRmsProp = SupportsGpuUpdate && 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) diff --git a/src/Optimizers/StochasticGradientDescentOptimizer.cs b/src/Optimizers/StochasticGradientDescentOptimizer.cs index 32eeb67df0..e7e071ef83 100644 --- a/src/Optimizers/StochasticGradientDescentOptimizer.cs +++ b/src/Optimizers/StochasticGradientDescentOptimizer.cs @@ -322,9 +322,8 @@ protected override string GenerateGradientCacheKey(IFullModel protected override void StepCore(TapeStepContext context) { - // GPU-resident step (AIDOTNET_GPU_ADAM=1); gated off, CPU fallback when not GPU-resident. + // GPU-resident step; CPU fallback 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) diff --git a/tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/GpuEagerTapeStepMatchesCpuTests.cs b/tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/GpuEagerTapeStepMatchesCpuTests.cs new file mode 100644 index 0000000000..a207c75175 --- /dev/null +++ b/tests/AiDotNet.Tests/IntegrationTests/NeuralNetworks/GpuEagerTapeStepMatchesCpuTests.cs @@ -0,0 +1,133 @@ +// Copyright (c) CrowdTrainer LLC. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.ActivationFunctions; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.Models.Options; +using AiDotNet.NeuralNetworks; +using AiDotNet.NeuralNetworks.Layers; +using AiDotNet.Optimizers; +using AiDotNet.Tensors.Engines; +using AiDotNet.Tensors.LinearAlgebra; +using Xunit; + +namespace AiDotNet.Tests.IntegrationTests.NeuralNetworks; + +/// +/// The eager tape step (TrainWithTape, the path a model takes when the compiled fused step does not apply) on the GPU +/// engine must follow the CPU engine's trajectory. Its parameters live in one ParameterBuffer, every parameter a view of +/// a flat array. Three bugs made the GPU step train on wrong values: the device weight cache served every view the +/// buffer cached for the whole flat array (the first parameter's data), the post-step invalidation dropped the device +/// optimizer's update, and the device optimizer was off unless an environment variable opted in. (Measured on a +/// 27M-parameter LM: GPU loss 10.91 -> 9.55 over five steps while the CPU engine went 10.91 -> 5.81.) +/// +[Collection("TrainingDiagnosticsSequential")] +public class GpuEagerTapeStepMatchesCpuTests +{ + private const int Inputs = 16, Hidden = 32, Outputs = 4, Batch = 32, Steps = 6; + + /// A network that always takes the eager tape step. + private sealed class EagerNetwork : NeuralNetwork + { + public EagerNetwork(NeuralNetworkArchitecture architecture, IGradientBasedOptimizer, Tensor> optimizer) + : base(architecture, optimizer) { } + + protected override bool SupportsFusedCompiledTraining => false; + } + + private static EagerNetwork BuildNetwork() + { + var layers = new List> + { + new DenseLayer(Hidden, (IActivationFunction)new ReLUActivation()), + // A norm gain read through the device weight path, at a nonzero offset of the flat parameter array. + new LayerNormalizationLayer(Hidden), + new DenseLayer(Outputs, (IActivationFunction?)null), + }; + var architecture = new NeuralNetworkArchitecture( + inputType: InputType.OneDimensional, taskType: NeuralNetworkTaskType.Regression, + inputSize: Inputs, outputSize: Outputs, layers: layers); + var options = new AdamWOptimizerOptions, Tensor> { InitialLearningRate = 1e-2 }; + return new EagerNetwork(architecture, new AdamWOptimizer, Tensor>(null, options)); + } + + private static (Tensor X, Tensor Y) Data() + { + var rng = AiDotNet.Tensors.Helpers.RandomHelper.CreateSeededRandom(20261102); + var x = new Tensor(new[] { Batch, Inputs }); + var y = new Tensor(new[] { Batch, Outputs }); + for (int r = 0; r < Batch; r++) + { + for (int i = 0; i < Inputs; i++) x[r, i] = (float)(rng.NextDouble() * 2 - 1); + for (int o = 0; o < Outputs; o++) y[r, o] = (float)(rng.NextDouble() * 2 - 1); + } + return (x, y); + } + + private static List Trajectory(IEngine engine, float[] initialParameters) + { + var previous = AiDotNetEngine.Current; + AiDotNetEngine.Current = engine; + try + { + var net = BuildNetwork(); + var (x, y) = Data(); + net.SetTrainingMode(true); + net.Predict(x); // materialize the lazy layers + net.SetTrainingMode(true); + net.UpdateParameters(new Vector((float[])initialParameters.Clone())); + var losses = new List(); + for (int s = 0; s < Steps; s++) + { + net.Train(x, y); + losses.Add(Convert.ToDouble(net.GetLastLoss())); + } + return losses; + } + finally + { + AiDotNetEngine.Current = previous; + } + } + + [SkippableFact] + public void Gpu_eager_step_follows_the_cpu_trajectory() + { + DirectGpuTensorEngine? gpu = null; + try { gpu = new DirectGpuTensorEngine(); } catch { /* no backend */ } + try + { + Skip.If(gpu is null || !gpu.SupportsGpu || !gpu.IsGpuAvailable, "No DirectGpuTensorEngine available."); + + // One shared initialization for both engines. + float[] init; + var previous = AiDotNetEngine.Current; + AiDotNetEngine.Current = new CpuEngine(); + try + { + var seedNet = BuildNetwork(); + seedNet.Predict(Data().X); + init = seedNet.GetParameters().ToArray(); + } + finally + { + AiDotNetEngine.Current = previous; + } + + var cpu = Trajectory(new CpuEngine(), init); + var gpuLosses = Trajectory(gpu!, init); + string Show(List l) => string.Join(" -> ", l.Select(v => v.ToString("G5"))); + Assert.True(cpu[^1] < 0.9 * cpu[0], $"the CPU reference did not learn: {Show(cpu)}"); + for (int s = 0; s < Steps; s++) + Assert.True(Math.Abs(gpuLosses[s] - cpu[s]) <= 2e-3 * Math.Max(1.0, Math.Abs(cpu[s])), + $"step {s}: GPU {Show(gpuLosses)} vs CPU {Show(cpu)}"); + } + finally + { + gpu?.Dispose(); + } + } +}