FlashOptim: Memory-Efficient PyTorch Optimizers

FlashOptim follows PyTorch's convention for weight decay styles:

• FlashAdam / FlashSGD: L2 regularization (coupled).
• FlashAdamW / FlashSGDW / FlashLion: Decoupled weight decay.

Standard Decoupled Weight Decay

For FlashAdamW, FlashSGDW, and FlashLion, setting weight_decay is equivalent to PyTorch's AdamW semantics: $\theta_t \leftarrow \theta_{t-1} \cdot (1 - \eta_t \cdot \lambda)$.

Fully LR-Decoupled Weight Decay

By setting decouple_lr=True, you enable a formulation where $\lambda$ is the absolute per-step decay rate, scaled by the LR ratio: $\theta_t \leftarrow \theta_{t-1} \cdot (1 - \lambda \cdot \frac{\eta_t}{\eta_0})$.

Warning: When using decouple_lr=True, you must use much smaller weight_decay values than in PyTorch. For example, torch.optim.AdamW(params, lr=1e-3, weight_decay=0.01) is equivalent to FlashAdamW(params, lr=1e-3, weight_decay=1e-5, decouple_lr=True).

FlashOptim provides drop-in replacements for PyTorch optimizers that reduce training memory by shrinking the footprint of optimizer states, master weights, and gradients. To use it, install the package, cast your model parameters to torch.bfloat16 using cast_model, and replace your standard PyTorch optimizer with a FlashOptim variant like FlashAdamW.

Supported optimizers include:

• FlashSGD
• FlashSGDW
• FlashAdam
• FlashAdamW
• FlashLion

import torch
from torch import nn
from flashoptim import FlashAdamW, cast_model

model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 10)).cuda()
# cast parameters to bf16
cast_model(model, dtype=torch.bfloat16)

# master_weight_bits=24 (default) means we have 24-bit parameter semantics
optimizer = FlashAdamW(model.parameters(), lr=1e-3)

x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)
loss = model(x).sum()
loss.backward()
optimizer.step()
optimizer.zero_grad()

FlashOptim allows you to reconstruct full-precision (FP32) weights from low-precision parameters and their associated error correction terms (stored in the optimizer state under the "error_bits" key).

To load an FP32 checkpoint, the model must first be initialized with the desired precision (e.g., via cast_model).

import torch
import torchvision
from flashoptim import FlashAdamW, cast_model

model = torchvision.models.resnet18().cuda()
cast_model(model, dtype=torch.bfloat16, full_precision_layers=["fc"])
optimizer = FlashAdamW(model.parameters(), lr=1e-3, master_weight_bits=24)

# ... training ...

# Save: reconstruct fp32 from bf16 + error bits
fp32_state_dict = optimizer.get_fp32_model_state_dict(model)
torch.save(fp32_state_dict, "checkpoint.pt")

# Load: restore fp32 weights into a bf16 model (error bits recomputed automatically)
fp32_state_dict = torch.load("checkpoint.pt", weights_only=True)
optimizer.set_fp32_model_state_dict(model, fp32_state_dict)

Gradient release updates parameters during the backward pass as soon as gradients are computed, reducing memory usage. This is implemented via post-backward hooks and must be explicitly enabled using enable_gradient_release.

Limitations: Gradient release is incompatible with:

• DDP: Due to custom communication hooks.
• Microbatch Accumulation: Parameters are updated immediately, preventing accumulation.
• Gradient Clipping: Global clipping cannot be applied.
• Gradient Scaling: torch.amp.GradScaler is not supported.

Gradient release is compatible with single-GPU training and FSDP2 (fully_shard).

from flashoptim import FlashAdamW, enable_gradient_release

optimizer = FlashAdamW(model.parameters(), lr=1e-3, master_weight_bits=24)
handle = enable_gradient_release(model, optimizer)

for x, y in dataloader:
    loss = loss_fn(model(x), y)
    loss.backward()
    # step() and zero_grad() are no-ops while gradient release is active;
    # parameters are updated during backward and gradients are freed immediately
    optimizer.step()
    optimizer.zero_grad()

# Call handle.remove() to restore normal optimizer behavior
handle.remove()

The FlashOptimizer is a base class for optimized PyTorch optimizers (like FlashAdam, FlashSGD, etc.) that implement several memory and speed optimizations.

Key optimizations include:

• Quantized optimizer states: Uses ~8.5 bits per parameter instead of 32 bits for momentum/variance. Requires CUDA.
• Compressed state dicts: Stores quantized states in checkpoints to save ~2-3 bytes per parameter.
• Fused optimizer steps: Runs the entire update in a single Triton kernel for speed. Requires contiguous parameters and gradients.
• Error correction bits (ECC): Allows using low-precision parameters (e.g., bf16) while maintaining higher effective precision (24-bit or 32-bit) via an additional correction tensor. This enables using 24-bit master weights to save 1 byte/param.
• Gradient release: Frees gradients immediately after the parameter update during the backward pass. This is incompatible with global gradient clipping.

Important Constraints:

• Parameters within the same param_group must have the same dtype.
• If using master_weight_bits (24 or 32), do not change parameter dtypes between saving and loading checkpoints.
• Fused kernels require contiguous tensors. If your parameters or gradients are non-contiguous (e.g., from tied weights), set fused=False.

# Example initialization (actual subclass usage depends on segment 2-4)
optimizer = FlashAdam(
    params=model.parameters(),
    lr=1e-3,
    quantize=True,
    master_weight_bits=24,
    fused=True,
    check_numerics=False
)

The compress_state_dict parameter controls how optimizer states (momentum and variance) are serialized in the state_dict().

• compress_state_dict=True (default): Optimizer states are quantized to int8 + fp16 scales. This produces checkpoints ~50% smaller than fp32. Note: These states are not loadable by vanilla PyTorch optimizers.
• compress_state_dict=False: Optimizer states are serialized as fp32. These are loadable by vanilla PyTorch optimizers. On load, FlashOptim pre-quantizes them to int8 to avoid lossy bf16 conversion.

Note: ECC (Error Correction Code) error correction bits are always preserved in the checkpoint regardless of this setting.

# Default: state_dict() saves states as quantized int8
optimizer = FlashAdamW(model.parameters(), lr=1e-3)
torch.save(optimizer.state_dict(), "checkpoint_int8.pt")

# Uncompressed: state_dict() saves states as fp32
optimizer = FlashAdamW(model.parameters(), lr=1e-3, compress_state_dict=False)
torch.save(optimizer.state_dict(), "checkpoint_fp32.pt")

Use enable_gradient_release to attach hooks to a model so that the optimizer steps each parameter as soon as its gradient is ready and immediately frees the gradient memory. This can significantly reduce peak memory usage during training.

Behavior:

• While active, calling optimizer.step() and optimizer.zero_grad() becomes a no-op (with a warning).
• Stepping and clearing are handled automatically by the hooks.

Compatibility & Restrictions:

• Supported: FSDP2 (fully_shard / FSDPModule).
• Unsupported (Raises TypeError): DistributedDataParallel (DDP) and FullyShardedDataParallel (FSDP1).
• Unsupported (Raises NotImplementedError): Providing a grad_scaler (integration is planned but not yet available).

Managing the lifecycle: enable_gradient_release returns a GradientReleaseHandle. Call its .remove() method to detach all hooks and restore normal (lazy) optimizer behavior.

from flashoptim import enable_gradient_release

handle = enable_gradient_release(model, optimizer)

# ... training loop ...

# Restore normal behavior
handle.remove()

Use cast_model to convert a model's parameters and buffers to a target dtype (e.g., torch.bfloat16).

Key features:

• Stability: By default, normalization layers with running statistics are kept in fp32 for training stability. Forward pre-hooks automatically upcast inputs to these fp32 modules.
• full_precision_layers: Specify layers to keep in fp32. Supports fnmatch patterns. Use "*.head" for nested modules (e.g., model.decoder.head).
• full_precision_recast_layers: Specify layers that should be kept in fp32 during computation, but whose outputs are recast back to the target dtype.

Note: Layer names are matched against the full dotted module name.

from flashoptim import cast_model

# Cast all parameters to bf16 (normalization layers kept in fp32 by default)
cast_model(model, dtype=torch.bfloat16)

# Terminal layers (e.g., lm_head) - kept fp32, output stays fp32
cast_model(model, dtype=torch.bfloat16, full_precision_layers=["lm_head", "*.head"])

# Middle layers - kept fp32 but output recast to bf16
cast_model(model, dtype=torch.bfloat16, full_precision_recast_layers=["target"])

# Module references work too
cast_model(model, full_precision_layers=[model.lm_head])

The master_weight_bits parameter in FlashOptim optimizers (like FlashAdamW) controls the precision of master weights maintained to save memory.

When using bf16/fp16 parameters, the optimizer quantizes moments to 8-bit and uses master_weight_bits to manage the precision of the master weights via a fused update kernel:

• master_weight_bits=24 (default): Uses 8-bit correction terms for 24-bit master weights. This is narrower than fp32 but preserves convergence.
• master_weight_bits=32: Uses 16-bit correction terms for full 32-bit master weight semantics.
• master_weight_bits=None: No master weight correction is used; parameters stay at their native precision (e.g., bf16), though optimizer states are still quantized to 8-bit.

When using fp32 parameters, master_weight_bits is not applicable because parameters are already full precision; however, optimizer states (moments) are still quantized to 8-bit to reduce memory.

from flashoptim import FlashAdamW

# Default: 24-bit master weights (bf16 param + 8-bit correction term)
optimizer = FlashAdamW(model.parameters(), lr=1e-3)

# 32-bit master weights (bf16 param + 16-bit correction term)
optimizer = FlashAdamW(model.parameters(), lr=1e-3, master_weight_bits=32)

# No master weight correction; parameters stay at native precision
optimizer = FlashAdamW(model.parameters(), lr=1e-3, master_weight_bits=None)