Repository navigation
feat: add 8-bit Adam optimizer for memory efficiency - #814
Conversation
Implements Adam optimizer with 8-bit quantized state storage, reducing optimizer memory usage by approximately 4x compared to standard Adam. Key features: - Block-wise quantization with configurable block size (default 2048) - Signed quantization for first moment (m), unsigned for second moment (v) - Dynamic quantization that adapts scales during training - Option to keep first moment in full precision (CompressBothMoments=false) - Stochastic rounding option for reduced quantization bias - Percentile-based scaling for outlier robustness - Full serialization/deserialization support - Memory usage tracking via GetMemoryUsage() This is particularly useful for: - Training large models where optimizer memory is a bottleneck - GPU training with limited VRAM - Distributed training with memory constraints Closes #314 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Summary by CodeRabbit
WalkthroughIntroduces an 8-bit quantized Adam optimizer with configurable options and comprehensive integration tests. Stores optimizer states in 8-bit quantized format with per-block scaling, dynamic quantization, and optional full-precision updates for memory efficiency. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Optimizer as Adam8BitOptimizer
participant State as Quantized State
participant Compute as Computation Engine
Client->>Optimizer: Initialize(parameters, options)
Optimizer->>State: InitializeQuantizedState()
State-->>Optimizer: _mQuantized, _vQuantized, scales
loop Each Batch
Client->>Optimizer: UpdateParameters(params, gradient)
Optimizer->>State: Dequantize(_mQuantized, _mScales)
State-->>Optimizer: m (full precision)
Optimizer->>State: Dequantize(_vQuantized, _vScales)
State-->>Optimizer: v (full precision)
Optimizer->>Compute: Adam Update with bias correction<br/>(m, v, gradient, β₁, β₂, α)
Compute-->>Optimizer: Updated m, v, parameters
Optimizer->>State: Quantize(m, blockSize)<br/>with optional stochastic rounding
State-->>Optimizer: _mQuantized, _mScales
Optimizer->>State: Quantize(v, blockSize)<br/>with optional stochastic rounding
State-->>Optimizer: _vQuantized, _vScales
Optimizer-->>Client: Updated parameters
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Pull request overview
This PR implements an 8-bit quantized Adam optimizer to reduce memory consumption for optimizer states. The implementation uses block-wise quantization to compress the momentum and variance estimates from full precision to 8-bit representation, achieving approximately 4x memory savings while maintaining optimization quality.
Changes:
- Added
Adam8BitOptimizer<T, TInput, TOutput>class with quantization/dequantization logic for optimizer states - Created
Adam8BitOptimizerOptions<T, TInput, TOutput>with configurable block size, quantization modes, and compression settings - Implemented comprehensive integration tests covering various optimization scenarios, edge cases, and configurations
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| Adam8BitOptimizer.cs | Core optimizer implementation with 8-bit quantization for Adam's moment estimates |
| Adam8BitOptimizerOptions.cs | Configuration options for block size, dynamic/static quantization, and compression modes |
| Adam8BitOptimizerIntegrationTests.cs | Integration tests verifying optimizer behavior across different scenarios and edge cases |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@src/Optimizers/Adam8BitOptimizer.cs`:
- Around line 667-676: Serialize currently writes v state only when _vQuantized
!= null but Deserialize unconditionally reads vLength, causing mismatch; modify
Serialize() to write an explicit null/hasValue marker (e.g., a bool or byte)
before writing v state when handling _vQuantized and _vScales, and update
Deserialize() to first read that marker and only read vLength, the _vQuantized
byte[] and the _vScales floats when the marker indicates presence; reference the
Serialize(), Deserialize(), _vQuantized, _vScales and the vLength read logic
when applying the change.
- Around line 461-467: UpdateParameters currently uses fixed
_options.Beta1/_options.Beta2 while UpdateSolution uses adaptive
_currentBeta1/_currentBeta2; change UpdateParameters to use the adaptive values
when UseAdaptiveBetas is true by replacing references to
_options.Beta1/_options.Beta2 with the current values
(_currentBeta1/_currentBeta2) when computing beta1, beta2, oneMinusBeta1,
oneMinusBeta2 and biasCorrection1/2 (compute biasCorrection using Math.Pow on
the underlying double adaptive betas or their converted values and _t), and fall
back to _options.* when adaptation is disabled so both methods behave
consistently (look for the block computing beta1, beta2, oneMinusBeta1,
oneMinusBeta2, epsilon, biasCorrection1, biasCorrection2).
🧹 Nitpick comments (6)
src/Models/Options/Adam8BitOptimizerOptions.cs (2)
46-46: Consider adding validation forBlockSize.A
BlockSizeof zero or negative would cause division-by-zero or invalid array allocations in the optimizer. Consider adding validation either here with a property backing field, or document that validation occurs in the optimizer.💡 Optional: Add validation in a backing field
+ private int _blockSize = 2048; + /// <summary> /// Gets or sets the block size for block-wise quantization. /// </summary> - public int BlockSize { get; set; } = 2048; + public int BlockSize + { + get => _blockSize; + set => _blockSize = value > 0 ? value : throw new ArgumentOutOfRangeException(nameof(value), "BlockSize must be positive."); + }
81-81: Consider validatingQuantizationPercentilerange.Values outside 0-100 could cause index-out-of-bounds errors in the percentile calculation logic.
tests/AiDotNet.Tests/IntegrationTests/Optimizers/Adam8BitOptimizerIntegrationTests.cs (1)
492-525: Good float type coverage.Consider adding tests for other numeric types if supported (e.g.,
decimal) to increase type coverage per the issue requirements.src/Optimizers/Adam8BitOptimizer.cs (3)
118-118: Fixed random seed may affect stochastic rounding reproducibility across instances.The fixed seed
42ensures reproducibility within a single optimizer instance, but multiple optimizers will share the same random sequence. Consider making the seed configurable or usingRandom.Sharedfor non-deterministic behavior whenUseStochasticRoundingis enabled.
706-725: Deserialization readscompressBothMomentsseparately from options.The serialized
compressBothMomentsflag (line 706) should match_options.CompressBothMoments(from deserialized options). If they ever diverge, the state restoration could be incorrect. Consider adding a consistency check.💡 Optional: Add consistency validation
// Deserialize first moment bool compressBothMoments = reader.ReadBoolean(); + if (compressBothMoments != _options.CompressBothMoments) + { + throw new InvalidOperationException( + $"Serialized state compression mode ({compressBothMoments}) does not match options ({_options.CompressBothMoments})."); + } if (compressBothMoments)
606-608: Type size assumption may not cover all numeric types.The code assumes
Tis eitherfloat(4 bytes) ordouble(8 bytes). Other numeric types likedecimal(16 bytes) orHalf(2 bytes) would report incorrect memory usage.
- Clarify comment about [-127, 127] to [1, 255] mapping (0 is unused) - Rename typeSize2 to bytesPerElement and compute once - Use adaptive betas (_currentBeta1, _currentBeta2) instead of fixed _options values - Fix serialization/deserialization mismatch by adding hasVQuantized marker - Document the alignment buffer constant in test Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
CodeQL found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
|
Summary
Motivation
Closes #314
Large models can have billions of parameters, and standard Adam stores 2 full-precision values per parameter (first and second moment estimates). This can require 16GB+ of optimizer memory alone. 8-bit Adam quantizes these states, reducing memory usage to approximately 2 bytes per parameter plus a small overhead for scaling factors.
Implementation Details
Adam8BitOptimizerOptions<T, TInput, TOutput>: Configuration class with options for:BlockSize(default 2048): Number of elements per quantization blockUseDynamicQuantization(default true): Whether to adapt scales during trainingCompressBothMoments(default true): Whether to quantize both m and vQuantizationPercentile(default 99.9): Percentile for outlier-aware scalingUseStochasticRounding(default false): Use probabilistic roundingAdam8BitOptimizer<T, TInput, TOutput>: Main optimizer implementation with:GetMemoryUsage()method for tracking memory savingsTest plan
🤖 Generated with Claude Code