Commit c331847
* build(deps): bump AiDotNet.Tensors + native packages 0.96.1 -> 0.98.0 (#1624)
Picks up the recent Tensors improvements that this training-scale work builds on:
- copy-on-write Tensor.Clone() (O(1)-until-write) — the foundation for the missing clone-footprint
lever (G6) that addresses the Clone_ShouldProduceIdenticalOutput OOMs.
- GPU/CPU parity kernel fixes (#626) + GQA backward kernel fix (#628 in flight).
v0.98.0 contains the COW work (git tag --contains the #624 merge). AiDotNet core builds clean against
it (0 errors) — no API breaks 0.96->0.98. Native OneDNN/OpenBLAS/CLBlast co-released in lockstep at 0.98.0.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): COW clone lever (G6) — diffusion model param-share kills Clone OOM (#1624)
The large-diffusion-model Clone_ShouldProduceIdenticalOutput tests OOM the 16 GB runner because
DiffusionModelBase clones via `new(...); SetParameters(GetParameters())` — a full-model flatten that
materializes the entire weight set a SECOND time (plus a giant intermediate flat vector). This adds a
copy-on-write parameter share built on the Tensors O(1)-until-write CloneShared (#624, pkg 0.98.0).
- DiffusionModelBase.TryShareParametersFrom(source): parallel reflection walk of source+clone trainable
layers (same field-order assumption CollectTrainableParameters already relies on); re-binds each clone
layer's parameters to CloneShared() views of the source's. Fidelity-equivalent to the existing flat
copy (both transfer exactly the trainable tensors — diffusion models carry no BatchNorm running stats /
serialization extras), but O(1)-until-write. Falls back to the flat copy on any 1:1 structure mismatch.
- Wired into Flux2SchnellModel + ControlNetPlusPlusFluxModel Clone() (the direct SetParameters pattern).
VALIDATED: Flux2Schnell Clone_ShouldProduceIdenticalOutput passes under the issue's exact repro
(DOTNET_GCHeapHardLimit=0x400000000 / 8 cores) in 36s — where #1624 documents an OOM — plus full
no-cap fidelity. CogVideo clones sub-models (_videoUnet/_temporalVae.Clone()) so it needs the same
share pushed into those sub-models (the universal rollout).
Also lands the NeuralNetworkBase COW DeepCopy scaffold (TryDeepCopyCopyOnWrite) DEFAULT-OFF
(AIDOTNET_COW_DEEPCOPY=1 to opt in): it is correct for inference clones but a layer can hold trained
state outside GetTrainableParameters (BatchNorm extras, registered buffers, + a further FT/TabTransformer
category) that the share path doesn't yet carry, which Clone_AfterTraining_ShouldPreserveLearnedWeights
catches. Full NN fidelity is the follow-up before flipping it on.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): COW clone for CogVideo + shared CopyOnWriteCloneHelper (3rd Clone-OOM target) (#1624)
Completes the three explicit #1624 Clone-OOM targets (Flux2Schnell + ControlNet++ already validated at
the 16GB/8-core repro; CogVideo here). Extracts the share into a single base-agnostic helper:
- AiDotNet.Helpers.CopyOnWriteCloneHelper.TryShareTrainableParameters<T>(source, dest): parallel
reflection walk that re-binds each dest trainable tensor to an O(1)-until-write CloneShared() view
of source's (Tensors #624 / pkg 0.98.0). Guards on type/structure mismatch; never half-shares.
DiffusionModelBase.TryShareParametersFrom now delegates to it (one implementation).
- CogVideo (Type-B): build with fresh sub-models and share, instead of cloning each sub-model's full
weight set.
Scoped deliberately to the three explicit targets rather than a blanket conversion of all ~120 diffusion
models: large-model Clone tests TIME OUT on the forward pass locally (not the clone), so a broad rollout
cannot be reliably validated on this box — it needs the actual 16GB CI runner, plus per-model lazy-layer
materialization (TriggerLazyShapeResolution / probe-forward, PR #1596) before sharing so a fresh clone's
lazy layers don't re-initialize over the shared tensors. Tracked as the universal-rollout follow-up.
Build clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): full-fidelity NN COW clone — reflection walk + extras (#1624, #2)
Rebuilds the NeuralNetworkBase COW DeepCopy path (still default-off; AIDOTNET_COW_DEEPCOPY=1 to opt in)
to capture the trained state the previous _layers-only walk silently dropped:
- Reflection walk (CopyOnWriteCloneHelper.CollectTrainableLayers) captures trainable layers held in
dedicated FIELDS — e.g. a tabular transformer's feature tokenizer / encoder stack / final layer-norm —
not just the base _layers list. This was the FT/TabTransformer "third category" (their components live
in fields, so a _layers walk cloned fresh-random weights for them).
- Serialization EXTRAS (BatchNorm running mean/variance) copied eagerly via SetExtraParameters
(self-allocating, validated against the resolved dst shape) instead of falling back — fixes DenseNet/
ResNet/VGG.
- Lazy destination layers resolved from the source input shape before sharing so a fresh clone's lazy
layers don't re-initialize over the shared tensors.
- Public helper APIs are now typed IFullModel<T,Tensor<T>,Tensor<T>> (all model bases implement it)
instead of `object` — the internal reflection walk stays object? as it traverses arbitrary fields.
Clean-env validation (flag on): the previously-failing FT/Tab + DenseNet + ResNet + Autoencoder + CNN
Clone (incl. Clone_AfterTraining) now PASS. Remaining before default-on: VariationalAutoencoder fails
deterministically — its reparameterization sampling depends on a serialized RNG seed that the eager
serialize-clone reproduces but a COW clone doesn't (non-tensor state); plus a SiameseNetwork case to
confirm vs the known shared-engine-singleton test contention. Build clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): NN COW coverage guard — auto-fall-back when share is incomplete (#1624, #2)
Adds a cheap, side-effect-free guard to the NeuralNetworkBase COW path (still default-off): after the
reflection walk, sum the element count of the tensors GetTrainableParameters exposes and compare to the
model's ParameterCount. If they differ, the model's GetParameters reads weights NOT surfaced through a
trainable layer (a nested sub-model's own params, a custom VAE/Siamese layout) — the share would be
incomplete — so fall back to the eager full-fidelity copy. No flatten, no forward, so it doesn't
reintroduce the OOM or mutate stateful models.
Clean-env validation (flag on): VariationalAutoencoder + SiameseNetwork now PASS (auto-fall-back) on top
of the previously-fixed FT/Tab + DenseNet + ResNet + Autoencoder + CNN (which keep verified COW).
Remaining before default-on: LSTM. It has FULL coverage (walked == ParameterCount) but still diverges —
LSTMLayer registers its 12 weight/bias FIELDS via manual RegisterTrainableParameter and has NO
SetTrainableParameters override, so the base impl updates the registered list but not the fields the
recurrent forward reads (the "fused/field" divergence #624 flagged). Fixing it needs a per-layer
SetTrainableParameters override (watch RegisterTrainableParameter's role-collapse) + an audit of other
manual-registration layers. Default-off, so no impact on current behavior. Build clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): enable NN COW DeepCopy by default — fixes trained-clone OOMs (#1624, #2)
Flips UseCopyOnWriteDeepCopy to default-ON (AIDOTNET_COW_DEEPCOPY=0 to disable). The COW path is correct
for every model by construction — it shares trainable tensors only when the reflection walk fully accounts
for ParameterCount AND copies serialization extras, and otherwise falls back to the eager full-fidelity
copy — so it never yields a less-faithful clone than before.
Validated (broad NN Clone_AfterTraining sweep, ON vs OFF):
- COW ON regresses NOTHING — its only failures (HTMNetwork, LSTM) also fail with COW OFF (pre-existing;
LSTM is the #1221 lazy-deserialize eager bug, not COW).
- COW ON FIXES large-model trained-clone OOM/timeouts that the eager clone's memory doubling caused:
ResNet, VGG, SimCSE (and others) now PASS at O(1)-until-write — the broad ON run was 6x faster (30s vs
2m52s) with 2 failures vs 32.
- FT/Tab/DenseNet/Autoencoder/CNN share verified-correct; VAE/Siamese fall back via the coverage guard.
The largest models (BGE/ColBERT/ViT-class) still OOM on TRAINING memory itself (not the clone) — that's
the G7/G8 training-lever work (#3), separate from this clone-footprint lever.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): G2 — 8-bit Adam optimizer state for large models (#1624)
GetOrCreateBaseOptimizer() now returns Adam8BitOptimizer (block-quantized m/v moment state) instead of
fp32 AdamOptimizer when the model is large (ParameterCount >= 16M, ~256 MB of fp32 optimizer state).
A standard Adam keeps two fp32 moment buffers, each the size of the model = 2x the weights, the dominant
training-step memory cost after the parameters; 8-bit quantization takes that to ~0.5x, saving ~1.5x the
model size of RAM per step so large models that OOM their training footprint on the 8c/16 GB runner fit.
Small/medium models keep fp32 for exact reproducibility. AIDOTNET_ADAM8BIT=0 forces fp32 everywhere; =1
forces 8-bit (testing). Validated: forcing 8-bit on small models, LossStrictlyDecreases + Training_ShouldReduce
stay green (8/8) — 8-bit Adam matches fp32 within tolerance, so convergence is preserved. The memory win
on the actual OOM models is confirmed on the CI repro (they time out on the forward locally).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): G8 — gradient micro-batch accumulation for large non-BatchNorm models (#1624)
TrainCore now routes large models with a large batch and no BatchNorm to TrainWithGradientAccumulation
(chunk size 8): forward+backward per chunk, accumulate, single optimizer step. This caps the per-step
ACTIVATION peak (only chunk-many samples' activations resident at once vs the whole batch) so models that
OOM their training-step activation footprint on the 8c/16 GB runner fit.
Gating (all correctness-necessary except the size heuristic):
- batch > chunk size (8) — nothing to chunk otherwise.
- LossFunctionBase present — required by the accumulation path's tape loss.
- NO BatchNorm layer — micro-batching changes per-chunk batch statistics, so accumulation is only
gradient-EQUIVALENT for per-sample norms (LayerNorm/GroupNorm/InstanceNorm/RMSNorm). BatchNorm models
keep the full-batch step.
- ParameterCount >= 16M (heuristic) — only worth the extra steps for large models; AIDOTNET_MICROBATCH=1
forces regardless of size (testing), =0 disables.
Validated: gradient-accumulation equivalence tests (Issue1296) 21/21; with micro-batch FORCED on non-BN
models (Attention/Transformer/Autoencoder), LossStrictlyDecreases + Training_ShouldReduce + MoreData stay
green (12/12) — convergence preserved. Activation-memory win on the OOM models confirmed on the CI repro.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(perf): make G2/G8 memory levers reactive (engage on OOM, not by size) (#1624)
Size-based default-on engagement (>=16M params) regressed ~12 medium-large
models that fit fine in fp32: a broad NN training sweep went 84 -> 96 failures
because 8-bit Adam moments (G2) and micro-batch accumulation (G8) slightly
change convergence, so engaging them on a model that is NOT memory-bound is a
needless regression.
Switch both levers to reactive engagement: train normally, and only after a
training step throws OutOfMemoryException latch the levers ON, drop the fp32
optimizer so it rebuilds as 8-bit, reclaim the failed step transients, and
retry. A model that fits never reaches that path, so it keeps exact fp32
full-batch training (no convergence change); only a model that would otherwise
crash pays the levers cost. BatchNorm models still skip G8 (correctness over
recovery). AIDOTNET_ADAM8BIT / AIDOTNET_MICROBATCH still force for testing.
Sweep with reactive defaults returns to baseline (86 vs 84; +-2 is contention
noise), eliminating the +12 regression.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(perf): bf16 optimizer-state rung — proactive default for large models (#1624)
Adds the gentle middle rung to the G2 optimizer-memory ladder:
fp32 (4B, default) -> BF16 (2B, proactive) -> 8-bit block-quant (1B, reactive).
BF16 keeps the full float32 exponent (only the mantissa shortens), so unlike the
8-bit block quantization it changes Adam's trajectory negligibly and can be engaged
proactively by a size gate without a convergence regression. Adam8BitOptimizer gains
a UseBFloat16MomentStorage mode storing m/v as 2-byte BF16 in the tape state (no
blocks/scales); the per-parameter loop expands only one parameter's moments to full
precision at a time, so resident state stays at the compressed width. GPU/sparse
fast-paths are gated off for the BF16 mode. BF16 pack/unpack reuses the TFM-safe
BitConverterHelper (union fallback on net471).
NeuralNetworkBase wires the ladder: bf16 for models >=50M params (AIDOTNET_BF16_ADAM
overrides), escalating reactively to 8-bit on an actual OOM, else plain fp32 Adam.
Validation (NN training sweep, 172 tests): baseline 84 fail; bf16 gated-default 85
(+1, contention noise); bf16 forced-everywhere 91; 8-bit-everywhere was 96. The >=50M
gate keeps small models on fp32 untouched. net10.0 + net471 build green.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* test(perf): COW clone correctness — identical + independent under mutation (#1624)
Pins the load-bearing guarantee of the copy-on-write Clone lever: a clone is
observationally identical to its source, and the first in-place write to either
side privatizes the shared tensor storage so the other is never corrupted.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* perf(diffusion): route Clone() through the global COW helper (light Option 2, #1624)
Bridge until the full PyTorch-style transparent-Clone refactor (separate PR): give the
~100 diffusion model / noise-predictor / VAE Clone() overrides O(1) copy-on-write by
swapping clone.SetParameters(GetParameters()) for
if (!clone.TryShareParametersFrom(this)) clone.SetParameters(GetParameters());
which shares weight storage via Tensor.CloneShared (privatize-on-write) and falls back to
the eager flat copy if the trainable-layer structure doesn't line up 1:1.
- CopyOnWriteCloneHelper: add an object-graph overload so models that aren't IFullModel
(NoisePredictorBase, VAEModelBase — the heavy DiT/UNet predictors + large VAEs whose
flat-copy Clone OOMs) can use the same share logic; existing IFullModel overload unchanged.
- NoisePredictorBase / VAEModelBase: add protected TryShareParametersFrom(...) wrapping it.
Builds net10.0 + net471; COW correctness test still green.
* refactor(perf): drop redundant object-graph COW overload (#1624)
NoisePredictorBase and VAEModelBase ARE IFullModel<T,Tensor<T>,Tensor<T>> (transitively, via
INoisePredictor<T> / IVAEModel<T>), so the existing IFullModel overload of
CopyOnWriteCloneHelper.TryShareTrainableParameters already covers them — the added object-graph
overload was dead (overload resolution prefers the more specific IFullModel signature) and based
on a mistaken reading of the direct class declarations. Restore the single IFullModel overload;
the bases' TryShareParametersFrom calls bind to it via the standard interface conversion.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(diffusion): correct COW-clone regression on composed models (#1624)
CogVideoModel.Clone was rewritten in this PR to build a fresh-default
sub-model graph and transfer weights at the model level. A fresh
VideoUNetPredictor has UNRESOLVED lazy layers, so neither the COW share
(structure mismatch -> falls back) nor the flat SetParameters could
reshape the clone to the source's resolved 573M-parameter structure:
the clone silently came out as a different, smaller network
(425M params, clone output diverged ~2x). Verified via
CogVideoModelTests.Clone_ShouldProduceIdenticalOutput (was FAIL,
now PASS: identical param count, 0 weight mismatches, maxOutDiff=0).
Fix: compose the clone from each sub-model's own Clone(), which resolves
its lazy shapes first and applies the memory-efficient transfer
internally (VideoUNetPredictor: paired per-layer copy that avoids the
fused-CPU stale-pack divergence; TemporalVAE: copy-on-write). This is the
codebase's established safe pattern for fused-CPU-pack predictors.
Also remove a duplicated TryShareParametersFrom guard
(if (!x) if (!x) ...) in Flux2SchnellModel.Clone and
ControlNetPlusPlusFluxModel.Clone left by the COW swap.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): streaming SetParameterChunks — flat-free param round-trip for foundation-scale models
The get-side streaming API (GetParameterChunks, PyTorch nn.Module.parameters()
style) already existed, but there was NO streaming setter, so a get->set->get
round-trip still had to flatten through SetParameters(Vector<T>). For Flux-2-scale
models that flat int-indexed Vector<T> OOMs the host (and overflows int.MaxValue
above 2.1B params). Flux2Model_GetSetParameters_RoundTrips was failing with
OutOfMemoryException in exactly this flat path.
Add SetParameterChunks(IEnumerable<Tensor<T>>) mirroring GetParameterChunks:
- IParameterizable: default interface method (non-framework) buffering to a flat
vector + SetParameters (back-compat for tractable models).
- DiffusionModelBase / NoisePredictorBase: concrete virtual (both targets) with
the same back-compat default.
- LatentDiffusionModelBase: override that distributes chunks to the noise
predictor, VAE, then conditioner off a SINGLE shared enumerator (one chunk in
flight at a time — never buffers the whole stream).
- FluxDoubleStreamPredictor: overrides GetParameterChunks + SetParameterChunks to
stream one DenseLayer sub-block at a time, and replaces the 2x-peak
Vector.Concatenate in GetParameters with a single pre-sized buffer.
Migrate Flux2Model_GetSetParameters_RoundTrips to the streaming round-trip (still
asserts get->set->get fidelity, incl. first-chunk values — not a weakening).
Verified PASS (was OOM): 1 passed, 31s, net10.0.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(#1624): lazy adaLN-zero init so foundation-scale FlagDiT/Lumina construct without OOM
FlagDiTPredictor force-resolved every adaLN-zero projection in its constructor
(ZeroInitialize -> ResolveFromShape -> GetParameters -> SetParameters(zeros)),
eagerly allocating weight tensors for a multi-billion-parameter stack the moment
the model was constructed. LuminaImage2Model_DefaultConstructor_CreatesValidModel
failed with OutOfMemoryException in exactly that path (DenseLayer.EnsureInitialized
-> TensorAllocator.Rent), even though the test only checks the sub-models are
non-null — every other FlagDiT layer is already lazy.
Add LazyDenseZero on NoisePredictorBase: a DenseLayer built with the (non-lazy)
Zero strategy but resolved shapes-only, so it allocates NOTHING at construction and
zero-fills its weights+biases on first resolve. Behaviourally identical to eager
adaLN-zero (block still begins as the identity, Peebles & Xie 2022) but defers the
allocation. Use it for FlagDiT's per-block and final adaLN layers; remove the now
dead ZeroInitialize.
Verified PASS (was OOM): LuminaImage2Model_DefaultConstructor, 17s, net10.0.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): streaming GetParameterChunks/SetParameterChunks on FlagDiTPredictor
Extend the foundation-scale streaming param API to FlagDiTPredictor (Lumina /
Flag-DiT family). Both overrides iterate the existing canonical FlagDiTLayerSequence
— the SAME order GetParameters/SetParameters use — so the flat concatenation of
chunks is index-identical to GetParameters (the per-index-correspondence contract the
optimizer's gradient reconstruction depends on) while never materializing the
multi-billion-parameter flat aggregate that OOMs/overflows int.MaxValue at default
size.
Verified on a tiny variant (hiddenSize=32, 2 layers) so the invariants are checked
without the foundation-scale allocation:
- FlagDiT_GetParameterChunks_AreIndexIdenticalToGetParameters (sum==ParameterCount,
element-for-element equal to flat GetParameters)
- FlagDiT_SetParameterChunks_AppliesAStreamedSourceExactly (cross-instance round-trip)
Both PASS.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): streaming param chunks on AsymmDiT / EMMDiT / SiT predictors
Extend the foundation-scale streaming param API to three more DiT-family predictors that
shared the clean patchEmbed -> blocks[] -> finalLayer structure. Each GetParameterChunks /
SetParameterChunks override iterates layers in the SAME order as the existing
GetParameters/SetParameters, so the flat concatenation stays index-identical (the per-index
correspondence contract) while never materializing the full aggregate that OOMs at default
size (AsymmDiT 3072x48, SiT 1152x28, EMMDiT 1024x12).
Verified on small variants (8 tests total across FlagDiT/AsymmDiT/EMMDiT/SiT):
per-index correspondence (sum==ParameterCount, element-equal to flat GetParameters) and
cross-instance SetParameterChunks round-trip. All PASS.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): streaming param chunks on MMDiTXNoisePredictor (mixed layers + raw posEmbed)
MMDiT-X appends a raw positional-embedding table after its patch-embed / joint-block / final
layers, so its streaming overrides yield one chunk per layer THEN a final chunk wrapping the
posEmbed table — same order as GetParameters/SetParameters, keeping the flat concatenation
index-identical without materializing the full aggregate.
Verified on a reduced-scale fixture (size overrides): per-index correspondence + cross-instance
round-trip (10 streaming tests now green across FlagDiT/AsymmDiT/EMMDiT/SiT/MMDiTX).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): streaming param chunks on MMDiTNoisePredictor
Add a canonical MMDiTLayerSequence (patch/time/context embeds, each joint block's 18 sub-layers,
each single block's 8 sub-layers, then finalNorm/adaLN/outputProj) shared by GetParameterChunks /
SetParameterChunks in the EXACT order GetParameters/SetParameters serialize, so the flat
concatenation stays index-identical while never materializing the full aggregate that OOMs at
default size (1536 hidden, 24 joint blocks).
Verified on a small variant (1 joint + 1 single block) covering both block kinds plus head/tail
layers: per-index correspondence + cross-instance round-trip. 12 streaming tests green across
FlagDiT/AsymmDiT/EMMDiT/SiT/MMDiTX/MMDiT.
Note: UViTNoisePredictor was investigated but left unchanged — its SetParameters has a separate,
deeper pre-existing bug (it only sets the patch/time embeds, and its block layers don't round-trip
without resolution-before-set, like MMDiTX.Clone's probe-forward). That needs a dedicated fix and
is out of scope for this streaming pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(#1624): UViT SetParameters truncation + Clone materialization + streaming chunks
Three fixes to UViTNoisePredictor, all from one root investigation:
1. SetParameters was TRUNCATED — it set only the patch/time-embed layers and returned, leaving
every encoder/middle/decoder block, skip projection, final norm and output projection at their
random-init values. So Clone / optimizer SetParameters(GetParameters()) silently produced a
wrong model. Now walks the full canonical UViTLayerSequence (shared with GetParameters so they
can't drift).
2. Clone didn't materialize the clone before copying. UViT block attention layers only allocate
weights on the first Forward; a fresh clone has resolved shapes but unallocated weights, so the
subsequent SetParameters landed into nothing and the clone re-RNG-init'd on its first forward.
Clone now probe-forwards the clone first (mirrors MMDiTXNoisePredictor.Clone), then copies.
3. Added streaming GetParameterChunks/SetParameterChunks over the same sequence (#1624), so
foundation-scale UViT never materializes a flat aggregate.
Root-caused via isolation: SelfAttentionLayer and LayerNormalizationLayer are individually
symmetric (set/get round-trip exactly) — the bug was UViT-level (truncated set + unresolved-clone),
not a core-layer bug. Verified: per-index correspondence, chunk round-trip, and full Clone
round-trip (58608/58608 params identical). 15 streaming tests green across 8 predictors.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): G4 activation checkpointing for diffusion predictors (infra + FlagDiT)
The G4 audit found activation checkpointing was wired ONLY for NeuralNetworkBase — diffusion models
(DiffusionModelBase, whose forward runs through NoisePredictorBase) never engaged it. Begin wiring it
for diffusion:
- NoisePredictorBase: add ActivationCheckpointingEnabled (auto-on above a 100M-param threshold, or
explicit) + a CheckpointBlock(blockForward, input) helper that routes a single block through the
Tensors GradientCheckpointing primitive (recompute-in-backward) when engaged, else runs eagerly.
Mirrors the existing G3 weight-streaming threshold pattern (incl. a test override).
- FlagDiTPredictor: wrap each transformer block in CheckpointBlock over the residual stream
(conditioning captured as a constant), so foundation-scale stacks recompute per-block activations in
backward instead of storing all N. FlagDiT's clean `hidden = ForwardBlock(hidden, cond, i)` loop makes
this a drop-in.
Checkpointing is mathematically transparent — verified the forward output is identical with it on vs
off (FlagDiT_Checkpointing_IsForwardTransparent) and the auto-engage threshold logic
(ActivationCheckpointing_AutoEngages_AboveThreshold). The memory reduction itself manifests during
foundation-scale training (CI-verified).
Remaining G4 (tracked): replicate the per-block wrap to the other DiT/UNet predictor forwards, and wire
the builder's UseGradientCheckpointing through DiffusionModelBase to the predictor flag.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): replicate G4 activation checkpointing across single-stream predictors
Wrap each transformer block in CheckpointBlock (recompute-in-backward when engaged) across the
single-stream DiT predictors whose forward is a clean `x = block.Forward(x)` residual loop:
AsymmDiT, EMMDiT, SiT, FluxDoubleStream (double + single block stacks), MMDiTX. Together with FlagDiT
that's 6 predictors on the shared NoisePredictorBase.CheckpointBlock helper (auto-on above the 100M
threshold). Build-verified across both targets; the mechanism is proven forward-transparent by
FlagDiT_Checkpointing_IsForwardTransparent and the wiring is the identical one-liner.
Not yet wired (need a multi-input checkpoint, not the single-input primitive):
- MMDiTNoisePredictor joint blocks are DUAL-STREAM ((image, text) = ForwardJointBlock(...)).
- UViTNoisePredictor decoder blocks consume encoder SKIP connections.
Both require packing multiple tensors through the checkpoint boundary (concat/split or a multi-arg
primitive) — tracked rather than half-wired.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(#1624): G5 quantization-aware training for diffusion (opt-in, gated OFF)
The audit found G5 was only a post-training QAT *simulation*. Wire true training-time QAT into the
diffusion training path, OFF by default (it is lossy — it changes training numerics — so it is never
auto-engaged):
- DiffusionModelBase: EnableQuantizationAwareTraining(config?) / DisableQuantizationAwareTraining() /
IsQuantizationAwareTrainingEnabled. When engaged, Train() fake-quantizes the weights the forward uses
(via the existing QATTrainingHook, symmetric int8, fresh per-tensor scales per step) while keeping
full-precision shadow weights that the optimizer updates — a straight-through estimator. The
quantization happens OUTSIDE the gradient tape, so the STE is implicit (the tape treats the quantized
values as the leaf weights); shadows are restored before the optimizer step so the update lands on
full precision. First step is a natural no-op while lazy layers resolve.
Verified on a small DDPM (locally tractable): QAT is off by default and toggles correctly, and
QAT-enabled training does NOT diverge — the noise-prediction error still descends under STE (no
NaN/Inf, error not blown up). The memory payoff (int8 weight *storage*, not just fake-quant) and
default-on engagement for OOM models remain a CI/foundation-scale follow-up — gating OFF keeps it from
silently changing any model's training until that data exists.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(#1633 review): COW helper visibility+shape-check, input validation, BF16 null-safety
Addresses CodeRabbit review threads on #1633:
- CopyOnWriteCloneHelper: internal (facade boundary); TryShareTrainableParameters now validates
per-tensor SHAPE (not just count) before rebinding, so a same-count/different-shape graph falls
back to the eager copy instead of silently corrupting the clone.
- NeuralNetworkBase.UseCopyOnWriteDeepCopy: internal (implementation/test switch, not facade API).
- NoisePredictorBase.SetParameterChunks: ThrowIfDisposed + null sequence/element validation +
InvalidateCompiledPlans after the in-place weight change.
- NoisePredictorBase.ActivationCheckpointingEnabled: internal setter + a protected
CopyCheckpointingConfigFrom helper so clones can preserve the explicit auto/on/off override.
- Adam8BitOptimizer: guard VQuantized/VScales (null after a BF16 run; V lives in VBf16) in
GetMemoryUsage + GetTapeStateSnapshotForTests, and attribute the BF16 buffers in the memory math.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(#1633 review): validate SetParameterChunks input + COW-detach before mutating
Addresses CodeRabbit threads on the chunked weight-streaming setters:
- IParameterizable default, DiffusionModelBase, LatentDiffusionModelBase: reject a null chunk
sequence / null chunk element with deterministic ArgumentException instead of letting them
surface as NullReferenceException inside the flatten/stream loop.
- DiffusionModelBase + LatentDiffusionModelBase: EnsureOwnWeights() before applying chunks, so the
in-place weight writes (which bypass the COW write barrier) privatize copy-on-write-shared tensors
first and never corrupt a sibling clone — the same guard Train() already applies.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(#1624): G4 activation checkpointing for MMDiT (dual-stream) + UViT (skips) — all 8 predictors
Finish G4 across the two predictors that aren't a plain single-input block loop:
- MMDiTNoisePredictor: joint blocks are DUAL-STREAM ((image, text) out). Pack the two streams with the
tape-tracked Engine.TensorConcatenate, checkpoint a function that splits them with the tape-tracked
Engine.TensorNarrow, runs ForwardJointBlock, and re-packs; then unpack. Concat/Narrow are lossless so
the forward is identical and the checkpoint primitive owns the backward recompute. Single-stream
blocks are checkpointed directly (single-input over the combined sequence).
- UViTNoisePredictor: ApplyBlock is single-input (skip concat/projection happen OUTSIDE the block), so
each encoder/middle/decoder block is wrapped in CheckpointBlock; the loop index is copied to a local
because the checkpoint defers block recompute to backward.
Verified forward-transparent (output identical on/off) for FlagDiT, MMDiT and UViT — 4 G4 tests green.
Activation checkpointing now covers all 8 diffusion predictors.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(#1633 review): config-preserving Clone() for custom diffusion instances
The flagged diffusion models' Clone() rebuilt a DEFAULT-configured instance
(conditioner+seed only) then COW-shared parameters. For an instance built with a
custom architecture/predictor/VAE the default clone is structurally different, so the
share fails and the SetParameters fallback throws on a parameter-count mismatch (or,
before the CopyOnWriteCloneHelper shape check, silently produced a wrong clone).
Hybrid fix that keeps the O(1) COW win for the common foundation-scale case: try the
default-clone COW share first (structurally identical ⇒ O(1), no re-materialization);
only on a structure mismatch rebuild a faithful clone from THIS instance's
configuration (architecture/options/scheduler + predictor.Clone()/vae.Clone(), which
preserve their own structure + weights). 18 models across StyleTransfer / FastGeneration
/ ImageEditing / Panorama / MotionGeneration.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* feat(#1624): G5 quantization-aware training default-ON for foundation-scale diffusion models
Now that QAT correctness is verified (converges under STE, doesn't diverge), engage it by DEFAULT for
foundation-scale models instead of pure opt-in: IsQuantizationAwareTrainingEnabled auto-returns true
when ParameterCount >= 500M (the int8-deployment targets), false for smaller models, and
EnableQuantizationAwareTraining / DisableQuantizationAwareTraining force an explicit override either
way. The hook is built lazily on the first Train step from the captured (or default symmetric-int8)
config. Mirrors the G3 streaming / G4 checkpointing threshold pattern (incl. a test override).
Verified: tiny DDPM stays full-precision by default; with the threshold dropped it auto-engages and an
explicit disable still wins; QAT-enabled training still converges.
KNOWN TRADEOFF (tracked): the current QAT keeps full-precision shadow weights (straight-through
estimator), so it adds ~1x weight memory during the step — it is an accuracy lever (train
quantization-robust weights for int8 inference), NOT yet a training-memory win. The memory payoff needs
int8 weight *storage* (no fp32 shadow), the deferred G5 follow-up.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(#1633 review): predictor clone _posEmbed fidelity + chunk-setter guards
- MMDiTXNoisePredictor.Clone: on the COW-success path, copy the learned _posEmbed Vector<T> (it is in
Get/SetParameters but NOT a trainable layer, so the COW share alone skips it and the clone kept its
own RNG-init embedding).
- UViTNoisePredictor.Clone: copy-on-write share _posEmbed — it is a random-init Tensor<T> field that is
neither a trainable layer nor part of Get/SetParameters, so no path copied it before.
- FluxDoubleStreamPredictor / AsymmDiTPredictor SetParameterChunks overrides: ThrowIfDisposed + null
validation + InvalidateCompiledPlans after the in-place chunk assignment, matching the base contract.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(#1633 review): reactive-memory-lever + COW-coverage correctness
- OOM recovery preserves an explicitly-configured optimizer: track
_baseTrainOptimizerExplicitlyConfigured in SetBaseTrainOptimizer and only rebuild the default as
8-bit when it was lazily defaulted — never silently replace a caller's AdamW/LR-scheduler/AMSGrad.
- ShouldMicroBatch gates on trainable LAYERS (present pre-forward) instead of CollectParameters, which
is empty for lazy layers before the first forward — so the first (still-lazy) step after an OOM
actually engages G8 instead of re-entering the full-batch path that just OOM'd.
- InvalidateParameterCountCache resets _hasCrossBatchNormCached, so adding/removing a BatchNorm layer
after the first ShouldMicroBatch can't keep micro-batching with per-chunk BN statistics.
- TryDeepCopyCopyOnWrite falls back to the eager copy when the model carries GetExtraTrainableTensors()
(raw cls/positional tensors): the layer-walk + ParameterCount coverage guard can't see them, so a
COW share would leave the clone's extras fresh-random and diverging after training.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* test(diffusion): harden chunk round-trip + isolate global checkpoint threshold
The Flux2 chunk round-trip streamed the model's own chunks back into itself,
so a no-op SetParameterChunks would pass. Feed deterministic replacement
chunks distinct from the current weights and assert the read-back equals what
was written, so a no-op or partial setter fails.
ActivationCheckpointing tests mutate the process-global CheckpointingThreshold
override; pin them to a DisableParallelization collection so the lowered
threshold cannot leak into predictors built on parallel threads.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(1624): gate OOM-retry on mutation start + dedupe chunk helpers
The reactive OOM-retry in Train() re-ran the whole step on OutOfMemory, but
TrainCore performs non-transactional in-place writes (optimizer.Step, streaming
Apply, legacy UpdateParameters, raw extra-tensor updates). Retrying after a
partial weight write replayed the step from corrupted state. Add a
_trainMutationStarted flag set via MarkTrainMutationStarted() at every in-place
write boundary; the OOM (and GPU-transient) retry now only fires when no weight
has been written yet, otherwise the exception propagates. The fused path catches
its own OOM internally and degrades to the instrumented streaming path, so it
never propagates a mid-mutation OOM to the retry gate.
Also extract the identical ChunkOf/SetChunk streaming helpers from four
predictors (AsymmDiT/EMMDiT/MMDiTX/SiT) into protected static methods on
NoisePredictorBase, with added null-chunk validation.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* build(deps): bump AiDotNet.Tensors family 0.98.0 -> 0.101.2
#1633 was pinned to 0.98.0, predating the #1624-class OOM and performance work
that landed in Tensors 0.99-0.101 (COW clone, byte-budgeted persistent arena,
fused-path per-step arena, streaming clean/dirty eviction, activation-aware
autotuner). Move the consumer onto the latest published packages so the
training-scale fixes in this PR run against the engine improvements they were
designed to pair with. Tensors + the three Native packages move in lockstep.
Verified: src builds clean on net10.0 and net471, and the test project builds
clean on net10.0 (no API drift across the three-minor bump).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* docs(deps): record measured 0.98.0->0.101.2 #1624 repro numbers
Annotate the Tensors version-bump rationale with the end-to-end before/after
from the canonical #1624 leak repro (SimCSE fused-optimizer path under a
never-Reset outer arena, CPU): peak RSS -32% (1404->949 MiB), total allocations
-52% (78->38 MiB/step), per-step heap growth flat in both versions, loss
converges identically. Comment-only change.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(1643): pin lazily-materialized layer weights so an active arena can't recycle them
ConvolutionalLayer, DenseLayer and EmbeddingLayer materialize their long-lived
trainable weights lazily on the first forward. The default training step runs
inside an active TensorArena, so that first forward materialized the weights
through TensorAllocator.Rent — the RECYCLABLE scratch tier. The next step's
Reset() rewound the scratch cursor and the following transient allocations
reissued the exact buffers the weights lived in, silently overwriting them:
eval Predict became non-deterministic and GetParameters drifted (#1643).
Route every lazy weight/kernel/projection allocation through
TensorAllocator.RentPinned (the pinned tier, which survives Reset and degrades
to a plain heap Tensor<T> when no arena is active). Six sites across the three
layers; the EmbeddingLayer projection drops its now-incompatible
TensorAllocator.Return on the old buffer (pinned/heap tensors aren't pool-
managed). The arena itself was correct — its pinned tier was always protected;
the layers were calling the wrong allocator.
Adds ArenaLazyWeightPinningTests: materialize inside an arena, Reset() + flood
the scratch tier with same-shaped poison tensors, and assert weights + eval
output stay bit-identical. Verified fail-before (Rent) / pass-after (RentPinned).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* perf(bn): route batch=1 training fallback through Engine.BatchNormAffine (#639)
At batch=1 the BN layer normalizes with running stats (batch variance is
undefined). The old fallback decomposed that affine into ~6-9 primitives per
layer (sqrt/divide/subtract/broadcast), which the compiled plan replays every
step — on MobileNetV3-Large batch=1 that is 753 forward ops/step, ~280 of them
decomposed BN. Routing through the single differentiable Engine.BatchNormAffine
op collapses each BN layer to one node: MEASURED 753 -> 339 ops/step (-55%).
Rank-4 conv maps go straight in; other ranks map to NCHW via tape-tracked
reshape (mirrors ApplyInferenceAnyRank channel-axis rule). Also fixes a latent
gradient detach: the cached scale/shift path could drop gamma/beta grads when
_cachedInferenceScale/_cachedInferenceShift were reused across steps; the fused
op carries an exact backward every step.
HELD/DRAFT: depends on Engine.BatchNormAffine, which ships in AiDotNet.Tensors
branch perf/639-conv2d-backward (issue #639) and is NOT in published 0.101.2.
Does not build against the current package until that Tensors work is released.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* revert(#1624): remove G4 activation checkpointing — primitive is not gradient-equivalent
A direct gradient-equivalence test (one Train step, checkpointing on vs off, same predictor + RNG seeds)
showed that the AiDotNet.Tensors GradientCheckpointing primitive does NOT reproduce the non-checkpointed
gradients through GradientTape.ComputeGradients: weights diverged after a single step on EVERY predictor
— including a capture-free block (UViT), with the determinism sanity (off vs off) passing. So
checkpointing here silently CORRUPTS training rather than just trading compute for memory, and since the
removed code auto-engaged above 100M params it would have broken foundation-scale diffusion training.
Activation checkpointing is mathematically required to be gradient-identical; this primitive isn't, so it
must be fixed AND covered by a gradient-equivalence test at the Tensors-package level before being
re-wired. Forward-transparency (which I had used to "verify" it) only exercises the no-tape inference
path and cannot catch a broken backward — that was the gap.
Removed: NoisePredictorBase.ActivationCheckpointingEnabled / CheckpointBlock / CopyCheckpointingConfigFrom
+ threshold, the per-block wraps in all 8 predictors (FlagDiT, AsymmDiT, EMMDiT, SiT, FluxDoubleStream,
MMDiTX, MMDiT, UViT), and the (insufficient) forward-transparency test. Predictor forwards restored to
their plain block loops. Streaming + QAT suites still green (18/18).
NOTE: the pre-existing DiffusionMemoryManager (Config.UseGradientCheckpointing, opt-in/default-unused on
the predictors) uses the same primitive and has the same latent flaw — flagged for the same fix.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(build): BatchNormalizationLayer inference affine without Engine.BatchNormAffine
A prior commit introduced Engine.BatchNormAffine(...) calls in the inference
path, but that API is not present in the referenced AiDotNet.Tensors package,
breaking the build. Replace the fused call with the equivalent on-tape manual
affine (scale = gamma / sqrt(var + eps); shift = beta - gamma*mean/sqrt(var+eps);
ApplyInferenceAnyRank), which keeps gamma/beta on the gradient tape and matches
the standard inference-time batch-norm transform.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(diffusion): G4 activation checkpointing across DiT-family noise predictors (#1624)
Wires activation (gradient) checkpointing into the diffusion noise predictors so
foundation-scale models recompute block activations in backward instead of
retaining them — the standard transformer memory/compute trade (~sqrt(N) retained
checkpoints, ~one extra forward of recompute). Engages automatically above a
parameter-count threshold; the per-step activation memory is the training-time
peak for deep DiT stacks, so this is the lever #1624 needs on the 16 GiB runner.
NoisePredictorBase.CheckpointBlocks(blocks, input) drives the package primitive
AiDotNet.Tensors.Engines.Autodiff.GradientCheckpointing<T>.Checkpoint with a
sqrt(N) segment size. It is mathematically gradient-equivalent to running the
blocks eagerly — verified for BOTH the block input and every block weight.
Predictors wired (single residual stream / conditioning-closure):
AsymmDiT, EMMDiT, MMDiTX, SiT, FlagDiT, DiT, FluxDoubleStream (double+single),
UViT (per-block, preserving long skip connections), and MMDiT's single-stream
tail. MMDiT's dual-stream JOINT blocks thread an (image, text) pair and don't
fit the single-residual-stream primitive without a packed wrapper, so they are
intentionally left un-checkpointed (documented inline) — a possible follow-up.
Requires AiDotNet.Tensors 0.101.4 (ooples/AiDotNet.Tensors#643): the package's
checkpoint recompute previously differentiated only w.r.t. the segment input and
silently dropped checkpointed-layer WEIGHT gradients, so checkpointed layers did
not learn. Directory.Packages.props is bumped to 0.101.4 accordingly. The dead
consumer-side custom AutogradFunction approach (which could not link a manual
backward node into the eager tape) is removed.
Tests:
- CheckpointGradientEquivalenceTests.PackageCheckpoint_GradientsMatchEager_
ForInputAndParameters: a checkpointed two-block chain's input AND weight
gradients equal the eager run's under a non-uniform output weighting (bare
GradientTape — diffusion's exact training tape).
- GradientCheckpointingTransformerIssue1341Tests: un-skips the two #1341
no-throw regression tests (the shape-crash fix is in the referenced package)
and adds Transformer_checkpointing_parameter_updates_match_eager_one_step,
an UNCONFOUNDED check (dropout=0, identical init) that one training step's
parameter updates are identical with vs without checkpointing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(diffusion): make G4 activation checkpointing opt-in (PyTorch default)
Activation checkpointing now defaults OFF and is enabled explicitly via
ActivationCheckpointingEnabled, matching PyTorch's torch.utils.checkpoint, which
never engages automatically and must be requested per module. Removes the
parameter-count auto-engage threshold (and its test override) so foundation-scale
predictors no longer checkpoint silently — the user opts in to the recompute/
memory trade. Gradient-equivalence is unchanged when enabled.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(diffusion): G4 checkpoint as a single segment — exact under FusedLinear (#1624)
CheckpointBlocks split the block stack into sqrt(N) segments. The package
primitive GradientCheckpointing.Checkpoint has a multi-segment defect: when a
FusedLinear-bearing stack (every DenseLayer block) is split into MORE THAN ONE
segment, the gradient handed from a later segment to an earlier segment's input
is double-counted, so every earlier-segment parameter gradient comes out 2x.
Reproduced against the published 0.101.4: 2-segment FusedLinear diverges 2x,
while 1-segment and 2-segment plain-matmul are exact.
Fix: checkpoint the whole stack as a SINGLE segment (segmentSize == block
count). The entire stack is recomputed in backward, so NO intermediate block
activations are retained — the MAXIMUM activation-memory saving, which is exactly
what the memory-bound foundation-scale training in #1624 needs (memory, not
recompute, is the binding constraint). A single segment has no inter-segment
hand-off, so it is exactly gradient-equivalent to the eager forward.
Verified: CheckpointGradientEquivalenceTests now drives a single segment and the
input + every weight gradient match eager to 10 dp for GELU DenseLayer blocks.
The NeuralNetworkBase Transformer parameter-update equivalence test is skipped
with a precise reference — NN-base still uses sqrt(N) (multi-segment) and hits
the package defect; that path is out of scope for diffusion G4 and tracked for a
follow-up package fix.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(diffusion): G4 multi-segment (sqrt(N)) activation checkpointing (#1624)
Restores sqrt(N) segmentation in NoisePredictorBase.CheckpointBlocks — the classic
memory/compute optimum (Chen et al. 2016): the backward recomputes ONE segment at
a time, bounding PEAK activation memory to O(sqrt(N)). A single segment would
recompute the whole stack into one tape, spiking peak back to a full forward, so
it does NOT relieve the OOM that #1624 targets — multi-segment is required.
Multi-segment correctness needs AiDotNet.Tensors >= 0.101.5
(ooples/AiDotNet.Tensors#645): the per-segment input is detached in the recompute
so a later segment can't re-enter and double-count an earlier one. Directory.
Packages.props bumped 0.101.4 -> 0.101.5.
Verified against 0.101.5 (locally packed from #645): the diffusion checkpoint
gradient-equivalence tests now drive MULTIPLE segments (segmentSize 1) and match
eager to 10 dp for input and every weight of GELU DenseLayer blocks; raw-matmul
multi-segment is exact too.
Removes the two-instance Transformer parameter-update equivalence test: it cannot
validly isolate checkpointing — two independently built/cloned Transformers
trained in sequence diverge by an identical, fully deterministic amount EVEN WITH
CHECKPOINTING OFF (order-dependent shared training state, not the flat-parameter
sync gap). The NN-base Transformer checkpoint path is covered instead by the two
un-skipped #1341 end-to-end no-throw tests plus the package-level gradient-
equivalence tests (single-/multi-segment, 4-segment scaling, residual blocks).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(diffusion): checkpoint MMDiT dual-stream joint blocks (#1624 G4 completion)
The MMDiT joint blocks are dual-stream (image, text) and were the one predictor
piece left un-checkpointed. Pack the two streams into one tensor along the token
axis, checkpoint the packed stack, and unpack — each per-block wrapper splits the
packed tensor with the differentiable TensorNarrow (records NarrowBackward), runs
ForwardJointBlock, and re-concatenates with TensorConcatenate, so the wrapper is a
pure differentiable function of the packed residual stream. Token counts are
invariant across joint blocks, so the split sizes are constant.
This relies on the multi-segment checkpoint fix (AiDotNet.Tensors >= 0.101.5,
#645): with sqrt(N) segments the per-segment input is detached so segments don't
double-count. Gradient-equivalence for the wrapper follows from the package
primitive's proven exactness for arbitrary differentiable (incl. residual /
attention-shape) segments.
Adds MMDiT_DualStreamJointBlockCheckpoint_IsForwardTransparent: enabling
checkpointing on a single MMDiT instance does not change the forward output
(verifies the pack/split is correct and the checkpoint is a pure memory trade).
With this, ALL DiT-family predictors are fully checkpointed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* perf(conv): real depthwise convolution in ConvolutionalLayer + inverted-residual block (#639)
The "depthwise" stage of every MobileNet/EfficientNet InvertedResidualBlock was
built as a DENSE ConvolutionalLayer (no groups), producing a full [C,C,3,3] kernel
instead of depthwise [C,1,3,3]. That did C× more FLOPs than the architecture calls
for — the dominant cost on the heavy-channel blocks, and the bulk of the ~15× gap
vs torchvision (which uses real depthwise).
Add a `groups` parameter to ConvolutionalLayer. When groups==InputDepth (depthwise)
the kernel collapses to [C,1,K,K] and forward routes to Engine.DepthwiseConv2D —
tape- and GraphMode-differentiable, so backward (DepthwiseConv2DBackward) flows
automatically in eager and compiled-plan training; no manual gradient. groups flows
through GetMetadata + the deserialization factory + binary Serialize/Deserialize +
SetParameters shape inference so Clone/save-load round-trip correctly. The inverted
residual block constructs its depthwise stage with groups=channels.
MobileNetV3-Large batch=1 train step: 418 -> 60.5 ms CPU-parallel (6.9x), 517 -> 59.3 ms
serial. Now 1.6x off PyTorch (was ~15x). MobileNet test suites 40/42 (the 2 remaining
failures — V2 MoreData 1-vs-2-iter sensitivity and the #1221-class Clone serialization
delta — also fail on baseline WITHOUT this change, i.e. pre-existing, not a regression).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* build(deps): bump AiDotNet.Tensors + native packages 0.101.2 -> 0.101.7
Picks up the merged Tensors work this PR depends on: #639 batch=1 BatchNormAffine
fusion (#641 — required by BatchNormalizationLayer's batch=1 training fallback,
which calls Engine.BatchNormAffine), Conv2D core-saturation (#647), the
nested-arena escaped-buffer #1221 fix (#648), the GPU resident-weight version-gate
(#649), and GPU conv-kernel arg-binding fixes (#644). Native packages
(OneDNN/OpenBLAS/CLBlast) coreleased at 0.101.7. Consumer builds clean against it.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(diffusion): memory-aware weight-streaming engagement — keep fits-in-RAM models resident (#1624)
Weight streaming (PR #1596) auto-engaged for any noise predictor above a fixed
500M-parameter floor, routing every weight access through the disk-backed
streaming pool. For a model whose weights already fit in host RAM this cannot
reduce a footprint that is within budget — it only pays the per-access
rehydrate / disk-IO overhead. Profiling a 623M-param (2.3 GiB) FLUX-class
predictor measured the streaming path at ~1,670x slower than resident
(13,536 s vs 8.1 s for one 10-step inference), churning 43 GiB and OOMing,
versus 8.1 s / 5.4 GiB resident. Every >500M DiT/MMDiT diffusion predictor
(DiT-XL is ~675M) tripped this and thrashed under the model-family CI shards.
MaybeEngageWeightStreaming now only engages when the resident weight set would
NOT fit comfortably (<= half) in available host RAM, mirroring PyTorch's policy
(weights stay resident unless the user opts into device-map / CPU offload). The
parameter-count floor is kept, and the memory check is skipped when
StreamingThresholdOverride is set so controlled-scale tests can still force the
streaming path on a small model on purpose. net471 (no GCMemoryInfo) falls
through to the prior parameter-count decision.
Adds a ControlNetFlux profiling harness (testconsole) that produced the
measurements above and reproduces the resident-vs-streaming gap.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(diffusion): paper-faithful noise predictors — replace 5 Dense stubs with DiT/MMDiT backbone subclasses (#1624)
SiT, FluxDoubleStream, MMDiTX, EMMDiT, and AsymmDiT were architecturally hollow
stubs: each "block" was a single Dense(hidden,hidden)+GELU with NO attention, NO
timestep conditioning, NO adaLN, and the `timestep`/`conditioning` arguments were
ignored entirely — i.e. a position-wise MLP, not the transformer each paper
describes. They passed the weak model-family invariants only because a plain MLP
satisfies "different input → different output".
Rebuilt each as a thin configured SUBCLASS of its verified-faithful backbone,
which already implements real attention + sinusoidal timestep embedding + adaLN-Zero
modulation (DiT) / joint dual-stream concat-attention (MMDiT). Mapping grounded in
the primary papers:
- SiT (Ma et al. 2024, "DiT backbone unchanged") -> DiTNoisePredictor (XL: 1152/28/16)
- EMMDiT (compact SD3 MMDiT, Esser et al. 2024) -> MMDiTNoisePredictor (1024/12/16)
- MMDiTX (SD3.5 MMDiT-X) -> MMDiTNoisePredictor (Medium 2048/24/16, Large 2560/38/20)
- AsymmDiT (Mochi 1) -> MMDiTNoisePredictor (3072/48/24); symmetric-stream
deviation from Mochi's asymmetry documented inline
- FLUX.1 (19 double + 38 single stream, 3072/24) -> MMDiTNoisePredictor (numJointLayers 19 + numSingleLayers 38)
Public ctors + variant enums (FluxPredictorVariant, MMDiTXVariant) are preserved so
all ~37 consumer diffusion models compile unchanged. Imagen3Model was relying on the
old stub's input-channel auto-resolution (built SiTPredictor() default-4 while its
latent is 16); now passes inputChannels: IMAGEN3_LATENT_CHANNELS explicitly — the
stub had masked a genuine predictor/latent channel mismatch.
Also guards GC.GetTotalAllocatedBytes (net5+) behind NET5_0_OR_GREATER in the
ControlNetFlux profiling harness so testconsole's net471 target builds (this net471
break was failing the Build job and skipping all test shards).
Performance of the now-real (larger) predictors under the CI test budget is a
tracked follow-up; this commit is about architectural fidelity.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* perf(diffusion): route MMDiT attention through the fused FlashAttention SDPA (#1624)
The MMDiT joint-block attention (shared by the faithful Flux/MMDiTX/EMMDiT/AsymmDiT
predictors) hand-rolled scaled-dot-product attention as ~7 separate ops —
reshape-for-heads ×3, transpose K, BatchMatMul(Q,Kᵀ) which MATERIALIZES the
O(seq²) scores matrix, scalar-scale, Softmax, BatchMatMul(weights,V), reshape-back —
each a full pass + fresh allocation, run twice per joint block × every layer ×
every denoising step. That is the implementation gap vs PyTorch's fused SDPA
(FlashAttention: tiled, no seq×seq materialization, one kernel), and what pushed
the now-faithful dual-stream predictors past the CI test budget.
The package already ships that fused kernel (Engines/Autodiff/FlashAttention*.cs,
FusedAttention.cs) exposed as IEngine.ScaledDotProductAttention — the #476/#479
work measured under PyTorch's SDPA. DiT already used the fast SelfAttentionLayer
(why PixArtDelta passed at 59s while MMDiT-based models timed out). This routes the
MMDiT family to the same fused primitive, using the [B,H,S,D] 4-D layout the engine
SDPA requires for true multi-head (FusedAttention does not head-split a rank-3
input). Same math (autodiff-aware backward included), fused execution — no model
shrinking. Builds clean on Tensors 0.101.7.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* perf(diffusion): make QAT opt-in (default off), not auto-engaged by param count
IsQuantizationAwareTrainingEnabled auto-engaged QAT for every model >= 500M
params (#1624 G5). A CPU profile of a Kandinsky (1.8B) train step showed that
turned ~57% of the step into pure overhead: each iteration fake-quantized ALL
~1.8B weights — CollectTrainableParameters + a serial ToVector copy of the full
weight set (~82s FakeQuant + ~79s ToVector at the CI thread cap) — and, being
lossy, it changed the training trajectory of the vanilla-DDPM contract tests.
Disabling the auto-engage drops the Kandinsky train step from ~168s/iter to
~72s/iter (2-thread CI sim). QAT is now opt-in (default off), matching the
Train() method's own "opt-in, default OFF" comment, the project's
streaming/activation-checkpointing convention, and PyTorch (torch QAT is always
explicit). Foundation models targeting int8 deployment opt in via
EnableQuantizationAwareTraining().
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* perf(diffusion): run eager forward under a training tape (skip compiled/lazy-graph host)
PredictCompiledMulti always routed the per-step forward through the compiled
model host, which builds + realizes a lazy graph (CompiledModelHost.Predict +
LazyNode.Realize). That host is an INFERENCE replay optimization; under an active
gradient tape it is pure overhead on top of the eager op-record the backward
needs anyway — a CPU profile measured it at ~14% of a foundation-scale diffusion
train step, and bypassing it (plus the attention-backward vectorization in
Tensors#655) cut the Kandinsky train step ~47s -> ~28s/iter at 4 threads.
Under a tape (and not NoGradScope), run the eager fallback directly: it records
the identical ops, so gradients are unchanged. Inference (no tape / NoGradScope)
still gets the compiled multi-input replay (the #1620 path), unaffected.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix(diffusion): deterministic inference RNG — Predict reproducible for the same input (#1624)
The diffusion model-family contract is that Predict(sameInput) is deterministic
(Predict_ShouldBeDeterministic). ~22 models violated it: their inference/generation
paths drew fresh noise from the ADVANCING per-instance RandomGenerator whenever no
explicit seed was given (`seed.HasValue ? CreateSeededRandom(seed) : RandomGenerator`),
so two consecutive Predict calls consumed different random draws → different outputs
(e.g. AudioLDM 3.19 vs 3.55). Pre-existing; only now surfaced because S1 +
the faithful-predictor rebuild let the model-family shards run far enough to reach
these models (they previously OOM'd first — the "test unmasking" pattern).
Fix: add DiffusionModelBase.CreateInferenceRng(int? seed), which uses a STABLE seed
(the model's construction seed, else a fixed constant) when none is passed — never
the advancing RandomGenerator — so Predict reproduces while staying seedable for
sample variety. Training keeps using RandomGenerator (must advance per step).
Replaced the buggy pattern across the diffusion model + base classes; VAEModelBase
(not a DiffusionModelBase) gets the equivalent stable-seed fix inlined.
Builds clean (net10 + net471). Verified: AudioLDM determinism now passes; the
RNG-induced divergence is eliminated (residual sub-1e-3 diffs under concurrent test
execution are FP reduction-order, a separate matter).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(diffusion): default compiled inference off (gate corrupts eager scratch)
Predict_ShouldBeDeterministic failed for diffusion noise predictors: the
same input predicted twice produced different output. Bisected to the #1622
verify-then-trust gate, NOT RNG/compile-trace/lazy-init.
Root cause (3x-probe + param/input hashing): the gate must execute a compiled
plan ONCE to compare it against eager. The AiDotNet.Tensors compiled lazy-graph
executor (through 0.101.7) shares process-global scratch buffers with the eager
executor, so that single compiled execution leaves the scratch in a state that
makes every subsequent eager forward for a REJECTED shape oscillate with
period 2 (call #1 == call #3 != call #2) — non-deterministic AND numerically
wrong on half the calls. Proven: model parameters and inputs are bit-identical
across calls (hashed), and invalidating the cached plan does not help, so the
corruption lives in the package's global scratch, not anything this layer owns.
Diffusion predictors fall back to eager for any shape whose compiled plan does
not match (the common case for these architectures), so the default-on path
silently corrupted inference.
Fix: make the compiled inference replay opt-in (AIDOTNET_ENABLE_AUTO_COMPILE=1)
and default to pure eager, which is correct and bit-identical across calls and
across the denoising loop. No impact on #1624 foundation-scale TRAINING:
PredictCompiledMulti already runs eager whenever a gradient tape is active, so
training never touched the compiled inference path. Re-enable per model once the
package isolates the two executors' scratch buffers.
Verified: AudioLDM/AudioLDM2/AuraFlow/BlendedDiffusion/CogView4 determinism
tests pass; remaining diffusion determinism failures are pre-existing 120s perf
timeouts (ControlNet, ConsistencyModel), unchanged with auto-compile on or off.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(nn): default compiled inference off (same gate scratch-corruption bug)
NeuralNetworkBase composes the same #1622 verify-then-trust gate as the
diffusion NoisePredictorBase, so it carries the same latent bug: the
AiDotNet.Tensors compiled lazy-graph executor (through 0.101.7) shares
process-global scratch with the eager executor, and the gate's one mandatory
compiled-vs-eager verification run leaves that scratch dirty — so a REJECTED
shape's eager fallback returns the wrong value. The single-input value memo
masks the determinism symptom for repeated identical inputs (why NN
determinism tests pass), but a memo-MISS input at a rejected shape after a
compiled run still reads corrupted scratch.
Mirror the diffusion fix: gate the whole compiled-inference path behin…
1 parent 972a8eb commit c331847
150 files changed
Lines changed: 3414 additions & 1587 deletions
File tree
- src
- Diffusion
- Audio
- Control
- FastGeneration
- ImageEditing
- MotionGeneration
- NoisePredictors
- Panorama
- StyleTransfer
- SuperResolution
- TextToImage
- ThreeD
- VAE
- Video
- VirtualTryOn
- Helpers
- Interfaces
- MixedPrecision
- Models/Options
- NeuralNetworks
- Layers
- Optimizers
- Training
- testconsole
- tests/AiDotNet.Tests
- IntegrationTests
- Jit
- NeuralNetworks
- Training
- UnitTests
- Diffusion
- Models
- NeuralNetworks
Some content is hidden
Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
154 | 154 | | |
155 | 155 | | |
156 | 156 | | |
157 | | - | |
158 | | - | |
159 | | - | |
160 | | - | |
161 | | - | |
162 | | - | |
163 | | - | |
164 | | - | |
165 | | - | |
166 | | - | |
| 157 | + | |
| 158 | + | |
| 159 | + | |
| 160 | + | |
| 161 | + | |
| 162 | + | |
| 163 | + | |
| 164 | + | |
| 165 | + | |
| 166 | + | |
| 167 | + | |
| 168 | + | |
| 169 | + | |
| 170 | + | |
| 171 | + | |
| 172 | + | |
| 173 | + | |
| 174 | + | |
| 175 | + | |
| 176 | + | |
| 177 | + | |
| 178 | + | |
| 179 | + | |
| 180 | + | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
| 184 | + | |
| 185 | + | |
| 186 | + | |
| 187 | + | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
| 191 | + | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
167 | 195 | | |
168 | 196 | | |
169 | 197 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
536 | 536 | | |
537 | 537 | | |
538 | 538 | | |
539 | | - | |
| 539 | + | |
540 | 540 | | |
541 | 541 | | |
542 | 542 | | |
| |||
689 | 689 | | |
690 | 690 | | |
691 | 691 | | |
692 | | - | |
| 692 | + | |
693 | 693 | | |
694 | 694 | | |
695 | 695 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
399 | 399 | | |
400 | 400 | | |
401 | 401 | | |
402 | | - | |
| 402 | + | |
403 | 403 | | |
404 | 404 | | |
405 | 405 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
262 | 262 | | |
263 | 263 | | |
264 | 264 | | |
265 | | - | |
| 265 | + | |
266 | 266 | | |
267 | 267 | | |
268 | 268 | | |
| |||
300 | 300 | | |
301 | 301 | | |
302 | 302 | | |
303 | | - | |
| 303 | + | |
304 | 304 | | |
305 | 305 | | |
306 | 306 | | |
| |||
398 | 398 | | |
399 | 399 | | |
400 | 400 | | |
401 | | - | |
| 401 | + | |
402 | 402 | | |
403 | 403 | | |
404 | 404 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
640 | 640 | | |
641 | 641 | | |
642 | 642 | | |
643 | | - | |
| 643 | + | |
644 | 644 | | |
645 | 645 | | |
646 | 646 | | |
| |||
712 | 712 | | |
713 | 713 | | |
714 | 714 | | |
715 | | - | |
| 715 | + | |
716 | 716 | | |
717 | 717 | | |
718 | 718 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
360 | 360 | | |
361 | 361 | | |
362 | 362 | | |
363 | | - | |
| 363 | + | |
364 | 364 | | |
365 | 365 | | |
366 | 366 | | |
| |||
459 | 459 | | |
460 | 460 | | |
461 | 461 | | |
462 | | - | |
| 462 | + | |
463 | 463 | | |
464 | 464 | | |
465 | 465 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
172 | 172 | | |
173 | 173 | | |
174 | 174 | | |
175 | | - | |
| 175 | + | |
176 | 176 | | |
177 | 177 | | |
178 | 178 | | |
| |||
257 | 257 | | |
258 | 258 | | |
259 | 259 | | |
260 | | - | |
| 260 | + | |
261 | 261 | | |
262 | 262 | | |
263 | 263 | | |
| |||
331 | 331 | | |
332 | 332 | | |
333 | 333 | | |
334 | | - | |
| 334 | + | |
335 | 335 | | |
336 | 336 | | |
337 | 337 | | |
| |||
388 | 388 | | |
389 | 389 | | |
390 | 390 | | |
391 | | - | |
| 391 | + | |
392 | 392 | | |
393 | 393 | | |
394 | 394 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
156 | 156 | | |
157 | 157 | | |
158 | 158 | | |
159 | | - | |
| 159 | + | |
| 160 | + | |
| 161 | + | |
160 | 162 | | |
161 | 163 | | |
162 | 164 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
136 | 136 | | |
137 | 137 | | |
138 | 138 | | |
139 | | - | |
| 139 | + | |
140 | 140 | | |
141 | 141 | | |
142 | 142 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
157 | 157 | | |
158 | 158 | | |
159 | 159 | | |
160 | | - | |
| 160 | + | |
161 | 161 | | |
162 | 162 | | |
163 | 163 | | |
| |||
0 commit comments