From d7181c6e4d64ee1a8d12ba54dd4369c6b3c77b07 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 08:44:14 -0400 Subject: [PATCH 01/15] feat(training): gPU-resident fused step for non-TS single-net models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kicks off the non-time-series GPU-residency sweep. Adds a shared helper for classes that don't inherit from NeuralNetworkBase or TimeSeriesModelBase but still want their Train() to route forward + backward + optimizer through the compiled fused plan (weights / activations / Adam moments resident on-device across the whole step). * src/Training/GpuResidentFusedStep.cs — shared helper. TryResolveOptimizerConfig maps a runtime IGradientBasedOptimizer to the fused-plan OptimizerType + hyperparameters (Adam / AdamW / SGD supported via case-insensitive class-name match + reflection over Options.InitialLearningRate / Beta1 / Beta2 / Epsilon / WeightDecay). TryStep is the one-shot entry. Wired the following single-net models through the helper (mirrors NeuralNetworkBase.TrainWithFusedStep): * GraphClassificationModel (NN base + GNN + pooling + cross-entropy) — routes tape training through the fused plan when float + DirectGpu + compilation are live; falls through to the existing eager tape+optimizer loop on any failure. * LinkPredictionModel (NN base + GNN + node embeddings + BCE) — same wire-up as GraphClassificationModel. * FourierNeuralOperator (lift → FourierLayers → project, hardcoded SGD + learning rate 0.001) — captures the whole spectral-conv chain in a fused SGD plan; falls through to the in-place SGD loop below. GAN-family generators (CTGAN, DPCTGAN, PATEGAN, ...), CLAPModel (learned scalar outside layer hierarchy), AutoDiffTabGenerator (per-sample variable shape defeats plan replay) and other non-conforming shapes are queued for follow-up commits. The shared helper is the common seam so they all wire the same way once their per-step structure fits the fused-plan contract (constant shape + Adam/AdamW/SGD optimizer + ITrainableLayer-carried params). Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- .../Tasks/Graph/GraphClassificationModel.cs | 24 ++++ .../Tasks/Graph/LinkPredictionModel.cs | 19 +++ .../NeuralOperators/FourierNeuralOperator.cs | 44 +++++++ src/Training/GpuResidentFusedStep.cs | 118 ++++++++++++++++++ 4 files changed, 205 insertions(+) create mode 100644 src/Training/GpuResidentFusedStep.cs 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/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs index 3f206cfdb7..ef3e57fa4b 100644 --- a/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs +++ b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs @@ -680,6 +680,50 @@ private void TapeTrainStep(Tensor input, Tensor expectedOutput) int[] spatialShape = input._shape.Skip(2).ToArray(); + // 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) + { + var allTrainable = new List>(); + foreach (var l in Layers) if (l is ITrainableLayer t) allTrainable.Add(t); + foreach (var fl in _fourierLayers) if (fl is ITrainableLayer t) allTrainable.Add(t); + if (allTrainable.Count > 0) + { + Tensor Forward(Tensor inp) + { + 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; + } + } + } + // Collect every trainable parameter tensor across Layers (lift + // project DenseLayers) and _fourierLayers. Cached between calls — // layer structure is stable after construction and parameter diff --git a/src/Training/GpuResidentFusedStep.cs b/src/Training/GpuResidentFusedStep.cs new file mode 100644 index 0000000000..85cf31218b --- /dev/null +++ b/src/Training/GpuResidentFusedStep.cs @@ -0,0 +1,118 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Interfaces; +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). +public 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) + { + lossValue = AiDotNet.Tensors.Helpers.MathHelper.GetNumericOperations().Zero; + if (layers is null || layers.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, input, target, forward, computeLoss, + type, lr, b1, b2, eps, wd, out lossValue, + maxGradNorm: maxGradNorm); + } +} From 1cf612fe0435e1563f1a5895b20e0b42f1ff0de7 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 09:02:11 -0400 Subject: [PATCH 02/15] feat(training): gPU-resident fused step for Finance foundation forecasters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires TFC, TOTEM, CSDI, MQCNN through the compiled fused SGD plan (all four use in-place SGD with lr=0.001 as their eager path). Extends the shared helper with an inline IsGpuResidentAvailable check so callers don't need to duplicate the availability gate. * GpuResidentFusedStep — adds IsGpuResidentAvailable static property (mirrors TimeSeriesModelBase.CanTrainOnGpu / NeuralNetworkBase.CanTrainOnGpu). TryStep short-circuits when unavailable so callers stay on their eager path. * TFC — supervised + contrastive branches fused into one closure; the fused plan captures both losses in a single backward. * TOTEM — reconstruction + VQ commitment terms fused; the recompute-loss closure re-derives the commitment loss on each replay for graph parity. * CSDI — denoising-score-matching: the forward closure re-samples timestep + noise each step (via ComputeDenoisingPairTape) so replay produces fresh training data even though the plan's captured graph shape is fixed. * MQCNN — multi-quantile pinball loss captured directly via ComputeMultiQuantilePinballLossTape. All four fall through to the existing in-place SGD path on any fused-plan failure (unsupported op, non-compilable graph, etc). Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- src/Finance/Forecasting/Foundation/CSDI.cs | 30 ++++++++++++++++ src/Finance/Forecasting/Foundation/TFC.cs | 40 +++++++++++++++++++++ src/Finance/Forecasting/Foundation/TOTEM.cs | 34 ++++++++++++++++++ src/Finance/Forecasting/Neural/MQCNN.cs | 16 +++++++++ src/Training/GpuResidentFusedStep.cs | 15 ++++++++ 5 files changed, 135 insertions(+) diff --git a/src/Finance/Forecasting/Foundation/CSDI.cs b/src/Finance/Forecasting/Foundation/CSDI.cs index e49938a7ad..d422d10f8d 100644 --- a/src/Finance/Forecasting/Foundation/CSDI.cs +++ b/src/Finance/Forecasting/Foundation/CSDI.cs @@ -287,6 +287,36 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + // GPU-RESIDENT fast path — fused SGD on the denoising-score-matching + // objective. The forward closure re-samples timestep + noise per step + // (via ComputeDenoisingPairTape), so replay produces fresh noise each + // step even though the compiled plan's captured graph shape is fixed. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + Tensor ForwardDenoise(Tensor inp) + { + var (pred, _) = ComputeDenoisingPairTape(inp, target); + return pred; + } + Tensor ComputeDenoiseLoss(Tensor pred, Tensor _) + { + // Recompute the noise target to match this replay's noise sample. + var (_, eps) = ComputeDenoisingPairTape(input, target); + return loss.ComputeTapeLoss(pred, eps); + } + if (AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( + trainableLayers, input, target, + forward: ForwardDenoise, computeLoss: ComputeDenoiseLoss, + 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 (epsilonPred, epsilonTarget) = ComputeDenoisingPairTape(input, target); diff --git a/src/Finance/Forecasting/Foundation/TFC.cs b/src/Finance/Forecasting/Foundation/TFC.cs index a4a6fc66d1..f4c44a0196 100644 --- a/src/Finance/Forecasting/Foundation/TFC.cs +++ b/src/Finance/Forecasting/Foundation/TFC.cs @@ -274,6 +274,46 @@ 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. Both branches run through + // the fused forward closure so gradients accumulate into shared + // params on-device. Falls through to the in-place SGD loop below + // when the fused path can't engage. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + Tensor ForwardCombined(Tensor inp) + { + // supervised forecast head (same alignment as eager path below). + var f = ForwardForTraining(inp); + return f; + } + 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 = ComputeContrastiveLossTape(input); + 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 diff --git a/src/Finance/Forecasting/Foundation/TOTEM.cs b/src/Finance/Forecasting/Foundation/TOTEM.cs index 3fc80e563a..67e007c716 100644 --- a/src/Finance/Forecasting/Foundation/TOTEM.cs +++ b/src/Finance/Forecasting/Foundation/TOTEM.cs @@ -330,6 +330,40 @@ 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. + // Falls through to the in-place SGD loop below when the fused path + // can't engage. + var trainableLayers = Layers.OfType>().ToList(); + if (trainableLayers.Count > 0) + { + Tensor ForwardCombined(Tensor inp) + { + var (fc, _) = ForwardNativeForTrainingWithCommitment(inp); + 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) = ForwardNativeForTrainingWithCommitment(input); + 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; + return; + } + } + using var tape = new GradientTape(); var (forecast, commitmentLoss) = ForwardNativeForTrainingWithCommitment(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/Training/GpuResidentFusedStep.cs b/src/Training/GpuResidentFusedStep.cs index 85cf31218b..aaee5ddf28 100644 --- a/src/Training/GpuResidentFusedStep.cs +++ b/src/Training/GpuResidentFusedStep.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Generic; using AiDotNet.Interfaces; +using AiDotNet.Tensors.Engines; using AiDotNet.Tensors.LinearAlgebra; using OptimizerType = AiDotNet.Tensors.Engines.Compilation.OptimizerType; @@ -107,6 +108,7 @@ public static bool TryStep( double maxGradNorm = 1.0) { lossValue = AiDotNet.Tensors.Helpers.MathHelper.GetNumericOperations().Zero; + if (!IsGpuResidentAvailable) return false; if (layers is null || layers.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; @@ -115,4 +117,17 @@ public static bool TryStep( type, lr, b1, b2, eps, wd, out lossValue, maxGradNorm: maxGradNorm); } + + /// + /// 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; } From 28bc7f7a06ab52463114e4bf5b08b2e4d115651d Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 09:26:07 -0400 Subject: [PATCH 03/15] feat(training): gPU-resident fused generator step for GAN-family synthetic generators + TVAE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the generator side of the dual-network training loop through the fused compiled plan for 9 SyntheticData generators. Each covers a distinct pattern: * PATEGAN — per-sample loop; fused plan compiles on first sample, replays across the batch with fresh noise per iteration. * MedSynth — batched noise -> decoder -> discriminator-frozen -> non-saturating log-sigmoid loss. * CausalGAN — batched noise -> generator -> optional causal-structure -> output activations -> disc-frozen -> -avgFake. * OCTGAN — per-sample loop; noise -> gen -> disc-frozen embedding -> SVDD dist². * TableGAN — noise + real batch -> gen -> composite loss (fake-scores + information-loss + optional classification). * CTGAN — the tricky one: pack (noise, cond, mask) into one persistent input tensor so the closure can slice cond and mask back out on replay; captures both -avgFake AND conditional cross-entropy on the fused plan. * DPCTGAN — same pattern as CTGAN but simpler (no mask, no CE). * CopulaGAN — same pattern as DPCTGAN. * TVAE — encoder + reparam + decoder + composite ELBO (recon + KL) all fused; Reparameterize re-samples inside the closure so training stays stochastic. Discriminator/critic/student layers are NOT passed to the fused step (their weights stay frozen on the gen step, matching the eager path semantics). Falls through to the eager tape+optimizer path on any failure. The discriminator STEPS of these GANs are not yet wired — they need a separate pass because their loss involves gradient penalty (WGAN-GP) or per-example DP clipping (DPCTGAN) which don't compose cleanly with the single-closure fused step yet. Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- .../SyntheticData/CTGANGenerator.cs | 60 +++++++++++++++++-- .../SyntheticData/CausalGANGenerator.cs | 31 +++++++++- .../SyntheticData/CopulaGANGenerator.cs | 35 ++++++++++- .../SyntheticData/DPCTGANGenerator.cs | 37 +++++++++++- .../SyntheticData/MedSynthGenerator.cs | 28 ++++++++- .../SyntheticData/OCTGANGenerator.cs | 25 ++++++++ .../SyntheticData/PATEGANGenerator.cs | 47 +++++++++++++++ .../SyntheticData/TVAEGenerator.cs | 29 +++++++++ .../SyntheticData/TableGANGenerator.cs | 22 ++++++- 9 files changed, 303 insertions(+), 11 deletions(-) diff --git a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 8980d3ea7f..90dd2e8b13 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -727,8 +727,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 +739,64 @@ 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 }); + 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]); + 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) + { + // Re-run gen forward on fused input to get fakeActivated, then + // compute the CE term. Redundant with Fwd but captured on the + // same fused plan and stays on-device. + var gi = Engine.TensorSlice(fusedInput, [0, 0], [totalSamples, embedDim + condDim]); + var faked = GeneratorForwardWithResidualBatched(gi); + var act = ApplyOutputActivationsBatched(faked); + var condFromInput = Engine.TensorSlice(fusedInput, [0, embedDim], [totalSamples, condDim]); + var maskFromInput = Engine.TensorSlice(fusedInput, [0, embedDim + condDim], [totalSamples, condDim]); + 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); diff --git a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs index a1bb5f0dc2..4927e8da58 100644 --- a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs @@ -648,13 +648,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); diff --git a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs index 39ba322071..59107cef68 100644 --- a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs @@ -714,7 +714,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 +726,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); diff --git a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs index 25c61a90aa..71a0eedee9 100644 --- a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs @@ -695,7 +695,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 +706,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); diff --git a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs index 9c0352b567..27a58f0e80 100644 --- a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs @@ -732,8 +732,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 +744,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/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/TableGANGenerator.cs b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs index d9d4f7e60e..632fa76e66 100644 --- a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs @@ -546,14 +546,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); From b22e1dac1087ac6bcdc59bb43d1d65ec6c44a4a2 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 09:35:11 -0400 Subject: [PATCH 04/15] feat(training): gPU-resident fused step for TabSyn VAE + MisGAN mask/data gens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two more SyntheticData generators covered — each covers a distinct pattern: * TabSyn.TrainVAEBatch — per-row VAE ELBO training. Fused plan compiles on first row, replays across the batch with refreshed input per iteration. Reparameterize re-samples inside the closure so training stays stochastic. * MisGAN.TrainDataGeneratorStep — dual-noise generator (data noise + mask noise) packed into a single persistent input tensor; closure slices them back out on replay. Loss = -E[D_x(fakeRow ⊙ fakeMask)]. * MisGAN.TrainMaskGeneratorStep — simpler single-noise version; per-sample loop with fused plan replay across the batch. The DataDiscriminator step (WGAN-style critic + weight clipping) and mask discriminator step aren't wired — they need separate plans and the ClipWeights post-step interacts poorly with the fused optimizer's in-place update. Follow-up. Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- .../SyntheticData/MisGANGenerator.cs | 60 +++++++++++++++++++ .../SyntheticData/TabSynGenerator.cs | 45 ++++++++++++-- 2 files changed, 100 insertions(+), 5 deletions(-) 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/TabSynGenerator.cs b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs index 7e951b3c8f..7bea43425f 100644 --- a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs @@ -755,6 +755,46 @@ 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) + { + Tensor Fwd(Tensor inp) + { + var enc = EncoderForwardOnTape(inp); + var (m, lv) = SplitEncoderOutput(enc); + var zz = Reparameterize(m, lv); + return DecoderForward(zz); + } + Tensor Loss(Tensor rawOutput, Tensor target) + { + var enc = EncoderForwardOnTape(target); + var (m, lv) = SplitEncoderOutput(enc); + return ComputeElboLossTape(rawOutput, target, m, lv); + } + bool fusedEngaged = false; + 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; continue; } + fusedEngaged = true; + } + if (fusedEngaged) return; + } + for (int row = startRow; row < endRow; row++) { var inputTensor = VectorToTensor(GetRow(transformedData, row)); @@ -766,11 +806,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); } } From b2115446b3b9eea2db791a141f8ce10d16f4a244 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 09:46:37 -0400 Subject: [PATCH 05/15] feat(training): gPU-resident fused step for TimeGAN P1/P2 + MedSynth non-DP disc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * TimeGAN Phase 1 (TrainReconstructionStepBatched) — embedder + recovery reconstruction MSE. Straight single-network fused-step conversion. * TimeGAN Phase 2 (TrainSupervisedStepBatched) — supervisor's next-step MSE in the embedded space, with the embedder frozen (its layers not in the trainable set for this step). Two-tensor closure (ht, htNext). * MedSynth non-DP disc step — pack (real, fake) into a single input tensor along axis 0; slice back in the loss for BCE-real + BCE-fake. Generator runs OUTSIDE the fused plan so its weights stay frozen on the critic step. TimeGAN Phase 3 (adversarial + supervised joint) and MedSynth DP-SGD paths still use their eager tape — the WGAN-GP gradient penalty and per-example DP clipping don't compose with the single-closure fused optimizer step. Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- .../SyntheticData/MedSynthGenerator.cs | 35 ++++++++++++ .../SyntheticData/TimeGANGenerator.cs | 54 +++++++++++++++++-- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs index 27a58f0e80..b50bd63436 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()); 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); From 062a4047c3b6b0646cd05efdb67c9ab872533ac9 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 10:18:09 -0400 Subject: [PATCH 06/15] feat(training): extra-tensors parameter on fused-step API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the fused compiled-plan training step to accept raw trainable tensors that aren't naturally carried by an ITrainableLayer (e.g. GOGGLE's soft adjacency A; CLAP's learned _logTemperature scalar). Solves the SOLID concerns raised in review of the ITrainableLayer-wrapper alternative: * ISP — no forcing raw tensors to implement the full ILayer surface with no-op Forward / SetTrainingMode / GetParameterGradients members * LSP — no risk of "layer" wrappers with identity Forward diverging from the layer contract's behavioral expectations elsewhere in the codebase Wire-up: * CompiledTapeTrainingStep.TryStepWithFusedOptimizer — new optional extraTensors param. Threaded into the dedup-aware parameter collection (CollectDeduplicatedParametersWithExtras) so extras get moment buffers, gradient accumulation, and GPU-residency in exactly the same code path that layer-carried params do. * GpuResidentFusedStep.TryStep — plumbs extras through. A callsite with only extras (no layers) is a valid config for models whose whole trainable surface is raw tensors. Dedup is by Tensor reference across BOTH sources, so an extra tensor that also happens to be layer-carried is registered exactly once — the same shared/tied-weight protection the layer-only collector provides. Enables Phase 4E (GOGGLE + CLAP wire-ups). Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- src/Training/CompiledTapeTrainingStep.cs | 38 ++++++++++++++++++++++-- src/Training/GpuResidentFusedStep.cs | 14 ++++++--- 2 files changed, 46 insertions(+), 6 deletions(-) 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/GpuResidentFusedStep.cs b/src/Training/GpuResidentFusedStep.cs index aaee5ddf28..400777edd0 100644 --- a/src/Training/GpuResidentFusedStep.cs +++ b/src/Training/GpuResidentFusedStep.cs @@ -105,17 +105,23 @@ public static bool TryStep( Func, Tensor, Tensor> computeLoss, object? optimizer, out T lossValue, - double maxGradNorm = 1.0) + double maxGradNorm = 1.0, + IReadOnlyList>? extraTensors = null) { lossValue = AiDotNet.Tensors.Helpers.MathHelper.GetNumericOperations().Zero; if (!IsGpuResidentAvailable) return false; - if (layers is null || layers.Count == 0) 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, input, target, forward, computeLoss, + layers ?? System.Array.Empty>(), + input, target, forward, computeLoss, type, lr, b1, b2, eps, wd, out lossValue, - maxGradNorm: maxGradNorm); + maxGradNorm: maxGradNorm, + extraTensors: extraTensors); } /// From 7a7e671b1c81a515836d7d3d55e5f45080319dbe Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 10:21:38 -0400 Subject: [PATCH 07/15] feat(training): wire GOGGLE + CLAP through fused-step extra-tensors API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consumes Phase 4A's extraTensors parameter to route models with raw trainable state (outside the ITrainableLayer contract) through the fused compiled plan. * GOGGLE — soft adjacency A (initialised in InitializeModel, projected after each optimizer step via ProjectAdjacencyConstraints) is passed through extraTensors. The fused optimizer allocates its moment buffer, accumulates gradients, and applies the update in-place. Reparameterize re-samples inside the forward closure so training stays stochastic; ProjectAdjacencyConstraints runs after the fused step to keep the adjacency on the valid soft-adjacency manifold. * CLAP — the learned _logTemperature scalar (Radford 2021 / Wu 2023 contrastive-alignment temperature) goes through extraTensors. Both audio and text encoders are in the layers list; the fused step handles all three parameter classes uniformly. Both fall through to the existing eager tape+optimizer path on any failure (unsupported optimizer, non-compilable graph, etc). Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- src/Audio/Fingerprinting/CLAPModel.cs | 46 +++++++++++++++++++ .../SyntheticData/GOGGLEGenerator.cs | 31 +++++++++++++ 2 files changed, 77 insertions(+) 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/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); From 572ff651c97c298815a2ef1ee6ec5e89ad711465 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 10:33:13 -0400 Subject: [PATCH 08/15] feat(diffusion): batched per-element timesteps for DDPM-canonical training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds industry-standard batched-per-element timestep API to the diffusion stack (Ho et al. 2020 §3, HuggingFace diffusers reference). Previously DiffusionModelBase.Train sampled ONE timestep for the whole batch — reducing the signal-to-noise diversity of the training gradient. Canonical DDPM training samples a distinct timestep per batch element. * INoiseScheduler.AddNoiseBatched(cleanBatch, noiseBatch, timesteps) — default interface implementation delegates to the scalar AddNoise per element for backward compatibility. Concrete schedulers can override with a fused batched implementation. * NoisePredictorBase.PredictNoiseBatched(noisyBatch, timesteps, conditioning) — virtual with default slice-then-call-scalar implementation. Subclasses that want a fused batched forward override this to keep the training loop on-device. * DiffusionModelBase.PredictNoiseBatched — model-level counterpart with the same default slice-then-call-scalar behavior. * DiffusionModelBase.Train detects batched vs rank-1 input and routes through the batched noise scheduler + batched predictor when input.Rank >= 2. Rank-1 unbatched inputs stay on the scalar path for backward compatibility. This is the foundation for the industry-exceeding "batched-per-element + fused-resident" diffusion training path — see follow-up commits that wire DiffusionModelBase.Train through the fused compiled plan. Co-Authored-By: Claude Opus 4.7 --- src/Diffusion/DiffusionModelBase.cs | 84 +++++++++++++++++-- .../NoisePredictors/NoisePredictorBase.cs | 44 ++++++++++ src/Interfaces/INoiseScheduler.cs | 53 ++++++++++++ 3 files changed, 173 insertions(+), 8 deletions(-) diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index f134bffe9e..1542494d12 100644 --- a/src/Diffusion/DiffusionModelBase.cs +++ b/src/Diffusion/DiffusionModelBase.cs @@ -802,6 +802,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 +1069,52 @@ 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()); + noisySampleTensor = _scheduler.AddNoiseBatched(input, noiseBatch, timesteps); + 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); + } 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) { diff --git a/src/Diffusion/NoisePredictors/NoisePredictorBase.cs b/src/Diffusion/NoisePredictors/NoisePredictorBase.cs index bde5f8c48a..b2ca0b2d75 100644 --- a/src/Diffusion/NoisePredictors/NoisePredictorBase.cs +++ b/src/Diffusion/NoisePredictors/NoisePredictorBase.cs @@ -1144,6 +1144,50 @@ 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(); + 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], conditioning); + 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/Interfaces/INoiseScheduler.cs b/src/Interfaces/INoiseScheduler.cs index 2a1c6c39d0..e88f77016d 100644 --- a/src/Interfaces/INoiseScheduler.cs +++ b/src/Interfaces/INoiseScheduler.cs @@ -164,6 +164,59 @@ public interface INoiseScheduler /// Vector AddNoise(Vector originalSample, Vector noise, int 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. and + /// are shape [B, ...]; + /// is a [B] int vector, one entry per batch element. + /// + /// + /// + /// Per-element timesteps decorrelate the noise-schedule signal across the batch, + /// which the DDPM / EDM reference implementations use as the canonical training + /// pattern (Ho et al. 2020, Karras et al. 2022). Applied to a single-element + /// batch (B = 1) this reduces to the scalar overload above. + /// + /// + /// Default implementation delegates to the scalar overload per-element for + /// backward compatibility with existing schedulers; concrete schedulers should + /// override with a fused batched implementation for on-device efficiency. + /// + /// + Tensor AddNoiseBatched(Tensor cleanBatch, Tensor noiseBatch, int[] timesteps) + { + int batchSize = cleanBatch.Shape[0]; + if (noiseBatch.Shape[0] != batchSize) + throw new System.ArgumentException( + $"noiseBatch batch size {noiseBatch.Shape[0]} does not match cleanBatch batch size {batchSize}.", + nameof(noiseBatch)); + if (timesteps.Length != batchSize) + throw new System.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; + } + /// /// Gets the scheduler configuration. /// From a6a270ab812e4aa528bf2e1984dd8a504fd6aa25 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 13:46:22 -0400 Subject: [PATCH 09/15] feat(training): wGAN-GP correctness fix + fused disc + DP-SGD wire-ups (4F/4G/4H) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consumes Tensors PR ooples/AiDotNet.Tensors#763 (compiled backward with createGraph=true, DP-SGD helper, multi-slot persistent inputs). ## 4G — WGAN-GP correctness fix + fused disc steps (5 files) Root cause fixed for AiDotNet issue #1844: every ComputeGradientPenalty implementation used the inner GradientTape's ComputeGradients WITHOUT createGraph=true. The inner backward's ops didn't record on the outer tape, so inputGradients had no GradFn chain back to the discriminator weights. Effect: the gradient penalty was included in the loss VALUE but NOT in the disc's gradient — WGAN-GP silently degraded to plain WGAN with no 1-Lipschitz enforcement. Fix (5 files: CTGAN, DPCTGAN, CopulaGAN, CausalGAN, TableGAN): add createGraph: true to the inner ComputeGradients so the outer tape can differentiate the penalty through the disc weights. Also wired 4 disc steps through the fused compiled plan (CausalGAN, CTGAN, CopulaGAN, TableGAN — DPCTGAN's disc is DP so goes via 4H). Packs (real, fake) into one persistent input along axis 0; loss closure splits scores back out and computes wasserstein + λ·GP. The fused path activates once Tensors PR #763 lands (its 4C change removes the !createGraph gate at GradientTape.cs:727); until then these wire-ups fall through to the (now-corrected) eager tape path. ## 4H — DP-SGD wire-ups (2 files) Local AiDotNet.Training.DpSgdStep — drop-in mirror of the same helper in Tensors PR #763 (Engines.Training.DpSgdStep). Enables the wire-ups to land in this PR without waiting on the Tensors NuGet publish. Both implementations enforce the Abadi 2016 §3 Algorithm 1 clip-BEFORE-aggregate order via their structure. * DPCTGAN.TrainDiscriminatorStepBatchedDP — routes the per-example WGAN-GP + DP-SGD critic through DpSgdStep. Objective: Wasserstein + λ·GP per example, clipped per-example, aggregated, noised, averaged. * MedSynth.TrainDiscriminatorStepPerExampleDPSGD — routes the per- example non-saturating BCE + DP-SGD critic through the same helper. Once Tensors PR #763 merges + NuGet publishes, the local mirror can be swapped for the Tensors version by changing the `using` — the API surface is identical by design. ## 4F — Batched-per-element diffusion foundation DiffusionModelBase.Train samples per-batch-element timesteps for rank ≥ 2 inputs (Ho et al. 2020 canonical pattern; HuggingFace diffusers reference), routing through NoiseSchedulerBase.AddNoiseBatched + DiffusionModelBase.PredictNoiseBatched. Rank-1 unbatched inputs keep the scalar path for backward compat. net471 doesn't support default interface implementations, so AddNoiseBatched lives on NoiseSchedulerBase (as virtual with a default per-element delegate) rather than on the INoiseScheduler interface. DiffusionModelBase's fused-training wire-up (via Tensors PR #763's PersistentInputRegistry) is queued for the NuGet-bump follow-up commit — the API foundation is already in place. Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- src/Diffusion/DiffusionModelBase.cs | 29 +++- .../Schedulers/NoiseSchedulerBase.cs | 49 +++++++ src/Interfaces/INoiseScheduler.cs | 53 -------- .../SyntheticData/CTGANGenerator.cs | 39 +++++- .../SyntheticData/CausalGANGenerator.cs | 49 ++++++- .../SyntheticData/CopulaGANGenerator.cs | 41 +++++- .../SyntheticData/DPCTGANGenerator.cs | 31 ++++- .../SyntheticData/MedSynthGenerator.cs | 121 ++++++----------- .../SyntheticData/TableGANGenerator.cs | 39 +++++- src/Training/DpSgdStep.cs | 124 ++++++++++++++++++ 10 files changed, 434 insertions(+), 141 deletions(-) create mode 100644 src/Training/DpSgdStep.cs diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index 1542494d12..9691eb0f00 100644 --- a/src/Diffusion/DiffusionModelBase.cs +++ b/src/Diffusion/DiffusionModelBase.cs @@ -1094,7 +1094,34 @@ public virtual void Train(Tensor input, Tensor expectedOutput) var noiseSpan = noiseBatch.AsWritableSpan(); for (int i = 0; i < noiseSpan.Length; i++) noiseSpan[i] = NumOps.FromDouble(RandomGenerator.NextGaussian()); - noisySampleTensor = _scheduler.AddNoiseBatched(input, noiseBatch, timesteps); + // 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 diff --git a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs index c542b7da47..a1d023b9e5 100644 --- a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs +++ b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs @@ -279,6 +279,55 @@ 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)); + + int batchSize = cleanBatch.Shape[0]; + if (noiseBatch.Shape[0] != batchSize) + throw new ArgumentException( + $"noiseBatch batch size {noiseBatch.Shape[0]} does not match cleanBatch batch size {batchSize}.", + 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/Interfaces/INoiseScheduler.cs b/src/Interfaces/INoiseScheduler.cs index e88f77016d..2a1c6c39d0 100644 --- a/src/Interfaces/INoiseScheduler.cs +++ b/src/Interfaces/INoiseScheduler.cs @@ -164,59 +164,6 @@ public interface INoiseScheduler /// Vector AddNoise(Vector originalSample, Vector noise, int 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. and - /// are shape [B, ...]; - /// is a [B] int vector, one entry per batch element. - /// - /// - /// - /// Per-element timesteps decorrelate the noise-schedule signal across the batch, - /// which the DDPM / EDM reference implementations use as the canonical training - /// pattern (Ho et al. 2020, Karras et al. 2022). Applied to a single-element - /// batch (B = 1) this reduces to the scalar overload above. - /// - /// - /// Default implementation delegates to the scalar overload per-element for - /// backward compatibility with existing schedulers; concrete schedulers should - /// override with a fused batched implementation for on-device efficiency. - /// - /// - Tensor AddNoiseBatched(Tensor cleanBatch, Tensor noiseBatch, int[] timesteps) - { - int batchSize = cleanBatch.Shape[0]; - if (noiseBatch.Shape[0] != batchSize) - throw new System.ArgumentException( - $"noiseBatch batch size {noiseBatch.Shape[0]} does not match cleanBatch batch size {batchSize}.", - nameof(noiseBatch)); - if (timesteps.Length != batchSize) - throw new System.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; - } - /// /// Gets the scheduler configuration. /// diff --git a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 90dd2e8b13..9240547ea4 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -614,6 +614,41 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); + // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + 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(); var discParams = TapeTrainingStep.CollectParameters(_discLayers); @@ -970,7 +1005,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 4927e8da58..87ab0a4dcc 100644 --- a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs @@ -602,6 +602,49 @@ 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(); + + // GPU-RESIDENT WGAN-GP disc via Tensors PR #763 (createGraph=true on + // compiled backward). The gradient-penalty inner tape's ops now record + // on the outer tape, so the fused plan can compile the full + // Wasserstein + λ·GP objective as one graph and backprop the penalty + // through the disc weights. + 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>()); @@ -1064,7 +1107,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 59107cef68..1482278083 100644 --- a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs @@ -668,6 +668,43 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); + // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. Same pack-real-fake pattern + // as CausalGAN; loss closure re-runs the gradient penalty on the outer tape + // via the fixed inner-tape createGraph=true (see AiDotNet #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(); var discParams = TapeTrainingStep.CollectParameters(_discLayers); @@ -896,7 +933,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 71a0eedee9..35feb29472 100644 --- a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs @@ -659,7 +659,32 @@ 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 shared + // Training.DpSgdStep helper (Phase 4H). Enforces clip-BEFORE-aggregate + // order via the helper's structure so a future refactor can't accidentally + // reverse it (which would break the Abadi 2016 privacy proof). + int exampleCount = Math.Max(1, realPacked.Shape[0]); + var noisedGrads = AiDotNet.Training.DpSgdStep.ComputeClippedAggregatedGradients( + batchSize: exampleCount, + perExampleLoss: exIdx => + { + var realEx = ExtractPackedExample(realPacked, exIdx); + var fakeEx = ExtractPackedExample(fakePacked, exIdx); + var rs = DiscriminatorForwardBatched(realEx, isTraining: true); + var fs = DiscriminatorForwardBatched(fakeEx, isTraining: true); + var axes = Enumerable.Range(0, rs.Shape.Length).ToArray(); + var w = Engine.TensorSubtract( + Engine.ReduceMean(fs, axes, keepDims: false), + Engine.ReduceMean(rs, axes, keepDims: false)); + var gpEx = ComputeGradientPenalty(realEx, fakeEx); + return Engine.TensorAdd( + w, + Engine.TensorMultiplyScalar(gpEx, NumOps.FromDouble(_options.GradientPenaltyWeight))); + }, + parameters: discParams, + clipNorm: _options.ClipNorm, + noiseMultiplier: _computedNoiseMultiplier, + rng: _random); T lossValue = lossTensor.Length > 0 ? lossTensor[0] : NumOps.Zero; // Replay must reproduce the full objective (Wasserstein + λ·GP). @@ -900,7 +925,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/MedSynthGenerator.cs b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs index b50bd63436..6134779b0c 100644 --- a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs @@ -643,90 +643,52 @@ 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) - { - var zero = new Tensor(p._shape); - zero.Fill(NumOps.Zero); - gradSum[p] = zero; - } - - double clipNorm = _options.ClipNorm; - double noiseStd = clipNorm * noiseMultiplier; - 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); - + // Route through the shared DpSgdStep helper (Phase 4H). The helper enforces + // the Abadi 2016 Algorithm 1 clip-BEFORE-aggregate order via its structure, + // so a future refactor cannot silently reverse it and break the privacy proof. + // + // We capture per-example (real, fake) tensors so the optimizer's replay + // closure can reconstruct the exact objective the noisedAvgGrads were + // computed against — replay noise != training noise 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 (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) - { - for (int i = 0; i < g.Length; i++) - { - double v = NumOps.ToDouble(g[i]); - globalSqSum += v * v; - } - } - 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) - { - var scaled = Engine.TensorMultiplyScalar(kvp.Value, clipFactorT); - gradSum[kvp.Key] = Engine.TensorAdd(gradSum[kvp.Key], scaled); - } + var (rb, fb) = BuildRealAndFakeBatches(data, row, row + 1); + perExampleReal.Add(rb); + perExampleFake.Add(fb); } - // 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++) + // Sum of per-example scalar losses (unclipped) — used to report avgLoss. + // Compute alongside the DP-SGD helper by running the same forward inside + // its perExampleLoss callback. + T lossSum = NumOps.Zero; + var noisedAvgGrads = AiDotNet.Training.DpSgdStep.ComputeClippedAggregatedGradients( + batchSize: batchSize, + perExampleLoss: exIdx => { - 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); - } - noisedAvgGrads[kvp.Key] = noisy; - } + var realBatch = perExampleReal[exIdx]; + var fakeBatch = perExampleFake[exIdx]; + 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]); + return lossTensor; + }, + parameters: discParams, + clipNorm: _options.ClipNorm, + noiseMultiplier: noiseMultiplier, + rng: _random); + 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 @@ -738,9 +700,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 _) diff --git a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs index 632fa76e66..b189080b7c 100644 --- a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs @@ -434,6 +434,41 @@ private void TrainDiscriminatorStepBatched(Tensor realBatch, Tensor noiseB var fakeBatch = GeneratorForwardBatched(noiseBatch); fakeBatch = ApplyOutputActivationsBatched(fakeBatch); + // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + 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(); var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); @@ -515,7 +550,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/Training/DpSgdStep.cs b/src/Training/DpSgdStep.cs new file mode 100644 index 0000000000..3cb75f558c --- /dev/null +++ b/src/Training/DpSgdStep.cs @@ -0,0 +1,124 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Tensors.Engines.Autodiff; +using AiDotNet.Tensors.Helpers; +using AiDotNet.Tensors.LinearAlgebra; + +namespace AiDotNet.Training; + +/// +/// Fused per-example DP-SGD step (Abadi et al. 2016 §3, Algorithm 1). Runs a batch +/// of examples through per-example forward+backward, clips each per-example gradient +/// against the GLOBAL parameter-vector L2 norm, aggregates the clipped gradients, +/// adds a single Gaussian noise draw N(0, σ² C² I), and averages by batch size. +/// +/// Note: this is an AiDotNet-local mirror of the same helper in Tensors PR #763 +/// (AiDotNet.Tensors.Engines.Training.DpSgdStep<T>). Once #763 merges and the +/// AiDotNet.Tensors NuGet publishes, callers should migrate to the Tensors version; this +/// local mirror is retained until then to unblock the AiDotNet-side wire-ups. Both are +/// drop-in equivalent — same public API, same clip-then-aggregate contract. +/// +/// The clip-BEFORE-aggregate order is the L2-sensitivity bound the DP proof +/// requires — reversing it (aggregate-then-clip) breaks the privacy guarantee. This +/// helper enforces the correct order so callers cannot regress it. +/// +/// Numeric type of the parameter tensors. +public static class DpSgdStep +{ + /// + /// Runs a DP-SGD training step over examples. + /// See the Tensors version for full documentation. + /// + public static Dictionary, Tensor> ComputeClippedAggregatedGradients( + int batchSize, + Func> perExampleLoss, + IReadOnlyList> parameters, + double clipNorm, + double noiseMultiplier, + Random rng) + { + if (batchSize <= 0) + throw new ArgumentOutOfRangeException(nameof(batchSize), batchSize, + "DP-SGD batch size must be positive."); + if (clipNorm <= 0) + throw new ArgumentOutOfRangeException(nameof(clipNorm), clipNorm, + "DP-SGD clip norm must be positive (defines the L2-sensitivity bound)."); + if (perExampleLoss is null) throw new ArgumentNullException(nameof(perExampleLoss)); + if (parameters is null) throw new ArgumentNullException(nameof(parameters)); + if (rng is null) throw new ArgumentNullException(nameof(rng)); + + var ops = MathHelper.GetNumericOperations(); + + var sums = new Dictionary, Tensor>(AiDotNet.Helpers.TensorReferenceComparer>.Instance); + foreach (var p in parameters) + { + if (p is null) + throw new ArgumentException("parameters must not contain null tensors.", nameof(parameters)); + sums[p] = new Tensor(p._shape); + } + + for (int example = 0; example < batchSize; example++) + { + using var tape = new GradientTape(); + var loss = perExampleLoss(example); + if (loss is null) + throw new InvalidOperationException( + $"perExampleLoss({example}) returned null; must return the scalar loss tensor."); + var grads = tape.ComputeGradients(loss, parameters); + + // Global L2 norm across all parameters — Abadi 2016 Algorithm 1 line 4. + double normSquared = 0.0; + foreach (var g in grads.Values) + { + if (g is null) continue; + var span = g.AsSpan(); + for (int i = 0; i < span.Length; i++) + { + double v = ops.ToDouble(span[i]); + normSquared += v * v; + } + } + double clipFactor = Math.Min(1.0, clipNorm / Math.Sqrt(normSquared + 1e-12)); + + // Accumulate clipped per-example gradient — Abadi 2016 Algorithm 1 line 5. + foreach (var p in parameters) + { + if (!grads.TryGetValue(p, out var g) || g is null) continue; + var sumSpan = sums[p].AsWritableSpan(); + var gSpan = g.AsSpan(); + for (int i = 0; i < gSpan.Length; i++) + { + double v = ops.ToDouble(sumSpan[i]) + ops.ToDouble(gSpan[i]) * clipFactor; + sumSpan[i] = ops.FromDouble(v); + } + } + } + + // Noise + average — Abadi 2016 Algorithm 1 line 6. + double invBatch = 1.0 / batchSize; + double noiseStd = clipNorm * noiseMultiplier * invBatch; + var result = new Dictionary, Tensor>(AiDotNet.Helpers.TensorReferenceComparer>.Instance); + foreach (var p in parameters) + { + var sum = sums[p]; + var averaged = new Tensor(p._shape); + var sumSpan = sum.AsSpan(); + var avgSpan = averaged.AsWritableSpan(); + for (int i = 0; i < sumSpan.Length; i++) + { + double noise = noiseStd > 0 ? SampleGaussian(rng) * noiseStd : 0.0; + avgSpan[i] = ops.FromDouble(ops.ToDouble(sumSpan[i]) * invBatch + noise); + } + result[p] = averaged; + } + return result; + } + + private static double SampleGaussian(Random rng) + { + double u1; + do { u1 = rng.NextDouble(); } while (u1 < 1e-300); + double u2 = rng.NextDouble(); + return Math.Sqrt(-2.0 * Math.Log(u1)) * Math.Cos(2.0 * Math.PI * u2); + } +} From bb64d10f4353f7a9a21f2ba9cca653ccf4b73f59 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 14:43:02 -0400 Subject: [PATCH 10/15] fix(review): resolve 12 CodeRabbit findings on PR #1843 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses every unresolved review comment on the non-TS GPU-residency PR. ## Correctness fixes * DiffusionModelBase.Train — RecomputeForward closure now dispatches isBatched ? PredictNoiseBatched(inp, timesteps) : PredictNoise(inp, timestep), so optimizer re-evaluation matches the recorded batched forward objective. * NoisePredictorBase.PredictNoiseBatched — detect batch-aligned conditioning (leading dim == batchSize) and slice per element; preserve shared conditioning when the leading dim doesn't match. Fixes shape-mismatch / conditioning-shared bugs in classifier-free-guidance-style predictors. * CSDI.Train — remove the fused-plan branch. ComputeDenoisingPairTape samples fresh (timestep, noise) each call, and calling it in both Fwd and Loss produces independent samples that don't match. The compiled plan can't refresh the RNG per replay either. Path stays on the eager tape until Tensors' PersistentInputRegistry (PR ooples/AiDotNet.Tensors#763) lands. * TFC.Train — ComputeContrastiveLossTape now runs INSIDE ForwardCombined so it consumes the current-step persistent input (`inp`), not the outer `input` which would freeze at compile time. Closure-captured local; Loss reads it with a null-guard covering the Fwd-then-Loss ordering invariant. * TOTEM.Train — capture the commitment tensor from ForwardNativeForTrainingWithCommitment's first call; reuse it in Loss instead of running the quantizer a second time. Without this, the compiled path performs an EMA SetCodebookValue update TWICE per step and diverges from the eager path. * NoiseSchedulerBase.AddNoiseBatched — validate full shape parity (rank + every dim), not only the leading batch dim. Prevents indexing beyond noiseBatch's span when a caller passes [B, smaller...] shapes. ## Correctness cleanup * CTGAN.TrainGeneratorStepBatched — capture (act, condFromInput, maskFromInput) from Fwd's single generator pass; reuse in Loss so the conditional-CE term doesn't re-run GeneratorForwardWithResidualBatched + ApplyOutputActivationsBatched (was doubling the per-step generator forward cost). * TabSynGenerator.TrainVAEBatch — capture (mean, logVar) from Fwd's single encoder pass; reuse in Loss so the encoder doesn't run twice per row. Also: if the fused step fails mid-batch after fusedEngaged=true, drop to the eager path for the remaining rows (no row is silently skipped). * FourierNeuralOperator.TapeTrainStep — remove the duplicate _fourierLayers loop in both the fused-path allTrainable collection AND the eager paramList. Fourier layers are already registered into Layers at construction, so iterating them again double-registered each Fourier parameter and drove the eager SGD to apply the update twice per step. ## API cleanup * GpuResidentFusedStep — public → internal. Aligns with CompiledTapeTrainingStep which is already internal; keeps in-assembly training plumbing out of the public API surface. Also fixed an inadvertent null-forgiving operator (`!`) usage that slipped into the review-fix edits. All new nullable annotations use explicit null-guards with descriptive InvalidOperationException on invariant violation (per CLAUDE.md's "never use null-forgiving operators" rule). INoiseScheduler default-interface issue (net471 incompatibility) was already fixed in an earlier commit (AddNoiseBatched moved to NoiseSchedulerBase). Builds green net471 / net8.0 / net10.0. Co-Authored-By: Claude Opus 4.7 --- src/Diffusion/DiffusionModelBase.cs | 5 ++- .../NoisePredictors/NoisePredictorBase.cs | 30 ++++++++++++++- .../Schedulers/NoiseSchedulerBase.cs | 19 +++++++++- src/Finance/Forecasting/Foundation/CSDI.cs | 38 +++++-------------- src/Finance/Forecasting/Foundation/TFC.cs | 17 ++++++++- src/Finance/Forecasting/Foundation/TOTEM.cs | 18 ++++++++- .../SyntheticData/CTGANGenerator.cs | 26 +++++++++---- .../SyntheticData/TabSynGenerator.cs | 36 ++++++++++++++++-- .../NeuralOperators/FourierNeuralOperator.cs | 29 +++++--------- src/Training/GpuResidentFusedStep.cs | 2 +- 10 files changed, 151 insertions(+), 69 deletions(-) diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index 9691eb0f00..e71a02f016 100644 --- a/src/Diffusion/DiffusionModelBase.cs +++ b/src/Diffusion/DiffusionModelBase.cs @@ -1245,7 +1245,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 b2ca0b2d75..90fea1db62 100644 --- a/src/Diffusion/NoisePredictors/NoisePredictorBase.cs +++ b/src/Diffusion/NoisePredictors/NoisePredictorBase.cs @@ -1174,13 +1174,41 @@ public virtual Tensor PredictNoiseBatched(Tensor noisyBatch, int[] timeste 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]; - var pred = PredictNoise(elem, timesteps[b], conditioning); + + 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]; diff --git a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs index a1d023b9e5..27bb6719f5 100644 --- a/src/Diffusion/Schedulers/NoiseSchedulerBase.cs +++ b/src/Diffusion/Schedulers/NoiseSchedulerBase.cs @@ -297,11 +297,26 @@ public virtual Tensor AddNoiseBatched(Tensor cleanBatch, Tensor noiseBa 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]; - if (noiseBatch.Shape[0] != batchSize) + + // 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 batch size {noiseBatch.Shape[0]} does not match cleanBatch batch size {batchSize}.", + $"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}.", diff --git a/src/Finance/Forecasting/Foundation/CSDI.cs b/src/Finance/Forecasting/Foundation/CSDI.cs index d422d10f8d..309e982117 100644 --- a/src/Finance/Forecasting/Foundation/CSDI.cs +++ b/src/Finance/Forecasting/Foundation/CSDI.cs @@ -287,35 +287,15 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); - // GPU-RESIDENT fast path — fused SGD on the denoising-score-matching - // objective. The forward closure re-samples timestep + noise per step - // (via ComputeDenoisingPairTape), so replay produces fresh noise each - // step even though the compiled plan's captured graph shape is fixed. - var trainableLayers = Layers.OfType>().ToList(); - if (trainableLayers.Count > 0) - { - Tensor ForwardDenoise(Tensor inp) - { - var (pred, _) = ComputeDenoisingPairTape(inp, target); - return pred; - } - Tensor ComputeDenoiseLoss(Tensor pred, Tensor _) - { - // Recompute the noise target to match this replay's noise sample. - var (_, eps) = ComputeDenoisingPairTape(input, target); - return loss.ComputeTapeLoss(pred, eps); - } - if (AiDotNet.Training.CompiledTapeTrainingStep.TryStepWithFusedOptimizer( - trainableLayers, input, target, - forward: ForwardDenoise, computeLoss: ComputeDenoiseLoss, - 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; - } - } + // CSDI's stochastic denoising step cannot be traced through the compiled + // fused plan on the current Tensors build. ComputeDenoisingPairTape samples + // a fresh (timestep, noise) pair each call, and calling it twice — once in + // the forward closure, once in the loss closure — produces two independent + // samples that don't correspond to each other. Since the compiled plan + // replays the SAME captured graph every step, we can't refresh the RNG + // sample per replay either. This path stays on the eager tape until + // Tensors' PersistentInputRegistry (PR #763) lands, at which point the + // (timestep, noise) pair can be sampled per step and passed in as data. using var tape = new GradientTape(); var (epsilonPred, epsilonTarget) = ComputeDenoisingPairTape(input, target); diff --git a/src/Finance/Forecasting/Foundation/TFC.cs b/src/Finance/Forecasting/Foundation/TFC.cs index f4c44a0196..6d8d8abf75 100644 --- a/src/Finance/Forecasting/Foundation/TFC.cs +++ b/src/Finance/Forecasting/Foundation/TFC.cs @@ -282,9 +282,17 @@ public override void Train(Tensor input, Tensor target) var trainableLayers = Layers.OfType>().ToList(); if (trainableLayers.Count > 0) { + // Closure-captured local: ComputeContrastiveLossTape runs INSIDE the + // forward closure so it consumes the CURRENT-step persistent input + // (`inp` — refreshed by CompiledTapeTrainingStep before every replay), + // not the outer `input` which would freeze at compile time. The loss + // closure reads the captured value with a null-guard for the initial + // trace pass. Fwd/Loss ordering is guaranteed by the fused-step + // contract (Fwd always runs before Loss on each step). + Tensor? capturedContrastive = null; Tensor ForwardCombined(Tensor inp) { - // supervised forecast head (same alignment as eager path below). + capturedContrastive = ComputeContrastiveLossTape(inp); var f = ForwardForTraining(inp); return f; } @@ -296,7 +304,12 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) 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 = ComputeContrastiveLossTape(input); + var contrastive = capturedContrastive; + if (contrastive is null) + 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); diff --git a/src/Finance/Forecasting/Foundation/TOTEM.cs b/src/Finance/Forecasting/Foundation/TOTEM.cs index 67e007c716..a26d72988d 100644 --- a/src/Finance/Forecasting/Foundation/TOTEM.cs +++ b/src/Finance/Forecasting/Foundation/TOTEM.cs @@ -336,9 +336,18 @@ public override void Train(Tensor input, Tensor target) var trainableLayers = Layers.OfType>().ToList(); if (trainableLayers.Count > 0) { + // Closure-captured commitment loss: ForwardNativeForTrainingWithCommitment + // internally calls VectorQuantize which performs an EMA SetCodebookValue + // update in training mode. Running it twice per step (once in Fwd, once + // in Loss) would update the codebook twice per replay and diverge from + // the eager path. Capture the commitment tensor from Fwd's single call + // and reuse it in Loss. Fwd/Loss ordering is guaranteed by the fused-step + // contract; a null-guard covers the invariant defensively. + Tensor? capturedCommitment = null; Tensor ForwardCombined(Tensor inp) { - var (fc, _) = ForwardNativeForTrainingWithCommitment(inp); + var (fc, commit) = ForwardNativeForTrainingWithCommitment(inp); + capturedCommitment = commit; return fc; } Tensor ComputeLossCombined(Tensor pred, Tensor tgt) @@ -349,7 +358,12 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) 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) = ForwardNativeForTrainingWithCommitment(input); + var commit = capturedCommitment; + if (commit is null) + 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( diff --git a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 9240547ea4..957a71e140 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -791,12 +791,22 @@ private void TrainGeneratorStepBatched(Matrix transformedData, int numPacks) // 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); @@ -807,14 +817,14 @@ Tensor Loss(Tensor scores, Tensor _) var lossT = Engine.TensorNegate(Engine.ReduceMean(scores, axes, keepDims: false)); if (_condWidth > 0 && _catOutputBlocks.Count > 0) { - // Re-run gen forward on fused input to get fakeActivated, then - // compute the CE term. Redundant with Fwd but captured on the - // same fused plan and stays on-device. - var gi = Engine.TensorSlice(fusedInput, [0, 0], [totalSamples, embedDim + condDim]); - var faked = GeneratorForwardWithResidualBatched(gi); - var act = ApplyOutputActivationsBatched(faked); - var condFromInput = Engine.TensorSlice(fusedInput, [0, embedDim], [totalSamples, condDim]); - var maskFromInput = Engine.TensorSlice(fusedInput, [0, embedDim + condDim], [totalSamples, condDim]); + 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); } diff --git a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs index 7bea43425f..8b852deadc 100644 --- a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs @@ -767,20 +767,38 @@ private void TrainVAEBatch(Matrix transformedData, int startRow, int endRow, 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 enc = EncoderForwardOnTape(target); - var (m, lv) = SplitEncoderOutput(enc); + 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)); @@ -789,10 +807,20 @@ Tensor Loss(Tensor rawOutput, Tensor target) forward: Fwd, computeLoss: Loss, optimizer: _optimizer, out T _); - if (!ran) { if (!fusedEngaged) break; continue; } + 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) return; + 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++) diff --git a/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs index ef3e57fa4b..d1e8bbbb69 100644 --- a/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs +++ b/src/PhysicsInformed/NeuralOperators/FourierNeuralOperator.cs @@ -687,9 +687,14 @@ private void TapeTrainStep(Tensor input, Tensor expectedOutput) // failure. Uses SGD to match this method's baseline behavior. if (CanTrainOnGpu) { + // _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); - foreach (var fl in _fourierLayers) if (fl is ITrainableLayer t) allTrainable.Add(t); if (allTrainable.Count > 0) { Tensor Forward(Tensor inp) @@ -724,10 +729,10 @@ Tensor ComputeLoss(Tensor pred, Tensor tgt) } } - // 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. + // 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) { @@ -743,20 +748,6 @@ Tensor ComputeLoss(Tensor pred, Tensor tgt) } } } - foreach (var fl in _fourierLayers) - { - if (fl is ITrainableLayer trainable) - { - var layerParams = trainable.GetTrainableParameters(); - if (layerParams is not null) - { - foreach (var p in layerParams) - { - if (p is not null && p.Length > 0) paramList.Add(p); - } - } - } - } var paramTensors = paramList.ToArray(); using var tape = new GradientTape(); diff --git a/src/Training/GpuResidentFusedStep.cs b/src/Training/GpuResidentFusedStep.cs index 400777edd0..429a75b041 100644 --- a/src/Training/GpuResidentFusedStep.cs +++ b/src/Training/GpuResidentFusedStep.cs @@ -30,7 +30,7 @@ namespace AiDotNet.Training; /// /// /// Numeric type (float supported end-to-end; other types fall back to eager). -public static class GpuResidentFusedStep +internal static class GpuResidentFusedStep { /// /// Maps the runtime optimizer instance to the fused-plan optimizer enum + hyperparameters. From 33a16f9becedec9992736e520c806a9d31c11704 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 16:10:40 -0400 Subject: [PATCH 11/15] feat(training): vectorized fused DP-SGD/WGAN-GP/multi-slot primitives + consumer wire-ups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three fused-plan training primitives (mirrors of AiDotNet.Tensors PR #763) and wires the DP-SGD consumers through them: * DpSgdFusedStep — Abadi 2016 §3 Algorithm 1 per-example clip-before-aggregate. Runs each per-example forward+backward through a compiled plan (LR=0 so weights don't drift between replays), then computes the global L2 norm, clips, noises, and aggregates via vectorized IEngine ops (TensorMultiply/ReduceSum for the norm, TensorMultiplyScalar/TensorAdd for the accumulator, TensorRandomNormalInto for on-device Gaussian noise). Returns aggregated gradients so the caller's configured optimizer (Adam/AdamW/SGD/...) applies the update — no hardcoded SGD step inside. * WganGpFusedStep — Gulrajani 2017 WGAN-GP critic step. Composes E[D(fake)] − E[D(real)] + λ·(‖∇_x̃ D(x̃)‖₂ − 1)² inside one compiled plan with the inner ∇_x̃ D(x̃) recorded via createGraph=true so it differentiates into disc weights (issue #1844 fix). OnesLike uses vectorized Engine.TensorFill. * MultiSlotFusedStep — N-slot persistent input mechanism with plan cache keyed by composite shape + parameter identity. Slots refreshed via AsSpan().CopyTo(AsWritableSpan()) — no per-element loops on the hot path. Consumer wire-ups (this PR): * DPCTGANGenerator.TrainDiscriminatorStepBatchedDP — primary path routes through DpSgdFusedStep.TryStep, falls back to the existing ComputePerExampleNoisedGradients when the fused path can't engage (non-GPU host / compilation disabled). * MedSynthGenerator.TrainDiscriminatorStepPerExampleDPSGD — same primary/fallback layering. * Both legacy fallback loops (ComputePerExampleNoisedGradients + MedSynth eager) are now themselves vectorized: global-L2-norm via ReduceSum(g·g), clipped accumulation via TensorMultiplyScalar+TensorAdd, on-device Gaussian noise via TensorRandomNormalInto+TensorAdd — no scalar per-element loops anywhere on the DP-SGD path. Codebase convention adopted: * Class-scope `private static IEngine Engine => AiDotNetEngine.Current;` and `private static readonly INumericOperations Ops = MathHelper.GetNumericOperations();` on each fused-step class (matches the ~20 activation/optimizer/etc. bases). No IEngine or INumericOperations threaded through method signatures. Skipped (follow-ups filed): * WGAN-GP consumer rewire → #1845 (blocked on IGradientBasedOptimizer config accessor extension; consumers already correct via GpuResidentFusedStep + local createGraph=true ComputeGradientPenalty). * Diffusion consumer rewire → #1846 (blocked on per-generator forward-refactor to move _timestepProjection inside compiled plan while treating raw timestep as a persistent slot). Verification: * net8.0 + net471 + net10.0 all build clean on both AiDotNet and AiDotNet.Tensors. * Old DpSgdStep.cs deleted (replaced by DpSgdFusedStep). Co-Authored-By: Claude Opus 4.7 --- .../SyntheticData/DPCTGANGenerator.cs | 106 +++--- .../SyntheticData/MedSynthGenerator.cs | 114 +++++- src/Training/DpSgdFusedStep.cs | 344 ++++++++++++++++++ src/Training/DpSgdStep.cs | 124 ------- src/Training/MultiSlotFusedStep.cs | 297 +++++++++++++++ src/Training/WganGpFusedStep.cs | 338 +++++++++++++++++ 6 files changed, 1138 insertions(+), 185 deletions(-) create mode 100644 src/Training/DpSgdFusedStep.cs delete mode 100644 src/Training/DpSgdStep.cs create mode 100644 src/Training/MultiSlotFusedStep.cs create mode 100644 src/Training/WganGpFusedStep.cs diff --git a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs index 35feb29472..226e95e5d5 100644 --- a/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/DPCTGANGenerator.cs @@ -659,32 +659,50 @@ private void TrainDiscriminatorStepBatchedDP(Matrix transformedData, int numP wasserstein, Engine.TensorMultiplyScalar(gp, NumOps.FromDouble(_options.GradientPenaltyWeight))); - // Route the per-example DP-SGD gradient computation through the shared - // Training.DpSgdStep helper (Phase 4H). Enforces clip-BEFORE-aggregate - // order via the helper's structure so a future refactor can't accidentally - // reverse it (which would break the Abadi 2016 privacy proof). + // 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]); - var noisedGrads = AiDotNet.Training.DpSgdStep.ComputeClippedAggregatedGradients( - batchSize: exampleCount, - perExampleLoss: exIdx => + using var dpSgdStep = new AiDotNet.Training.DpSgdFusedStep(); + var gpWeightConst = _options.GradientPenaltyWeight; + bool dpFusedRan = dpSgdStep.TryStep( + parameters: discParams, + perExampleSlotData: exIdx => new[] { - var realEx = ExtractPackedExample(realPacked, exIdx); - var fakeEx = ExtractPackedExample(fakePacked, exIdx); - var rs = DiscriminatorForwardBatched(realEx, isTraining: true); - var fs = DiscriminatorForwardBatched(fakeEx, isTraining: true); - var axes = Enumerable.Range(0, rs.Shape.Length).ToArray(); + 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(fs, axes, keepDims: false), - Engine.ReduceMean(rs, axes, keepDims: false)); - var gpEx = ComputeGradientPenalty(realEx, fakeEx); + 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(_options.GradientPenaltyWeight))); + Engine.TensorMultiplyScalar(gpEx, NumOps.FromDouble(gpWeightConst))); }, - parameters: discParams, + batchSize: exampleCount, clipNorm: _options.ClipNorm, noiseMultiplier: _computedNoiseMultiplier, - rng: _random); + 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). @@ -811,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++) { @@ -840,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; diff --git a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs index 6134779b0c..79d486adec 100644 --- a/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/MedSynthGenerator.cs @@ -643,13 +643,9 @@ private void TrainDiscriminatorStepPerExampleDPSGD(Matrix data, int startRow, var discLayerList = BuildDiscLayerList(); var discParams = TapeTrainingStep.CollectParameters(discLayerList); - // Route through the shared DpSgdStep helper (Phase 4H). The helper enforces - // the Abadi 2016 Algorithm 1 clip-BEFORE-aggregate order via its structure, - // so a future refactor cannot silently reverse it and break the privacy proof. - // - // We capture per-example (real, fake) tensors so the optimizer's replay - // closure can reconstruct the exact objective the noisedAvgGrads were - // computed against — replay noise != training noise would train against a + // 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); @@ -660,31 +656,109 @@ private void TrainDiscriminatorStepPerExampleDPSGD(Matrix data, int startRow, perExampleFake.Add(fb); } - // Sum of per-example scalar losses (unclipped) — used to report avgLoss. - // Compute alongside the DP-SGD helper by running the same forward inside - // its perExampleLoss callback. + // 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; - var noisedAvgGrads = AiDotNet.Training.DpSgdStep.ComputeClippedAggregatedGradients( - batchSize: batchSize, - perExampleLoss: exIdx => + using var dpSgdStep = new AiDotNet.Training.DpSgdFusedStep(); + bool dpFusedRan = dpSgdStep.TryStep( + parameters: discParams, + perExampleSlotData: exIdx => new[] { - var realBatch = perExampleReal[exIdx]; - var fakeBatch = perExampleFake[exIdx]; - var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); + 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]); + if (lossTensor.Length > 0) lossSum = NumOps.Add(lossSum, lossTensor[0]); return lossTensor; }, - parameters: discParams, + batchSize: batchSize, clipNorm: _options.ClipNorm, noiseMultiplier: noiseMultiplier, - rng: _random); + 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) + { + noisedAvgGrads = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); + var gradSum = new Dictionary, Tensor>(TensorReferenceComparer>.Instance); + foreach (var p in discParams) + { + var zero = new Tensor(p._shape); + zero.Fill(NumOps.Zero); + gradSum[p] = zero; + } + for (int row = startRow; row < endRow; row++) + { + 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); + } + } + + // 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) + { + 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; + } + } + } var stackedReal = Engine.TensorConcatenate([.. perExampleReal], axis: 0); var stackedFake = Engine.TensorConcatenate([.. perExampleFake], axis: 0); 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/DpSgdStep.cs b/src/Training/DpSgdStep.cs deleted file mode 100644 index 3cb75f558c..0000000000 --- a/src/Training/DpSgdStep.cs +++ /dev/null @@ -1,124 +0,0 @@ -using System; -using System.Collections.Generic; -using AiDotNet.Tensors.Engines.Autodiff; -using AiDotNet.Tensors.Helpers; -using AiDotNet.Tensors.LinearAlgebra; - -namespace AiDotNet.Training; - -/// -/// Fused per-example DP-SGD step (Abadi et al. 2016 §3, Algorithm 1). Runs a batch -/// of examples through per-example forward+backward, clips each per-example gradient -/// against the GLOBAL parameter-vector L2 norm, aggregates the clipped gradients, -/// adds a single Gaussian noise draw N(0, σ² C² I), and averages by batch size. -/// -/// Note: this is an AiDotNet-local mirror of the same helper in Tensors PR #763 -/// (AiDotNet.Tensors.Engines.Training.DpSgdStep<T>). Once #763 merges and the -/// AiDotNet.Tensors NuGet publishes, callers should migrate to the Tensors version; this -/// local mirror is retained until then to unblock the AiDotNet-side wire-ups. Both are -/// drop-in equivalent — same public API, same clip-then-aggregate contract. -/// -/// The clip-BEFORE-aggregate order is the L2-sensitivity bound the DP proof -/// requires — reversing it (aggregate-then-clip) breaks the privacy guarantee. This -/// helper enforces the correct order so callers cannot regress it. -/// -/// Numeric type of the parameter tensors. -public static class DpSgdStep -{ - /// - /// Runs a DP-SGD training step over examples. - /// See the Tensors version for full documentation. - /// - public static Dictionary, Tensor> ComputeClippedAggregatedGradients( - int batchSize, - Func> perExampleLoss, - IReadOnlyList> parameters, - double clipNorm, - double noiseMultiplier, - Random rng) - { - if (batchSize <= 0) - throw new ArgumentOutOfRangeException(nameof(batchSize), batchSize, - "DP-SGD batch size must be positive."); - if (clipNorm <= 0) - throw new ArgumentOutOfRangeException(nameof(clipNorm), clipNorm, - "DP-SGD clip norm must be positive (defines the L2-sensitivity bound)."); - if (perExampleLoss is null) throw new ArgumentNullException(nameof(perExampleLoss)); - if (parameters is null) throw new ArgumentNullException(nameof(parameters)); - if (rng is null) throw new ArgumentNullException(nameof(rng)); - - var ops = MathHelper.GetNumericOperations(); - - var sums = new Dictionary, Tensor>(AiDotNet.Helpers.TensorReferenceComparer>.Instance); - foreach (var p in parameters) - { - if (p is null) - throw new ArgumentException("parameters must not contain null tensors.", nameof(parameters)); - sums[p] = new Tensor(p._shape); - } - - for (int example = 0; example < batchSize; example++) - { - using var tape = new GradientTape(); - var loss = perExampleLoss(example); - if (loss is null) - throw new InvalidOperationException( - $"perExampleLoss({example}) returned null; must return the scalar loss tensor."); - var grads = tape.ComputeGradients(loss, parameters); - - // Global L2 norm across all parameters — Abadi 2016 Algorithm 1 line 4. - double normSquared = 0.0; - foreach (var g in grads.Values) - { - if (g is null) continue; - var span = g.AsSpan(); - for (int i = 0; i < span.Length; i++) - { - double v = ops.ToDouble(span[i]); - normSquared += v * v; - } - } - double clipFactor = Math.Min(1.0, clipNorm / Math.Sqrt(normSquared + 1e-12)); - - // Accumulate clipped per-example gradient — Abadi 2016 Algorithm 1 line 5. - foreach (var p in parameters) - { - if (!grads.TryGetValue(p, out var g) || g is null) continue; - var sumSpan = sums[p].AsWritableSpan(); - var gSpan = g.AsSpan(); - for (int i = 0; i < gSpan.Length; i++) - { - double v = ops.ToDouble(sumSpan[i]) + ops.ToDouble(gSpan[i]) * clipFactor; - sumSpan[i] = ops.FromDouble(v); - } - } - } - - // Noise + average — Abadi 2016 Algorithm 1 line 6. - double invBatch = 1.0 / batchSize; - double noiseStd = clipNorm * noiseMultiplier * invBatch; - var result = new Dictionary, Tensor>(AiDotNet.Helpers.TensorReferenceComparer>.Instance); - foreach (var p in parameters) - { - var sum = sums[p]; - var averaged = new Tensor(p._shape); - var sumSpan = sum.AsSpan(); - var avgSpan = averaged.AsWritableSpan(); - for (int i = 0; i < sumSpan.Length; i++) - { - double noise = noiseStd > 0 ? SampleGaussian(rng) * noiseStd : 0.0; - avgSpan[i] = ops.FromDouble(ops.ToDouble(sumSpan[i]) * invBatch + noise); - } - result[p] = averaged; - } - return result; - } - - private static double SampleGaussian(Random rng) - { - double u1; - do { u1 = rng.NextDouble(); } while (u1 < 1e-300); - double u2 = rng.NextDouble(); - return Math.Sqrt(-2.0 * Math.Log(u1)) * Math.Cos(2.0 * Math.PI * u2); - } -} 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)); + } +} From 87900aa14ada68f5afb59357334f3a6a986b4cae Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 16:48:38 -0400 Subject: [PATCH 12/15] fix(review): tFC + TOTEM preprocessing rewritten with traceable engine ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the two BLOCKING CodeRabbit findings on PR #1843: both fused-plan training paths were freezing preprocessing into the compiled plan because their inner ops used host-side .Data.Span loops, so later replays reused the first batch's normalized values / spectrum / argmin decisions instead of recomputing per batch. TFC — ApplyInstanceNormalization (RevIN) + ComputeFrequencyRepresentation: * ApplyInstanceNormalization now delegates to a new stateless NormalizeWithStats that returns (normalized, mean, std) as tensors. Under the hood: ReduceMean + ReduceVariance + TensorSqrt + TensorBroadcastSubtract + TensorBroadcastDivide. All ops record on the tape and re-execute per replay under the compiled plan. * DenormalizeForecast likewise delegates to DenormalizeForecastWithStats, which takes mean/std as explicit tensor parameters. ForwardNative threads them through as locals so the compile-mode replay uses CURRENT-step stats instead of frozen trace-time values. * _revinMean/_revinStd (Vector scalars) → _revinMeanTensor/_revinStdTensor (Tensor? nullable) — kept for the abstract override's external callers, but the fused path never reads them. * ComputeFrequencyRepresentation now uses Engine.RFFT for batched real FFT → reshape to [B, halfN, 2] pairs → TensorMultiply + ReduceSum(axis=2) for magSquared → TensorSqrt + TensorMultiplyScalar(1/n) for the one-sided magnitude spectrum → TensorSlice + TensorFlip + TensorConcatenate to mirror bins [1..n-halfN] into the tail. Handles even and odd n identically to the old scalar impl. * Fused-plan fast path in TFC.Train restored — now safe because both preprocessing methods trace correctly. TOTEM — VectorQuantize (VQ-VAE) traceable rewrite + post-Step EMA: * New VectorQuantizeTraceable returns (quantized, commitmentLoss, argmin, head) using Engine ops end-to-end: TensorBroadcastSubtract + TensorMultiply + ReduceSum for distances, TensorArgMin along the codebookSize axis, per-c TensorSliceAxis + TensorIndexSelectDiff + TensorStack for the gather, Engine.StopGradient for the straight-through estimator, ReduceSum-based commitment loss weighted by β/totalLen. * EMA moved OUT of the compiled forward into a new UpdateCodebookEMA(head, argmin). Called POST-Step by the fused path with the trace-time graph-node references — their .Data reflects the LAST replay so the update lands exactly once per batch (matches CodeRabbit's "EMA must execute exactly once per batch" contract). Under the compiled plan, argmin/head are refreshed by every _plan.Step() so post-Step reads see the current batch's values. * UpdateCodebookEMA expresses the per-codebook scatter as: current codebook slice + TensorScatterAdd((1-decay)·(head - gathered), argmin) → new slice, then TensorConcatenate across codebook axis into the full [numCodebooks, codebookSize, codebookDim] tensor, then Engine.TensorCopy back into the _codebooks tensor object to preserve identity (future reads via the same reference see the update). * Legacy VectorQuantize is now a thin adapter around VectorQuantizeTraceable for callers that don't need the extras. * ForwardNativeForTrainingWithCommitment delegates to a new ForwardNativeForTrainingWithVQExtras that exposes the argmin/head; the original (forecast, commitmentLoss) contract is preserved for callers that don't need EMA state. * Fused-plan fast path in TOTEM.Train restored — now safe because VQ is fully traceable AND the EMA runs exactly once per batch in post-Step eager code. No new Tensors primitives were needed — the engine already has RFFT, ReduceMean/Variance, TensorSqrt, TensorBroadcastSubtract/Divide, TensorFlip, TensorConcatenate, TensorArgMin, TensorIndexSelectDiff, TensorSliceAxis, TensorStack, StopGradient, TensorScatterAdd, and TensorCopy, all in the autodiff/compile registry per OpRegistry.cs. Verification: * net8.0 + net471 + net10.0 all build clean. Co-Authored-By: Claude Opus 4.7 --- src/Finance/Forecasting/Foundation/TFC.cs | 231 ++++++++------ src/Finance/Forecasting/Foundation/TOTEM.cs | 335 ++++++++++++++------ 2 files changed, 368 insertions(+), 198 deletions(-) diff --git a/src/Finance/Forecasting/Foundation/TFC.cs b/src/Finance/Forecasting/Foundation/TFC.cs index 6d8d8abf75..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,27 +281,24 @@ 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. Both branches run through - // the fused forward closure so gradients accumulate into shared - // params on-device. Falls through to the in-place SGD loop below - // when the fused path can't engage. + // 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 local: ComputeContrastiveLossTape runs INSIDE the - // forward closure so it consumes the CURRENT-step persistent input - // (`inp` — refreshed by CompiledTapeTrainingStep before every replay), - // not the outer `input` which would freeze at compile time. The loss - // closure reads the captured value with a null-guard for the initial - // trace pass. Fwd/Loss ordering is guaranteed by the fused-step - // contract (Fwd always runs before Loss on each step). + // 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); - var f = ForwardForTraining(inp); - return f; + return ForwardForTraining(inp); } Tensor ComputeLossCombined(Tensor pred, Tensor tgt) { @@ -304,9 +308,8 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) 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; - if (contrastive is null) - throw new InvalidOperationException( + 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."); @@ -675,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; } @@ -765,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; @@ -815,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] }); @@ -876,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 a26d72988d..64059726ce 100644 --- a/src/Finance/Forecasting/Foundation/TOTEM.cs +++ b/src/Finance/Forecasting/Foundation/TOTEM.cs @@ -330,24 +330,26 @@ 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. - // Falls through to the in-place SGD loop below when the fused path - // can't engage. + // 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) { - // Closure-captured commitment loss: ForwardNativeForTrainingWithCommitment - // internally calls VectorQuantize which performs an EMA SetCodebookValue - // update in training mode. Running it twice per step (once in Fwd, once - // in Loss) would update the codebook twice per replay and diverge from - // the eager path. Capture the commitment tensor from Fwd's single call - // and reuse it in Loss. Fwd/Loss ordering is guaranteed by the fused-step - // contract; a null-guard covers the invariant defensively. Tensor? capturedCommitment = null; + Tensor? capturedArgmin = null; + Tensor? capturedHead = null; Tensor ForwardCombined(Tensor inp) { - var (fc, commit) = ForwardNativeForTrainingWithCommitment(inp); + var (fc, commit, argmin, head) = ForwardNativeForTrainingWithVQExtras(inp); capturedCommitment = commit; + capturedArgmin = argmin; + capturedHead = head; return fc; } Tensor ComputeLossCombined(Tensor pred, Tensor tgt) @@ -358,9 +360,8 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) 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; - if (commit is null) - throw new InvalidOperationException( + 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."); @@ -374,6 +375,8 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) out T fusedLoss)) { LastLoss = fusedLoss; + if (IsTrainingMode && capturedArgmin is not null && capturedHead is not null) + UpdateCodebookEMA(capturedHead, capturedArgmin); return; } } @@ -423,6 +426,24 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) /// 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. @@ -436,30 +457,9 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) 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) @@ -475,7 +475,7 @@ Tensor ComputeLossCombined(Tensor pred, Tensor tgt) // RevIN reverse: train against the input-scale forecast. decoded = DenormalizeForecast(decoded); - return (decoded, commitmentLoss); + return (decoded, commitmentLoss, argmin, head); } /// @@ -792,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) From ec718f3954c0c8b163c2140232c5c48deedeb699 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Fri, 10 Jul 2026 19:15:03 -0400 Subject: [PATCH 13/15] feat(training): centralize WGAN-GP + diffusion consumers through fused primitives (#1847) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1845 and #1846. Routes all four WGAN-GP critics through WganGpFusedStep (the fused-plan primitive from PR #1843) and adds MultiSlotFusedStep wire-ups to the three tabular diffusion consumers plus a base-class opt-in hook for DiffusionModelBase subclasses. ## Optimizer config plumbing Zero new API surface: IFusedOptimizerSpec.TryGetFusedOptimizerConfig already exposes (OptimizerType, LR, Beta1, Beta2, Epsilon, WeightDecay, Schedule, UseBf16Moments). Widened NeuralNetworkBase.TryMapToFusedOptimizerConfig from private to internal so the sibling generators in the same assembly can reuse the existing helper. ## WGAN-GP consumers (#1845) * CTGANGenerator, CopulaGANGenerator, TableGANGenerator, CausalGANGenerator — each critic training method now attempts WganGpFusedStep.TryStep FIRST with the discriminator's optimizer hyperparameters extracted via TryMapToFusedOptimizerConfig, falls back to the existing GpuResidentFusedStep path (secondary fused), then the eager tape (final fallback). The ε ∈ [0, 1]^B epsilon sampler uses Engine.TensorRandomUniformRange to match each critic's local ComputeGradientPenalty behavior. * Non-Adam optimizers (Lion, LBFGS) that don't implement IFusedOptimizerSpec cleanly fall through — TryMapToFusedOptimizerConfig returns false and the code path skips to GpuResidentFusedStep as before. ## Diffusion consumers (#1846) * TabDDPMGenerator — refactored to expose a slot-based forward (BuildTabDDPMSlots + DenoiserForwardFromTensors + ComputeDiffusionLossTapeFromTensors). Per-row TrainBatch loop now attempts MultiSlotFusedStep with (numNoisy, actualNoise, catNoisy, catClean, rawSinusoidalTimeEmbed) as persistent slots. The learnable _timestepProjection stays INSIDE the compiled forward closure so its weights participate in the backward pass. Plan is compiled once on the first row and replayed via slot-data refresh for subsequent rows. * TabSynGenerator — TrainDiffusionBatch's per-row loop wired with MultiSlotFusedStep on (noisyLatent, actualNoise, projectedTimeEmbed). Matches the existing eager path's semantic that _timestepProjection is NOT in _diffMLPLayers (kept detached in the eager path too), so the projected embedding is precomputed host-side per row and passed as slot data. * Finance/Forecasting/Foundation/CSDI — - ApplyInstanceNormalization rewritten with traceable engine ops (ReduceMean + ReduceVariance + TensorSqrt + broadcast subtract/divide) — same pattern as the TFC RevIN fix. The previous `.Data.Span` per-batch loop froze at trace time. - New BuildCsdiSlots + DenoiserForwardFromSlots express the DDPM x_t formation and packed denoising input via TensorConcatenate + engine scalar multiplies. Replaces the `.Data.Span[i] = xt[0, i]` fill that baked the trace batch's x_t into the compiled plan. - Train() attempts MultiSlotFusedStep first, falls back to the existing eager ComputeDenoisingPairTape path when the fused path can't engage. * Diffusion/DiffusionModelBase — - New opt-in `protected virtual bool SupportsFusedDenoising => false;` property. Base default is false so no existing subclass changes behavior. - Train() attempts MultiSlotFusedStep when SupportsFusedDenoising is true AND the training optimizer maps cleanly to a fused config. Slots: (noisySample, noise). Loss = MSE(pred, noise). QAT shadow restoration is preserved on the fused-success path. - Subclasses with fully-traceable PredictNoise / PredictNoiseBatched (e.g. after auditing to remove `.Data.Span` host loops) can opt in via a single-line override; no infrastructure changes needed elsewhere. ## Verification * net8.0, net471, net10.0 all build clean. * No API surface changes on IGradientBasedOptimizer — the existing IFusedOptimizerSpec interface (already implemented by all fuse-able optimizers) provided everything needed. * All consumers preserve eager fallback path for non-fuse-able optimizers and non-GPU hosts. Co-authored-by: franklinic Co-authored-by: Claude Opus 4.7 --- src/Diffusion/DiffusionModelBase.cs | 87 +++++++ src/Finance/Forecasting/Foundation/CSDI.cs | 190 ++++++++++++++-- src/NeuralNetworks/NeuralNetworkBase.cs | 2 +- .../SyntheticData/CTGANGenerator.cs | 44 +++- .../SyntheticData/CausalGANGenerator.cs | 41 +++- .../SyntheticData/CopulaGANGenerator.cs | 39 +++- .../SyntheticData/TabDDPMGenerator.cs | 213 +++++++++++++++++- .../SyntheticData/TabSynGenerator.cs | 91 +++++++- .../SyntheticData/TableGANGenerator.cs | 37 ++- 9 files changed, 684 insertions(+), 60 deletions(-) diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index e71a02f016..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. /// @@ -1132,6 +1155,70 @@ public virtual void Train(Tensor input, Tensor expectedOutput) 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 diff --git a/src/Finance/Forecasting/Foundation/CSDI.cs b/src/Finance/Forecasting/Foundation/CSDI.cs index 309e982117..dfd3f08e64 100644 --- a/src/Finance/Forecasting/Foundation/CSDI.cs +++ b/src/Finance/Forecasting/Foundation/CSDI.cs @@ -287,16 +287,51 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); - // CSDI's stochastic denoising step cannot be traced through the compiled - // fused plan on the current Tensors build. ComputeDenoisingPairTape samples - // a fresh (timestep, noise) pair each call, and calling it twice — once in - // the forward closure, once in the loss closure — produces two independent - // samples that don't correspond to each other. Since the compiled plan - // replays the SAME captured graph every step, we can't refresh the RNG - // sample per replay either. This path stays on the eager tape until - // Tensors' PersistentInputRegistry (PR #763) lands, at which point the - // (timestep, noise) pair can be sampled per step and passed in as data. + // 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); @@ -352,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, @@ -535,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/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/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 957a71e140..46ce9adb6f 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -614,7 +614,46 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + // 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) { @@ -650,7 +689,8 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers); + // 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); diff --git a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs index 87ab0a4dcc..130fa68710 100644 --- a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs @@ -604,11 +604,40 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int batchS var (realBatch, fakeBatch) = BuildRealAndFakeBatches(transformedData, batchSize); var discLayerList = _discLayers.Cast>().ToList(); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763 (createGraph=true on - // compiled backward). The gradient-penalty inner tape's ops now record - // on the outer tape, so the fused plan can compile the full - // Wasserstein + λ·GP objective as one graph and backprop the penalty - // through the disc weights. + // 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) { @@ -647,7 +676,7 @@ Tensor Loss(Tensor allScores, Tensor _) } 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); diff --git a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs index 1482278083..137907eab9 100644 --- a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs @@ -668,9 +668,40 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. Same pack-real-fake pattern - // as CausalGAN; loss closure re-runs the gradient penalty on the outer tape - // via the fixed inner-tape createGraph=true (see AiDotNet #1844). + // 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) { @@ -706,7 +737,7 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true); 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 8b852deadc..39326a3f0e 100644 --- a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs @@ -941,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 b189080b7c..57000481fa 100644 --- a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs @@ -434,7 +434,40 @@ private void TrainDiscriminatorStepBatched(Tensor realBatch, Tensor noiseB var fakeBatch = GeneratorForwardBatched(noiseBatch); fakeBatch = ApplyOutputActivationsBatched(fakeBatch); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + // 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) { @@ -470,7 +503,7 @@ Tensor Loss(Tensor allScores, Tensor _) } 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); From 618795deabb838a148396a0d2d592ea9b60a90b3 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Fri, 10 Jul 2026 19:55:54 -0400 Subject: [PATCH 14/15] feat(training): wire additional WGAN-GP + diffusion consumers through fused primitives (#1848) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends PR #1847 with three additional consumer wire-ups discovered during a comprehensive audit against the primitives added in PR #1843. Each consumer attempts the fused path first via the existing IFusedOptimizerSpec / TryMapToFusedOptimizerConfig plumbing and falls back to eager tape training when the fused path can't engage. ## WGAN-GP consumers (additional beyond #1845) * NeuralNetworks/WGANGP.cs — standalone WGAN-GP class. Previously used the legacy flat-vector round-trip (three Critic.Predict calls per step, GetParameterGradients + host-side vector combine + UpdateCriticWithOptimizer). TrainCriticBatchWithGP now attempts WganGpFusedStep.TryStep first, using Critic.Layers as the parameter source and a per-batch ε ~ U(0, 1) sampler matching each critic's local ComputeGradientPenalty behavior. Falls back to the legacy path when the critic's optimizer has no fused-kernel mapping. ## Diffusion consumers (additional beyond #1846) * NeuralNetworks/SyntheticData/FinDiffGenerator.cs — per-row DDPM training (Sattarov et al. 2023). TrainBatch caches a MultiSlotFusedStep across rows so the compiled plan is built once and replayed via slot-data refresh for subsequent rows. Slots: (packedDenoiserInput, targetNoise). Falls back to the existing Train(input, targetNoise) call on miss. * NeuralNetworks/SyntheticData/AutoDiffTabGenerator.cs — per-row DDPM training with a custom TapeStepOver optimizer step. Same MultiSlotFusedStep caching pattern as FinDiff. Slots: (denoiserInput, targetNoise). Falls back to the existing tape-based TapeStepOver path on miss. ## Explicitly not wired in this PR (documented) * NeuralNetworks/GenerativeAdversarialNetwork.cs — uses BCE loss with optional GP regularization (a separate auxiliary optimizer step), NOT Wasserstein + GP. Wiring WganGpFusedStep would change training semantics from BCE to Wasserstein. Requires a separate design decision. * Finance/Forecasting/Foundation/{CCDM,MGTSD,TSDiff,TimeDiff,TimeGrad}.cs and Finance/Probabilistic/DiffusionTS.cs — each uses .Data.Span host-side packs for the denoising input (same trace-freeze issue TFC/CSDI hit in PR #1843). Each needs a TFC/CSDI-scale traceable rewrite BEFORE fused wiring is safe. Deferred to individual per-consumer PRs. * MetaLearning/Algorithms/{MetaDDPMAlgorithm,MetaDMAlgorithm,MetaDiffAlgorithm}.cs — hand-rolled Vector<T> params with index arithmetic, NOT the ILayer/Engine/tape infrastructure. Would need full rewrite to fit MultiSlotFusedStep. * DiffusionModelBase subclasses (DDPMModel, DiffWaveModel, LatentDiffusionModelBase) — the SupportsFusedDenoising opt-in hook (added in PR #1847) is available, but flipping it safely requires per-class audit of PredictNoise → UNet traceability. Left as follow-up work; the mechanism is in place. ## Verification * net8.0, net471, net10.0 all build clean. * No API surface changes. Co-authored-by: franklinic Co-authored-by: Claude Opus 4.7 --- .../SyntheticData/AutoDiffTabGenerator.cs | 82 ++++++++++--- .../SyntheticData/FinDiffGenerator.cs | 115 ++++++++++++++---- src/NeuralNetworks/WGANGP.cs | 42 +++++++ 3 files changed, 198 insertions(+), 41 deletions(-) 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/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/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)); From f8fb2c96d675624818a5f5cf5623b7f04c75c73b Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Fri, 10 Jul 2026 20:01:37 -0400 Subject: [PATCH 15/15] chore(deps): bump AiDotNet.Tensors 0.111.2 -> 0.113.0 (activates #763 for fused paths) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 0.113.0 publishes Tensors PR #763 — this branch's gating dependency. #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.111.2 the fused GPU-resident WGAN-GP path would have silently degraded to plain WGAN. The src/Training fused-primitive mirrors (WganGpFusedStep, MultiSlotFusedStep, DpSgdFusedStep) stay as the consumer-side primitives that #1847/#1848 centralize every consumer through; they call the public engine API, so the bump alone routes them onto #763's fixed compiled backward. Also picks up #765 (compiled ReduceMax axis fill) and #764 (resident-param fused-Adam fix). Native OneDNN/OpenBLAS/CLBlast bumped in lockstep (all published). Builds green on net8.0. Co-Authored-By: Claude Opus 4.8 --- Directory.Packages.props | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/Directory.Packages.props b/Directory.Packages.props index def86e2a9f..dee116c705 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -229,11 +229,25 @@ activations from a prior training step on the same thread) leaked straight into the concat result — an intermittent, pool-state-dependent NaN, exactly the failure mode behind the model-family shard flakes this PR targets. 0.111.2 is PUBLISHED to NuGet and is a superset of - 0.111.1, so all rationale above is retained. --> - - - - + 0.111.1, so all rationale above is retained. + + Bumped 0.111.2 -> 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.111.2 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 + #764 resident-param fused-Adam fix), so all + rationale above is retained. --> + + + +