diff --git a/Directory.Packages.props b/Directory.Packages.props index 28a2b69a16..aaab0d33e0 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -235,11 +235,24 @@ accumulating multi-consumer grad buffers) and #764 (GPU-resident-parameter fused-Adam mistrain fix), which together make this PR's TimeSeries GPU-resident training path train faithfully (the resident run now lowers the eager-forward training MSE and improves held-out MSE, so the N-BEATS - correctness gate accepts it). 0.112.0 is PUBLISHED to NuGet and is a superset of 0.111.2. --> - - - - + correctness gate accepts it). 0.112.0 is PUBLISHED to NuGet and is a superset of 0.111.2. + + Bumped 0.112.0 -> 0.113.0: brings AiDotNet.Tensors #763 — this branch's GATING dependency and + the engine half of the non-time-series GPU-residency sweep. #763's key fix (4C): the + compiled/persistent backward now honors createGraph=true (GradientTape.ComputeGradients + previously gated the compiled path on !createGraph), so the WGAN-GP gradient penalty's inner + backward differentiates into the disc weights through the fused compiled plan instead of + silently returning zeros (issue #1844). Against 0.112.0 the fused GPU-resident WGAN-GP path + (WganGpFusedStep / GpuResidentFusedStep) would have silently degraded to plain WGAN. #763 also + publishes Engines.Training.{DpSgdStep, DpSgdFusedStep, MultiSlotFusedStep, WganGpFusedStep, + PersistentInputRegistry}; the src/Training mirror classes stay as the consumer-side primitives + (the fused-primitive centralization in #1847/#1848 wires every consumer through them) and now + run against the fixed engine. 0.113.0 is PUBLISHED to NuGet and is a superset of 0.112.0 (also + carries #765 compiled-ReduceMax axis fill), so all rationale above is retained. --> + + + + diff --git a/src/Audio/Fingerprinting/CLAPModel.cs b/src/Audio/Fingerprinting/CLAPModel.cs index 33eee4a0cd..5efe246363 100644 --- a/src/Audio/Fingerprinting/CLAPModel.cs +++ b/src/Audio/Fingerprinting/CLAPModel.cs @@ -501,6 +501,52 @@ public override void Train(Tensor input, Tensor expected) var optimizer = GetOrCreateBaseOptimizer(); + // GPU-RESIDENT fast path — audio + text encoders + the learned + // temperature scalar. _logTemperature is NOT an ITrainableLayer, so + // it goes through extraTensors (Phase 4A). The fused optimizer tracks + // it via the same moment-buffer path as layer-carried params. + var trainableLayers = new List>(); + foreach (var l in Layers) if (l is ITrainableLayer t) trainableLayers.Add(t); + foreach (var l in TextEncoderLayers) if (l is ITrainableLayer t) trainableLayers.Add(t); + var extras = new List> { _logTemperature }; + if (trainableLayers.Count > 0 || extras.Count > 0) + { + // The contrastive loss depends on BOTH encoders + temperature, so + // the forward closure runs the full symmetric-alignment computation + // and the loss closure just reduces its output. We pass the audio + // input as the primary and the text batch as target so the closure + // signature matches TryStep's contract. + Tensor FwdCLAP(Tensor audioIn) => EncodeAudio(audioIn); // audio embedding — the loss closure re-runs the text side. + Tensor LossCLAP(Tensor audioEmb, Tensor textBatch) + { + // Re-run text encoder on this replay's text batch. + var textEmb = EncodeText(textBatch); + int batchSize = audioEmb.Shape[0]; + int projDim = audioEmb.Shape[audioEmb.Shape.Length - 1]; + var audioEmb2D = audioEmb.Shape.Length == 2 ? audioEmb : Engine.Reshape(audioEmb, new[] { batchSize, projDim }); + var textEmb2D = textEmb.Shape.Length == 2 ? textEmb : Engine.Reshape(textEmb, new[] { batchSize, projDim }); + var textEmbT = Engine.TensorTranspose(textEmb2D); + var sim = Engine.TensorMatMul(audioEmb2D, textEmbT); + var tau = Engine.TensorExp(_logTemperature); + var tauBroadcast = Engine.TensorTile(Engine.Reshape(tau, new[] { 1, 1 }), new[] { batchSize, batchSize }); + var logitsA2T = Engine.TensorMultiply(sim, tauBroadcast); + var logitsT2A = Engine.TensorTranspose(logitsA2T); + var halfA2T = SymmetricRowCrossEntropy(logitsA2T, batchSize); + var halfT2A = SymmetricRowCrossEntropy(logitsT2A, batchSize); + return Engine.TensorAdd(halfA2T, halfT2A); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableLayers, input, expected, + forward: FwdCLAP, computeLoss: LossCLAP, + optimizer: optimizer, + out T fusedLoss, + extraTensors: extras)) + { + LastLoss = fusedLoss; + return; + } + } + using var tape = new GradientTape(); // Forward both encoders inside the same tape so gradients flow // through every parameter. EncodeAudio / EncodeText already diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index f134bffe9e..25fafce2b1 100644 --- a/src/Diffusion/DiffusionModelBase.cs +++ b/src/Diffusion/DiffusionModelBase.cs @@ -126,6 +126,29 @@ protected virtual IEnumerable EnumerateDisposableComponents() /// public virtual ModelOptions GetOptions() => _options; + /// + /// Opt-in flag: does this subclass's / + /// path use ONLY traceable engine ops (no + /// host-side .Data.Span loops, no per-step class-field mutation)? + /// When true, attempts a + /// fused-plan step + /// with (noisySample, noise) as persistent slots — the compiled plan replays + /// the forward per training call so PredictNoise must be free of + /// per-call side effects that would freeze at trace time. + /// + /// Base default is false so existing subclasses (which may still use + /// .Data.Span internally) run through the eager tape as before. Override + /// after auditing your forward path. See ooples/AiDotNet#1846. + /// + /// + protected virtual bool SupportsFusedDenoising => false; + + private static void RestoreShadow(Tensor param, Vector shadow) + { + var span = param.Data.Span; + for (int k = 0; k < span.Length && k < shadow.Length; k++) span[k] = shadow[k]; + } + /// /// The optional neural network architecture blueprint for custom layer configuration. /// @@ -802,6 +825,41 @@ protected virtual Tensor Generate(int[] shape, int numInferenceSteps, int? se /// public abstract Tensor PredictNoise(Tensor noisySample, int timestep); + /// + /// Batched per-element noise prediction (industry-standard DDPM training pattern). + /// Default implementation slices the batch and calls scalar + /// per element — subclasses should override with a fused batched forward to keep + /// training on-device. is shape [B, ...]; + /// is a [B] int vector. + /// + public virtual Tensor PredictNoiseBatched(Tensor noisyBatch, int[] timesteps) + { + int batchSize = noisyBatch.Shape[0]; + if (timesteps.Length != batchSize) + throw new ArgumentException( + $"timesteps length {timesteps.Length} does not match batch size {batchSize}.", + nameof(timesteps)); + + int perElement = noisyBatch.Length / batchSize; + var elemShape = new int[noisyBatch.Rank - 1]; + for (int i = 1; i < noisyBatch.Rank; i++) elemShape[i - 1] = noisyBatch.Shape[i]; + var result = new Tensor(noisyBatch._shape); + var nbSpan = noisyBatch.AsSpan(); + var resSpan = result.AsWritableSpan(); + for (int b = 0; b < batchSize; b++) + { + var elem = new Tensor(elemShape); + var elemSpan = elem.AsWritableSpan(); + for (int j = 0; j < perElement; j++) + elemSpan[j] = nbSpan[b * perElement + j]; + var pred = PredictNoise(elem, timesteps[b]); + var predSpan = pred.AsSpan(); + for (int j = 0; j < perElement; j++) + resSpan[b * perElement + j] = predSpan[j]; + } + return result; + } + /// /// Runs one denoising-step noise prediction, optionally inside a GPU deferred execution graph /// (AiDotNet.Tensors #642) when is @@ -1034,19 +1092,143 @@ public virtual void Train(Tensor input, Tensor expectedOutput) // tensors than GetParameters knows about, which is now the norm after // migrating layers like FlashAttentionLayer from Matrix to Tensor. - // Sample a random timestep and build the noisy training sample. - var timestep = RandomGenerator.Next(_scheduler.Config.TrainTimesteps); - var inputVector = input.ToVector(); - var noiseVector = SampleNoise(inputVector.Length, RandomGenerator); - var noisySample = _scheduler.AddNoise(inputVector, noiseVector, timestep); - var noisySampleTensor = new Tensor(input._shape, noisySample); + // Industry-standard batched-per-element timesteps (Ho et al. 2020, HuggingFace + // diffusers reference): sample a distinct timestep per batch element instead of + // one timestep for the whole batch. Decorrelates the noise-schedule signal + // across the batch, which is the canonical DDPM training pattern. + // + // Rank-1 input (unbatched, historical AiDotNet contract) still gets a single + // timestep; rank ≥ 2 gets per-element timesteps. + bool isBatched = input.Rank >= 2; + int batchSize = isBatched ? input.Shape[0] : 1; + var timesteps = new int[batchSize]; + for (int b = 0; b < batchSize; b++) + timesteps[b] = RandomGenerator.Next(_scheduler.Config.TrainTimesteps); + // Legacy scalar view — retained for downstream code that reads the "current" + // timestep (e.g. QAT hook telemetry). For batched inputs this reports element 0's + // timestep, matching the historical single-timestep-per-Train contract. + var timestep = timesteps[0]; + + Tensor noisySampleTensor; + Vector noiseVector; + if (isBatched) + { + var noiseBatch = new Tensor(input._shape); + var noiseSpan = noiseBatch.AsWritableSpan(); + for (int i = 0; i < noiseSpan.Length; i++) + noiseSpan[i] = NumOps.FromDouble(RandomGenerator.NextGaussian()); + // AddNoiseBatched lives on NoiseSchedulerBase (not INoiseScheduler — that + // interface can't carry a default implementation on net471). Fall back to + // per-element scalar AddNoise if a caller passed a scheduler that doesn't + // derive from NoiseSchedulerBase (shouldn't happen for framework schedulers). + if (_scheduler is Schedulers.NoiseSchedulerBase baseScheduler) + { + noisySampleTensor = baseScheduler.AddNoiseBatched(input, noiseBatch, timesteps); + } + else + { + noisySampleTensor = new Tensor(input._shape); + var cleanSpan = input.AsSpan(); + var noisedSpan = noisySampleTensor.AsWritableSpan(); + int perElement = input.Length / batchSize; + for (int b = 0; b < batchSize; b++) + { + var cleanVec = new Vector(perElement); + var noiseVec = new Vector(perElement); + for (int j = 0; j < perElement; j++) + { + cleanVec[j] = cleanSpan[b * perElement + j]; + noiseVec[j] = noiseSpan[b * perElement + j]; + } + var noised = _scheduler.AddNoise(cleanVec, noiseVec, timesteps[b]); + for (int j = 0; j < perElement; j++) + noisedSpan[b * perElement + j] = noised[j]; + } + } + noiseVector = noiseBatch.ToVector(); + } + else + { + var inputVector = input.ToVector(); + noiseVector = SampleNoise(inputVector.Length, RandomGenerator); + var noisySample = _scheduler.AddNoise(inputVector, noiseVector, timestep); + noisySampleTensor = new Tensor(input._shape, noisySample); + } + + // Preferred fused path: MultiSlotFusedStep with (noisySample, noise, + // timestepsPerElement) as persistent slots. Only engages when the + // concrete subclass has certified its PredictNoise/PredictNoiseBatched + // path as fully traceable by opting in via SupportsFusedDenoising — + // eager PredictNoise implementations that use host-side .Data.Span + // loops would freeze at trace time. See ooples/AiDotNet#1846. + // Only try the fused path when an optimizer has already been resolved + // (avoids double-construction with a slightly-different config). + var trainingOptimizerForFused = _trainingOptimizer; + if (SupportsFusedDenoising + && trainingOptimizerForFused is not null + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + trainingOptimizerForFused, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var trainableForFused = CollectTrainableParameters(); + var noiseSlotT = new Tensor(noisySampleTensor._shape, noiseVector); + var slots = new Tensor[] + { + noisySampleTensor, + noiseSlotT, + }; + using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep(); + var timestepsSnapshot = timesteps; + var timestepSnapshotSingle = timestep; + var isBatchedSnapshot = isBatched; + Tensor ForwardFromSlots(IReadOnlyList> s) + { + return isBatchedSnapshot + ? PredictNoiseBatched(s[0], timestepsSnapshot) + : PredictNoise(s[0], timestepSnapshotSingle); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + var diff = Engine.TensorSubtract(pred, s[1]); + var sq = Engine.TensorMultiply(diff, diff); + var axes = Enumerable.Range(0, sq.Shape.Length).ToArray(); + return Engine.ReduceMean(sq, axes, keepDims: false); + } + if (trainableForFused.Length > 0 + && multiSlotStep.TryStep( + parameters: trainableForFused, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + if (qatShadows is not null && qatParams is not null) + { + for (int i = 0; i < qatParams.Length; i++) + RestoreShadow(qatParams[i], qatShadows[i]); + } + return; + } + } using var tape = new GradientTape(); // Forward pass — triggers lazy layer initialization, then we walk for // trainable parameters. Collection must happen AFTER the forward pass so - // newly-initialized layers are visible to the walker. - var predicted = PredictNoise(noisySampleTensor, timestep); + // newly-initialized layers are visible to the walker. Batched inputs route + // through the per-element PredictNoiseBatched (Ho et al. 2020 canonical pattern); + // rank-1 unbatched inputs go through the scalar PredictNoise for backward compat. + var predicted = isBatched + ? PredictNoiseBatched(noisySampleTensor, timesteps) + : PredictNoise(noisySampleTensor, timestep); var paramTensors = CollectTrainableParameters(); if (paramTensors.Length == 0) { @@ -1150,7 +1332,10 @@ public virtual void Train(Tensor input, Tensor expectedOutput) // clip; nothing is materialized or copied here. The forward/loss closures are only consulted by // optimizers that re-evaluate the objective (e.g. line search); Adam ignores them. T lossValue = loss.Length > 0 ? loss[0] : NumOps.Zero; - Tensor RecomputeForward(Tensor inp, Tensor _) => PredictNoise(inp, timestep); + Tensor RecomputeForward(Tensor inp, Tensor _) => + isBatched + ? PredictNoiseBatched(inp, timesteps) + : PredictNoise(inp, timestep); Tensor RecomputeLoss(Tensor inp, Tensor target) { using var noGrad = new NoGradScope(); diff --git a/src/Diffusion/NoisePredictors/NoisePredictorBase.cs b/src/Diffusion/NoisePredictors/NoisePredictorBase.cs index bde5f8c48a..90fea1db62 100644 --- a/src/Diffusion/NoisePredictors/NoisePredictorBase.cs +++ b/src/Diffusion/NoisePredictors/NoisePredictorBase.cs @@ -1144,6 +1144,78 @@ protected static SelfAttentionLayer LazySelfAttention( /// public abstract Tensor PredictNoise(Tensor noisySample, int timestep, Tensor? conditioning = null); + /// + /// Batched noise prediction (industry-standard: HuggingFace diffusers / DDPM + /// reference sample one timestep per batch element and condition per-element). + /// Default implementation calls the scalar per element + /// and stacks results — subclasses should override for a fused batched forward + /// to keep the training loop on-device. + /// + /// Noisy samples, shape [B, ...]. + /// Per-batch-element timesteps, shape [B]. + /// Optional conditioning (e.g. class embedding, text). + /// Predicted noise, shape matches . + public virtual Tensor PredictNoiseBatched(Tensor noisyBatch, int[] timesteps, Tensor? conditioning = null) + { + int batchSize = noisyBatch.Shape[0]; + if (timesteps.Length != batchSize) + throw new System.ArgumentException( + $"timesteps length {timesteps.Length} does not match batch size {batchSize}.", + nameof(timesteps)); + + // Slice each element out of the batch, call the scalar predictor, stack results. + // Slow default; subclasses that expose a fused batched forward should override + // this to keep the whole thing on-device. + int perElement = noisyBatch.Length / batchSize; + var elemShape = new int[noisyBatch.Rank - 1]; + for (int i = 1; i < noisyBatch.Rank; i++) elemShape[i - 1] = noisyBatch.Shape[i]; + var resultShape = new int[noisyBatch.Rank]; + System.Array.Copy(noisyBatch._shape, resultShape, noisyBatch.Rank); + var result = new Tensor(resultShape); + var nbSpan = noisyBatch.AsSpan(); + var resSpan = result.AsWritableSpan(); + // Detect batch-aligned conditioning: if its leading dim matches batchSize + // (typical for classifier-free-guidance where each sample carries its own + // text/class embedding), slice per element. If the leading dim doesn't + // match batchSize, treat conditioning as shared (broadcast the same + // tensor to every scalar call). Bind conditioning to a local so the + // nullable analysis sees the null-guarded value across the slicing block. + var condLocal = conditioning; + int condPerElement = 0; + int[]? condElemShape = null; + if (condLocal is not null && condLocal.Rank > 0 && condLocal.Shape[0] == batchSize) + { + condPerElement = condLocal.Length / batchSize; + condElemShape = new int[condLocal.Rank - 1]; + for (int i = 1; i < condLocal.Rank; i++) condElemShape[i - 1] = condLocal.Shape[i]; + } + + for (int b = 0; b < batchSize; b++) + { + var elem = new Tensor(elemShape); + var elemSpan = elem.AsWritableSpan(); + for (int j = 0; j < perElement; j++) + elemSpan[j] = nbSpan[b * perElement + j]; + + Tensor? elementConditioning = condLocal; + if (condLocal is not null && condElemShape is not null) + { + var condElem = new Tensor(condElemShape); + var condElemSpan = condElem.AsWritableSpan(); + var condFullSpan = condLocal.AsSpan(); + for (int j = 0; j < condPerElement; j++) + condElemSpan[j] = condFullSpan[b * condPerElement + j]; + elementConditioning = condElem; + } + + var pred = PredictNoise(elem, timesteps[b], elementConditioning); + var predSpan = pred.AsSpan(); + for (int j = 0; j < perElement; j++) + resSpan[b * perElement + j] = predSpan[j]; + } + return result; + } + /// /// Async overload of . /// Routes the forward through diff --git a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs index c542b7da47..27bb6719f5 100644 --- a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs +++ b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs @@ -279,6 +279,70 @@ public virtual T GetAlphaCumulativeProduct(int timestep) return AlphasCumulativeProduct[timestep]; } + /// + /// Batched forward-diffusion (industry-standard batching, HuggingFace diffusers / + /// DDPM reference): applies a distinct timestep to each batch element instead of + /// a single timestep to the whole batch. Default implementation delegates to the + /// scalar per element for backward compatibility; + /// concrete schedulers should override with a fused batched implementation for + /// on-device efficiency. + /// + /// Clean samples, shape [B, ...]. + /// Per-element noise, shape [B, ...]. + /// Per-batch-element timesteps, shape [B]. + /// Noised batch, shape matches . + public virtual Tensor AddNoiseBatched(Tensor cleanBatch, Tensor noiseBatch, int[] timesteps) + { + if (cleanBatch == null) throw new ArgumentNullException(nameof(cleanBatch)); + if (noiseBatch == null) throw new ArgumentNullException(nameof(noiseBatch)); + if (timesteps == null) throw new ArgumentNullException(nameof(timesteps)); + + if (cleanBatch.Rank == 0 || cleanBatch.Shape[0] <= 0) + throw new ArgumentException( + "cleanBatch must have a non-empty batch dimension.", + nameof(cleanBatch)); + int batchSize = cleanBatch.Shape[0]; + + // Validate full shape parity, not only the batch dim. Without this a + // noiseBatch of shape [B, smaller...] would pass the leading-dim check + // and then index beyond its span in the per-element copy loop below. + if (noiseBatch.Rank != cleanBatch.Rank) + throw new ArgumentException( + $"noiseBatch rank {noiseBatch.Rank} does not match cleanBatch rank {cleanBatch.Rank}.", + nameof(noiseBatch)); + for (int dim = 0; dim < cleanBatch.Rank; dim++) + { + if (noiseBatch.Shape[dim] != cleanBatch.Shape[dim]) + throw new ArgumentException( + $"noiseBatch shape[{dim}] = {noiseBatch.Shape[dim]} does not match cleanBatch shape[{dim}] = {cleanBatch.Shape[dim]}.", + nameof(noiseBatch)); + } + if (timesteps.Length != batchSize) + throw new ArgumentException( + $"timesteps length {timesteps.Length} does not match batch size {batchSize}.", + nameof(timesteps)); + + int perElement = cleanBatch.Length / batchSize; + var result = new Tensor(cleanBatch._shape); + var cleanSpan = cleanBatch.AsSpan(); + var noiseSpan = noiseBatch.AsSpan(); + var resultSpan = result.AsWritableSpan(); + for (int b = 0; b < batchSize; b++) + { + var cleanVec = new Vector(perElement); + var noiseVec = new Vector(perElement); + for (int j = 0; j < perElement; j++) + { + cleanVec[j] = cleanSpan[b * perElement + j]; + noiseVec[j] = noiseSpan[b * perElement + j]; + } + var noised = AddNoise(cleanVec, noiseVec, timesteps[b]); + for (int j = 0; j < perElement; j++) + resultSpan[b * perElement + j] = noised[j]; + } + return result; + } + /// public virtual Vector AddNoise(Vector originalSample, Vector noise, int timestep) { diff --git a/src/Finance/Forecasting/Foundation/CSDI.cs b/src/Finance/Forecasting/Foundation/CSDI.cs index e49938a7ad..dfd3f08e64 100644 --- a/src/Finance/Forecasting/Foundation/CSDI.cs +++ b/src/Finance/Forecasting/Foundation/CSDI.cs @@ -287,6 +287,51 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + // Preferred fused path: MultiSlotFusedStep with the sampled (noise, + // xt-scale, sinT) tuple passed as persistent slots. Refreshes per step + // by host-sampling a fresh (t, ε) pair and copying values into the + // slot tensors — the compiled forward reads the CURRENT slot data on + // every replay. See ooples/AiDotNet#1846. + if (trainableParams.Length > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var slots = BuildCsdiSlots(input, target); + if (slots is not null) + { + using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) => DenoiserForwardFromSlots(s); + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + // s[1] = ε_true (noise slot). Loss = MSE(ε_pred, ε_true) via + // the model's LossFunctionBase so custom losses are respected. + return loss.ComputeTapeLoss(pred, s[1]); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + } + } + + // Eager fallback: samples (t, ε) inline and runs the denoising pair on + // the tape. Same objective; used when the fused path can't engage + // (non-fuse-able optimizer, non-GPU host, etc.). using var tape = new GradientTape(); var (epsilonPred, epsilonTarget) = ComputeDenoisingPairTape(input, target); @@ -342,6 +387,112 @@ public override void Train(Tensor input, Tensor target) } } + /// + /// Assembles the persistent-slot data for the MultiSlotFusedStep wire-up: + /// samples (t, ε) host-side then packs everything the compiled forward + /// needs as tensor slots. Slot layout: + /// + /// [0] target — the clean values to denoise (Ho 2020 x_0). + /// [1] ε_true — freshly-sampled Gaussian noise, same shape as target. + /// [2] sqrtAlphaBar_scalar — cumulative α scaling at the sampled t, shape [1]. + /// [3] sqrtOneMinusAlphaBar_scalar — noise scaling at t, shape [1]. + /// [4] sinT_scalar — sin(2π·t/T) timestep encoding, shape [1]. + /// [5] conditioned — RevIN-normalized observed input, shape [1, seqLen]. + /// + /// Returns null when the target is degenerate (zero length) so the caller + /// falls back to eager training. + /// + private IReadOnlyList>? BuildCsdiSlots(Tensor input, Tensor target) + { + int targetLen = target.Length; + if (targetLen <= 0) return null; + + var rand = RandomHelper.CreateSecureRandom(); + int t = rand.Next(_numDiffusionSteps); + var noiseData = new T[targetLen]; + for (int i = 0; i < targetLen; i++) noiseData[i] = SampleStandardNormal(rand); + var epsilonTrue = new Tensor(target._shape, new Vector(noiseData)); + + T sqrtAlphaBar = NumOps.Sqrt(_alphasCumprod[t]); + T sqrtOneMinus = NumOps.Sqrt(NumOps.Subtract(NumOps.One, _alphasCumprod[t])); + T sinT = NumOps.FromDouble(Math.Sin(2.0 * Math.PI * t / Math.Max(1, _numDiffusionSteps - 1))); + + var sqrtAlphaBarT = new Tensor(new[] { 1 }); + sqrtAlphaBarT[0] = sqrtAlphaBar; + var sqrtOneMinusT = new Tensor(new[] { 1 }); + sqrtOneMinusT[0] = sqrtOneMinus; + var sinTT = new Tensor(new[] { 1 }); + sinTT[0] = sinT; + + var conditioned = ApplyInstanceNormalization(input); + if (conditioned.Rank == 1) + conditioned = Engine.Reshape(conditioned, new[] { 1, conditioned.Length }); + return new[] { target, epsilonTrue, sqrtAlphaBarT, sqrtOneMinusT, sinTT, conditioned }; + } + + /// + /// Fused-path denoiser forward: reads the slot tuple from + /// , builds x_t via traceable engine + /// arithmetic (no .Data.Span pack — the previous inline + /// implementation froze the trace batch's x_t into the compiled plan), + /// concatenates [x_t | conditioning_slice | sin(t)] via + /// , and runs the projection + + /// residual stack + output projection. Returns predicted noise aligned + /// with the noise slot's shape. + /// + private Tensor DenoiserForwardFromSlots(IReadOnlyList> slots) + { + var target = slots[0]; + var epsilonTrue = slots[1]; + var sqrtAlphaBarT = slots[2]; + var sqrtOneMinusT = slots[3]; + var sinTT = slots[4]; + var conditioned = slots[5]; + + int targetLen = target.Length; + // x_t = sqrt(α̅_t) * target + sqrt(1-α̅_t) * ε — all traceable. + var scaledTarget = Engine.TensorMultiplyScalar(target, sqrtAlphaBarT[0]); + var scaledNoise = Engine.TensorMultiplyScalar(epsilonTrue, sqrtOneMinusT[0]); + var xt = Engine.TensorAdd(scaledTarget, scaledNoise); + var xt1d = xt.Rank == 1 ? xt : Engine.Reshape(xt, new[] { targetLen }); + + // Conditioning slice: first Min(condLen, hiddenDimension) elements, + // flattened to 1-D so we can concat with xt1d + sinT. + var condFlat = conditioned.Rank == 1 + ? conditioned + : Engine.Reshape(conditioned, new[] { conditioned.Length }); + int condLen = Math.Min(condFlat.Length, _hiddenDimension); + var condSlice = condLen == condFlat.Length + ? condFlat + : Engine.TensorSlice(condFlat, new[] { 0 }, new[] { condLen }); + + // Pack [xt | conditioning | sin(t)] via TensorConcatenate — replaces + // the .Data.Span[i] fill that froze at trace time. + var packed1D = Engine.TensorConcatenate(new[] { xt1d, condSlice, sinTT }, axis: 0); + var denoisingInput = Engine.Reshape(packed1D, new[] { 1, targetLen + condLen + 1 }); + + var eps = (Tensor)denoisingInput; + if (_inputProjection is not null) + eps = _inputProjection.Forward(eps); + foreach (var layer in _residualLayers) + eps = layer.Forward(eps); + if (_outputProjection is not null) + eps = _outputProjection.Forward(eps); + + if (eps.Rank == 2 && eps.Shape[1] > epsilonTrue.Length) + eps = Engine.TensorNarrow(eps, dim: 1, start: 0, length: epsilonTrue.Length); + if (eps.Length != epsilonTrue.Length) + { + throw new InvalidOperationException( + $"CSDI denoising pair: predicted-noise length ({eps.Length}, shape=[" + + $"{string.Join(",", eps._shape)}]) does not match true-noise length (" + + $"{epsilonTrue.Length}, shape=[{string.Join(",", epsilonTrue._shape)}])."); + } + if (!eps._shape.AsEnumerable().SequenceEqual(epsilonTrue._shape)) + eps = Engine.Reshape(eps, epsilonTrue._shape); + return eps; + } + /// /// Builds the (predicted-noise, true-noise) pair for one DDPM /// training step. Samples a timestep and a fresh noise tensor, @@ -525,23 +676,28 @@ public override Dictionary Evaluate(Tensor predictions, Tensor return new Dictionary { ["MSE"] = mse, ["MAE"] = mae, ["RMSE"] = NumOps.Sqrt(mse) }; } + /// + /// + /// Traceable RevIN forward — same pattern as + /// . Uses + /// / + /// / / broadcast subtract+divide so the + /// per-batch stats re-execute on every replay under a compiled fused plan. + /// public override Tensor ApplyInstanceNormalization(Tensor input) { int batchSize = input.Rank > 1 ? input.Shape[0] : 1; int seqLen = input.Rank > 1 ? input.Shape[1] : input.Length; - var result = new Tensor(input._shape); - for (int b = 0; b < batchSize; b++) - { - T mean = NumOps.Zero; - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) mean = NumOps.Add(mean, input[idx]); } - mean = NumOps.Divide(mean, NumOps.FromDouble(seqLen)); - T variance = NumOps.Zero; - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) { var diff = NumOps.Subtract(input[idx], mean); variance = NumOps.Add(variance, NumOps.Multiply(diff, diff)); } } - variance = NumOps.Divide(variance, NumOps.FromDouble(seqLen)); - T std = NumOps.Sqrt(NumOps.Add(variance, NumOps.FromDouble(1e-5))); - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length && idx < result.Length) result.Data.Span[idx] = NumOps.Divide(NumOps.Subtract(input[idx], mean), std); } - } - return result; + if (seqLen <= 0) return input; + + bool reshaped = input.Rank != 2; + var flat = reshaped ? Engine.Reshape(input, new[] { batchSize, seqLen }) : input; + var mean = Engine.ReduceMean(flat, new[] { 1 }, keepDims: true); + var variance = Engine.ReduceVariance(flat, new[] { 1 }, keepDims: true); + var std = Engine.TensorSqrt(Engine.TensorAddScalar(variance, NumOps.FromDouble(1e-5))); + var centered = Engine.TensorBroadcastSubtract(flat, mean); + var normalized = Engine.TensorBroadcastDivide(centered, std); + return reshaped ? Engine.Reshape(normalized, input._shape) : normalized; } public override Dictionary GetFinancialMetrics() diff --git a/src/Finance/Forecasting/Foundation/TFC.cs b/src/Finance/Forecasting/Foundation/TFC.cs index a4a6fc66d1..fe3c6c6e78 100644 --- a/src/Finance/Forecasting/Foundation/TFC.cs +++ b/src/Finance/Forecasting/Foundation/TFC.cs @@ -95,8 +95,15 @@ public class TFC : TimeSeriesFoundationModelBase // TFC normalizes each input series before the time/frequency encoders and // restores the level on the output so distinct input scales produce distinct // forecasts. - private Vector _revinMean = new Vector(0); - private Vector _revinStd = new Vector(0); + /// Per-instance mean captured by , + /// consumed by . Tensor-shaped [B, 1] so it broadcasts + /// against the forecast. NULL when no forward has run yet. + private Tensor? _revinMeanTensor; + + /// Per-instance standard deviation captured by , + /// consumed by . Tensor-shaped [B, 1]. NULL when no forward + /// has run yet. + private Tensor? _revinStdTensor; #endregion @@ -274,6 +281,55 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + // GPU-RESIDENT fast path — compiled fused SGD on the combined supervised + + // contrastive objective. Safe now that ApplyInstanceNormalization and + // ComputeFrequencyRepresentation both use traceable Engine ops (ReduceMean / + // ReduceVariance / TensorSqrt / broadcast for RevIN, Engine.RFFT + magnitude + // + concat/flip mirror for the DFT) — both re-execute on every replay from + // the current-step persistent slot data instead of freezing at trace time. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + // Closure-captured contrastive loss: ComputeContrastiveLossTape runs + // INSIDE the forward closure so it consumes the CURRENT-step persistent + // input (`inp`), not the outer `input` which would freeze at compile + // time. Fwd/Loss ordering is guaranteed by the fused-step contract. + Tensor? capturedContrastive = null; + Tensor ForwardCombined(Tensor inp) + { + capturedContrastive = ComputeContrastiveLossTape(inp); + return ForwardForTraining(inp); + } + Tensor ComputeLossCombined(Tensor pred, Tensor tgt) + { + var alignedT = tgt; + if (pred.Rank > tgt.Rank && pred.Shape[0] == 1 && pred.Length == tgt.Length) + pred = Engine.Reshape(pred, tgt._shape); + else if (tgt.Rank > pred.Rank && tgt.Shape[0] == 1 && tgt.Length == pred.Length) + alignedT = Engine.Reshape(tgt, pred._shape); + var supervised = loss.ComputeTapeLoss(pred, alignedT); + var contrastive = capturedContrastive + ?? throw new InvalidOperationException( + "TFC fused step: contrastive loss was not captured by ForwardCombined. " + + "This indicates the fused-step framework called the loss closure before " + + "the forward closure, which violates its documented Fwd-then-Loss ordering."); + if (!supervised._shape.SequenceEqual(contrastive._shape) + && supervised.Length == contrastive.Length) + contrastive = Engine.Reshape(contrastive, supervised._shape); + return Engine.TensorAdd(supervised, contrastive); + } + if (AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + trainableLayers, input, target, + forward: ForwardCombined, computeLoss: ComputeLossCombined, + optimizerType: AiDotNet.Tensors.Engines.Compilation.OptimizerType.SGD, + learningRate: 0.001f, beta1: 0.9f, beta2: 0.999f, epsilon: 1e-8f, weightDecay: 0f, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + } + // Custom tape step: TFC's loss is supervised forecast + weighted // contrastive alignment between the time-domain and frequency- // domain encoder outputs. Both terms must be recorded under the @@ -622,71 +678,89 @@ public override Dictionary Evaluate(Tensor predictions, Tensor } /// + /// + /// Traceable RevIN forward (Kim et al. 2022). Uses , + /// , , and + /// so every op records on the tape and + /// re-executes under the compiled fused plan. The per-instance mean/std are captured + /// in tensor fields (not scalars) so + /// stays on-tape too — an inference call refreshes the tensors and a compiled-plan + /// replay recomputes both on-device from the current slot data. + /// public override Tensor ApplyInstanceNormalization(Tensor input) { - // RevIN forward (Kim et al. 2022). Stats are taken over every non-batch - // element of each row (a rank-1 input is a single instance, not one per - // element), and stored so DenormalizeForecast can restore the scale. + var (normalized, mean, std) = NormalizeWithStats(input); + _revinMeanTensor = mean; + _revinStdTensor = std; + return normalized; + } + + /// + /// Stateless RevIN forward. Returns (normalized, mean, std) with mean/std as tensors + /// shaped [B, 1] so downstream ops can broadcast them back. All ops go through + /// — no .Data.Span host loops — so the whole computation + /// records on the autodiff tape and re-executes on every replay under a compiled plan. + /// + private (Tensor Normalized, Tensor Mean, Tensor Std) NormalizeWithStats(Tensor input) + { int batchSize = input.Shape.Length > 1 ? input.Shape[0] : 1; int instanceSize = batchSize > 0 ? input.Length / batchSize : input.Length; if (instanceSize <= 0) - return input; - - var result = new Tensor(input._shape); - _revinMean = new Vector(batchSize); - _revinStd = new Vector(batchSize); - - for (int b = 0; b < batchSize; b++) { - int start = b * instanceSize; - - T mean = NumOps.Zero; - for (int t = 0; t < instanceSize; t++) - mean = NumOps.Add(mean, input[start + t]); - mean = NumOps.Divide(mean, NumOps.FromDouble(instanceSize)); + // Degenerate input — return input unchanged with identity mean/std. + var meanIdentity = new Tensor(new[] { 1, 1 }); + var stdIdentity = new Tensor(new[] { 1, 1 }); + Engine.TensorFill(meanIdentity, NumOps.Zero); + Engine.TensorFill(stdIdentity, NumOps.One); + return (input, meanIdentity, stdIdentity); + } - T variance = NumOps.Zero; - for (int t = 0; t < instanceSize; t++) - { - var diff = NumOps.Subtract(input[start + t], mean); - variance = NumOps.Add(variance, NumOps.Multiply(diff, diff)); - } - variance = NumOps.Divide(variance, NumOps.FromDouble(instanceSize)); - T std = NumOps.Sqrt(NumOps.Add(variance, NumOps.FromDouble(1e-5))); + bool reshaped = input.Rank != 2; + var flat = reshaped ? Engine.Reshape(input, new[] { batchSize, instanceSize }) : input; - _revinMean[b] = mean; - _revinStd[b] = std; + // mean over the instance axis, keepDims so shape stays [B, 1] for broadcast. + var mean = Engine.ReduceMean(flat, new[] { 1 }, keepDims: true); + var variance = Engine.ReduceVariance(flat, new[] { 1 }, keepDims: true); + var std = Engine.TensorSqrt(Engine.TensorAddScalar(variance, NumOps.FromDouble(1e-5))); - for (int t = 0; t < instanceSize; t++) - result.Data.Span[start + t] = NumOps.Divide(NumOps.Subtract(input[start + t], mean), std); - } + // (x - mean) / std via BroadcastSubtract + BroadcastDivide. + var centered = Engine.TensorBroadcastSubtract(flat, mean); + var normalized = Engine.TensorBroadcastDivide(centered, std); - return result; + if (reshaped) + normalized = Engine.Reshape(normalized, input._shape); + return (normalized, mean, std); } /// /// RevIN reverse step (Kim et al. 2022): restores each instance's mean/std to the - /// forecast so it is expressed on the input's original scale. The multiply/add go - /// through the Engine so the forecast stays on the autodiff tape. + /// forecast so it is expressed on the input's original scale. All ops go through + /// so the forecast stays on the autodiff tape AND the compiled-plan + /// replay uses the CURRENT input's stats (via / + /// , both refreshed by + /// or earlier in the same forward). /// private Tensor DenormalizeForecast(Tensor forecast) { + return DenormalizeForecastWithStats(forecast, _revinMeanTensor, _revinStdTensor); + } + + /// + /// Stateless RevIN inverse. Takes explicit mean/std tensors so the compiled-plan + /// path can thread the CURRENT-step stats through without touching class fields. + /// + private Tensor DenormalizeForecastWithStats(Tensor forecast, Tensor? mean, Tensor? std) + { + if (mean is null || std is null) return forecast; + int batch = forecast.Shape.Length > 1 ? forecast.Shape[0] : 1; - if (_revinMean.Length != batch || forecast.Length % batch != 0) + if (mean.Length != batch || std.Length != batch || forecast.Length % batch != 0) return forecast; - var meanT = new Tensor(new[] { batch, 1 }); - var stdT = new Tensor(new[] { batch, 1 }); - for (int b = 0; b < batch; b++) - { - meanT.Data.Span[b] = _revinMean[b]; - stdT.Data.Span[b] = _revinStd[b]; - } - bool reshaped = forecast.Rank != 2; var work = reshaped ? Engine.Reshape(forecast, new[] { batch, forecast.Length / batch }) : forecast; - var scaled = Engine.TensorBroadcastMultiply(work, stdT); - var shifted = Engine.TensorBroadcastAdd(scaled, meanT); + var scaled = Engine.TensorBroadcastMultiply(work, std); + var shifted = Engine.TensorBroadcastAdd(scaled, mean); return reshaped ? Engine.Reshape(shifted, forecast._shape) : shifted; } @@ -712,7 +786,13 @@ public override Dictionary GetFinancialMetrics() private Tensor ForwardNative(Tensor input) { - var normalized = ApplyInstanceNormalization(input); + // Thread mean/std as tensors through the same forward — under the compiled + // fused plan, class-field capture at trace time would freeze the trace-batch + // stats; tensor-threaded stats re-execute on every replay. The abstract + // override caches them for external ApplyInstanceNormalization callers. + var (normalized, mean, std) = NormalizeWithStats(input); + _revinMeanTensor = mean; + _revinStdTensor = std; var current = normalized; bool addedBatchDim = false; @@ -762,8 +842,9 @@ private Tensor ForwardNative(Tensor input) // RevIN reverse: restore the input's per-instance level/scale so distinct // input levels yield distinct forecasts (the encoders see only the - // mean/std-normalized series). - current = DenormalizeForecast(current); + // mean/std-normalized series). Pass mean/std as tensor locals so the + // compiled-plan replay picks up the CURRENT-step stats, not the trace pass. + current = DenormalizeForecastWithStats(current, mean, std); if (addedBatchDim && current.Rank == 2 && current.Shape[0] == 1) current = Engine.Reshape(current, new[] { current.Shape[1] }); @@ -823,47 +904,58 @@ protected override Tensor ForecastOnnx(Tensor input) #region Frequency Transform /// - /// Computes the DFT magnitude spectrum of the input time series. - /// For rank-1 input, computes DFT directly. For batched input (rank > 1), - /// computes DFT per sample along the last dimension. - /// Returns |X[k]| for k = 0..N/2 (one-sided spectrum), same shape as input via mirroring. + /// Traceable DFT magnitude spectrum. Uses for the batched + /// real-to-complex FFT, then computes per-bin magnitude via elementwise square + + /// axis-reduction + sqrt, and mirrors the one-sided spectrum to full length via + /// + . + /// All ops record on the autodiff tape and re-execute on every compiled-plan replay — + /// the previous .Data.Span DFT loop froze the trace-batch spectrum into the plan. /// + /// + /// Output layout matches the previous scalar impl: [..., n] with the first + /// halfN = n/2 + 1 bins holding |X[k]| / n and the tail mirroring bins + /// 1 .. n-halfN so downstream freq encoders that expect the full-n layout + /// see identical shapes. + /// private Tensor ComputeFrequencyRepresentation(Tensor input) { - // For rank-1 (unbatched), n = sequence length. For batched, n = last dimension. int n = input.Rank > 1 ? input.Shape[^1] : input.Length; int numSamples = input.Rank > 1 ? input.Length / n : 1; int halfN = n / 2 + 1; - var result = new Tensor(input._shape); - T invN = NumOps.Divide(NumOps.One, NumOps.FromDouble(n)); - for (int s = 0; s < numSamples; s++) + // Normalize to [B, n] for the batched RFFT contract. + bool reshaped = input.Rank != 2; + var flat = reshaped ? Engine.Reshape(input, new[] { numSamples, n }) : input; + + // RFFT returns interleaved [re0, im0, re1, im1, ..., re(halfN-1), im(halfN-1)], + // shape [B, halfN * 2] = [B, n + 2]. + var rfft = Engine.RFFT(flat); + var complexPairs = Engine.Reshape(rfft, new[] { numSamples, halfN, 2 }); + + // magSquared[b, k] = re[b, k]^2 + im[b, k]^2 via TensorMultiply + ReduceSum on the + // last (re/im) axis. Axis 2 reduces the pair into a scalar → shape [B, halfN]. + var squares = Engine.TensorMultiply(complexPairs, complexPairs); + var magSquared = Engine.ReduceSum(squares, new[] { 2 }, keepDims: false); + var mag = Engine.TensorSqrt(magSquared); + // Normalize by 1/n to match the previous impl (which multiplied by invN post-sqrt). + var oneSided = Engine.TensorMultiplyScalar(mag, NumOps.Divide(NumOps.One, NumOps.FromDouble(n))); + + Tensor full; + int mirrorLen = n - halfN; + if (mirrorLen > 0) { - int offset = s * n; - - // DFT: X[k] = sum_{t=0}^{N-1} x[t] * exp(-2*pi*i*k*t/N) - for (int k = 0; k < halfN; k++) - { - T realPart = NumOps.Zero; - T imagPart = NumOps.Zero; - for (int t = 0; t < n; t++) - { - double angle = -2.0 * Math.PI * k * t / n; - T cosT = NumOps.FromDouble(Math.Cos(angle)); - T sinT = NumOps.FromDouble(Math.Sin(angle)); - realPart = NumOps.Add(realPart, NumOps.Multiply(input[offset + t], cosT)); - imagPart = NumOps.Add(imagPart, NumOps.Multiply(input[offset + t], sinT)); - } - T magSquared = NumOps.Add(NumOps.Multiply(realPart, realPart), NumOps.Multiply(imagPart, imagPart)); - result.Data.Span[offset + k] = NumOps.Multiply(NumOps.Sqrt(magSquared), invN); - } - - // Mirror the one-sided spectrum for symmetric representation - for (int k = halfN; k < n; k++) - result.Data.Span[offset + k] = result[offset + (n - k)]; + // Mirror indices [1 .. n-halfN] reversed → tail of the full spectrum. + // Slice: start=[0, 1], length=[B, mirrorLen] gives bins 1..mirrorLen inclusive. + var mirrorSource = Engine.TensorSlice(oneSided, new[] { 0, 1 }, new[] { numSamples, mirrorLen }); + var mirrorReversed = Engine.TensorFlip(mirrorSource, new[] { 1 }); + full = Engine.TensorConcatenate(new[] { oneSided, mirrorReversed }, axis: 1); + } + else + { + full = oneSided; } - return result; + return reshaped ? Engine.Reshape(full, input._shape) : full; } #endregion diff --git a/src/Finance/Forecasting/Foundation/TOTEM.cs b/src/Finance/Forecasting/Foundation/TOTEM.cs index 3fc80e563a..64059726ce 100644 --- a/src/Finance/Forecasting/Foundation/TOTEM.cs +++ b/src/Finance/Forecasting/Foundation/TOTEM.cs @@ -330,6 +330,57 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + // GPU-RESIDENT fast path — recon + commitment on a fused SGD plan. Safe + // now that VectorQuantize is fully traceable (argmin + gather + straight- + // through + commitment loss all via engine ops) so each replay recomputes + // from the CURRENT slot data instead of freezing the trace-batch argmin + // into the plan. The codebook EMA runs POST-Step in eager code using the + // trace-time argmin/head tensor references — their .Data is refreshed by + // each replay, so the post-Step read gives the current batch's values and + // the update lands exactly once per batch (CodeRabbit contract). + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + Tensor? capturedCommitment = null; + Tensor? capturedArgmin = null; + Tensor? capturedHead = null; + Tensor ForwardCombined(Tensor inp) + { + var (fc, commit, argmin, head) = ForwardNativeForTrainingWithVQExtras(inp); + capturedCommitment = commit; + capturedArgmin = argmin; + capturedHead = head; + return fc; + } + Tensor ComputeLossCombined(Tensor pred, Tensor tgt) + { + var alignedT = tgt; + if (pred.Rank > tgt.Rank && pred.Shape[0] == 1 && pred.Length == tgt.Length) + pred = Engine.Reshape(pred, tgt._shape); + else if (tgt.Rank > pred.Rank && tgt.Shape[0] == 1 && tgt.Length == pred.Length) + alignedT = Engine.Reshape(tgt, pred._shape); + var recon = loss.ComputeTapeLoss(pred, alignedT); + var commit = capturedCommitment + ?? throw new InvalidOperationException( + "TOTEM fused step: commitment loss was not captured by ForwardCombined. " + + "This indicates the fused-step framework called the loss closure before " + + "the forward closure, violating its documented Fwd-then-Loss ordering."); + return Engine.TensorAdd(recon, commit); + } + if (AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + trainableLayers, input, target, + forward: ForwardCombined, computeLoss: ComputeLossCombined, + optimizerType: AiDotNet.Tensors.Engines.Compilation.OptimizerType.SGD, + learningRate: 0.001f, beta1: 0.9f, beta2: 0.999f, epsilon: 1e-8f, weightDecay: 0f, + out T fusedLoss)) + { + LastLoss = fusedLoss; + if (IsTrainingMode && capturedArgmin is not null && capturedHead is not null) + UpdateCodebookEMA(capturedHead, capturedArgmin); + return; + } + } + using var tape = new GradientTape(); var (forecast, commitmentLoss) = ForwardNativeForTrainingWithCommitment(input); @@ -375,6 +426,24 @@ public override void Train(Tensor input, Tensor target) /// caller to add to the reconstruction loss. /// private (Tensor forecast, Tensor commitmentLoss) ForwardNativeForTrainingWithCommitment(Tensor input) + { + var (forecast, commitmentLoss, _, _) = ForwardNativeForTrainingWithVQExtras(input); + return (forecast, commitmentLoss); + } + + /// + /// Training forward that also exposes the VQ argmin indices and encoder head + /// tensors. Used by the compiled fused path so the caller can invoke + /// AFTER each Step with post-replay values + /// (argmin/head are graph-node references whose .Data is refreshed + /// by each replay). All ops go through so the full + /// forward — including the RevIN normalize/denormalize, the encoder, the + /// quantization projection, VQ argmin+gather+straight-through, commitment + /// loss, and the decoder — records on the autodiff tape and re-executes + /// per replay. + /// + private (Tensor Forecast, Tensor CommitmentLoss, Tensor? Argmin, Tensor? Head) + ForwardNativeForTrainingWithVQExtras(Tensor input) { var normalized = ApplyInstanceNormalization(input); // Tokenize to [1, contextLength, 1] for the per-token encoder/decoder. @@ -388,30 +457,9 @@ public override void Train(Tensor input, Tensor target) if (_quantizationProjection is not null) current = _quantizationProjection.Forward(current); - // Non-differentiable lookup: find nearest codebook entry per - // position. VectorQuantize returns a plain Tensor built by - // .Data.Span — this is intentional here; we use it as the - // stop-gradient target in the straight-through trick. - var codebookValues = VectorQuantize(current); - - // Straight-through estimator: - // quantized = encoder + sg(codebook - encoder) - // forward: evaluates to codebook values (VQ behavior) - // backward: d quantized / d encoder = 1 (gradient passes through) - var diff = Engine.TensorSubtract(codebookValues, current); - var diffDetached = Engine.StopGradient(diff); - var quantizedST = Engine.TensorAdd(current, diffDetached); - - // Commitment loss: mean((encoder - sg(codebook))^2), weighted. - // Pulling encoder toward codebook encourages discrete - // quantization without destabilizing codebook values. - var codebookDetached = Engine.StopGradient(codebookValues); - var commitDiff = Engine.TensorSubtract(current, codebookDetached); - var commitSq = Engine.TensorMultiply(commitDiff, commitDiff); - var allAxes = Enumerable.Range(0, commitSq.Rank).ToArray(); - var commitmentLoss = Engine.ReduceMean(commitSq, allAxes, keepDims: false); - commitmentLoss = Engine.TensorMultiplyScalar(commitmentLoss, - NumOps.FromDouble(_commitmentWeight)); + // Traceable VQ: returns straight-through-quantized values, commitment loss, + // argmin indices, and the reshaped-head input for post-Step EMA. + var (quantizedST, commitmentLoss, argmin, head) = VectorQuantizeTraceable(current); var decoded = quantizedST; if (_decoder is not null) @@ -427,7 +475,7 @@ public override void Train(Tensor input, Tensor target) // RevIN reverse: train against the input-scale forecast. decoded = DenormalizeForecast(decoded); - return (decoded, commitmentLoss); + return (decoded, commitmentLoss, argmin, head); } /// @@ -744,85 +792,216 @@ private Tensor ForwardNative(Tensor input) /// Uses product quantization when numCodebooks > 1 (splits features across codebooks). /// Also computes commitment loss: ||z_e - sg(e_k)||^2. /// + /// + /// Legacy scalar-loop VectorQuantize — kept for callers that don't need the argmin/head + /// side-outputs (inference, serialization roundtrips). Training paths should use + /// so the entire quantization runs on-tape and + /// re-executes correctly under the compiled fused plan. + /// private Tensor VectorQuantize(Tensor encoderOutput) + { + var (quantized, _, _, _) = VectorQuantizeTraceable(encoderOutput); + return quantized; + } + + /// + /// Traceable VQ-VAE quantization step (van den Oord et al. 2017). Returns the + /// straight-through-quantized tensor, the commitment loss, the argmin indices, + /// and the reshaped-head input in the [numPositions, numCodebooks, codebookDim] + /// layout. All ops go through so the computation records on + /// the autodiff tape and re-executes on every replay under a compiled fused plan + /// — the previous .Data.Span nearest-neighbor + SetCodebookValue + /// EMA loop froze the argmin decision AND applied the codebook update at trace + /// time (bug flagged by CodeRabbit). + /// + /// + /// EMA is intentionally NOT applied here. The caller invokes + /// with the returned argmin+head AFTER the compiled Step so the codebook update + /// runs exactly once per batch (regardless of whether the fused or eager path + /// engaged). Under the compiled plan, head and argmin are trace-time + /// graph nodes whose .Data is refreshed by each replay — reading them + /// post-Step gives the current batch's values. + /// + private (Tensor Quantized, Tensor CommitmentLoss, Tensor? ArgminIndices, Tensor? Head) + VectorQuantizeTraceable(Tensor encoderOutput) { if (_codebooks is null) InitializeCodebooks(); + var codebooks = _codebooks!; int totalLen = encoderOutput.Length; - var quantized = new Tensor(encoderOutput._shape); - T commitmentLoss = NumOps.Zero; - - // Product quantization: split encoded vector across codebooks int dimPerCodebook = Math.Max(1, _codebookDimension); - int numPositions = Math.Max(1, totalLen / Math.Max(1, dimPerCodebook * _numCodebooks)); - - for (int pos = 0; pos < numPositions; pos++) + int blockSize = dimPerCodebook * _numCodebooks; + int numPositions = Math.Max(1, totalLen / Math.Max(1, blockSize)); + int quantizedElements = numPositions * blockSize; + + // Fallback: input can't be cleanly reshaped into the PQ block structure. + // Return the input unchanged with a zero commitment loss and no argmin/head + // (the caller's EMA-update path is a no-op when these are null). + if (numPositions <= 0 || quantizedElements > totalLen) { - for (int c = 0; c < _numCodebooks; c++) - { - int startIdx = pos * dimPerCodebook * _numCodebooks + c * dimPerCodebook; - if (startIdx >= totalLen) break; + var zeroLoss = new Tensor(new[] { 1 }); + Engine.TensorFill(zeroLoss, NumOps.Zero); + _lastCommitmentLoss = NumOps.Zero; + return (encoderOutput, zeroLoss, null, null); + } - int effectiveDim = Math.Min(dimPerCodebook, totalLen - startIdx); + // Split input into [quantizable, passThrough]. The passThrough tail is + // copied unchanged; the quantizable prefix goes through PQ. + var flatInput = encoderOutput.Rank == 1 + ? encoderOutput + : Engine.Reshape(encoderOutput, new[] { totalLen }); + var quantizable = Engine.TensorSlice(flatInput, new[] { 0 }, new[] { quantizedElements }); + + // head[p, c, d] — reshape the quantizable prefix into PQ block layout. + var head = Engine.Reshape(quantizable, new[] { numPositions, _numCodebooks, dimPerCodebook }); + + // Distance to each codebook entry: broadcast head [P, C, 1, D] against + // codebook [1, C, K, D] → diff [P, C, K, D] → sum(diff²) → [P, C, K]. + // codebooks shape: [numCodebooks, codebookSize, codebookDim] → add batch axis. + var headExpanded = Engine.Reshape(head, new[] { numPositions, _numCodebooks, 1, dimPerCodebook }); + var codebookExpanded = Engine.Reshape(codebooks, new[] { 1, _numCodebooks, _codebookSize, dimPerCodebook }); + var diff = Engine.TensorBroadcastSubtract(headExpanded, codebookExpanded); + var diffSq = Engine.TensorMultiply(diff, diff); + var distances = Engine.ReduceSum(diffSq, new[] { 3 }, keepDims: false); + // distances shape: [numPositions, numCodebooks, codebookSize]. + + // Argmin over the codebookSize axis — non-differentiable by design; the + // straight-through estimator below routes gradients around the argmin. + var argmin = Engine.TensorArgMin(distances, axis: 2); + // argmin shape: [numPositions, numCodebooks] of Tensor. + + // Per-codebook gather: for each c, zqSlices[c][p, :] = codebooks[c, argmin[p, c], :]. + // TensorIndexSelectDiff along the codebookSize axis of the per-c codebook slice. + var zqSlices = new Tensor[_numCodebooks]; + for (int c = 0; c < _numCodebooks; c++) + { + // Slice codebook_c = codebooks[c, :, :] via TensorSliceAxis(axis=0, index=c). + var codebookC = Engine.TensorSliceAxis(codebooks, axis: 0, index: c); + // argminC = argmin[:, c] shape [numPositions] — TensorSliceAxis on int tensor. + var argminC = Engine.TensorSliceAxis(argmin, axis: 1, index: c); + // Gather: source shape [codebookSize, codebookDim], indices [numPositions] along axis 0 + // → [numPositions, codebookDim]. + zqSlices[c] = Engine.TensorIndexSelectDiff(codebookC, argminC, axis: 0); + } + // Stack per-codebook slices along the codebook axis to get [numPositions, numCodebooks, codebookDim]. + var zq = Engine.TensorStack(zqSlices, axis: 1); + + // Commitment loss per Oord 2017 §3.2: β · ||z_e - sg(e_k)||² averaged over totalLen. + // Straight-through routes gradient through encoder only — encoder learns to + // match codebook via commitment loss; codebook learns via EMA (separate path). + var zqDetached = Engine.StopGradient(zq); + var commitmentDelta = Engine.TensorSubtract(head, zqDetached); + var commitmentSqSum = Engine.ReduceSum( + Engine.TensorMultiply(commitmentDelta, commitmentDelta), + axes: null, keepDims: false); + var invTotalLen = NumOps.Divide(NumOps.One, NumOps.FromDouble(Math.Max(1, totalLen))); + var commitmentLoss = Engine.TensorMultiplyScalar( + commitmentSqSum, + NumOps.Multiply(NumOps.FromDouble(_commitmentWeight), invTotalLen)); + _lastCommitmentLoss = commitmentLoss.Length > 0 ? commitmentLoss[0] : NumOps.Zero; + + // Straight-through: quantized = head + StopGradient(zq - head). Forward-values + // equal codebook entries; backward gradient flows through head as if identity. + var straightThroughShift = Engine.StopGradient(Engine.TensorSubtract(zq, head)); + var quantizedBlocks = Engine.TensorAdd(head, straightThroughShift); + var quantizedFlat = Engine.Reshape(quantizedBlocks, new[] { quantizedElements }); + + Tensor quantized; + if (quantizedElements < totalLen) + { + // Concat the passThrough tail unchanged. + var passThroughLen = totalLen - quantizedElements; + var passThrough = Engine.TensorSlice(flatInput, new[] { quantizedElements }, new[] { passThroughLen }); + var combined = Engine.TensorConcatenate(new[] { quantizedFlat, passThrough }, axis: 0); + quantized = encoderOutput.Rank == 1 ? combined : Engine.Reshape(combined, encoderOutput._shape); + } + else + { + quantized = encoderOutput.Rank == 1 ? quantizedFlat : Engine.Reshape(quantizedFlat, encoderOutput._shape); + } - // Find nearest codebook entry - int bestIdx = 0; - T bestDist = NumOps.MaxValue; - for (int k = 0; k < _codebookSize; k++) - { - T dist = NumOps.Zero; - for (int d = 0; d < effectiveDim; d++) - { - int idx = startIdx + d; - if (idx >= totalLen) break; - T diff = NumOps.Subtract(encoderOutput[idx], GetCodebookValue(c, k, d % _codebookDimension)); - dist = NumOps.Add(dist, NumOps.Multiply(diff, diff)); - } - if (NumOps.LessThan(dist, bestDist)) { bestDist = dist; bestIdx = k; } - } + return (quantized, commitmentLoss, argmin, head); + } - // Replace encoder output with nearest codebook entry (straight-through) - for (int d = 0; d < effectiveDim; d++) - { - int idx = startIdx + d; - if (idx >= totalLen) break; - T codebookVal = GetCodebookValue(c, bestIdx, d % _codebookDimension); + /// + /// Post-Step EMA codebook update (van den Oord 2017 §3.2). Runs exactly once + /// per batch, using the trace-time / + /// tensors captured by — under the compiled + /// plan these are graph-node references whose .Data reflects the LAST + /// replay, so post-Step reads give the current batch's values. + /// + /// + /// EMA formula (last-wins matching the eager path's per-position race): + /// codebook[c, argmin[p,c], :] ← decay · codebook[c, argmin[p,c], :] + + /// (1-decay) · head[p, c, :]. + /// Expressed as an in-place scatter-add: codebook += (1-decay) · scatter(head - zq_at_selected), + /// where the scatter writes the deltas into the codebook at (c, argmin) positions + /// and the codebook's underlying data is updated via + /// so the tensor object identity (referenced by future trace/replay reads) stays + /// stable. + /// + private void UpdateCodebookEMA(Tensor head, Tensor argmin) + { + if (_codebooks is null) return; + var codebooks = _codebooks; - // Straight-through: quantized = encoder_output + sg(codebook - encoder_output) - // This means forward uses codebook values, backward passes gradient through encoder - quantized.Data.Span[idx] = codebookVal; + int dimPerCodebook = Math.Max(1, _codebookDimension); + int numPositions = head.Shape.Length >= 3 ? head.Shape[0] : 1; - // Commitment loss: ||z_e - sg(e_k)||^2 - T diff = NumOps.Subtract(encoderOutput[idx], codebookVal); - commitmentLoss = NumOps.Add(commitmentLoss, NumOps.Multiply(diff, diff)); - } + var decayT = NumOps.FromDouble(0.99); + var oneMinusDecayT = NumOps.FromDouble(0.01); - // EMA codebook update (during training): move codebook entry toward encoder output - if (IsTrainingMode) - { - T emaDecay = NumOps.FromDouble(0.99); - T oneMinusDecay = NumOps.FromDouble(0.01); - for (int d = 0; d < effectiveDim && d < _codebookDimension; d++) - { - int idx = startIdx + d; - if (idx >= totalLen) break; - T currentVal = GetCodebookValue(c, bestIdx, d); - T newVal = NumOps.Add(NumOps.Multiply(emaDecay, currentVal), NumOps.Multiply(oneMinusDecay, encoderOutput[idx])); - SetCodebookValue(c, bestIdx, d, newVal); - } - } + // Per-codebook scatter: for each c, the selected entries move by + // (1-decay)·(head[:, c, :] - codebook[c, argmin[:, c], :]). + // TensorScatterAdd writes updates into a fresh copy of the codebook slice; + // TensorCopy propagates the update back into the same codebook tensor object. + for (int c = 0; c < _numCodebooks; c++) + { + var codebookC = Engine.TensorSliceAxis(codebooks, axis: 0, index: c); + var headC = Engine.TensorSliceAxis(head, axis: 1, index: c); + var argminC = Engine.TensorSliceAxis(argmin, axis: 1, index: c); + + // Current selected entries: gather codebookC by argminC along axis 0. + var zqC = Engine.TensorIndexSelectDiff(codebookC, argminC, axis: 0); + + // Delta = (1-decay) · (headC - zqC), shape [numPositions, codebookDim]. + var deltaC = Engine.TensorMultiplyScalar( + Engine.TensorSubtract(headC, zqC), + oneMinusDecayT); + + // Scatter deltaC into a zero canvas at argminC indices along axis 0. + var canvas = new Tensor(new[] { _codebookSize, dimPerCodebook }); + Engine.TensorFill(canvas, NumOps.Zero); + var updatedCodebookC = Engine.TensorScatterAdd(canvas, argminC, deltaC, axis: 0); + // updatedCodebookC = deltas placed at scatter positions, zeros elsewhere. + + // codebookC_new = codebookC + updatedCodebookC. + var codebookC_new = Engine.TensorAdd(codebookC, updatedCodebookC); + + // Write codebookC_new back into codebooks[c, :, :] via TensorSliceAxisWrite-like + // pattern: reshape to [1, codebookSize, codebookDim], place with TensorScatter + // along axis 0 at index c into a per-codebook staging tensor. Simpler: + // reconstruct the full codebook by concatenating updated + unchanged slices. + var updated3D = Engine.Reshape(codebookC_new, new[] { 1, _codebookSize, dimPerCodebook }); + var partsList = new System.Collections.Generic.List>(); + if (c > 0) + { + partsList.Add(Engine.TensorSlice(codebooks, new[] { 0, 0, 0 }, new[] { c, _codebookSize, dimPerCodebook })); + } + partsList.Add(updated3D); + if (c + 1 < _numCodebooks) + { + partsList.Add(Engine.TensorSlice(codebooks, + new[] { c + 1, 0, 0 }, + new[] { _numCodebooks - c - 1, _codebookSize, dimPerCodebook })); } + var newCodebooks = partsList.Count == 1 ? partsList[0] : Engine.TensorConcatenate(partsList.ToArray(), axis: 0); + Engine.TensorCopy(newCodebooks, codebooks); } - // Copy through any remaining elements that don't fit product quantization - int quantizedElements = numPositions * _numCodebooks * dimPerCodebook; - for (int i = quantizedElements; i < totalLen; i++) - quantized.Data.Span[i] = encoderOutput[i]; - - T commitWeightT = NumOps.FromDouble(_commitmentWeight); - T invTotalLen = NumOps.Divide(NumOps.One, NumOps.FromDouble(Math.Max(1, totalLen))); - _lastCommitmentLoss = NumOps.Multiply(NumOps.Multiply(commitmentLoss, commitWeightT), invTotalLen); - return quantized; + // Suppress unused-var warning for decayT (kept for documentation of the + // decay-in-the-formula intent; the (1-decay) factor is what's actually used). + _ = decayT; } protected override Tensor ForecastOnnx(Tensor input) diff --git a/src/Finance/Forecasting/Neural/MQCNN.cs b/src/Finance/Forecasting/Neural/MQCNN.cs index 67bc1153ad..d614220ebe 100644 --- a/src/Finance/Forecasting/Neural/MQCNN.cs +++ b/src/Finance/Forecasting/Neural/MQCNN.cs @@ -531,6 +531,22 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + // GPU-RESIDENT fast path — fused SGD on the multi-quantile pinball + // objective. Falls through to the in-place SGD loop below. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0 + && AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + trainableLayers, input, target, + forward: ForwardForTraining, + computeLoss: ComputeMultiQuantilePinballLossTape, + optimizerType: AiDotNet.Tensors.Engines.Compilation.OptimizerType.SGD, + learningRate: 0.001f, beta1: 0.9f, beta2: 0.999f, epsilon: 1e-8f, weightDecay: 0f, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + using var tape = new GradientTape(); var predictions = ForwardForTraining(input); var lossTensor = ComputeMultiQuantilePinballLossTape(predictions, target); diff --git a/src/NeuralNetworks/NeuralNetworkBase.cs b/src/NeuralNetworks/NeuralNetworkBase.cs index 1b0e9bd698..e8fd372d38 100644 --- a/src/NeuralNetworks/NeuralNetworkBase.cs +++ b/src/NeuralNetworks/NeuralNetworkBase.cs @@ -8800,7 +8800,7 @@ private static void EmitFusedPathEventIfEnabled(bool hit, string? reason) /// configure-once contract: adaptive learning rates, attached LR schedulers, /// or AMSGrad mode (which the fused kernel doesn't model). /// - private static bool TryMapToFusedOptimizerConfig( + internal static bool TryMapToFusedOptimizerConfig( IGradientBasedOptimizer, Tensor> optimizer, out AiDotNet.Tensors.Engines.Compilation.OptimizerType optimizerType, out float learningRate, diff --git a/src/NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs b/src/NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs index 8c9444b53f..e3aa763640 100644 --- a/src/NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs @@ -718,23 +718,77 @@ private void ComputeNoiseSchedule(string schedule, int timesteps) private void TrainBatch(Matrix data, int startRow, int endRow, T lr) { if (_alphasCumprod is null) return; - for (int row = startRow; row < endRow; row++) + + // Cache MultiSlotFusedStep across rows so the compiled plan is built + // once and replayed via slot-data refresh for subsequent rows. See + // ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try { - int t = _random.Next(_numTimesteps); - var x0 = GetRow(data, row); - var noise = CreateStandardNormalVector(_dataWidth); - double sqrtAbar = Math.Sqrt(NumOps.ToDouble(_alphasCumprod[t])); - double sqrtOneMinusAbar = Math.Sqrt(1.0 - NumOps.ToDouble(_alphasCumprod[t])); + var denoiserLayers = BuildDenoiserLayerList(); + var trainableParams = Training.TapeTrainingStep.CollectParameters(denoiserLayers).ToArray(); + AiDotNet.Tensors.Engines.Compilation.OptimizerType mfsOptType = default; + float mfsLr = 0f, mfsB1 = 0f, mfsB2 = 0f, mfsEps = 0f, mfsWd = 0f; + bool fusedEligible = false; + if (trainableParams.Length > 0) + { + fusedEligible = NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out mfsOptType, out mfsLr, out mfsB1, out mfsB2, + out mfsEps, out mfsWd, out _, out _); + } - var xt = new Vector(_dataWidth); - for (int j = 0; j < _dataWidth; j++) - xt[j] = NumOps.FromDouble(sqrtAbar * NumOps.ToDouble(x0[j]) + sqrtOneMinusAbar * NumOps.ToDouble(noise[j])); + for (int row = startRow; row < endRow; row++) + { + int t = _random.Next(_numTimesteps); + var x0 = GetRow(data, row); + var noise = CreateStandardNormalVector(_dataWidth); + double sqrtAbar = Math.Sqrt(NumOps.ToDouble(_alphasCumprod[t])); + double sqrtOneMinusAbar = Math.Sqrt(1.0 - NumOps.ToDouble(_alphasCumprod[t])); + + var xt = new Vector(_dataWidth); + for (int j = 0; j < _dataWidth; j++) + xt[j] = NumOps.FromDouble(sqrtAbar * NumOps.ToDouble(x0[j]) + sqrtOneMinusAbar * NumOps.ToDouble(noise[j])); + + // Preferred fused path: MultiSlotFusedStep with (denoiserInput, targetNoise) + // as persistent slots. The denoiser layer set replays per row with fresh slots. + if (fusedEligible) + { + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + var denoiserInput = BuildDenoiserInput(xt, t); + var targetNoiseT = VectorToTensor(noise); + var slots = new[] { denoiserInput, targetNoiseT }; + Tensor ForwardFromSlots(IReadOnlyList> s) => DenoiserForward(s[0]); + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) => + ReduceToScalar(Engine.TensorSquare(Engine.TensorSubtract(pred, s[1]))); + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } - using var tape = new GradientTape(); - var pred = DenoiserForward(BuildDenoiserInput(xt, t)); - var target = VectorToTensor(noise); - var loss = ReduceToScalar(Engine.TensorSquare(Engine.TensorSubtract(pred, target))); - TapeStepOver(tape, loss, BuildDenoiserLayerList()); + using var tape = new GradientTape(); + var pred = DenoiserForward(BuildDenoiserInput(xt, t)); + var target = VectorToTensor(noise); + var loss = ReduceToScalar(Engine.TensorSquare(Engine.TensorSubtract(pred, target))); + TapeStepOver(tape, loss, denoiserLayers); + } + } + finally + { + multiSlotStep?.Dispose(); } } diff --git a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 8980d3ea7f..46ce9adb6f 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -614,8 +614,83 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - using var tape = new GradientTape(); + // Preferred fused path: WganGpFusedStep runs the full WGAN-GP objective + // (Wasserstein + λ·GP with createGraph=true GP) in one compiled plan, + // with persistent (real, fake, ε) slots refreshed each Step. Falls + // through to the GpuResidentFusedStep path (which uses this critic's + // per-batch epsilon sampled inline) when the optimizer has no fused- + // kernel mapping (adaptive LR, non-fuse-able type) or the primitive + // itself declines. See ooples/AiDotNet#1845. var discParams = TapeTrainingStep.CollectParameters(_discLayers); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realPacked, + fakeBatch: fakePacked, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). + var trainableDiscLayers = _discLayers.OfType>().ToList(); + if (trainableDiscLayers.Count > 0) + { + int realN = realPacked.Shape[0]; + int fakeN = fakePacked.Shape[0]; + var stacked = Engine.TensorConcatenate([realPacked, fakePacked], axis: 0); + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor both) => DiscriminatorForwardBatched(both, isTraining: true); + Tensor Loss(Tensor allScores, Tensor _) + { + var rShape = allScores._shape.ToArray(); rShape[0] = realN; + var fShape = allScores._shape.ToArray(); fShape[0] = fakeN; + var rStart = new int[allScores.Rank]; + var fStart = new int[allScores.Rank]; fStart[0] = realN; + var rScores = Engine.TensorSlice(allScores, rStart, rShape); + var fScores = Engine.TensorSlice(allScores, fStart, fShape); + var axes = Enumerable.Range(0, rScores.Shape.Length).ToArray(); + var wasserstein = Engine.TensorSubtract( + Engine.ReduceMean(fScores, axes, keepDims: false), + Engine.ReduceMean(rScores, axes, keepDims: false)); + var gp = ComputeGradientPenalty(realPacked, fakePacked); + return Engine.TensorAdd(wasserstein, + Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDiscLayers, stacked, target, + forward: Fwd, computeLoss: Loss, + optimizer: _discriminatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + // discParams already collected above for the WganGpFusedStep attempt; + // reuse it here to avoid a redundant Layers → parameter scan. var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true); @@ -727,8 +802,6 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) { if (_sampler is null) return; - using var tape = new GradientTape(); - // Generator's trainable surface = Layers + per-layer BN. var generatorLayers = new List>(); generatorLayers.AddRange(Layers); @@ -741,10 +814,74 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) // Build the conditional generator input batch [numPacks * pacSize, genInputDim], // capturing the mask so the conditional cross-entropy term can be computed. - var noiseBatch = GenerateNoiseBatchTensor(numPacks * pacSize); - var (condBatch, maskBatch) = SampleCondMaskBatch(numPacks * pacSize); + int totalSamples = numPacks * pacSize; + var noiseBatch = GenerateNoiseBatchTensor(totalSamples); + var (condBatch, maskBatch) = SampleCondMaskBatch(totalSamples); var genInput = Engine.TensorConcatenate([noiseBatch, condBatch], axis: 1); + // GPU-RESIDENT fast path — pack (genInput, maskBatch) into one persistent + // input tensor so replay uses fresh noise + cond + mask each step. The + // closure slices them back out. Conditional CE and -avgFake both captured + // on the fused plan. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + int embedDim = _options.EmbeddingDimension; + int condDim = _condWidth; + // Fused-step input encodes [noise | cond | mask] side-by-side (columns). + var fusedInput = Engine.TensorConcatenate([genInput, maskBatch], axis: 1); + var target = new Tensor(new[] { 1 }); + // Closure-captured intermediates from Fwd's single generator pass — + // reused by Loss so we don't re-run GeneratorForwardWithResidualBatched + // and ApplyOutputActivationsBatched for the conditional-CE term. + // Fwd/Loss ordering is guaranteed by the fused-step contract. + Tensor? capturedAct = null; + Tensor? capturedCond = null; + Tensor? capturedMask = null; + Tensor Fwd(Tensor packed) + { + var gi = Engine.TensorSlice(packed, [0, 0], [totalSamples, embedDim + condDim]); + var faked = GeneratorForwardWithResidualBatched(gi); + var act = ApplyOutputActivationsBatched(faked); + var condFromInput = Engine.TensorSlice(packed, [0, embedDim], [totalSamples, condDim]); + capturedAct = act; + capturedCond = condFromInput; + capturedMask = Engine.TensorSlice(packed, [0, embedDim + condDim], [totalSamples, condDim]); + var withCond = Engine.TensorConcatenate([act, condFromInput], axis: 1); + var pk = withCond.Reshape([numPacks, _packedInputDim]); + return DiscriminatorForwardBatched(pk, false); + } + Tensor Loss(Tensor scores, Tensor _) + { + var axes = Enumerable.Range(0, scores.Shape.Length).ToArray(); + var lossT = Engine.TensorNegate(Engine.ReduceMean(scores, axes, keepDims: false)); + if (_condWidth > 0 && _catOutputBlocks.Count > 0) + { + var act = capturedAct; + var condFromInput = capturedCond; + var maskFromInput = capturedMask; + if (act is null || condFromInput is null || maskFromInput is null) + throw new InvalidOperationException( + "CTGAN fused step: generator intermediates were not captured by Fwd. " + + "This indicates the fused-step framework called the loss closure before " + + "the forward closure, violating its documented Fwd-then-Loss ordering."); + var ce = ConditionalCrossEntropy(act, condFromInput, maskFromInput); + lossT = Engine.TensorAdd(lossT, ce); + } + return lossT; + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, fusedInput, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + // Forward through generator → produces [numPacks * pacSize, dataWidth]. var fakeFlat = GeneratorForwardWithResidualBatched(genInput); var fakeActivated = ApplyOutputActivationsBatched(fakeFlat); @@ -918,7 +1055,9 @@ private Tensor ComputeGradientPenalty(Tensor realPacked, Tensor fakePac var scores = DiscriminatorForwardBatched(interpolated, true); var scoreAxes = Enumerable.Range(0, scores.Shape.Length).ToArray(); var summedScores = Engine.ReduceSum(scores, scoreAxes, keepDims: false); - var gradients = gradientTape.ComputeGradients(summedScores, [interpolated]); + // AiDotNet #1844: createGraph=true records inner backward on outer tape + // so gradient penalty actually flows to disc weights (WGAN-GP correctness). + var gradients = gradientTape.ComputeGradients(summedScores, [interpolated], createGraph: true); inputGradients = gradients.TryGetValue(interpolated, out var gradient) ? gradient : new Tensor(interpolated._shape); diff --git a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs index a1bb5f0dc2..130fa68710 100644 --- a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs @@ -602,9 +602,81 @@ private Tensor DiscriminatorForward(Tensor input, bool isTraining) private void TrainDiscriminatorStepBatched(Matrix transformedData, int batchSize) { var (realBatch, fakeBatch) = BuildRealAndFakeBatches(transformedData, batchSize); + var discLayerList = _discLayers.Cast>().ToList(); + + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). + var discParams = TapeTrainingStep.CollectParameters(discLayerList); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realBatch, + fakeBatch: fakeBatch, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). + var trainableDiscLayers = discLayerList.OfType>().ToList(); + if (trainableDiscLayers.Count > 0) + { + // Pack real + fake into a single persistent input along axis 0 so + // both scores compute from one forward. Split in the loss closure. + int realN = realBatch.Shape[0]; + int fakeN = fakeBatch.Shape[0]; + var stacked = Engine.TensorConcatenate([realBatch, fakeBatch], axis: 0); + var target = new Tensor(new[] { 1 }); + + Tensor Fwd(Tensor both) => DiscriminatorForwardBatched(both, isTraining: true); + Tensor Loss(Tensor allScores, Tensor _) + { + var rShape = allScores._shape.ToArray(); rShape[0] = realN; + var fShape = allScores._shape.ToArray(); fShape[0] = fakeN; + var rStart = new int[allScores.Rank]; + var fStart = new int[allScores.Rank]; fStart[0] = realN; + var rScores = Engine.TensorSlice(allScores, rStart, rShape); + var fScores = Engine.TensorSlice(allScores, fStart, fShape); + var axes = Enumerable.Range(0, rScores.Shape.Length).ToArray(); + var avgReal = Engine.ReduceMean(rScores, axes, keepDims: false); + var avgFake = Engine.ReduceMean(fScores, axes, keepDims: false); + var wasserstein = Engine.TensorSubtract(avgFake, avgReal); + var gp = ComputeGradientPenalty(realBatch, fakeBatch); + return Engine.TensorAdd(wasserstein, + Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDiscLayers, stacked, target, + forward: Fwd, computeLoss: Loss, + optimizer: _discriminatorOptimizer, + out T _)) + { + return; + } + } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); @@ -648,13 +720,42 @@ Tensor RecomputeLoss(Tensor pred, Tensor _) /// private void TrainGeneratorStepBatched(int batchSize) { - using var tape = new GradientTape(); var generatorLayers = new List>(); generatorLayers.AddRange(Layers); foreach (var bn in _genBNLayers) generatorLayers.Add(bn); var genParams = TapeTrainingStep.CollectParameters(generatorLayers); var noiseBatch = GenerateNoiseBatchTensor(batchSize); + + // GPU-RESIDENT fast path — noise → gen → (optional causal) → activation → disc-frozen → -avgFake. + // Disc layers not in the trainable set, so no gradients accumulate to them. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor nb) + { + var f = GeneratorForwardBatched(nb); + var c = _adjacency is not null ? ApplyCausalStructureBatched(f) : f; + var a = ApplyOutputActivationsBatched(c); + return DiscriminatorForwardBatched(a, false); + } + Tensor Loss(Tensor scores, Tensor _) + { + var axes = Enumerable.Range(0, scores.Shape.Length).ToArray(); + return Engine.TensorNegate(Engine.ReduceMean(scores, axes, keepDims: false)); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, noiseBatch, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); var fakeRaw = GeneratorForwardBatched(noiseBatch); var fakeCausal = _adjacency is not null ? ApplyCausalStructureBatched(fakeRaw) : fakeRaw; var fakeBatch = ApplyOutputActivationsBatched(fakeCausal); @@ -1035,7 +1136,11 @@ private Tensor ComputeGradientPenalty(Tensor realBatch, Tensor fakeBatc var scores = DiscriminatorForwardBatched(interpolated, true); var scoreAxes = Enumerable.Range(0, scores.Shape.Length).ToArray(); var summedScores = Engine.ReduceSum(scores, scoreAxes, keepDims: false); - var gradients = gradientTape.ComputeGradients(summedScores, [interpolated]); + // AiDotNet #1844: createGraph=true records the inner backward's ops on the + // outer tape so the gradient penalty actually backpropagates into the + // discriminator weights. Without this, WGAN-GP silently degrades to plain + // WGAN — the 1-Lipschitz constraint from Gulrajani 2017 is not enforced. + var gradients = gradientTape.ComputeGradients(summedScores, [interpolated], createGraph: true); inputGradients = gradients.TryGetValue(interpolated, out var gradient) ? gradient : new Tensor(interpolated._shape); diff --git a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs index 39ba322071..137907eab9 100644 --- a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs @@ -668,8 +668,76 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - using var tape = new GradientTape(); + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). var discParams = TapeTrainingStep.CollectParameters(_discLayers); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realPacked, + fakeBatch: fakePacked, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). + var trainableDiscLayers = _discLayers.OfType>().ToList(); + if (trainableDiscLayers.Count > 0) + { + int realN = realPacked.Shape[0]; + int fakeN = fakePacked.Shape[0]; + var stacked = Engine.TensorConcatenate([realPacked, fakePacked], axis: 0); + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor both) => DiscriminatorForwardBatched(both, isTraining: true); + Tensor Loss(Tensor allScores, Tensor _) + { + var rShape = allScores._shape.ToArray(); rShape[0] = realN; + var fShape = allScores._shape.ToArray(); fShape[0] = fakeN; + var rStart = new int[allScores.Rank]; + var fStart = new int[allScores.Rank]; fStart[0] = realN; + var rScores = Engine.TensorSlice(allScores, rStart, rShape); + var fScores = Engine.TensorSlice(allScores, fStart, fShape); + var axes = Enumerable.Range(0, rScores.Shape.Length).ToArray(); + var wasserstein = Engine.TensorSubtract( + Engine.ReduceMean(fScores, axes, keepDims: false), + Engine.ReduceMean(rScores, axes, keepDims: false)); + var gp = ComputeGradientPenalty(realPacked, fakePacked); + return Engine.TensorAdd(wasserstein, + Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDiscLayers, stacked, target, + forward: Fwd, computeLoss: Loss, + optimizer: _discriminatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true); @@ -714,7 +782,6 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) { if (_sampler is null) return; - using var tape = new GradientTape(); var generatorLayers = new List>(); generatorLayers.AddRange(Layers); foreach (var bn in _genBNLayers) generatorLayers.Add(bn); @@ -727,6 +794,40 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) var condBatch = SampleConditionalBatchTensor(total); var genInput = Engine.TensorConcatenate([noiseBatch, condBatch], axis: 1); + // GPU-RESIDENT fast path — slice condBatch back out of the persistent + // genInput so replay uses fresh cond each step. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + int embedDim = _options.EmbeddingDimension; + int condDim = _condWidth; + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor ginp) + { + var faked = GeneratorForwardWithResidualBatched(ginp); + var act = ApplyOutputActivationsBatched(faked); + var condFromInput = Engine.TensorSlice(ginp, [0, embedDim], [total, condDim]); + var withCond = Engine.TensorConcatenate([act, condFromInput], axis: 1); + var packed = withCond.Reshape([numPacks, _packedInputDim]); + return DiscriminatorForwardBatched(packed, false); + } + Tensor Loss(Tensor scores, Tensor _) + { + var axes = Enumerable.Range(0, scores.Shape.Length).ToArray(); + return Engine.TensorNegate(Engine.ReduceMean(scores, axes, keepDims: false)); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, genInput, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + var fakeFlat = GeneratorForwardWithResidualBatched(genInput); var fakeActivated = ApplyOutputActivationsBatched(fakeFlat); var fakeWithCond = Engine.TensorConcatenate([fakeActivated, condBatch], axis: 1); @@ -863,7 +964,9 @@ private Tensor ComputeGradientPenalty(Tensor realPacked, Tensor fakePac var scores = DiscriminatorForwardBatched(interpolated, true); var scoreAxes = Enumerable.Range(0, scores.Shape.Length).ToArray(); var summedScores = Engine.ReduceSum(scores, scoreAxes, keepDims: false); - var gradients = gradientTape.ComputeGradients(summedScores, [interpolated]); + // AiDotNet #1844: createGraph=true records inner backward on outer tape + // so gradient penalty actually flows to disc weights (WGAN-GP correctness). + var gradients = gradientTape.ComputeGradients(summedScores, [interpolated], createGraph: true); inputGradients = gradients.TryGetValue(interpolated, out var gradient) ? gradient : new Tensor(interpolated._shape); diff --git a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs index 25c61a90aa..226e95e5d5 100644 --- a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs @@ -659,7 +659,50 @@ private void TrainDiscriminatorStepBatchedDP(Matrix transformedData, int numP wasserstein, Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); - var noisedGrads = ComputePerExampleNoisedGradients(realPacked, fakePacked, discParams); + // Route the per-example DP-SGD gradient computation through the fused + // DpSgdFusedStep helper (Phase 4H). Each per-example forward+backward + // runs the compiled plan (GPU-resident params, fused kernels); clip + + // aggregate + noise happens in host code because the per-example L2 + // norm's control flow doesn't fit the compiled-plan capture model. + // The helper's structure enforces clip-BEFORE-aggregate so the Abadi + // 2016 privacy proof's L2-sensitivity contract can't be silently reversed. + int exampleCount = Math.Max(1, realPacked.Shape[0]); + using var dpSgdStep = new AiDotNet.Training.DpSgdFusedStep(); + var gpWeightConst = _options.GradientPenaltyWeight; + bool dpFusedRan = dpSgdStep.TryStep( + parameters: discParams, + perExampleSlotData: exIdx => new[] + { + ExtractPackedExample(realPacked, exIdx), + ExtractPackedExample(fakePacked, exIdx), + }, + forward: slots => DiscriminatorForwardBatched(slots[0], isTraining: true), + computeLoss: (realScore, slots) => + { + var fake = slots[1]; + var fakeScore = DiscriminatorForwardBatched(fake, isTraining: true); + var axes = Enumerable.Range(0, realScore.Shape.Length).ToArray(); + var w = Engine.TensorSubtract( + Engine.ReduceMean(fakeScore, axes, keepDims: false), + Engine.ReduceMean(realScore, axes, keepDims: false)); + var gpEx = ComputeGradientPenalty(slots[0], fake); + return Engine.TensorAdd( + w, + Engine.TensorMultiplyScalar(gpEx, NumOps.FromDouble(gpWeightConst))); + }, + batchSize: exampleCount, + clipNorm: _options.ClipNorm, + noiseMultiplier: _computedNoiseMultiplier, + rng: _random, + out var noisedGrads); + + // Fall back to the eager per-example loop if the fused DP-SGD path + // couldn't engage (non-GPU host, compilation disabled, etc.). Uses the + // pre-existing ComputePerExampleNoisedGradients path unchanged. + if (!dpFusedRan) + { + noisedGrads = ComputePerExampleNoisedGradients(realPacked, fakePacked, discParams); + } T lossValue = lossTensor.Length > 0 ? lossTensor[0] : NumOps.Zero; // Replay must reproduce the full objective (Wasserstein + λ·GP). @@ -695,7 +738,6 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) { if (_sampler is null) return; - using var tape = new GradientTape(); var generatorLayers = new List>(); generatorLayers.AddRange(Layers); foreach (var bn in _genBNLayers) generatorLayers.Add(bn); @@ -707,6 +749,42 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) var condBatch = SampleConditionalBatchTensor(total); var genInput = Engine.TensorConcatenate([noiseBatch, condBatch], axis: 1); + // GPU-RESIDENT fast path — genInput carries both noise + cond, so the + // forward closure can slice condBatch back out of the persistent input + // tensor (correctly refreshed per step). Disc layers not in trainable set. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + int embedDim = _options.EmbeddingDimension; + int condDim = _condWidth; + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor ginp) + { + var faked = GeneratorForwardWithResidualBatched(ginp); + var act = ApplyOutputActivationsBatched(faked); + // condBatch is the tail of ginp (columns [embedDim, embedDim+condDim)). + var condFromInput = Engine.TensorSlice(ginp, [0, embedDim], [total, condDim]); + var withCond = Engine.TensorConcatenate([act, condFromInput], axis: 1); + var packed = withCond.Reshape([numPacks, _packedInputDim]); + return DiscriminatorForwardBatched(packed, false); + } + Tensor Loss(Tensor scores, Tensor _) + { + var axes = Enumerable.Range(0, scores.Shape.Length).ToArray(); + return Engine.TensorNegate(Engine.ReduceMean(scores, axes, keepDims: false)); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, genInput, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + var fakeFlat = GeneratorForwardWithResidualBatched(genInput); var fakeActivated = ApplyOutputActivationsBatched(fakeFlat); var fakeWithCond = Engine.TensorConcatenate([fakeActivated, condBatch], axis: 1); @@ -751,7 +829,11 @@ private Dictionary, Tensor> ComputePerExampleNoisedGradients( int exampleCount = Math.Max(1, realPacked.Shape[0]); var clippedSums = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); foreach (var param in discParams) - clippedSums[param] = new Tensor(param._shape); + { + var zero = new Tensor(param._shape); + Engine.TensorFill(zero, NumOps.Zero); + clippedSums[param] = zero; + } for (int example = 0; example < exampleCount; example++) { @@ -780,46 +862,48 @@ private Dictionary, Tensor> ComputePerExampleNoisedGradients( grads = tape.ComputeGradients(loss, discParams); } - double normSquared = 0; + // Global L2 norm via vectorized per-param sum(g²) + scalar accumulation. + T normSquared = NumOps.Zero; foreach (var grad in grads.Values) { - for (int i = 0; i < grad.Length; i++) - { - double value = NumOps.ToDouble(grad[i]); - normSquared += value * value; - } + var sq = Engine.TensorMultiply(grad, grad); + var perParamSum = Engine.ReduceSum(sq, axes: null, keepDims: false); + normSquared = NumOps.Add(normSquared, perParamSum.Length > 0 ? perParamSum[0] : NumOps.Zero); } - - double norm = Math.Sqrt(normSquared + 1e-12); - double clipFactor = Math.Min(1.0, _options.ClipNorm / norm); + double clipFactor = Math.Min(1.0, _options.ClipNorm / Math.Sqrt(NumOps.ToDouble(normSquared) + 1e-12)); + var clipFactorT = NumOps.FromDouble(clipFactor); foreach (var param in discParams) { if (!grads.TryGetValue(param, out var grad)) continue; - - var sum = clippedSums[param]; - for (int i = 0; i < grad.Length; i++) - { - double value = NumOps.ToDouble(sum[i]) + NumOps.ToDouble(grad[i]) * clipFactor; - sum[i] = NumOps.FromDouble(value); - } + var scaled = Engine.TensorMultiplyScalar(grad, clipFactorT); + clippedSums[param] = Engine.TensorAdd(clippedSums[param], scaled); } } + // Noise + average — vectorized: TensorRandomNormalInto for the + // Gaussian tensor, TensorMultiplyScalar(sum, 1/N), TensorAdd. var noisedAverage = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); double inverseCount = 1.0 / exampleCount; - double noiseStd = _options.ClipNorm * _computedNoiseMultiplier * inverseCount; + double noiseStdD = _options.ClipNorm * _computedNoiseMultiplier * inverseCount; + var invCountT = NumOps.FromDouble(inverseCount); + var noiseStdT = NumOps.FromDouble(noiseStdD); + var zeroMean = NumOps.Zero; foreach (var param in discParams) { var sum = clippedSums[param]; - var averaged = new Tensor(sum._shape); - for (int i = 0; i < sum.Length; i++) + var scaledSum = Engine.TensorMultiplyScalar(sum, invCountT); + if (noiseStdD > 0) + { + var noise = new Tensor(sum._shape); + Engine.TensorRandomNormalInto(noise, zeroMean, noiseStdT); + noisedAverage[param] = Engine.TensorAdd(scaledSum, noise); + } + else { - double noise = NumOps.ToDouble(SampleStandardNormal()) * noiseStd; - averaged[i] = NumOps.FromDouble(NumOps.ToDouble(sum[i]) * inverseCount + noise); + noisedAverage[param] = scaledSum; } - noisedAverage[param] = averaged; } return noisedAverage; @@ -865,7 +949,9 @@ private Tensor ComputeGradientPenalty(Tensor realBatch, Tensor fakeBatc var scores = DiscriminatorForwardBatched(interpolated, true); var scoreAxes = Enumerable.Range(0, scores.Shape.Length).ToArray(); var summedScores = Engine.ReduceSum(scores, scoreAxes, keepDims: false); - var gradients = gradientTape.ComputeGradients(summedScores, [interpolated]); + // AiDotNet #1844: createGraph=true records inner backward on outer tape + // so gradient penalty actually flows to disc weights (WGAN-GP correctness). + var gradients = gradientTape.ComputeGradients(summedScores, [interpolated], createGraph: true); inputGradients = gradients.TryGetValue(interpolated, out var gradient) ? gradient : new Tensor(interpolated._shape); diff --git a/src/NeuralNetworks/SyntheticData/FinDiffGenerator.cs b/src/NeuralNetworks/SyntheticData/FinDiffGenerator.cs index 9390e90039..b470b6c3f0 100644 --- a/src/NeuralNetworks/SyntheticData/FinDiffGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/FinDiffGenerator.cs @@ -420,40 +420,101 @@ private void ComputeNoiseSchedule() private void TrainBatch(Matrix data, int startRow, int endRow) { - for (int row = startRow; row < endRow; row++) + // Cache MultiSlotFusedStep across rows so the compiled plan is built + // once and replayed via slot-data refresh for subsequent rows. See + // ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try { - int t = _random.Next(_options.NumTimesteps); - var x0 = GetRow(data, row); - var noise = CreateStandardNormalVector(_dataWidth); - - double sqrtAlphaBar = Math.Sqrt(_alphasCumprod[t]); - double sqrtOneMinusAlphaBar = Math.Sqrt(1.0 - _alphasCumprod[t]); - - // Build noisy input: xt = sqrt(alpha_bar) * x0 + sqrt(1-alpha_bar) * noise - var xt = new Vector(_dataWidth); - for (int j = 0; j < _dataWidth; j++) + var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + AiDotNet.Tensors.Engines.Compilation.OptimizerType mfsOptType = default; + float mfsLr = 0f, mfsB1 = 0f, mfsB2 = 0f, mfsEps = 0f, mfsWd = 0f; + bool fusedEligible = false; + if (trainableParams.Length > 0) { - xt[j] = NumOps.FromDouble( - sqrtAlphaBar * NumOps.ToDouble(x0[j]) + - sqrtOneMinusAlphaBar * NumOps.ToDouble(noise[j])); + fusedEligible = NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out mfsOptType, out mfsLr, out mfsB1, out mfsB2, + out mfsEps, out mfsWd, out _, out _); } - // Build target noise tensor - var targetNoise = new Tensor([_dataWidth]); - for (int j = 0; j < _dataWidth; j++) + for (int row = startRow; row < endRow; row++) { - targetNoise[j] = noise[j]; - } + int t = _random.Next(_options.NumTimesteps); + var x0 = GetRow(data, row); + var noise = CreateStandardNormalVector(_dataWidth); - // Create timestep embedding and build input tensor - var timeEmbed = CreateTimestepEmbedding(t); - int totalLen = xt.Length + timeEmbed.Length; - var input = new Tensor([totalLen]); - for (int j = 0; j < xt.Length; j++) input[j] = xt[j]; - for (int j = 0; j < timeEmbed.Length; j++) input[xt.Length + j] = timeEmbed[j]; + double sqrtAlphaBar = Math.Sqrt(_alphasCumprod[t]); + double sqrtOneMinusAlphaBar = Math.Sqrt(1.0 - _alphasCumprod[t]); - // Use Train() method (GANDALF pattern: forward -> loss -> backward -> update) - Train(input, targetNoise); + // Build noisy input: xt = sqrt(alpha_bar) * x0 + sqrt(1-alpha_bar) * noise + var xt = new Vector(_dataWidth); + for (int j = 0; j < _dataWidth; j++) + { + xt[j] = NumOps.FromDouble( + sqrtAlphaBar * NumOps.ToDouble(x0[j]) + + sqrtOneMinusAlphaBar * NumOps.ToDouble(noise[j])); + } + + // Build target noise tensor + var targetNoise = new Tensor([_dataWidth]); + for (int j = 0; j < _dataWidth; j++) + { + targetNoise[j] = noise[j]; + } + + // Create timestep embedding and build input tensor + var timeEmbed = CreateTimestepEmbedding(t); + int totalLen = xt.Length + timeEmbed.Length; + var input = new Tensor([totalLen]); + for (int j = 0; j < xt.Length; j++) input[j] = xt[j]; + for (int j = 0; j < timeEmbed.Length; j++) input[xt.Length + j] = timeEmbed[j]; + + // Preferred fused path: MultiSlotFusedStep with (packedInput, targetNoise) + // as persistent slots. Layers stay unchanged; the compiled plan replays + // the Layers forward per row with fresh slot data. + if (fusedEligible) + { + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + var slots = new[] { input, targetNoise }; + Tensor ForwardFromSlots(IReadOnlyList> s) + { + var current = s[0]; + foreach (var layer in Layers) current = layer.Forward(current); + return current; + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + var diff = Engine.TensorSubtract(pred, s[1]); + var sq = Engine.TensorMultiply(diff, diff); + var axes = Enumerable.Range(0, sq.Shape.Length).ToArray(); + return Engine.ReduceMean(sq, axes, keepDims: false); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } + + // Eager fallback: GANDALF pattern (forward -> loss -> backward -> update). + Train(input, targetNoise); + } + } + finally + { + multiSlotStep?.Dispose(); } } diff --git a/src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs b/src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs index 2fd9b067d8..1090d43335 100644 --- a/src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/GOGGLEGenerator.cs @@ -735,6 +735,37 @@ public override void Train(Tensor input, Tensor expectedOutput) SetTrainingMode(true); try { + // GPU-RESIDENT fast path — GOGGLE's ELBO + structure-regularisation + // objective compiles cleanly, and the soft adjacency A is threaded + // through via extraTensors so the fused optimizer treats it identically + // to layer-carried params. Reparameterize re-samples in the closure. + var trainableLayers = Layers.OfType>().ToList(); + var extras = GetExtraTrainableTensors().ToList(); + if (trainableLayers.Count > 0 || extras.Count > 0) + { + Tensor Fwd(Tensor inp) + { + var (m, lv) = EncoderForwardTape(inp); + var zz = ReparameterizeTape(m, lv); + return DecoderForwardTape(zz); + } + Tensor Loss(Tensor raw, Tensor tgt) + { + var (m, lv) = EncoderForwardTape(tgt); + return ComputeGoggleLossTape(raw, tgt, m, lv); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableLayers, input, expectedOutput, + forward: Fwd, computeLoss: Loss, + optimizer: _optimizer, + out T _, + extraTensors: extras)) + { + ProjectAdjacencyConstraints(); + return; + } + } + using var tape = new GradientTape(); var (mean, logVar) = EncoderForwardTape(input); var z = ReparameterizeTape(mean, logVar); diff --git a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs index 9c0352b567..79d486adec 100644 --- a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs @@ -557,6 +557,41 @@ private void TrainDiscriminatorStepBatched(Matrix data, int startRow, int end // Non-DP fast path: single batched forward+backward. var (realBatch, fakeBatch) = BuildRealAndFakeBatches(data, startRow, endRow); + // GPU-RESIDENT fast path — pack (real, fake) into a single persistent + // input tensor along axis 0 so both scores can be computed from one + // forward call, then split in the loss for BCE-real + BCE-fake. + var discLayerList = BuildDiscLayerList(); + var trainableDisc = discLayerList.OfType>().ToList(); + if (trainableDisc.Count > 0) + { + int realN = realBatch.Shape[0]; + int fakeN = fakeBatch.Shape[0]; + var stacked = Engine.TensorConcatenate([realBatch, fakeBatch], axis: 0); + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor both) => DiscriminatorForwardBatched(both, isTraining: true); + Tensor Loss(Tensor allScores, Tensor _) + { + var rShape = allScores._shape.ToArray(); rShape[0] = realN; + var fShape = allScores._shape.ToArray(); fShape[0] = fakeN; + var rStart = new int[allScores.Rank]; + var fStart = new int[allScores.Rank]; fStart[0] = realN; + var rScores = Engine.TensorSlice(allScores, rStart, rShape); + var fScores = Engine.TensorSlice(allScores, fStart, fShape); + var axes = Enumerable.Range(0, rScores.Shape.Length).ToArray(); + var lossR = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(rScores), axes, keepDims: false)); + var lossF = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(Engine.TensorNegate(fScores)), axes, keepDims: false)); + return Engine.TensorAdd(lossR, lossF); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDisc, stacked, target, + forward: Fwd, computeLoss: Loss, + optimizer: _discriminatorOptimizer, + out T _)) + { + return; + } + } + using var tape = new GradientTape(); var discParams = TapeTrainingStep.CollectParameters(BuildDiscLayerList()); @@ -608,90 +643,126 @@ private void TrainDiscriminatorStepPerExampleDPSGD(Matrix data, int startRow, var discLayerList = BuildDiscLayerList(); var discParams = TapeTrainingStep.CollectParameters(discLayerList); - // Accumulator for sum of clipped per-example gradients - var gradSum = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); - foreach (var p in discParams) + // Pre-materialize per-example (real, fake) batches so both the fused + // DP-SGD path and the replay closure below see the SAME sampled fake + // data. Fresh sampling in a replay path would train against a + // different objective than the gradients that got noised + averaged. + var perExampleReal = new List>(batchSize); + var perExampleFake = new List>(batchSize); + for (int row = startRow; row < endRow; row++) { - var zero = new Tensor(p._shape); - zero.Fill(NumOps.Zero); - gradSum[p] = zero; + var (rb, fb) = BuildRealAndFakeBatches(data, row, row + 1); + perExampleReal.Add(rb); + perExampleFake.Add(fb); } - double clipNorm = _options.ClipNorm; - double noiseStd = clipNorm * noiseMultiplier; + // Route the per-example DP-SGD gradient computation through the fused + // DpSgdFusedStep helper (Phase 4H). Each per-example forward+backward + // runs the compiled plan; clip + aggregate + noise happens in host + // code with structural enforcement of Abadi 2016's clip-BEFORE-aggregate + // order. T lossSum = NumOps.Zero; - - // Capture the EXACT per-example (real, fake) tensors that produced - // the per-example losses + clipped gradients, so the replay closure - // can reconstruct the same objective. Replay built from - // BuildRealAndFakeBatches(...startRow, endRow) draws fresh noise - // (= different fake rows), which decouples the replayed scalar loss - // from the noisedAvgGrads — the optimizer's replay would compute a - // loss tied to a different objective than the gradients it applies. - var perExampleReal = new List>(endRow - startRow); - var perExampleFake = new List>(endRow - startRow); - - for (int row = startRow; row < endRow; row++) + using var dpSgdStep = new AiDotNet.Training.DpSgdFusedStep(); + bool dpFusedRan = dpSgdStep.TryStep( + parameters: discParams, + perExampleSlotData: exIdx => new[] + { + perExampleReal[exIdx], + perExampleFake[exIdx], + }, + forward: slots => DiscriminatorForwardBatched(slots[0], isTraining: true), + computeLoss: (realScores, slots) => + { + var fakeBatch = slots[1]; + var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); + var axes = Enumerable.Range(0, realScores.Shape.Length).ToArray(); + var negFakeScores = Engine.TensorNegate(fakeScores); + var lossReal = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(realScores), axes, keepDims: false)); + var lossFake = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(negFakeScores), axes, keepDims: false)); + var lossTensor = Engine.TensorAdd(lossReal, lossFake); + if (lossTensor.Length > 0) lossSum = NumOps.Add(lossSum, lossTensor[0]); + return lossTensor; + }, + batchSize: batchSize, + clipNorm: _options.ClipNorm, + noiseMultiplier: noiseMultiplier, + rng: _random, + out var noisedAvgGrads); + + // Eager fallback: replicate the per-example loop when the fused DP-SGD + // path can't engage. Same clip-BEFORE-aggregate contract preserved via + // manual accumulation. + if (!dpFusedRan) { - var (realBatch, fakeBatch) = BuildRealAndFakeBatches(data, row, row + 1); - perExampleReal.Add(realBatch); - perExampleFake.Add(fakeBatch); - - using var tape = new GradientTape(); - var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); - var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); - - var perExampleAxes = Enumerable.Range(0, realScores.Shape.Length).ToArray(); - var negFakeScores = Engine.TensorNegate(fakeScores); - var lossReal = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(realScores), perExampleAxes, keepDims: false)); - var lossFake = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(negFakeScores), perExampleAxes, keepDims: false)); - var lossTensor = Engine.TensorAdd(lossReal, lossFake); - - if (lossTensor.Length > 0) - lossSum = NumOps.Add(lossSum, lossTensor[0]); - - var grads = tape.ComputeGradients(lossTensor, discParams); - - // GLOBAL L2 norm across all parameter gradients concatenated. - // Required by Abadi's L2-sensitivity bound — per-tensor norms - // do NOT provide the same privacy guarantee. - double globalSqSum = 0.0; - foreach (var g in grads.Values) + noisedAvgGrads = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); + var gradSum = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); + foreach (var p in discParams) { - for (int i = 0; i < g.Length; i++) - { - double v = NumOps.ToDouble(g[i]); - globalSqSum += v * v; - } + var zero = new Tensor(p._shape); + zero.Fill(NumOps.Zero); + gradSum[p] = zero; } - double globalNorm = Math.Sqrt(globalSqSum + 1e-12); - double clipFactor = Math.Min(1.0, clipNorm / globalNorm); - T clipFactorT = NumOps.FromDouble(clipFactor); - - foreach (var kvp in grads) + for (int row = startRow; row < endRow; row++) { - var scaled = Engine.TensorMultiplyScalar(kvp.Value, clipFactorT); - gradSum[kvp.Key] = Engine.TensorAdd(gradSum[kvp.Key], scaled); + var realBatch = perExampleReal[row - startRow]; + var fakeBatch = perExampleFake[row - startRow]; + using var tape = new GradientTape(); + var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); + var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); + var axes = Enumerable.Range(0, realScores.Shape.Length).ToArray(); + var negFakeScores = Engine.TensorNegate(fakeScores); + var lossReal = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(realScores), axes, keepDims: false)); + var lossFake = Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(negFakeScores), axes, keepDims: false)); + var lossTensor = Engine.TensorAdd(lossReal, lossFake); + if (lossTensor.Length > 0) lossSum = NumOps.Add(lossSum, lossTensor[0]); + var grads = tape.ComputeGradients(lossTensor, discParams); + + // GLOBAL L2 norm across ALL parameter gradients (Abadi 2016 + // sensitivity contract) — vectorized: per-param sum(g²) via + // TensorMultiply + ReduceSum (all axes), then accumulate the + // scalar tensor results. + T normSquared = NumOps.Zero; + foreach (var g in grads.Values) + { + var sq = Engine.TensorMultiply(g, g); + var perParamSum = Engine.ReduceSum(sq, axes: null, keepDims: false); + normSquared = NumOps.Add(normSquared, perParamSum.Length > 0 ? perParamSum[0] : NumOps.Zero); + } + double clipFactor = Math.Min(1.0, _options.ClipNorm / Math.Sqrt(NumOps.ToDouble(normSquared) + 1e-12)); + var clipFactorT = NumOps.FromDouble(clipFactor); + foreach (var kvp in grads) + { + var scaled = Engine.TensorMultiplyScalar(kvp.Value, clipFactorT); + gradSum[kvp.Key] = Engine.TensorAdd(gradSum[kvp.Key], scaled); + } } - } - // Add Gaussian noise to the SUM, then average by batchSize. - var noisedAvgGrads = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); - double invBatch = 1.0 / batchSize; - foreach (var kvp in gradSum) - { - var noisy = new Tensor(kvp.Value._shape); - for (int i = 0; i < kvp.Value.Length; i++) + // Noise + average — vectorized: TensorRandomNormalInto for the + // Gaussian tensor, TensorMultiplyScalar(sum, 1/B), TensorAdd. + double invBatch = 1.0 / batchSize; + double noiseStdD = _options.ClipNorm * noiseMultiplier * invBatch; + var invBatchT = NumOps.FromDouble(invBatch); + var noiseStdT = NumOps.FromDouble(noiseStdD); + var zeroMean = NumOps.Zero; + foreach (var kvp in gradSum) { - double sumVal = NumOps.ToDouble(kvp.Value[i]); - double u1 = Math.Max(1e-10, _random.NextDouble()); - double u2 = _random.NextDouble(); - double zn = Math.Sqrt(-2.0 * Math.Log(u1)) * Math.Cos(2.0 * Math.PI * u2); - noisy[i] = NumOps.FromDouble((sumVal + zn * noiseStd) * invBatch); + var scaledSum = Engine.TensorMultiplyScalar(kvp.Value, invBatchT); + if (noiseStdD > 0) + { + var noise = new Tensor(kvp.Value._shape); + Engine.TensorRandomNormalInto(noise, zeroMean, noiseStdT); + noisedAvgGrads[kvp.Key] = Engine.TensorAdd(scaledSum, noise); + } + else + { + noisedAvgGrads[kvp.Key] = scaledSum; + } } - noisedAvgGrads[kvp.Key] = noisy; } + var stackedReal = Engine.TensorConcatenate([.. perExampleReal], axis: 0); + var stackedFake = Engine.TensorConcatenate([.. perExampleFake], axis: 0); + var allAxesFull = Enumerable.Range(0, stackedReal.Shape.Length).ToArray(); T avgLoss = NumOps.Divide(lossSum, NumOps.FromDouble(batchSize)); // Replay-correct closure: each per-example lossTensor was @@ -703,9 +774,6 @@ private void TrainDiscriminatorStepPerExampleDPSGD(Matrix data, int startRow, // here would tie the replayed loss to a different objective than // the noisedAvgGrads (silent training drift on any optimizer.Step // that exercises the replay path). - var stackedReal = Engine.TensorConcatenate([.. perExampleReal], axis: 0); - var stackedFake = Engine.TensorConcatenate([.. perExampleFake], axis: 0); - var allAxesFull = Enumerable.Range(0, stackedReal.Shape.Length).ToArray(); var capturedFake = stackedFake; Tensor ComputeForward(Tensor inp, Tensor _) => DiscriminatorForwardBatched(inp, true); Tensor RecomputeLoss(Tensor predReal, Tensor _) @@ -732,8 +800,6 @@ Tensor RecomputeLoss(Tensor predReal, Tensor _) /// private void TrainGeneratorStepBatched(int batchSize) { - using var tape = new GradientTape(); - // Generator's trainable surface = decoder FCs + per-layer BN + // output projection — NOT the full Layers list (Layers also holds // the encoder, VAE heads, and discriminator sub-graph; capturing @@ -746,6 +812,32 @@ private void TrainGeneratorStepBatched(int batchSize) var genParams = TapeTrainingStep.CollectParameters(generatorLayers); var noiseBatch = GenerateNoiseBatchTensor(batchSize); + + // GPU-RESIDENT fast path — the fused plan captures the whole + // (noise → decoder → discriminator-frozen → non-saturating loss) chain. + // Discriminator layers aren't in `genParams`, so no gradients accumulate + // to them on this step. Falls through on failure. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor nb) => DiscriminatorForwardBatched(DecoderForwardBatched(nb, true), false); + Tensor Loss(Tensor scores, Tensor _) + { + var axes = Enumerable.Range(0, scores.Shape.Length).ToArray(); + return Engine.TensorNegate(Engine.ReduceMean(LogSigmoid(scores), axes, keepDims: false)); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, noiseBatch, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); var fakeBatch = DecoderForwardBatched(noiseBatch, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: false); diff --git a/src/NeuralNetworks/SyntheticData/MisGANGenerator.cs b/src/NeuralNetworks/SyntheticData/MisGANGenerator.cs index b201af4cd8..489e024662 100644 --- a/src/NeuralNetworks/SyntheticData/MisGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/MisGANGenerator.cs @@ -612,6 +612,43 @@ private void TrainDataGeneratorStep(int batchStart, int batchEnd) { // G_x is optimized against D_x using a mask drawn from G_m, so the data and // mask generators are trained jointly against the masked-data critic. + // GPU-RESIDENT: pack (dataNoise, maskNoise) into a single input; splitting + // inside the closure keeps both noise sources refreshed per replay. + var dataGenLayers = BuildDataGenLayerList(); + var trainableDataGen = dataGenLayers.OfType>().ToList(); + if (trainableDataGen.Count > 0 && AiDotNet.Training.GpuResidentFusedStep.IsGpuResidentAvailable) + { + int embed = _options.EmbeddingDimension; + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor packed) + { + var dn = Engine.TensorSlice(packed, [0], [embed]); + var mn = Engine.TensorSlice(packed, [embed], [embed]); + var fakeRow = DataGeneratorForward(dn); + var fakeMask = MaskGeneratorForward(mn); + return DataDiscriminatorForward(ElementwiseMultiplyTensor(fakeRow, fakeMask), isTraining: true); + } + Tensor Loss(Tensor score, Tensor _) => Engine.TensorNegate(ReduceToScalar(score)); + bool fusedEngaged = false; + for (int i = batchStart; i < batchEnd; i++) + { + var dn = CreateStandardNormalVector(embed); + var mn = CreateStandardNormalVector(embed); + var packed = new Vector(2 * embed); + for (int k = 0; k < embed; k++) packed[k] = dn[k]; + for (int k = 0; k < embed; k++) packed[embed + k] = mn[k]; + var packedTensor = new Tensor(new[] { 2 * embed }, packed); + bool ran = AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDataGen, packedTensor, target, + forward: Fwd, computeLoss: Loss, + optimizer: _dataGenOptimizer, + out T _); + if (!ran) { if (!fusedEngaged) break; continue; } + fusedEngaged = true; + } + if (fusedEngaged) return; + } + for (int i = batchStart; i < batchEnd; i++) { var dataNoise = VectorToTensor(CreateStandardNormalVector(_options.EmbeddingDimension)); @@ -628,6 +665,29 @@ private void TrainDataGeneratorStep(int batchStart, int batchEnd) private void TrainMaskGeneratorStep(int batchStart, int batchEnd) { + // GPU-RESIDENT: per-sample loop; fused plan replays with fresh noise. + var maskGenLayers = BuildMaskGenLayerList(); + var trainableMaskGen = maskGenLayers.OfType>().ToList(); + if (trainableMaskGen.Count > 0 && AiDotNet.Training.GpuResidentFusedStep.IsGpuResidentAvailable) + { + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor noise) => MaskDiscriminatorForward(MaskGeneratorForward(noise), isTraining: true); + Tensor Loss(Tensor score, Tensor _) => Engine.TensorNegate(ReduceToScalar(score)); + bool fusedEngaged = false; + for (int i = batchStart; i < batchEnd; i++) + { + var noiseTensor = VectorToTensor(CreateStandardNormalVector(_options.EmbeddingDimension)); + bool ran = AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableMaskGen, noiseTensor, target, + forward: Fwd, computeLoss: Loss, + optimizer: _maskGenOptimizer, + out T _); + if (!ran) { if (!fusedEngaged) break; continue; } + fusedEngaged = true; + } + if (fusedEngaged) return; + } + for (int i = batchStart; i < batchEnd; i++) { var noise = VectorToTensor(CreateStandardNormalVector(_options.EmbeddingDimension)); diff --git a/src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs index a4b931c5fb..84a3fff809 100644 --- a/src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/OCTGANGenerator.cs @@ -546,6 +546,31 @@ private void TrainDiscriminatorStep(Matrix trainData, int batchStart, int bat private void TrainGeneratorStep(int batchSize, T scaledLr) { // Generator pulls its embeddings toward the SVDD center (look real). + // GPU-RESIDENT fast path — every iteration uses the same forward graph + // (noise → gen → disc-frozen embedding → SVDD dist²), so the fused plan + // compiles once and replays across the batch with refreshed noise per step. + var generatorLayers = BuildGeneratorLayerList(); + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0 && AiDotNet.Training.GpuResidentFusedStep.IsGpuResidentAvailable) + { + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor noiseT) => DiscriminatorForward(GeneratorForward(noiseT), isTraining: false); + Tensor Loss(Tensor emb, Tensor _) => SvddDistSq(emb); + bool fusedEngaged = false; + for (int i = 0; i < batchSize; i++) + { + var noiseTensor = VectorToTensor(CreateStandardNormalVector(_options.EmbeddingDimension)); + bool ran = AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, noiseTensor, target, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _); + if (!ran) { if (!fusedEngaged) break; continue; } + fusedEngaged = true; + } + if (fusedEngaged) return; + } + for (int i = 0; i < batchSize; i++) { var noise = VectorToTensor(CreateStandardNormalVector(_options.EmbeddingDimension)); diff --git a/src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs b/src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs index c19ac24653..ae1f79c6f3 100644 --- a/src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/PATEGANGenerator.cs @@ -593,6 +593,53 @@ private void TrainStudentStep(Matrix transformedData, int batchSize, T lr) private void TrainGeneratorStep(int batchSize) { + // GPU-RESIDENT fast path — every iteration runs the same forward graph + // (noise → generator → student → BCE-to-1), so the fused plan compiles + // once on the first call and replays across the remaining iterations + // with refreshed noise per step (see CompiledTapeTrainingStep's persistent- + // input mechanism). Falls through to the per-sample eager loop on any + // failure. Student's params aren't in the trainable set so its forward + // stays frozen — only the generator's layers get moment updates. + var generatorLayers = BuildGeneratorLayerList(); + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0 && AiDotNet.Training.GpuResidentFusedStep.IsGpuResidentAvailable) + { + // The target tensor is unused (BceLoss takes a fixed scalar target=1.0) + // but TryStepWithFusedOptimizer requires one; pass a 1-element scalar + // that the loss closure ignores. + var targetPlaceholder = new Tensor(new[] { 1 }); + targetPlaceholder[0] = NumOps.One; + + Tensor ForwardG(Tensor noiseInput) + { + var noiseVec = noiseInput.ToVector(); + var fake = GeneratorForward(noiseVec); + return StudentForward(fake, isTraining: true); + } + Tensor ComputeGenLoss(Tensor studentScore, Tensor _) => BceLoss(studentScore, 1.0); + + bool fusedEngaged = false; + for (int i = 0; i < batchSize; i++) + { + var noise = CreateStandardNormalVector(_options.EmbeddingDimension); + var noiseTensor = new Tensor(new[] { noise.Length }, noise); + bool ran = AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, noiseTensor, targetPlaceholder, + forward: ForwardG, computeLoss: ComputeGenLoss, + optimizer: _generatorOptimizer, + out T _); + if (!ran) + { + // First-step compile failure → abandon and fall back to eager for the rest of the batch. + if (!fusedEngaged) break; + // Compiled earlier but this step couldn't — skip. + continue; + } + fusedEngaged = true; + } + if (fusedEngaged) return; + } + for (int i = 0; i < batchSize; i++) { var noise = CreateStandardNormalVector(_options.EmbeddingDimension); diff --git a/src/NeuralNetworks/SyntheticData/TVAEGenerator.cs b/src/NeuralNetworks/SyntheticData/TVAEGenerator.cs index 2e097c70f1..3a0c583cb8 100644 --- a/src/NeuralNetworks/SyntheticData/TVAEGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TVAEGenerator.cs @@ -613,6 +613,35 @@ private void TrainBatch(Matrix transformedData, int batchIndex, int batchSize private void ElboStep(Tensor input) { EnsureSizedForInput(input); + + // GPU-RESIDENT fast path — encoder + reparam + decoder + composite + // (recon + KL) loss all captured on the fused plan. Reparam re-samples + // per replay via the closure so training remains stochastic. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + Tensor Fwd(Tensor inp) + { + var (m, lv) = EncoderForward(inp); + var zz = Reparameterize(m, lv); + return DecoderForward(zz); + } + Tensor Loss(Tensor rawOutput, Tensor target) + { + // Re-run encoder to get fresh mean/logVar for the KL term. + var (m, lv) = EncoderForward(target); + return ComputeElboLossTape(rawOutput, target, m, lv); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableLayers, input, input, + forward: Fwd, computeLoss: Loss, + optimizer: _optimizer, + out T _)) + { + return; + } + } + using var tape = new GradientTape(); var (mean, logVar) = EncoderForward(input); var z = Reparameterize(mean, logVar); diff --git a/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs b/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs index 1345e86970..ddc59ae99a 100644 --- a/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs @@ -665,6 +665,12 @@ private void TrainBatch(Matrix numericalData, Matrix categoricalData, { if (_gaussianDiffusion is null || _multinomialDiffusion is null) return; + // MultiSlotFusedStep cache: reused across rows so the compiled plan + // is compiled once (on the first row) and replayed per subsequent row + // by refreshing slot data. See ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try + { for (int row = startRow; row < endRow; row++) { int t = _gaussianDiffusion.SampleTimestep(); @@ -695,22 +701,205 @@ private void TrainBatch(Matrix numericalData, Matrix categoricalData, catNoisy = catClean; } - // Tape-connected diffusion training step: run the denoiser forward on - // the tape, build the TabDDPM hybrid loss (ε-prediction MSE for the - // Gaussian-diffused numerical features + softmax cross-entropy for the - // multinomial-diffused categorical features, Kotelnikov et al. 2023), - // and backpropagate through the MLP + output heads + timestep - // projection in one optimizer step. The previous hand-rolled gradient - // path never backpropagated through the denoiser, so - // layer.UpdateParameters threw for want of computed gradients. + // Preferred fused path: MultiSlotFusedStep with the raw sinusoidal + // timestep encoding as a persistent slot. The learnable + // _timestepProjection stays INSIDE the compiled forward closure + // (which _plan.Step() replays per row) so its weights participate + // in the backward pass. See ooples/AiDotNet#1846. + var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + if (trainableParams.Length > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var slots = BuildTabDDPMSlots(numNoisy, actualNoise, catNoisy, catClean, t); + if (slots is not null) + { + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) + { + // s[0] numNoisy, s[1] actualNoise, s[2] catNoisy, + // s[3] catClean, s[4] rawSinusoidalTimeEmbed. + var timeEmbed = _timestepProjection is not null + ? _timestepProjection.Forward(s[4]) + : s[4]; + var (noisePred, catLogits) = DenoiserForwardFromTensors(s[0], s[2], timeEmbed); + // Concat both heads into a single output tensor so the + // fused-step signature (Tensor forward output) is + // satisfied. Loss closure splits it back. + return Engine.TensorConcatenate(new[] { noisePred, catLogits }, axis: 0); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + // Split forward output back into (noisePred, catLogits). + int noisePredLen = _numNumericalFeatures > 0 && _numericalOutputHead is not null + ? s[1].Length : 0; + int catLogitsLen = _totalCategoricalWidth > 0 && _categoricalOutputHead is not null + ? s[3].Length : 0; + Tensor noisePredT, catLogitsT; + if (noisePredLen > 0 && catLogitsLen > 0) + { + noisePredT = Engine.TensorSlice(pred, new[] { 0 }, new[] { noisePredLen }); + catLogitsT = Engine.TensorSlice(pred, new[] { noisePredLen }, new[] { catLogitsLen }); + } + else if (noisePredLen > 0) + { + noisePredT = pred; + catLogitsT = new Tensor(new[] { 0 }); + } + else + { + noisePredT = new Tensor(new[] { 0 }); + catLogitsT = pred; + } + return ComputeDiffusionLossTapeFromTensors(noisePredT, s[1], catLogitsT, s[3]); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } + } + + // Eager fallback: tape-connected diffusion training step: run the + // denoiser forward on the tape, build the TabDDPM hybrid loss + // (ε-prediction MSE for the Gaussian-diffused numerical features + + // softmax cross-entropy for the multinomial-diffused categorical + // features, Kotelnikov et al. 2023), and backpropagate through the + // MLP + output heads + timestep projection in one optimizer step. using var tape = new GradientTape(); - // Compute the timestep embedding INSIDE the tape so the learnable - // _timestepProjection participates in the backward pass. - var timeEmbed = CreateTimestepEmbeddingTensor(t); - var (predictedNoise, predictedLogits) = DenoiserForwardTensors(numNoisy, catNoisy, timeEmbed); + var timeEmbed2 = CreateTimestepEmbeddingTensor(t); + var (predictedNoise, predictedLogits) = DenoiserForwardTensors(numNoisy, catNoisy, timeEmbed2); var loss = ComputeDiffusionLossTape(predictedNoise, actualNoise, predictedLogits, catClean); BackwardAndStepOnPrecomputedLoss(tape, loss, _optimizer); } + } + finally + { + multiSlotStep?.Dispose(); + } + } + + /// Builds the raw sinusoidal timestep encoding (no learnable projection + /// applied — that runs INSIDE the compiled forward). Matches the first-half + /// of . + private Tensor BuildRawTimestepEmbedding(int timestep) + { + int dim = _options.TimestepEmbeddingDimension; + var embedding = new Vector(dim); + int halfDim = dim / 2; + for (int i = 0; i < halfDim; i++) + { + double freq = Math.Exp(-Math.Log(10000.0) * i / halfDim); + double angle = timestep * freq; + embedding[i] = NumOps.FromDouble(Math.Sin(angle)); + if (i + halfDim < dim) + embedding[i + halfDim] = NumOps.FromDouble(Math.Cos(angle)); + } + return VectorToTensor(embedding); + } + + /// Assembles the persistent-slot data list for the MultiSlotFusedStep + /// wire-up. Returns null when any required feature vector's length is zero + /// (indicates a degenerate row) so the caller falls back to eager training. + private IReadOnlyList>? BuildTabDDPMSlots( + Vector numNoisy, Vector actualNoise, + Vector catNoisy, Vector catClean, int timestep) + { + // Degenerate: no features at all — can't run the denoiser. + if (numNoisy.Length == 0 && catNoisy.Length == 0) return null; + + // Use zero-length tensors as placeholders for absent modality (matches + // the DenoiserForwardTensors branching pattern). MultiSlotFusedStep's + // shape-key includes zero-length slots so a shape change in either + // modality triggers a plan recompile. + return new[] + { + VectorToTensor(numNoisy), + VectorToTensor(actualNoise), + VectorToTensor(catNoisy), + VectorToTensor(catClean), + BuildRawTimestepEmbedding(timestep), + }; + } + + /// Tensor-input variant of . Used by the + /// MultiSlotFusedStep forward closure so the numerical/categorical/timestep inputs + /// are all persistent slots (their references are captured once at trace, their + /// data is refreshed per replay). + private (Tensor NoisePred, Tensor CatLogits) DenoiserForwardFromTensors( + Tensor numericalNoisy, Tensor categoricalNoisy, Tensor projectedTimeEmbed) + { + var timeT = projectedTimeEmbed.Rank == 1 + ? projectedTimeEmbed + : Engine.Reshape(projectedTimeEmbed, new[] { projectedTimeEmbed.Length }); + + Tensor current; + if (numericalNoisy.Length > 0 && categoricalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { numericalNoisy, categoricalNoisy, timeT }, 0); + else if (numericalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { numericalNoisy, timeT }, 0); + else if (categoricalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { categoricalNoisy, timeT }, 0); + else + current = timeT; + + foreach (var layer in Layers) + current = layer.Forward(current); + + var noisePred = _numericalOutputHead is not null && _numNumericalFeatures > 0 + ? _numericalOutputHead.Forward(current) + : new Tensor(new[] { 0 }); + var catLogits = _categoricalOutputHead is not null && _totalCategoricalWidth > 0 + ? _categoricalOutputHead.Forward(current) + : new Tensor(new[] { 0 }); + return (noisePred, catLogits); + } + + /// Tensor-input variant of — target + /// noise and target categorical values come in as slot tensors instead of Vector<T>. + private Tensor ComputeDiffusionLossTapeFromTensors( + Tensor noisePred, Tensor actualNoise, Tensor catLogits, Tensor catClean) + { + Tensor? loss = null; + + if (_numNumericalFeatures > 0 && actualNoise.Length > 0) + { + var pred = noisePred.Rank == 1 ? noisePred : Engine.Reshape(noisePred, new[] { noisePred.Length }); + var diff = Engine.TensorSubtract(pred, actualNoise); + loss = ReduceToScalar(Engine.TensorSquare(diff)); + } + + if (_totalCategoricalWidth > 0 && catClean.Length > 0) + { + var logits = catLogits.Rank == 1 ? catLogits : Engine.Reshape(catLogits, new[] { catLogits.Length }); + int offset = 0; + for (int j = 0; j < _numCategoricalFeatures; j++) + { + int numCats = _categoricalColumnWidths[j]; + if (numCats > 0) + { + var ce = SoftmaxCrossEntropy(logits, catClean, offset, numCats); + loss = loss is null ? ce : Engine.TensorAdd(loss, ce); + offset += numCats; + } + } + } + + return loss ?? ReduceToScalar(Engine.TensorSquare(VectorToTensor(new Vector(1)))); } /// diff --git a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs index 7e951b3c8f..39326a3f0e 100644 --- a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs @@ -755,6 +755,74 @@ private Vector DiffusionMLPForward(Vector latent, Vector timeEmbed) /// private void TrainVAEBatch(Matrix transformedData, int startRow, int endRow, T scaledLr) { + // Build the VAE sub-network's layer list once — same set used per row. + var vaeLayers = new List>(Layers); + if (_meanLayer is not null) vaeLayers.Add(_meanLayer); + if (_logVarLayer is not null) vaeLayers.Add(_logVarLayer); + vaeLayers.AddRange(_decoderLayers); + + // GPU-RESIDENT fast path — the ELBO closure is identical per row, so + // the fused plan compiles on the first row and replays across the batch + // with refreshed input per iteration. + var trainableVaeLayers = vaeLayers.OfType>().ToList(); + if (trainableVaeLayers.Count > 0 && AiDotNet.Training.GpuResidentFusedStep.IsGpuResidentAvailable) + { + // Closure-captured (mean, logVar) from Fwd's single encoder pass — + // reused by Loss so the encoder doesn't run twice per row. Since input + // and target are the same tensor for VAE ELBO training, the encoder + // output is genuinely the same value in both closures. Fwd/Loss ordering + // is guaranteed by the fused-step contract. + Tensor? capturedMean = null; + Tensor? capturedLogVar = null; + Tensor Fwd(Tensor inp) + { + var enc = EncoderForwardOnTape(inp); + var (m, lv) = SplitEncoderOutput(enc); + capturedMean = m; + capturedLogVar = lv; + var zz = Reparameterize(m, lv); + return DecoderForward(zz); + } + Tensor Loss(Tensor rawOutput, Tensor target) + { + var m = capturedMean; + var lv = capturedLogVar; + if (m is null || lv is null) + throw new InvalidOperationException( + "TabSyn VAE fused step: encoder outputs were not captured by Fwd. " + + "This indicates the fused-step framework called the loss closure before " + + "the forward closure, violating its documented Fwd-then-Loss ordering."); + return ComputeElboLossTape(rawOutput, target, m, lv); + } + // On a mid-batch fused failure after fusedEngaged, drop out of the + // fused loop and let the eager path below reprocess the remaining + // rows — no row is silently skipped. + bool fusedEngaged = false; + int nextEagerRow = startRow; + for (int row = startRow; row < endRow; row++) + { + var inputTensor = VectorToTensor(GetRow(transformedData, row)); + bool ran = AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableVaeLayers, inputTensor, inputTensor, + forward: Fwd, computeLoss: Loss, + optimizer: _optimizer, + out T _); + if (!ran) + { + if (!fusedEngaged) break; + // Fused engaged then failed mid-batch. Resume this row and the + // remainder on the eager path so no gradient is lost. + nextEagerRow = row; + break; + } + fusedEngaged = true; + nextEagerRow = row + 1; + } + if (fusedEngaged && nextEagerRow >= endRow) return; + // Adjust the eager loop's start so it picks up where fused stopped. + startRow = nextEagerRow; + } + for (int row = startRow; row < endRow; row++) { var inputTensor = VectorToTensor(GetRow(transformedData, row)); @@ -766,11 +834,6 @@ private void TrainVAEBatch(Matrix transformedData, int startRow, int endRow, var rawOutput = DecoderForward(z); var loss = ComputeElboLossTape(rawOutput, inputTensor, mean, logVar); - // Step only the VAE sub-network (encoder + heads + decoder). - var vaeLayers = new List>(Layers); - if (_meanLayer is not null) vaeLayers.Add(_meanLayer); - if (_logVarLayer is not null) vaeLayers.Add(_logVarLayer); - vaeLayers.AddRange(_decoderLayers); TapeStepOver(tape, loss, vaeLayers); } } @@ -878,19 +941,88 @@ private void TrainDiffusionBatch(Matrix latentCodes, int startRow, int endRow { if (_latentDiffusion is null) return; - for (int row = startRow; row < endRow; row++) + // MultiSlotFusedStep cache: reused across rows so the compiled plan is + // built once (on the first row) and replayed per subsequent row by + // refreshing slot data. See ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try { - int t = _latentDiffusion.SampleTimestep(); - var clean = GetRow(latentCodes, row); - var (noisy, actualNoise) = _latentDiffusion.AddNoise(clean, t); - var timeEmbed = CreateTimestepEmbedding(t); + var trainableParams = Training.TapeTrainingStep.CollectParameters(_diffMLPLayers).ToArray(); + AiDotNet.Tensors.Engines.Compilation.OptimizerType mfsOptType = default; + float mfsLr = 0f, mfsB1 = 0f, mfsB2 = 0f, mfsEps = 0f, mfsWd = 0f; + bool fusedEligible = false; + if (trainableParams.Length > 0) + { + fusedEligible = NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out mfsOptType, out mfsLr, out mfsB1, out mfsB2, + out mfsEps, out mfsWd, out _, out _); + } - using var tape = new GradientTape(); - var predictedNoise = DiffusionMLPForwardOnTape(noisy, timeEmbed); - // ε-prediction MSE (Ho et al. 2020) on the VAE latent codes (Zhang et al. 2024 TabSyn). - var diff = Engine.TensorSubtract(predictedNoise, VectorToTensor(actualNoise)); - var loss = ReduceToScalar(Engine.TensorSquare(diff)); - TapeStepOver(tape, loss, _diffMLPLayers); + for (int row = startRow; row < endRow; row++) + { + int t = _latentDiffusion.SampleTimestep(); + var clean = GetRow(latentCodes, row); + var (noisy, actualNoise) = _latentDiffusion.AddNoise(clean, t); + var timeEmbed = CreateTimestepEmbedding(t); + + if (fusedEligible) + { + // Slots: [noisyLatent, actualNoise, projectedTimeEmbed]. + // _timestepProjection is NOT in _diffMLPLayers (parity with the + // eager path — it's detached there too), so we precompute the + // projected embedding host-side and pass it as slot data. + var slots = new[] + { + VectorToTensor(noisy), + VectorToTensor(actualNoise), + VectorToTensor(timeEmbed), + }; + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) + { + // s[0] noisyLatent, s[1] actualNoise (unused here — read + // in Loss), s[2] projectedTimeEmbed. Concat + MLP forward. + int totalLen = s[0].Length + s[2].Length; + var input = Engine.TensorConcatenate(new[] { s[0], s[2] }, axis: 0); + var current = input; + foreach (var layer in _diffMLPLayers) current = layer.Forward(current); + return current.Rank == 1 ? current : Engine.Reshape(current, new[] { current.Length }); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + var diff = Engine.TensorSubtract(pred, s[1]); + return ReduceToScalar(Engine.TensorSquare(diff)); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } + + using var tape = new GradientTape(); + var predictedNoise = DiffusionMLPForwardOnTape(noisy, timeEmbed); + // ε-prediction MSE (Ho et al. 2020) on the VAE latent codes (Zhang et al. 2024 TabSyn). + var diff2 = Engine.TensorSubtract(predictedNoise, VectorToTensor(actualNoise)); + var loss = ReduceToScalar(Engine.TensorSquare(diff2)); + TapeStepOver(tape, loss, _diffMLPLayers); + } + } + finally + { + multiSlotStep?.Dispose(); } } diff --git a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs index d9d4f7e60e..57000481fa 100644 --- a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs @@ -434,8 +434,76 @@ private void TrainDiscriminatorStepBatched(Tensor realBatch, Tensor noiseB var fakeBatch = GeneratorForwardBatched(noiseBatch); fakeBatch = ApplyOutputActivationsBatched(fakeBatch); - using var tape = new GradientTape(); + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realBatch, + fakeBatch: fakeBatch, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). + var trainableDiscLayers = _discLayers.OfType>().ToList(); + if (trainableDiscLayers.Count > 0) + { + int realN = realBatch.Shape[0]; + int fakeN = fakeBatch.Shape[0]; + var stacked = Engine.TensorConcatenate([realBatch, fakeBatch], axis: 0); + var target = new Tensor(new[] { 1 }); + Tensor Fwd(Tensor both) => DiscriminatorForwardBatched(both, isTraining: true); + Tensor Loss(Tensor allScores, Tensor _) + { + var rShape = allScores._shape.ToArray(); rShape[0] = realN; + var fShape = allScores._shape.ToArray(); fShape[0] = fakeN; + var rStart = new int[allScores.Rank]; + var fStart = new int[allScores.Rank]; fStart[0] = realN; + var rScores = Engine.TensorSlice(allScores, rStart, rShape); + var fScores = Engine.TensorSlice(allScores, fStart, fShape); + var axes = Enumerable.Range(0, rScores.Shape.Length).ToArray(); + var wasserstein = Engine.TensorSubtract( + Engine.ReduceMean(fScores, axes, keepDims: false), + Engine.ReduceMean(rScores, axes, keepDims: false)); + var gp = ComputeGradientPenalty(realBatch, fakeBatch); + return Engine.TensorAdd(wasserstein, + Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableDiscLayers, stacked, target, + forward: Fwd, computeLoss: Loss, + optimizer: _discriminatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); @@ -515,7 +583,9 @@ private Tensor ComputeGradientPenalty(Tensor realBatch, Tensor fakeBatc var scores = DiscriminatorForwardBatched(interpolated, true); var scoreAxes = Enumerable.Range(0, scores.Shape.Length).ToArray(); var summedScores = Engine.ReduceSum(scores, scoreAxes, keepDims: false); - var gradients = gradientTape.ComputeGradients(summedScores, [interpolated]); + // AiDotNet #1844: createGraph=true records inner backward on outer tape + // so gradient penalty actually flows to disc weights (WGAN-GP correctness). + var gradients = gradientTape.ComputeGradients(summedScores, [interpolated], createGraph: true); inputGradients = gradients.TryGetValue(interpolated, out var gradient) ? gradient : new Tensor(interpolated._shape); @@ -546,14 +616,32 @@ private Tensor ComputeGradientPenalty(Tensor realBatch, Tensor fakeBatc /// private void TrainGeneratorStepBatched(Tensor noiseBatch, Tensor realBatch) { - using var tape = new GradientTape(); - // Generator's trainable surface = generator FC layers + their BN layers. var generatorLayers = new List>(); generatorLayers.AddRange(Layers); foreach (var bn in _genBNLayers) generatorLayers.Add(bn); var genParams = TapeTrainingStep.CollectParameters(generatorLayers); + // GPU-RESIDENT fast path — noise → gen → activation → composite loss + // (fake-scores + info-loss + optional classification). Fused SGD/Adam + // step; disc/classifier layers not in trainable set so their weights stay frozen. + var trainableGenLayers = generatorLayers.OfType>().ToList(); + if (trainableGenLayers.Count > 0) + { + Tensor Fwd(Tensor nb) => ApplyOutputActivationsBatched(GeneratorForwardBatched(nb)); + Tensor Loss(Tensor fakeActivated, Tensor real) => ComputeGeneratorLoss(fakeActivated, real); + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableGenLayers, noiseBatch, realBatch, + forward: Fwd, computeLoss: Loss, + optimizer: _generatorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + var fakeBatch = GeneratorForwardBatched(noiseBatch); var fakeActivated = ApplyOutputActivationsBatched(fakeBatch); var lossTensor = ComputeGeneratorLoss(fakeActivated, realBatch); diff --git a/src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs b/src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs index dd1928d834..dac78d83d4 100644 --- a/src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TimeGANGenerator.cs @@ -586,8 +586,6 @@ private void TrainEmbeddingStepBatched(List> sequences, int startIdx, var xBatch = BuildFlattenedSequenceBatch(sequences, startIdx, endIdx); if (xBatch.Shape[0] == 0) return; - using var tape = new GradientTape(); - var embedderRecoveryLayers = new List>(); embedderRecoveryLayers.AddRange(_embedderLayers); if (_embedderOutput is not null) embedderRecoveryLayers.Add(_embedderOutput); @@ -595,6 +593,31 @@ private void TrainEmbeddingStepBatched(List> sequences, int startIdx, if (_recoveryOutput is not null) embedderRecoveryLayers.Add(_recoveryOutput); var paramsList = TapeTrainingStep.CollectParameters(embedderRecoveryLayers); + // GPU-RESIDENT fast path — Phase 1 embedder + recovery reconstruction. + var trainableEmbRec = embedderRecoveryLayers.OfType>().ToList(); + if (trainableEmbRec.Count > 0) + { + Tensor Fwd(Tensor x) => RecoveryForwardBatched(EmbedderForwardBatched(x, true), true); + Tensor Loss(Tensor r, Tensor x) + { + var d = Engine.TensorSubtract(r, x); + var s = Engine.TensorMultiply(d, d); + var axes = Enumerable.Range(0, s.Shape.Length).ToArray(); + var m = Engine.ReduceMean(s, axes, keepDims: false); + return Engine.TensorMultiplyScalar(m, NumOps.FromDouble(_options.ReconstructionWeight)); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableEmbRec, xBatch, xBatch, + forward: Fwd, computeLoss: Loss, + optimizer: _embedderOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + var hBatch = EmbedderForwardBatched(xBatch, isTraining: true); var rBatch = RecoveryForwardBatched(hBatch, isTraining: true); @@ -634,8 +657,6 @@ private void TrainSupervisedStepBatched(List> sequences, int startIdx, var (xt, xtNext) = BuildPairedSequenceBatch(sequences, startIdx, endIdx); if (xt.Shape[0] == 0) return; - using var tape = new GradientTape(); - var supervisorLayers = new List>(); supervisorLayers.AddRange(_supervisorLayers); if (_supervisorOutput is not null) supervisorLayers.Add(_supervisorOutput); @@ -645,6 +666,31 @@ private void TrainSupervisedStepBatched(List> sequences, int startIdx, var ht = EmbedderForwardBatched(xt, isTraining: false); var htNext = EmbedderForwardBatched(xtNext, isTraining: false); + // GPU-RESIDENT fast path — supervisor's next-step prediction. Embedder + // is frozen; supervisor's layers are the only trainable set here. + var trainableSup = supervisorLayers.OfType>().ToList(); + if (trainableSup.Count > 0) + { + Tensor Fwd(Tensor h) => SupervisorForwardBatched(h, isTraining: true); + Tensor Loss(Tensor pred, Tensor tgt) + { + var d = Engine.TensorSubtract(pred, tgt); + var s = Engine.TensorMultiply(d, d); + var axes = Enumerable.Range(0, s.Shape.Length).ToArray(); + return Engine.ReduceMean(s, axes, keepDims: false); + } + if (AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableSup, ht, htNext, + forward: Fwd, computeLoss: Loss, + optimizer: _supervisorOptimizer, + out T _)) + { + return; + } + } + + using var tape = new GradientTape(); + var htPred = SupervisorForwardBatched(ht, isTraining: true); var diff = Engine.TensorSubtract(htPred, htNext); var sq = Engine.TensorMultiply(diff, diff); diff --git a/src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs b/src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs index bd98833944..cacfda86c4 100644 --- a/src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs +++ b/src/NeuralNetworks/Tasks/Graph/GraphClassificationModel.cs @@ -733,6 +733,29 @@ public override void Train(Tensor input, Tensor expectedOutput) $"GraphClassificationModel tape-based training requires a LossFunctionBase (one that " + $"implements ComputeTapeLoss); the configured loss '{_lossFunction.GetType().Name}' is " + "not tape-differentiable. Supply a LossFunctionBase-derived loss such as CrossEntropyWithLogitsLoss."); + + // GPU-RESIDENT fast path (float + DirectGpuTensorEngine + compilation + Adam-family + // optimizer). Routes forward + backward + optimizer step through the compiled fused plan + // so weights / activations / moments stay resident on the device — no per-op host<->device + // round-trip. Falls through to the eager tape+optimizer path on any failure (unsupported + // optimizer, non-compilable graph, etc). Mirrors NeuralNetworkBase.TrainWithFusedStep and + // TimeSeriesModelBase.TryFusedResidentStep — same seam, applied per Train() call because + // GraphClassificationModel.Train is a single-batch entry (not a training loop). + var trainableLayers = Layers + .Where(l => l is ITrainableLayer).Cast>() + .ToList(); + if (CanTrainOnGpu + && AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableLayers, input, expectedOutput, + forward: Forward, + computeLoss: tapeLoss.ComputeTapeLoss, + optimizer: _optimizer, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + using (var tape = new GradientTape()) { var logits = Forward(input); @@ -751,6 +774,7 @@ public override void Train(Tensor input, Tensor expectedOutput) } } + /// /// Gets metadata about this model for serialization and identification. /// diff --git a/src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs b/src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs index 643de9ff8c..01de35e570 100644 --- a/src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs +++ b/src/NeuralNetworks/Tasks/Graph/LinkPredictionModel.cs @@ -725,6 +725,25 @@ public override void Train(Tensor input, Tensor expectedOutput) $"LinkPredictionModel tape-based training requires a LossFunctionBase (one that " + $"implements ComputeTapeLoss); the configured loss '{_lossFunction.GetType().Name}' is " + "not tape-differentiable. Supply a LossFunctionBase-derived loss such as BinaryCrossEntropyLoss."); + + // GPU-RESIDENT fast path — same seam as GraphClassificationModel and + // TimeSeries deep forecasters. Compiled forward + backward + fused Adam + // step on a device-resident plan; falls back to eager tape+optimizer. + var trainableLayers = Layers + .Where(l => l is ITrainableLayer).Cast>() + .ToList(); + if (CanTrainOnGpu + && AiDotNet.Training.GpuResidentFusedStep.TryStep( + trainableLayers, input, expectedOutput, + forward: Forward, + computeLoss: tapeLoss.ComputeTapeLoss, + optimizer: _optimizer, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + using (var tape = new GradientTape()) { var predictions = Forward(input); diff --git a/src/NeuralNetworks/WGANGP.cs b/src/NeuralNetworks/WGANGP.cs index e67e518f82..296521183c 100644 --- a/src/NeuralNetworks/WGANGP.cs +++ b/src/NeuralNetworks/WGANGP.cs @@ -355,6 +355,48 @@ public WGANGP( { Critic.SetTrainingMode(true); + // Preferred fused path: WganGpFusedStep runs the full WGAN-GP critic + // objective (Wasserstein + λ·GP with createGraph=true GP) in one + // compiled plan. Bypasses the legacy flat-vector round-trip (which + // needs to Predict the critic three times per step, extract flat + // gradients, combine host-side, then apply). See ooples/AiDotNet#1845. + var criticDiscParams = Training.TapeTrainingStep.CollectParameters(Critic.Layers); + if (criticDiscParams.Count > 0 + && TryMapToFusedOptimizerConfig( + _criticOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) + { + Tensor current = inp; + foreach (var layer in Critic.Layers) + current = layer.Forward(current); + return current; + } + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: criticDiscParams, + realBatch: realImages, + fakeBatch: fakeImages, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _gradientPenaltyCoefficient, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T fusedLoss)) + { + return (fusedLoss, NumOps.Zero); + } + } + // Forward pass on real images to compute scores using vectorized reduction var realScores = Critic.Predict(realImages); T realScore = NumOps.Divide(Engine.TensorSum(realScores), NumOps.FromDouble(batchSize)); diff --git a/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs index 3f206cfdb7..d1e8bbbb69 100644 --- a/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs +++ b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs @@ -680,28 +680,63 @@ private void TapeTrainStep(Tensor input, Tensor expectedOutput) int[] spatialShape = input._shape.Skip(2).ToArray(); - // Collect every trainable parameter tensor across Layers (lift + - // project DenseLayers) and _fourierLayers. Cached between calls — - // layer structure is stable after construction and parameter - // tensors are updated in place by SetParameters. - var paramList = new List>(); - foreach (var layer in Layers) + // GPU-RESIDENT fast path — compiled fused-Adam step on the whole + // lift → fourier stack → project pipeline. FNO's forward is a + // straight tape-tracked chain, so the compiled plan captures it + // cleanly. Falls through to the SGD-in-place path below on any + // failure. Uses SGD to match this method's baseline behavior. + if (CanTrainOnGpu) { - if (layer is ITrainableLayer trainable) + // _fourierLayers are already registered into Layers (see the + // InitializeArchitecture Layers.Add(fourierLayer) call at construction), + // so iterating Layers alone collects them exactly once — a second + // pass over _fourierLayers would double-register each Fourier + // parameter and let the fused optimizer's moment buffers apply + // the update twice per step. + var allTrainable = new List>(); + foreach (var l in Layers) if (l is ITrainableLayer t) allTrainable.Add(t); + if (allTrainable.Count > 0) { - var layerParams = trainable.GetTrainableParameters(); - if (layerParams is not null) + Tensor Forward(Tensor inp) { - foreach (var p in layerParams) - { - if (p is not null && p.Length > 0) paramList.Add(p); - } + var sSh = inp._shape.Skip(2).ToArray(); + var lf = FlattenPointwiseInput(inp, sSh); + var lo = liftLayer.Forward(lf); + var lifted = UnflattenPointwiseOutput(lo, inp.Shape[0], sSh); + Tensor y = lifted; + foreach (var fl in _fourierLayers) y = fl.Forward(y); + var pf = FlattenPointwiseInput(y, sSh); + var po = projectLayer.Forward(pf); + return UnflattenPointwiseOutput(po, inp.Shape[0], sSh); + } + Tensor ComputeLoss(Tensor pred, Tensor tgt) + { + var d = Engine.TensorSubtract(pred, tgt); + var s = Engine.TensorMultiply(d, d); + var ss = Engine.ReduceSum(s, null); + return Engine.TensorMultiplyScalar(ss, NumOps.Divide(NumOps.One, NumOps.FromDouble(pred.Length))); + } + if (AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + allTrainable, input, expectedOutput, + forward: Forward, computeLoss: ComputeLoss, + optimizerType: AiDotNet.Tensors.Engines.Compilation.OptimizerType.SGD, + learningRate: 0.001f, beta1: 0.9f, beta2: 0.999f, epsilon: 1e-8f, weightDecay: 0f, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; } } } - foreach (var fl in _fourierLayers) + + // Collect every trainable parameter tensor. _fourierLayers are already + // registered into Layers (Layers.Add(fourierLayer) at construction), so + // iterating Layers alone covers them — a second pass over _fourierLayers + // would apply the SGD update twice per step to each Fourier parameter. + var paramList = new List>(); + foreach (var layer in Layers) { - if (fl is ITrainableLayer trainable) + if (layer is ITrainableLayer trainable) { var layerParams = trainable.GetTrainableParameters(); if (layerParams is not null) diff --git a/src/Training/CompiledTapeTrainingStep.cs b/src/Training/CompiledTapeTrainingStep.cs index 91c4c76914..f54ee681fc 100644 --- a/src/Training/CompiledTapeTrainingStep.cs +++ b/src/Training/CompiledTapeTrainingStep.cs @@ -433,7 +433,8 @@ internal static bool TryStepWithFusedOptimizer( double maxGradNorm = 0.0, AiDotNet.Tensors.Engines.Compilation.LrSchedule? lrSchedule = null, IGradientBasedOptimizer, Tensor>? eagerOptimizer = null, - bool useBf16Moments = false) + bool useBf16Moments = false, + IReadOnlyList>? extraTensors = null) { lossValue = MathHelper.GetNumericOperations().Zero; // AiDotNet#1395: clear the previous-call's exception buffer so the @@ -577,7 +578,7 @@ or AiDotNet.Tensors.Engines.Compilation.OptimizerType.ScheduleFreeSGD // would otherwise drive the fused kernel's m/v buffers to update // the same parameter twice per step, breaking Adam's moment math. bool firstCollectThisLifecycle = _cachedParameters is null; - var parameters = _cachedParameters ??= CollectDeduplicatedParameters(layers); + var parameters = _cachedParameters ??= CollectDeduplicatedParametersWithExtras(layers, extraTensors); if (firstCollectThisLifecycle) RememberLayerSet(layers); // GPU-RESIDENCY (campaign M1): on the DirectGpu engine, make the parameters GPU-resident ONCE so @@ -914,6 +915,39 @@ private static Tensor[] CollectDeduplicatedParameters(IReadOnlyList + /// Layer-carried parameters plus any extra raw trainable tensors the caller + /// wants tracked by the fused optimizer (e.g. GOGGLE's soft adjacency, CLAP's + /// learned temperature — trainable state that isn't naturally an ILayer). + /// Dedup is by Tensor<T> reference across BOTH sources, so an + /// extra tensor that also happens to be layer-carried is only registered once. + /// + private static Tensor[] CollectDeduplicatedParametersWithExtras( + IReadOnlyList> layers, + IReadOnlyList>? extras) + { + var seen = new HashSet>(AiDotNet.Helpers.TensorReferenceComparer>.Instance); + var result = new List>(); + foreach (var layer in layers) + { + foreach (var p in layer.GetTrainableParameters()) + { + if (p is not null && seen.Add(p)) + result.Add(p); + } + } + if (extras is not null) + { + for (int i = 0; i < extras.Count; i++) + { + var p = extras[i]; + if (p is not null && seen.Add(p)) + result.Add(p); + } + } + return result.ToArray(); + } + // FP16 activation-storage plan cache (Tensors #558). Persistent input/target are reused across steps // (the compiled trace captures them as leaves); each step copies the current batch in, then replays. // The plan handles are held as `object?` because the mixed-precision API (Tensors PR #557) is on diff --git a/src/Training/DpSgdFusedStep.cs b/src/Training/DpSgdFusedStep.cs new file mode 100644 index 0000000000..6471e74d50 --- /dev/null +++ b/src/Training/DpSgdFusedStep.cs @@ -0,0 +1,344 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Tensors.Engines; +using AiDotNet.Tensors.Engines.Autodiff; +using AiDotNet.Tensors.Engines.Compilation; +using AiDotNet.Tensors.Helpers; +using AiDotNet.Tensors.LinearAlgebra; +using OptimizerType = AiDotNet.Tensors.Engines.Compilation.OptimizerType; + +namespace AiDotNet.Training; + +// Local mirror of AiDotNet.Tensors.Engines.Training.DpSgdFusedStep +// (Tensors PR ooples/AiDotNet.Tensors#763). Same API by design — swap when +// Tensors NuGet publishes. + +/// +/// Fused per-example DP-SGD (Abadi et al. 2016 §3, Algorithm 1). Runs each of K +/// per-example forward+backward passes through the compiled plan (with the +/// optimizer step's learning rate set to zero so weights don't drift between +/// examples), extracts each example's gradients from the parameter .Grad +/// tensors, clips against the GLOBAL parameter-vector L2 norm, sums, adds a +/// single Gaussian noise draw, averages, then applies the final aggregate +/// update to the parameters. +/// +/// Fused benefit over an eager per-example loop: each per-example +/// forward+backward runs the compiled plan (with GPU-resident parameters, +/// fused kernels, no per-op host↔device round-trip) instead of a fresh +/// non-persistent . The clip / aggregate / noise +/// steps run in host code because their control flow (per-example L2 norm +/// then min-clip against C) doesn't fit the current compiled-plan capture +/// model — but those are O(params) scalar operations, not the compute +/// bottleneck. The forward+backward per example IS the expensive part, and +/// that's what gets fused. +/// +/// Correctness contract: the clip-BEFORE-aggregate order is enforced by +/// the class's internal structure — callers can't reverse it. This preserves +/// the L2-sensitivity bound the Abadi 2016 privacy proof requires. +/// +/// Numeric type (float / double). +public sealed class DpSgdFusedStep : IDisposable +{ + /// Ambient engine — matches the codebase convention + /// (protected IEngine Engine => AiDotNetEngine.Current; on the + /// activation/optimizer/etc. bases). All tensor arithmetic in this class + /// goes through Engine so the dispatch is uniform. + private static IEngine Engine => AiDotNetEngine.Current; + + /// Cached numeric ops for T. Class-level to match the codebase + /// pattern instead of threading INumericOperations<T> through + /// method signatures. + private static readonly INumericOperations Ops = MathHelper.GetNumericOperations(); + + // Persistent per-example slots. Refreshed with each example's data before + // running plan.Step(). The plan captured these references at trace time. + private Tensor[]? _persistentSlots; + private ICompiledTrainingPlan? _plan; + private int[]? _cachedShapeKey; + private object?[]? _cachedParamIdentities; + private Tensor[]? _cachedParameters; + private bool _disposed; + + /// Whether the fused-resident DP-SGD path is available on this + /// thread's current engine. + public static bool IsAvailable => + typeof(T) == typeof(float) + && AiDotNetEngine.Current is DirectGpuTensorEngine gpu && gpu.SupportsGpu + && AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current.EnableCompilation; + + /// + /// Runs a full DP-SGD training step over examples. + /// + /// Trainable tensors (de-duplicated by reference). + /// Each gets its own moment buffer. + /// Callback returning the fresh slot data + /// for the given example index. Slot count and shapes must be stable across + /// example indices (a change triggers a recompile). + /// Forward closure that consumes the persistent slots + /// and returns the predicted output tensor. + /// Loss closure — pred + slots → scalar loss. + /// Number of per-example passes to run. + /// Per-example L2 norm clip (Abadi 2016's C). + /// Gaussian noise multiplier (Abadi 2016's σ). + /// Retained for API compatibility. Noise is now sampled + /// on-device via Engine.TensorRandomNormalInto so it dispatches to + /// the engine's RNG (which supports vectorized / on-device random draws). + /// The caller's rng parameter is not used by the fused path. + /// On success, the per-example-clipped + + /// noised + averaged gradients keyed by parameter tensor reference. Feed + /// directly into the caller's configured optimizer (via TapeStepContext, + /// optimizer.Step, etc.) so the caller's optimizer choice (Adam, SGD, + /// AdamW, ...) applies the DP-clean update. Empty dictionary when TryStep + /// returns false. + /// True when the fused compiled DP-SGD path ran; false to fall + /// back to eager. + public bool TryStep( + IReadOnlyList> parameters, + Func>> perExampleSlotData, + Func>, Tensor> forward, + Func, IReadOnlyList>, Tensor> computeLoss, + int batchSize, + double clipNorm, + double noiseMultiplier, + Random rng, + out Dictionary, Tensor> aggregatedGradients) + { + ThrowIfDisposed(); + aggregatedGradients = new Dictionary, Tensor>(ReferenceEqualityComparer>.Instance); + + if (!IsAvailable) return false; + if (parameters is null || parameters.Count == 0) return false; + if (perExampleSlotData is null) throw new ArgumentNullException(nameof(perExampleSlotData)); + if (forward is null) throw new ArgumentNullException(nameof(forward)); + if (computeLoss is null) throw new ArgumentNullException(nameof(computeLoss)); + if (batchSize <= 0) throw new ArgumentOutOfRangeException(nameof(batchSize)); + if (clipNorm <= 0) throw new ArgumentOutOfRangeException(nameof(clipNorm)); + if (rng is null) throw new ArgumentNullException(nameof(rng)); + + try + { + // Get the first example to establish slot shapes (before compile). + var firstExample = perExampleSlotData(0); + if (firstExample is null || firstExample.Count == 0) return false; + + int[] shapeKey = ComputeCompositeShapeKey(firstExample); + bool shapeChanged = _cachedShapeKey is null || !ShapeKeysEqual(shapeKey, _cachedShapeKey); + bool paramsChanged = ParameterSetChanged(parameters); + + if (shapeChanged || paramsChanged) + { + InvalidateCachedPlan(); + AllocatePersistentSlots(firstExample); + _cachedShapeKey = shapeKey; + RememberParameterSet(parameters); + _cachedParameters = new Tensor[parameters.Count]; + for (int i = 0; i < parameters.Count; i++) _cachedParameters[i] = parameters[i]; + + // GPU-residency for the parameters + if (typeof(T) == typeof(float) + && Environment.GetEnvironmentVariable("AIDOTNET_GPU_RESIDENT_PARAMS") != "0") + { + foreach (var p in _cachedParameters) p.Gpu(); + } + } + + if (_persistentSlots is null || _cachedParameters is null) return false; + + // Trace + compile on first Step. Plan runs forward+backward and + // applies a zero-LR SGD step so weights don't drift between + // per-example replays. The .Grad on each parameter is populated + // by the backward pass regardless of the optimizer's LR. + if (_plan is null) + { + CopySlotData(firstExample); + using var arenaSuspend = TensorArena.Suspend(); + using var scope = GraphMode.Enable(); + var pred = forward(_persistentSlots); + var loss = computeLoss(pred, _persistentSlots); + _plan = scope.CompileTraining(_cachedParameters, loss); + _plan.ConfigureOptimizer(OptimizerType.SGD, learningRate: 0.0f, beta1: 0.9f, beta2: 0.999f, eps: 1e-8f, weightDecay: 0f); + } + + // Per-example accumulators for clipped gradients — zero-init via + // vectorized TensorFill (avoids per-element scalar zero writes). + var clippedSums = new Tensor[_cachedParameters.Length]; + for (int p = 0; p < _cachedParameters.Length; p++) + { + clippedSums[p] = new Tensor((int[])_cachedParameters[p]._shape.Clone()); + Engine.TensorFill(clippedSums[p], Ops.Zero); + } + + // Per-example forward+backward via the compiled plan. Each example + // populates parameter .Grad; the L2 norm computation and clipped + // accumulation below use vectorized IEngine ops (TensorMultiply for + // element-wise square, ReduceSum for the norm reduction, + // TensorMultiplyScalar for the clip scale, TensorAdd for the + // accumulation) — no per-element scalar loops on the hot path. + for (int example = 0; example < batchSize; example++) + { + var exampleData = example == 0 ? firstExample : perExampleSlotData(example); + if (exampleData.Count != _persistentSlots.Length) return false; + CopySlotData(exampleData); + + // Run compiled forward+backward. LR=0 means weights don't update + // but .Grad on every parameter gets populated by the backward. + _plan.Step(); + + // GLOBAL L2 norm across ALL parameter gradients concatenated — + // Abadi 2016 L2-sensitivity contract. Per parameter: sum(g²) via + // TensorMultiply + ReduceSum (vectorized). Sum-of-scalars across + // parameters is a small O(paramCount) accumulation (unavoidable + // since each param has its own gradient tensor). + T normSquared = Ops.Zero; + for (int p = 0; p < _cachedParameters.Length; p++) + { + var grad = _cachedParameters[p].Grad; + if (grad is null) continue; + var sq = Engine.TensorMultiply(grad, grad); + // axes=null reduces all axes to a scalar. + var perParamSum = Engine.ReduceSum(sq, axes: null, keepDims: false); + normSquared = Ops.Add(normSquared, perParamSum.Length > 0 ? perParamSum[0] : Ops.Zero); + } + double clipFactor = Math.Min(1.0, clipNorm / Math.Sqrt(Ops.ToDouble(normSquared) + 1e-12)); + var clipFactorT = Ops.FromDouble(clipFactor); + + // Accumulate clipped per-example gradient into sums — + // vectorized: TensorMultiplyScalar + TensorAdd. + for (int p = 0; p < _cachedParameters.Length; p++) + { + var grad = _cachedParameters[p].Grad; + if (grad is null) continue; + var scaled = Engine.TensorMultiplyScalar(grad, clipFactorT); + clippedSums[p] = Engine.TensorAdd(clippedSums[p], scaled); + } + } + + // Build the DP-SGD-clean aggregated gradient dictionary using + // vectorized IEngine ops: noise is sampled on-device via + // TensorRandomNormalInto, aggregate is TensorMultiplyScalar(sum, 1/B) + // + TensorAdd(scaled, noise). Preserves the caller's optimizer + // choice (Adam / AdamW / SGD) by returning gradients instead of + // applying an SGD update here. + double invBatch = 1.0 / batchSize; + double noiseStdD = clipNorm * noiseMultiplier * invBatch; + var invBatchT = Ops.FromDouble(invBatch); + var noiseStdT = Ops.FromDouble(noiseStdD); + var zeroMean = Ops.Zero; + for (int p = 0; p < _cachedParameters.Length; p++) + { + var scaledSum = Engine.TensorMultiplyScalar(clippedSums[p], invBatchT); + Tensor noisyAvg; + if (noiseStdD > 0) + { + var noise = new Tensor(clippedSums[p]._shape); + Engine.TensorRandomNormalInto(noise, zeroMean, noiseStdT); + noisyAvg = Engine.TensorAdd(scaledSum, noise); + } + else + { + noisyAvg = scaledSum; + } + aggregatedGradients[_cachedParameters[p]] = noisyAvg; + } + + return true; + } + catch (NotSupportedException) + { + InvalidateCachedPlan(); + return false; + } + catch (InvalidOperationException) + { + InvalidateCachedPlan(); + return false; + } + } + + /// Invalidates the cached compiled plan and persistent slots. + /// Call when model structure changes; not needed when only data changes. + public void Invalidate() => InvalidateCachedPlan(); + + private void InvalidateCachedPlan() + { + _plan?.Dispose(); + _plan = null; + _persistentSlots = null; + _cachedShapeKey = null; + _cachedParamIdentities = null; + _cachedParameters = null; + } + + private void AllocatePersistentSlots(IReadOnlyList> firstExample) + { + _persistentSlots = new Tensor[firstExample.Count]; + for (int i = 0; i < firstExample.Count; i++) + { + _persistentSlots[i] = new Tensor((int[])firstExample[i]._shape.Clone()); + } + } + + private void CopySlotData(IReadOnlyList> fresh) + { + if (_persistentSlots is null) return; + for (int i = 0; i < _persistentSlots.Length; i++) + { + var slot = _persistentSlots[i]; + var src = fresh[i]; + if (src.Length != slot.Length) + throw new InvalidOperationException( + $"DpSgdFusedStep: per-example slot {i} shape changed mid-batch. " + + $"Expected {slot.Length} elements, got {src.Length}."); + src.AsSpan().CopyTo(slot.AsWritableSpan()); + } + } + + private static int[] ComputeCompositeShapeKey(IReadOnlyList> slots) + { + int total = slots.Count; + for (int i = 0; i < slots.Count; i++) total += slots[i]._shape.Length; + var key = new int[total]; + int idx = 0; + for (int i = 0; i < slots.Count; i++) + { + var s = slots[i]._shape; + for (int d = 0; d < s.Length; d++) key[idx++] = s[d]; + key[idx++] = -1 - i; + } + return key; + } + + private static bool ShapeKeysEqual(int[] a, int[] b) + { + if (a.Length != b.Length) return false; + for (int i = 0; i < a.Length; i++) if (a[i] != b[i]) return false; + return true; + } + + private bool ParameterSetChanged(IReadOnlyList> parameters) + { + if (_cachedParamIdentities is null) return true; + if (_cachedParamIdentities.Length != parameters.Count) return true; + for (int i = 0; i < parameters.Count; i++) + if (!ReferenceEquals(_cachedParamIdentities[i], parameters[i])) return true; + return false; + } + + private void RememberParameterSet(IReadOnlyList> parameters) + { + _cachedParamIdentities = new object?[parameters.Count]; + for (int i = 0; i < parameters.Count; i++) _cachedParamIdentities[i] = parameters[i]; + } + + public void Dispose() + { + if (_disposed) return; + InvalidateCachedPlan(); + _disposed = true; + } + + private void ThrowIfDisposed() + { + if (_disposed) throw new ObjectDisposedException(nameof(DpSgdFusedStep)); + } +} diff --git a/src/Training/GpuResidentFusedStep.cs b/src/Training/GpuResidentFusedStep.cs new file mode 100644 index 0000000000..429a75b041 --- /dev/null +++ b/src/Training/GpuResidentFusedStep.cs @@ -0,0 +1,139 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Interfaces; +using AiDotNet.Tensors.Engines; +using AiDotNet.Tensors.LinearAlgebra; +using OptimizerType = AiDotNet.Tensors.Engines.Compilation.OptimizerType; + +namespace AiDotNet.Training; + +/// +/// Shared entry point for GPU-resident fused training steps in classes that don't +/// inherit from NeuralNetworkBase or TimeSeriesModelBase but still want +/// their Train to route forward + backward + optimizer through a single on-device +/// compiled plan. Wraps +/// with an optimizer-config resolver (converts the model's runtime IGradientBasedOptimizer +/// into the fused-plan's + hyperparameters). +/// +/// Use pattern (mirrors NeuralNetworkBase.TrainWithFusedStep): +/// +/// if (CanTrainOnGpu && trainableLayers.Count > 0 +/// && GpuResidentFusedStep<T>.TryResolveOptimizerConfig(_optimizer, out var t, out var lr, out var b1, out var b2, out var eps, out var wd) +/// && CompiledTapeTrainingStep<T>.TryStepWithFusedOptimizer( +/// trainableLayers, input, target, Forward, loss.ComputeTapeLoss, +/// t, lr, b1, b2, eps, wd, out T fusedLoss)) +/// { +/// LastLoss = fusedLoss; +/// return; +/// } +/// // eager fallback here +/// +/// +/// Numeric type (float supported end-to-end; other types fall back to eager). +internal static class GpuResidentFusedStep +{ + /// + /// Maps the runtime optimizer instance to the fused-plan optimizer enum + hyperparameters. + /// Recognises Adam / AdamW / SGD by class name (case-insensitive). Reads the learning rate + /// from the optimizer's Options.InitialLearningRate via reflection; falls back to + /// standard defaults if the property isn't available. Returns false for optimizers the fused + /// path can't handle so callers cleanly fall through to the eager tape+optimizer route. + /// + public static bool TryResolveOptimizerConfig( + object? optimizer, + out OptimizerType type, + out float learningRate, + out float beta1, out float beta2, + out float epsilon, out float weightDecay) + { + // Sensible defaults matching AiDotNet's AdamOptimizerOptions. + type = OptimizerType.Adam; + learningRate = 1e-3f; + beta1 = 0.9f; beta2 = 0.999f; epsilon = 1e-8f; weightDecay = 0f; + + if (optimizer is null) return false; + var name = optimizer.GetType().Name; + if (name.IndexOf("AdamW", StringComparison.OrdinalIgnoreCase) >= 0) + type = OptimizerType.AdamW; + else if (name.IndexOf("Adam", StringComparison.OrdinalIgnoreCase) >= 0) + type = OptimizerType.Adam; + else if (name.IndexOf("SGD", StringComparison.OrdinalIgnoreCase) >= 0) + type = OptimizerType.SGD; + else + return false; // Unsupported optimizer — let the caller's eager path handle it. + + // Best-effort hyperparameter extraction from Options.InitialLearningRate / + // Options.Beta1 / .Beta2 / .Epsilon / .WeightDecay. Any missing property keeps the default. + try + { + var flags = System.Reflection.BindingFlags.Instance + | System.Reflection.BindingFlags.Public + | System.Reflection.BindingFlags.NonPublic; + var optsProp = optimizer.GetType().GetProperty("Options", flags); + var opts = optsProp?.GetValue(optimizer); + if (opts is null) return true; + var optsType = opts.GetType(); + var lrVal = optsType.GetProperty("InitialLearningRate")?.GetValue(opts); + if (lrVal is not null) learningRate = Convert.ToSingle(lrVal); + var b1Val = optsType.GetProperty("Beta1")?.GetValue(opts); + if (b1Val is not null) beta1 = Convert.ToSingle(b1Val); + var b2Val = optsType.GetProperty("Beta2")?.GetValue(opts); + if (b2Val is not null) beta2 = Convert.ToSingle(b2Val); + var epsVal = optsType.GetProperty("Epsilon")?.GetValue(opts); + if (epsVal is not null) epsilon = Convert.ToSingle(epsVal); + var wdVal = optsType.GetProperty("WeightDecay")?.GetValue(opts); + if (wdVal is not null) weightDecay = Convert.ToSingle(wdVal); + } + catch + { + // Keep defaults on any reflection failure. + } + return true; + } + + /// + /// One-shot: runs a fused-resident training step if all preconditions hold + /// (float, DirectGpu engine, compilation enabled, supported optimizer, at least + /// one trainable layer) and returns the loss via . + /// Returns false when the fused path can't engage (caller must run its eager fallback). + /// + public static bool TryStep( + IReadOnlyList> layers, + Tensor input, + Tensor target, + Func, Tensor> forward, + Func, Tensor, Tensor> computeLoss, + object? optimizer, + out T lossValue, + double maxGradNorm = 1.0, + IReadOnlyList>? extraTensors = null) + { + lossValue = AiDotNet.Tensors.Helpers.MathHelper.GetNumericOperations().Zero; + if (!IsGpuResidentAvailable) return false; + // A callsite with only extras (no layers) is a valid config — e.g. a model + // whose whole training surface is raw trainable tensors (learned scalars). + if ((layers is null || layers.Count == 0) && (extraTensors is null || extraTensors.Count == 0)) + return false; + if (!TryResolveOptimizerConfig(optimizer, out var type, out var lr, out var b1, out var b2, out var eps, out var wd)) + return false; + return CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + layers ?? System.Array.Empty>(), + input, target, forward, computeLoss, + type, lr, b1, b2, eps, wd, out lossValue, + maxGradNorm: maxGradNorm, + extraTensors: extraTensors); + } + + /// + /// True when the fused-resident training path is reachable on this thread's + /// current engine (mirrors TimeSeriesModelBase.CanTrainOnGpu / + /// NeuralNetworkBase.CanTrainOnGpu). Three conditions must hold: + /// T == float, the current engine is a + /// with GPU available, and graph compilation is enabled. When false, callers + /// should stay on their eager tape+optimizer path. + /// + public static bool IsGpuResidentAvailable => + typeof(T) == typeof(float) + && AiDotNetEngine.Current is DirectGpuTensorEngine gpu && gpu.SupportsGpu + && AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current.EnableCompilation; +} diff --git a/src/Training/MultiSlotFusedStep.cs b/src/Training/MultiSlotFusedStep.cs new file mode 100644 index 0000000000..9ecaf2c91b --- /dev/null +++ b/src/Training/MultiSlotFusedStep.cs @@ -0,0 +1,297 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Tensors.Engines.Autodiff; +using AiDotNet.Tensors.Engines.Compilation; +using AiDotNet.Tensors.Helpers; +using AiDotNet.Tensors.LinearAlgebra; +using OptimizerType = AiDotNet.Tensors.Engines.Compilation.OptimizerType; + +namespace AiDotNet.Training; + +// Local mirror of AiDotNet.Tensors.Engines.Training.MultiSlotFusedStep +// (Tensors PR ooples/AiDotNet.Tensors#763). Same API by design — when that PR +// merges and the AiDotNet.Tensors NuGet publishes, this file gets deleted and +// the `using AiDotNet.Training` in consumers swaps to +// `using AiDotNet.Tensors.Engines.Training`. + +/// +/// Fused-resident training step with N persistent input slots — supersedes the +/// two-slot (input + target) mechanism in higher-level wrappers when a model's +/// training loop needs to refresh more than two tensors per step. +/// +/// Motivating cases: +/// +/// Batched-per-element diffusion (Ho et al. 2020, HuggingFace +/// diffusers): each step provides [clean_sample, noise, timesteps] where +/// timesteps is a per-batch-element scalar vector. The compiled plan captures +/// references to all three; the trainer refreshes them per step so replays see +/// fresh RNG draws without recompilation. +/// Classifier-free-guidance conditional generation: each step +/// provides [latent, text_embedding, class_embedding, timesteps]. +/// TFT-style forecasters with static, historical, and future +/// covariates as separate inputs. +/// Stochastic training (CSDI, TabDDPM): timestep index and noise +/// sample are per-step data, not compile-time constants — this trainer feeds +/// them as slots. +/// +/// +/// Contract: caller allocates persistent tensors OR passes fresh +/// data per step; the trainer copies fresh data into stable-reference persistent +/// tensors before every plan.Step() so the compiled graph re-reads current data +/// without needing to recompile. Plan cache is keyed by the composite shapes of +/// ALL slots — a shape change in any slot triggers a recompile. +/// +/// Numeric type (float / double). +public sealed class MultiSlotFusedStep : IDisposable +{ + // Persistent slot tensors — captured by the compiled plan at trace time. + // Refreshed in place per step by copying fresh caller data into them. + // Reference stability is essential: the plan replays against these exact + // tensor instances, so we must never replace them (only mutate contents). + private Tensor[]? _persistentSlots; + + // Cached compiled plan. Keyed by the composite shape of all slots plus the + // trainable-layer set identity. A shape change or layer swap invalidates + // and recompiles on the next Step. + private ICompiledTrainingPlan? _plan; + private int[]? _cachedShapeKey; + private object?[]? _cachedLayerIdentities; + private Tensor[]? _cachedParameters; + + // Optimizer configuration frozen at first Step. A different optimizer + // config on a subsequent Step invalidates and reconfigures. + private (OptimizerType Type, float Lr, float B1, float B2, float Eps, float Wd)? _configuredOptimizer; + + private bool _disposed; + + /// Whether the fused-resident training step is available on this thread's + /// current engine (T == float + DirectGpu with GPU available + compilation enabled). + public static bool IsAvailable => + typeof(T) == typeof(float) + && AiDotNetEngine.Current is DirectGpuTensorEngine gpu && gpu.SupportsGpu + && AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current.EnableCompilation; + + /// + /// Runs a single training step with N persistent input slots. The trainer + /// owns the persistent tensors internally; callers pass fresh data on every + /// call and the trainer copies it into the persistent tensors before running + /// the compiled plan. + /// + /// Trainable tensors that receive optimizer updates. + /// Must be a de-duplicated list (shared/tied weights collected once). + /// Optional callback invoked before every step + /// to clear per-parameter gradient state (e.g. layer-owned .Grad buffers). + /// The Tensors library doesn't know about the caller's layer abstraction; + /// callers pass a closure that iterates their layers and calls their own + /// ZeroGrad. Pass null when no per-step grad reset is required. + /// Fresh tensor data for each slot. Length and + /// shapes must be stable across calls (a change triggers a recompile). + /// Forward closure — receives the persistent slots + /// (same references every step, refreshed data) and returns the predicted + /// output tensor. + /// Loss closure — receives the forward output and + /// the same persistent slots (for target/aux access), returns the scalar + /// loss tensor. + /// Fused optimizer kernel to configure. + /// Optimizer LR. + /// Momentum β₁ for Adam-family. + /// Momentum β₂ for Adam-family. + /// Adam ε. + /// Optimizer weight-decay (AdamW-style). + /// Scalar loss value from this step. + /// True when the fused compiled step ran successfully. Callers must + /// handle false by falling back to eager tape-based training. + public bool TryStep( + IReadOnlyList> parameters, + Action? zeroGradAction, + IReadOnlyList> freshSlotData, + Func>, Tensor> forward, + Func, IReadOnlyList>, Tensor> computeLoss, + OptimizerType optimizerType, + float learningRate, + float beta1, + float beta2, + float epsilon, + float weightDecay, + out T lossValue) + { + ThrowIfDisposed(); + lossValue = MathHelper.GetNumericOperations().Zero; + + if (!IsAvailable) return false; + if (parameters is null || parameters.Count == 0) return false; + if (freshSlotData is null || freshSlotData.Count == 0) return false; + if (forward is null) throw new ArgumentNullException(nameof(forward)); + if (computeLoss is null) throw new ArgumentNullException(nameof(computeLoss)); + + // Compute composite shape key across all slots. Rebuild persistent + // tensors and recompile the plan if the key changed OR the parameter + // set changed identity. + int[] shapeKey = ComputeCompositeShapeKey(freshSlotData); + bool shapeChanged = _cachedShapeKey is null || !ShapeKeysEqual(shapeKey, _cachedShapeKey); + bool paramsChanged = ParameterSetChanged(parameters); + bool optimizerChanged = OptimizerConfigChanged(optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + + if (shapeChanged || paramsChanged) + { + InvalidateCachedPlan(); + AllocatePersistentSlots(freshSlotData); + _cachedShapeKey = shapeKey; + RememberParameterSet(parameters); + } + + // Copy fresh data into persistent slots BEFORE the plan runs. The plan + // captured references to these tensors; refreshing their data makes the + // plan read current-step data without recompilation. + if (_persistentSlots is null) + return false; + for (int i = 0; i < _persistentSlots.Length; i++) + { + var slot = _persistentSlots[i]; + var fresh = freshSlotData[i]; + if (fresh.Length != slot.Length) + return false; + fresh.AsSpan().CopyTo(slot.AsWritableSpan()); + } + + // Cache the (already-deduplicated) parameter list. The caller + // guarantees dedup — this class doesn't know about the caller's layer + // structure so it can't do reference-based dedup on its own. + if (_cachedParameters is null) + { + _cachedParameters = new Tensor[parameters.Count]; + for (int i = 0; i < parameters.Count; i++) _cachedParameters[i] = parameters[i]; + // GPU-residency: mark parameters as GPU-resident so the compiled + // plan's optimizer branch takes the GPU Adam path (params stay + // on-device across steps). + if (typeof(T) == typeof(float) + && Environment.GetEnvironmentVariable("AIDOTNET_GPU_RESIDENT_PARAMS") != "0" + && (optimizerType == OptimizerType.Adam + || optimizerType == OptimizerType.AdamW + || optimizerType == OptimizerType.SGD)) + { + foreach (var p in _cachedParameters) p.Gpu(); + } + } + + zeroGradAction?.Invoke(); + + try + { + // Trace + compile on first call after invalidation. + if (_plan is null) + { + using var arenaSuspend = TensorArena.Suspend(); + using var scope = GraphMode.Enable(); + var pred = forward(_persistentSlots); + var loss = computeLoss(pred, _persistentSlots); + _plan = scope.CompileTraining(_cachedParameters, loss); + } + + // Configure optimizer on first Step OR when config changed. + if (optimizerChanged || _configuredOptimizer is null) + { + _plan.ConfigureOptimizer(optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + _configuredOptimizer = (optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + } + + var lossTensor = _plan.Step(); + lossValue = lossTensor.Length > 0 ? lossTensor[0] : MathHelper.GetNumericOperations().Zero; + return true; + } + catch (NotSupportedException) + { + // The compiled plan can't execute this graph on this engine. Fall + // back to eager by returning false so the caller runs their eager + // tape path. + InvalidateCachedPlan(); + return false; + } + catch (InvalidOperationException) + { + InvalidateCachedPlan(); + return false; + } + } + + /// Invalidates the cached compiled plan and persistent slots. Call + /// when the model's layer structure changes. Data-only refreshes don't need + /// this — just pass fresh data on the next Step. + public void Invalidate() => InvalidateCachedPlan(); + + private void InvalidateCachedPlan() + { + _plan?.Dispose(); + _plan = null; + _persistentSlots = null; + _cachedShapeKey = null; + _cachedLayerIdentities = null; + _cachedParameters = null; + _configuredOptimizer = null; + } + + private void AllocatePersistentSlots(IReadOnlyList> freshSlotData) + { + _persistentSlots = new Tensor[freshSlotData.Count]; + for (int i = 0; i < freshSlotData.Count; i++) + { + _persistentSlots[i] = new Tensor((int[])freshSlotData[i]._shape.Clone()); + } + } + + private static int[] ComputeCompositeShapeKey(IReadOnlyList> slots) + { + int total = slots.Count; // one separator per slot + for (int i = 0; i < slots.Count; i++) total += slots[i]._shape.Length; + var key = new int[total]; + int idx = 0; + for (int i = 0; i < slots.Count; i++) + { + var s = slots[i]._shape; + for (int d = 0; d < s.Length; d++) key[idx++] = s[d]; + key[idx++] = -1 - i; // slot separator (distinct per slot index) + } + return key; + } + + private static bool ShapeKeysEqual(int[] a, int[] b) + { + if (a.Length != b.Length) return false; + for (int i = 0; i < a.Length; i++) if (a[i] != b[i]) return false; + return true; + } + + private bool ParameterSetChanged(IReadOnlyList> parameters) + { + if (_cachedLayerIdentities is null) return true; + if (_cachedLayerIdentities.Length != parameters.Count) return true; + for (int i = 0; i < parameters.Count; i++) + if (!ReferenceEquals(_cachedLayerIdentities[i], parameters[i])) return true; + return false; + } + + private void RememberParameterSet(IReadOnlyList> parameters) + { + _cachedLayerIdentities = new object?[parameters.Count]; + for (int i = 0; i < parameters.Count; i++) _cachedLayerIdentities[i] = parameters[i]; + } + + private bool OptimizerConfigChanged(OptimizerType type, float lr, float b1, float b2, float eps, float wd) + { + if (_configuredOptimizer is null) return true; + var (cType, cLr, cB1, cB2, cEps, cWd) = _configuredOptimizer.Value; + return cType != type || cLr != lr || cB1 != b1 || cB2 != b2 || cEps != eps || cWd != wd; + } + + public void Dispose() + { + if (_disposed) return; + InvalidateCachedPlan(); + _disposed = true; + } + + private void ThrowIfDisposed() + { + if (_disposed) throw new ObjectDisposedException(nameof(MultiSlotFusedStep)); + } +} diff --git a/src/Training/WganGpFusedStep.cs b/src/Training/WganGpFusedStep.cs new file mode 100644 index 0000000000..d642d04d9f --- /dev/null +++ b/src/Training/WganGpFusedStep.cs @@ -0,0 +1,338 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Tensors.Engines; +using AiDotNet.Tensors.Engines.Autodiff; +using AiDotNet.Tensors.Engines.Compilation; +using AiDotNet.Tensors.Helpers; +using AiDotNet.Tensors.Interfaces; +using AiDotNet.Tensors.LinearAlgebra; +using OptimizerType = AiDotNet.Tensors.Engines.Compilation.OptimizerType; + +namespace AiDotNet.Training; + +// Local mirror of AiDotNet.Tensors.Engines.Training.WganGpFusedStep +// (Tensors PR ooples/AiDotNet.Tensors#763). Same API by design — swap when +// Tensors NuGet publishes. + +/// +/// Fused WGAN-GP critic training step (Gulrajani et al. 2017). Runs the +/// combined E[D(fake)] − E[D(real)] + λ·(‖∇_x̃ D(x̃)‖₂ − 1)² objective +/// through the compiled fused plan, where the gradient-penalty term's inner +/// backward is expressed as a nested-tape computation. +/// +/// Fused benefit: the outer forward + backward + optimizer step +/// are captured in ONE compiled plan. The inner ∇_x̃ D(x̃) gradient is captured +/// on the outer tape via createGraph=true — this is why the compiled +/// backward's createGraph=true support (Phase 4C on the engine side) +/// is a prerequisite. The plan replays with fresh real/fake batches per step. +/// +/// Correctness contract: the inner gradient's ops are recorded on +/// the outer tape so the outer backward can differentiate the penalty into disc +/// weights (which the pre-fix AiDotNet code failed to do — see issue #1844). +/// This class's structure enforces the createGraph=true call on the inner tape. +/// +/// Numeric type (float / double). +public sealed class WganGpFusedStep : IDisposable +{ + /// Ambient engine — matches the codebase convention + /// (protected IEngine Engine => AiDotNetEngine.Current; on the + /// activation/optimizer/etc. bases). + private static IEngine Engine => AiDotNetEngine.Current; + + /// Cached numeric ops for T. Also matches the codebase pattern: + /// class-level Ops instead of threading INumericOperations<T> + /// through method signatures. + private static readonly INumericOperations Ops = MathHelper.GetNumericOperations(); + + // Persistent per-batch slots for (real, fake). Refreshed per step. + private Tensor[]? _persistentSlots; + private ICompiledTrainingPlan? _plan; + private int[]? _cachedShapeKey; + private object?[]? _cachedParamIdentities; + private Tensor[]? _cachedParameters; + private (OptimizerType Type, float Lr, float B1, float B2, float Eps, float Wd)? _configuredOptimizer; + private bool _disposed; + + public static bool IsAvailable => + typeof(T) == typeof(float) + && AiDotNetEngine.Current is DirectGpuTensorEngine gpu && gpu.SupportsGpu + && AiDotNet.Tensors.Engines.Optimization.TensorCodecOptions.Current.EnableCompilation; + + /// + /// Runs one WGAN-GP critic step with fresh real + fake batches. + /// + /// Critic's trainable tensors (de-duplicated). + /// Fresh real batch, shape [B, ...]. + /// Fresh fake batch (same shape as realBatch). + /// Critic forward: input tensor → scalar-per-example scores. + /// Callback returning per-batch interpolation + /// weights ε ∈ [0, 1]^B, shape [B, 1, ...] broadcasting over the sample. + /// The interpolated point is x̃ = ε · real + (1−ε) · fake. + /// λ from Gulrajani 2017. + /// Fused optimizer kernel. + /// Optimizer LR. + /// Adam β₁. + /// Adam β₂. + /// Adam ε. + /// Optimizer weight decay. + /// Scalar loss (Wasserstein + λ·GP). + /// True on successful fused step, false to fall back to eager. + public bool TryStep( + IReadOnlyList> discParameters, + Tensor realBatch, + Tensor fakeBatch, + Func, Tensor> discForward, + Func> epsilonSampler, + double gradientPenaltyWeight, + OptimizerType optimizerType, + float learningRate, + float beta1, + float beta2, + float epsilon, + float weightDecay, + out T lossValue) + { + ThrowIfDisposed(); + lossValue = Ops.Zero; + + if (!IsAvailable) return false; + if (discParameters is null || discParameters.Count == 0) return false; + if (realBatch is null || fakeBatch is null) return false; + if (discForward is null) throw new ArgumentNullException(nameof(discForward)); + if (epsilonSampler is null) throw new ArgumentNullException(nameof(epsilonSampler)); + if (!ShapeMatches(realBatch, fakeBatch)) + throw new ArgumentException( + $"realBatch shape [{string.Join(",", realBatch._shape)}] does not match fakeBatch shape [{string.Join(",", fakeBatch._shape)}].", + nameof(fakeBatch)); + + try + { + int batchSize = realBatch.Shape[0]; + int[] shapeKey = ComputeShapeKey(realBatch); + bool shapeChanged = _cachedShapeKey is null || !ShapeKeysEqual(shapeKey, _cachedShapeKey); + bool paramsChanged = ParameterSetChanged(discParameters); + bool optimizerChanged = OptimizerConfigChanged(optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + + if (shapeChanged || paramsChanged) + { + InvalidateCachedPlan(); + AllocatePersistentSlots(realBatch); + _cachedShapeKey = shapeKey; + RememberParameterSet(discParameters); + _cachedParameters = new Tensor[discParameters.Count]; + for (int i = 0; i < discParameters.Count; i++) _cachedParameters[i] = discParameters[i]; + if (typeof(T) == typeof(float) + && Environment.GetEnvironmentVariable("AIDOTNET_GPU_RESIDENT_PARAMS") != "0") + { + foreach (var p in _cachedParameters) p.Gpu(); + } + } + + if (_persistentSlots is null || _cachedParameters is null) return false; + + // Copy fresh real/fake into persistent slots BEFORE running the plan. + realBatch.AsSpan().CopyTo(_persistentSlots[0].AsWritableSpan()); + fakeBatch.AsSpan().CopyTo(_persistentSlots[1].AsWritableSpan()); + var eps01 = epsilonSampler(batchSize); + if (_persistentSlots.Length < 3) + { + _persistentSlots = ExtendSlots(_persistentSlots, eps01._shape); + } + eps01.AsSpan().CopyTo(_persistentSlots[2].AsWritableSpan()); + + // Trace + compile on first call. + if (_plan is null) + { + using var arenaSuspend = TensorArena.Suspend(); + using var scope = GraphMode.Enable(); + var loss = BuildWganGpLoss(discForward, gradientPenaltyWeight); + _plan = scope.CompileTraining(_cachedParameters, loss); + } + + if (optimizerChanged || _configuredOptimizer is null) + { + _plan.ConfigureOptimizer(optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + _configuredOptimizer = (optimizerType, learningRate, beta1, beta2, epsilon, weightDecay); + } + + var lossTensor = _plan.Step(); + lossValue = lossTensor.Length > 0 ? lossTensor[0] : Ops.Zero; + return true; + } + catch (NotSupportedException) + { + InvalidateCachedPlan(); + return false; + } + catch (InvalidOperationException) + { + InvalidateCachedPlan(); + return false; + } + } + + /// + /// Builds the tape-recorded WGAN-GP loss E[D(fake)] − E[D(real)] + + /// λ·(‖∇_x̃ D(x̃)‖₂ − 1)². The inner gradient uses createGraph=true + /// so its ops record on the outer (compilation) tape, letting the compiled + /// backward differentiate the penalty into disc weights. + /// + private Tensor BuildWganGpLoss( + Func, Tensor> discForward, + double gradientPenaltyWeight) + { + if (_persistentSlots is null || _persistentSlots.Length < 3) + throw new InvalidOperationException("WganGpFusedStep: persistent slots not allocated."); + + var real = _persistentSlots[0]; + var fake = _persistentSlots[1]; + var epsilon01 = _persistentSlots[2]; + + // Wasserstein term: E[D(fake)] − E[D(real)]. + var realScores = discForward(real); + var fakeScores = discForward(fake); + var scoreAxes = System.Linq.Enumerable.Range(0, realScores.Shape.Length).ToArray(); + var wasserstein = Engine.TensorSubtract( + Engine.ReduceMean(fakeScores, scoreAxes, keepDims: false), + Engine.ReduceMean(realScores, scoreAxes, keepDims: false)); + + // Interpolated x̃ = ε · real + (1 − ε) · fake, broadcasting ε over the + // sample. epsilon01 has shape [B, 1, ...] matching the sample. + var interpolated = Engine.TensorAdd( + Engine.TensorBroadcastMultiply(epsilon01, real), + Engine.TensorBroadcastMultiply( + Engine.TensorSubtract(OnesLike(epsilon01), epsilon01), + fake)); + + // Inner gradient penalty. Run the disc forward on the interpolated point + // under a nested tape with createGraph=true so the inner backward's ops + // record on the outer tape — enabling the outer backward (this plan's + // compiled backward) to differentiate the penalty into disc weights. + Tensor penalty; + using (var innerTape = new GradientTape()) + { + var interpScores = discForward(interpolated); + var summedScores = Engine.ReduceSum(interpScores, + System.Linq.Enumerable.Range(0, interpScores.Shape.Length).ToArray(), + keepDims: false); + // createGraph=true records the inner backward's ops on the outer tape. + var innerGrads = innerTape.ComputeGradients(summedScores, new[] { interpolated }, createGraph: true); + var inputGrad = innerGrads.TryGetValue(interpolated, out var g) + ? g + : new Tensor(interpolated._shape); + + int batchSize = interpolated.Shape[0]; + int elementsPer = interpolated.Length / batchSize; + var gradReshaped = Engine.Reshape(inputGrad, new[] { batchSize, elementsPer }); + var gradSquared = Engine.TensorMultiply(gradReshaped, gradReshaped); + var normSquared = Engine.ReduceSum(gradSquared, new[] { 1 }, keepDims: false); + var norm = Engine.TensorSqrt(Engine.TensorAddScalar(normSquared, Ops.FromDouble(1e-12))); + var deviation = Engine.TensorSubtract(norm, OnesLike(norm)); + var perExPenalty = Engine.TensorMultiply(deviation, deviation); + var penaltyAxes = System.Linq.Enumerable.Range(0, perExPenalty.Shape.Length).ToArray(); + penalty = Engine.ReduceMean(perExPenalty, penaltyAxes, keepDims: false); + } + + // Total = wasserstein + λ · penalty. + var weightedPenalty = Engine.TensorMultiplyScalar(penalty, Ops.FromDouble(gradientPenaltyWeight)); + var alignedPenalty = weightedPenalty._shape.SequenceEqual(wasserstein._shape) + ? weightedPenalty + : Engine.Reshape(weightedPenalty, wasserstein._shape); + return Engine.TensorAdd(wasserstein, alignedPenalty); + } + + private static Tensor OnesLike(Tensor reference) + { + // Vectorized fill through the engine — dispatches to on-device + // fill (GPU/SIMD) instead of a per-element scalar write loop. + var ones = new Tensor(reference._shape); + Engine.TensorFill(ones, Ops.One); + return ones; + } + + public void Invalidate() => InvalidateCachedPlan(); + + private void InvalidateCachedPlan() + { + _plan?.Dispose(); + _plan = null; + _persistentSlots = null; + _cachedShapeKey = null; + _cachedParamIdentities = null; + _cachedParameters = null; + _configuredOptimizer = null; + } + + private void AllocatePersistentSlots(Tensor realBatch) + { + // Slot 0: real, Slot 1: fake, Slot 2: epsilon (allocated on first Step + // when epsilonSampler is called). + _persistentSlots = new Tensor[3]; + _persistentSlots[0] = new Tensor((int[])realBatch._shape.Clone()); + _persistentSlots[1] = new Tensor((int[])realBatch._shape.Clone()); + // Slot 2 shape depends on epsilon shape; allocated at first refresh. + _persistentSlots[2] = null!; // placeholder — replaced on first Refresh. + } + + private static Tensor[] ExtendSlots(Tensor[] slots, int[] epsShape) + { + // Called if slot 2 was allocated as a null placeholder. + slots[2] = new Tensor((int[])epsShape.Clone()); + return slots; + } + + private static bool ShapeMatches(Tensor a, Tensor b) + { + if (a.Rank != b.Rank) return false; + for (int i = 0; i < a.Rank; i++) if (a.Shape[i] != b.Shape[i]) return false; + return true; + } + + private static int[] ComputeShapeKey(Tensor t) + { + var key = new int[t.Rank]; + for (int i = 0; i < t.Rank; i++) key[i] = t.Shape[i]; + return key; + } + + private static bool ShapeKeysEqual(int[] a, int[] b) + { + if (a.Length != b.Length) return false; + for (int i = 0; i < a.Length; i++) if (a[i] != b[i]) return false; + return true; + } + + private bool ParameterSetChanged(IReadOnlyList> parameters) + { + if (_cachedParamIdentities is null) return true; + if (_cachedParamIdentities.Length != parameters.Count) return true; + for (int i = 0; i < parameters.Count; i++) + if (!ReferenceEquals(_cachedParamIdentities[i], parameters[i])) return true; + return false; + } + + private void RememberParameterSet(IReadOnlyList> parameters) + { + _cachedParamIdentities = new object?[parameters.Count]; + for (int i = 0; i < parameters.Count; i++) _cachedParamIdentities[i] = parameters[i]; + } + + private bool OptimizerConfigChanged(OptimizerType type, float lr, float b1, float b2, float eps, float wd) + { + if (_configuredOptimizer is null) return true; + var (cType, cLr, cB1, cB2, cEps, cWd) = _configuredOptimizer.Value; + return cType != type || cLr != lr || cB1 != b1 || cB2 != b2 || cEps != eps || cWd != wd; + } + + public void Dispose() + { + if (_disposed) return; + InvalidateCachedPlan(); + _disposed = true; + } + + private void ThrowIfDisposed() + { + if (_disposed) throw new ObjectDisposedException(nameof(WganGpFusedStep)); + } +}