mixle.models.memory_efficient_training module¶
Memory efficiency for training state (roadmap F6): fp8 hardening + compressed optimizer moments + a per-block selective activation-recompute policy.
Three pieces, per the roadmap spec:
fp8 hardening (
fp8_cast_with_guard()) – the existing fp8 mention in this codebase (mixle/utils/parallel/torch_neural.py’sprecisiondocstring:"fp32"|"bf16" (fp8 = torchao, vendored)) is, honestly, just a comment: no fp8 code path actually exists there yet. “Hardening” here means building the real thing with real edge-case handling – using torch’s NATIVEfloat8_e4m3fn/float8_e5m2dtypes (no torchao dependency needed for the guard logic itself) plus explicit overflow detection (values exceeding the format’s representable range, which fp8 hardware silently clamps to +/-inf rather than raising), underflow detection (fp8’s tiny dynamic range flushing a large fraction of small-but-nonzero values to zero), and a graceful fallback to a wider dtype (bf16/fp32) whenever either guard fires – rather than accepting a silently-corrupted fp8 tensor.Optimizer-state compression (
CompressedOptimizerState,CompressedAdam) – Adam-style optimizers keep two moment buffers (m,v) at the SAME size as the parameters they track (states are 2x params fp32, per the F6 spec). This directly reuses G4 (mixle.models.sorted_profile_quantizer.fit_sorted_profile/reconstruct), which was explicitly built with “optimizer states (F6)” as its first named use case, as one compression path, PLUS a simpler/cheaper 8-bit blockwise quantization path (the standard bitsandbytes-style 8-bit-Adam technique) as a fast alternative when G4’s fuller (and more expensive to fit) sorted- profile machinery is not worth its cost.choose_compression_method()picks between them per tensor using a goodness-of-fit-based rule (reusing G4’s own KS-statistic receipt for the G4 path, and a real reconstruction-error check for the int8 path), with a DENSE fallback – mirroring G4’s own “receipt-driven dense fallback” pattern – when neither compressed representation is trustworthy for a given (possibly adversarial) tensor.Selective activation-recompute policy (
SelectiveRecomputePolicy) – extendsmixle.models.transformer.CausalLM’s previously all-or-nothinggradient_checkpointingbool flag to a PER-BLOCK decision, using a real cost model: a block’s stored-activation memory footprint (the benefit of recomputing it instead) versus its recompute FLOP cost. This mirrors the cost/benefit-tradeoff SHAPE of D6’s compile economics (mixle.inference.backend_respecialization.estimate_compile_cost/estimate_compile_benefit, PR #153) – a fixed/proxy-unit cost estimate compared against a fixed/proxy-unit benefit estimate, with a net-benefit-positive gate – applied here to a different decision (recompute vs. store, not eager vs. compiled). D6 itself lives on a separate, not-yet-merged branch (backend-respecialization) so it is not imported here; the pattern is reapplied standalone.
- fp8_cast_with_guard(tensor, fp8_dtype='float8_e4m3fn', fallback_dtype=None, underflow_fraction_threshold=_DEFAULT_UNDERFLOW_FRACTION_THRESHOLD)[source]
Attempt an fp8 cast of
tensor, guarded against the two ways fp8 silently corrupts data.- Parameters:
tensor (Any) – A torch tensor (any float dtype).
fp8_dtype (str) –
"float8_e4m3fn"(higher precision, smaller range – the default, matching common fp8-training practice for weights/activations) or"float8_e5m2"(wider range, coarser precision).fallback_dtype (Any) – Dtype to fall back to when a guard fires. Defaults to
torch.bfloat16(the codebase’s existing wide-range training dtype, permixle.utils.parallel.torch_neural).underflow_fraction_threshold (float) – Maximum tolerable fraction of nonzero input values that are allowed to flush to exactly zero under the fp8 round-trip before the underflow guard fires.
- Returns:
always carries the real, computed guard statistics, whether or not the fp8 cast was ultimately accepted – consistent with this codebase’s receipt-over-silent- assumption convention (see
mixle.models.sorted_profile_quantizer’s goodness-of-fit receipt).- Return type:
Fp8CastResult
- class Fp8CastResult(tensor, used_fp8, reason, max_abs, underflow_fraction)[source]
Bases:
objectReceipt of one
fp8_cast_with_guard()call – never silently swallowed.- tensor
The output tensor – either the fp8-cast tensor (
used_fp8=True) or thefallback_dtypecast (used_fp8=False).- Type:
Any
- used_fp8
Whether the fp8 cast was accepted.
- Type:
- reason
Human-readable reason for the decision (acceptance or the specific guard that fired).
- Type:
- max_abs
The input tensor’s max absolute value (0.0 for an empty tensor) – the statistic the overflow guard checks.
- Type:
- underflow_fraction
Fraction of the input’s nonzero values that flushed to exactly zero under the fp8 round-trip – the statistic the underflow guard checks.
0.0when the overflow guard fired first (the round-trip was never attempted).- Type:
- quantize_int8_blockwise(flat, block_size=_DEFAULT_INT8_BLOCK_SIZE)[source]
Dynamic per-block symmetric int8 quantization (bitsandbytes-style 8-bit-Adam technique).
Each contiguous block of
block_sizeelements gets its own scale (absmax / 127), so one extreme value only degrades the resolution of ITS OWN block rather than the whole tensor.
- dequantize_int8_blockwise(codes, scales, block_size=_DEFAULT_INT8_BLOCK_SIZE)[source]
Invert
quantize_int8_blockwise(). Returns afloat32array the same length ascodes.
- class CompressedMomentEncoding(method, shape, g4_encoding=None, int8_codes=None, int8_scales=None, int8_block_size=2048, dense_values=None)[source]
Bases:
objectStorage format for ONE moment tensor (
morv), compressed by exactly one of the three available methods – only the fields formethodare populated, mirroringSortedProfileEncoding’s single-active-branch convention.- Parameters:
- method
"g4"(sorted-profile,mixle.models.sorted_profile_quantizer),"int8"(blockwise quantization), or"dense"(fallback – neither compressed representation was trustworthy for this tensor).- Type:
- g4_encoding
Populated iff
method == "g4".- Type:
mixle.models.sorted_profile_quantizer.SortedProfileEncoding | None
- int8_codes / int8_scales / int8_block_size
Populated iff
method == "int8".
- dense_values
Populated iff
method == "dense".- Type:
numpy.ndarray | None
- choose_compression_method(tensor, *, min_size_for_g4=_DEFAULT_MIN_SIZE_FOR_G4, g4_top_k_fraction=_DEFAULT_G4_TOP_K_FRACTION, g4_gof_threshold=DEFAULT_GOF_THRESHOLD, tail_family=None)[source]
Goodness-of-fit-based per-tensor method picker (the “one level down” method-picker pattern I1 uses for its own picker –
mixle.task.banditis a reasonable fit for a LEARNED, reward- driven picker, but the choice here is cheaper and just as principled as a fixed rule: G4 already computes a real KS goodness-of-fit receipt as part of fitting, so reusing THAT receipt directly is simpler than bolting on a bandit that would need to learn what the receipt already tells us for free).Rule: for tensors at or above
min_size_for_g4(below which G4’s fixed per-tensor fitting overhead is not worth it), attempt a G4 fit; if it does not dense-fall-back on its own goodness-of-fit receipt, use G4 (the more accurate, more expensive path). Otherwise, use int8 (the cheap path) – callers should still checkcompress_moment()’s returnedmethodfor a possible further downgrade to"dense"if int8 itself proves untrustworthy for this specific tensor (see_DEFAULT_INT8_ADVERSARIAL_RELATIVE_ERROR).
- compress_moment(tensor, method='auto', *, min_size_for_g4=_DEFAULT_MIN_SIZE_FOR_G4, g4_top_k_fraction=_DEFAULT_G4_TOP_K_FRACTION, g4_gof_threshold=DEFAULT_GOF_THRESHOLD, int8_block_size=_DEFAULT_INT8_BLOCK_SIZE, int8_adversarial_relative_error=_DEFAULT_INT8_ADVERSARIAL_RELATIVE_ERROR, int8_flushed_fraction_threshold=_DEFAULT_INT8_FLUSHED_FRACTION_THRESHOLD, tail_family=None)[source]
Compress one Adam moment tensor (
morv) via G4, int8, or dense storage.- Parameters:
tensor (Any) – A torch tensor or numpy array (one Adam moment buffer, any shape).
method (str) –
"auto"(usechoose_compression_method()),"g4","int8", or"dense"to force a specific path. A forced"g4"/"int8"still honestly downgrades to"dense"if the chosen method’s own receipt (G4’s KS statistic, or int8’s reconstruction error) rejects the fit – this function never returns a silently bad compressed representation.min_size_for_g4 (int)
g4_top_k_fraction (float)
g4_gof_threshold (float)
int8_block_size (int)
int8_adversarial_relative_error (float)
int8_flushed_fraction_threshold (float)
tail_family (Any)
- Returns:
CompressedMomentEncoding
- Return type:
CompressedMomentEncoding
- decompress_moment(encoding)[source]
Invert
compress_moment(). Returns afloat32array reshaped toencoding.shape.- Parameters:
encoding (CompressedMomentEncoding)
- Return type:
- class CompressedOptimizerState(shape, method='auto', **compress_kwargs)[source]
Bases:
objectAdam-style
m/vmoment storage for ONE parameter tensor, held COMPRESSED between optimizer steps (rather than as two dense fp32 buffers).setcompresses;getdecompresses back to plain tensors for the optimizer math to use.v(Adam’s second moment, an EMA of squared gradients) is stored as its COMPRESSED SQUARE ROOT rather than compressed directly. This is a real, load-bearing design choice, not cosmetic:vroutinely spans many orders of magnitude within one parameter tensor (a few large-gradient elements next to many near-zero ones), which is exactly the dynamic range that defeats both compression paths – linear int8 quantization sets one block-wide scale from the block’s max, so any element more than ~127x smaller than that max quantizes to LITERALLY zero; G4’s Gaussian tail fit is symmetric and can reconstruct small true values as slightly negative. Either failure, fed straight into Adam’s1/(sqrt(v_hat) + eps)denominator, produces a step blown up by orders of magnitude (empirically confirmed: an early version of this module, compressingvdirectly with int8, diverged to a >100x loss spike within ~5 steps on the tiny transformer this module’s own test suite trains).sqrt(v)roughly halves the dynamic range in log-space (Adam’s own second-moment buffer already tracks squared gradients FOR this reason –sqrt(v)is the RMS gradient-magnitude scale, the quantity Adam’s update actually normalizes by), and squaring the reconstructed value back onget()is a nonnegative projection for free – it can never reconstruct a negativev, closing the negative-v/NaN failure mode without a separate clamp needing to paper over it (a defensiveclamp_min(0.0)is still applied as a last line of defense inCompressedAdam, since callers ofCompressedOptimizerStatedirectly should not have to know this internal detail to stay safe).- nbytes()[source]
Total measured compressed footprint of both moment buffers, in bytes.
- Return type:
- class RecomputeDecision(block_index, should_recompute, estimated_cost, estimated_benefit, activation_bytes, recompute_flops, rationale)[source]
Bases:
objectThe cost/benefit tradeoff and chosen action for ONE block’s activation-recompute policy – mirrors D6’s
RespecializationDecisionshape: a flat, inspectable dataclass carrying both the raw estimates and the derived decision, not just a boolean.- Parameters:
- property net_benefit: float
estimated_benefit - estimated_cost– positive iff recomputing this block is worth it.
- class SelectiveRecomputePolicy(memory_value_per_byte=_DEFAULT_MEMORY_VALUE_PER_BYTE, flop_cost_per_unit=_DEFAULT_FLOP_COST_PER_UNIT)[source]
Bases:
objectPer-block, cost-model-driven activation-checkpointing decision – extends
mixle.models.transformer.CausalLM’s previously all-or-nothinggradient_checkpointingbool flag (see that module’sforward, which now also accepts a per-block list) to a PER-BLOCK decision.A block is recommended for recompute when the value of the memory freed (its stored- activation footprint, valued at
memory_value_per_byte) exceeds the cost of the extra compute spent recomputing it (its recompute FLOPs, valued atflop_cost_per_unit) – the same cost-vs-benefit tradeoff SHAPE as D6’s compile economics, applied to this different decision (recompute-vs-store, not eager-vs-compiled).- decide_model(lm, batch, seq_len, dtype_bytes=4)[source]
Decide per-block recompute for every block of a
CausalLM(mixle.models.transformer.build_causal_lm), using its ownd_model/n_layer.
- apply_to_model(lm, batch, seq_len, dtype_bytes=4)[source]
Compute the per-block decisions and set them directly as
lm.gradient_checkpointing(a per-block bool list – seemixle.models.transformer.CausalLM.forward, which accepts either a single bool for all blocks or a per-block list).
- estimate_block_activation_bytes(batch, seq_len, d_model, dtype_bytes=4)[source]
Estimate the memory footprint of ONE transformer block’s stored output activation (
mixle.models.transformer.Block’s output, shape(batch, seq_len, d_model)) – the memory that activation checkpointing (recomputing instead of storing) frees.
- estimate_block_recompute_flops(batch, seq_len, d_model)[source]
Estimate the FLOP cost of recomputing ONE transformer block’s forward pass, tied directly to
mixle.models.transformer.Block’s actual layer shapes:qkv(d -> 3d),proj(d -> d), and the two MLP linears (d -> 4d -> d) give3d^2 + d^2 + 4d^2 + 4d^2 = 12*d_model^2linear-layer parameters per block; the standard “2 FLOPs per parameter per token” forward-pass heuristic turns that into a FLOP estimate, plus the attention score/value matmuls (QK^Tandattn @ V, each~2*batch*seq_len^2*d_modelFLOPs) that scale withseq_len^2rather than with parameter count.
- estimate_recompute_benefit(activation_bytes, memory_value_per_byte=_DEFAULT_MEMORY_VALUE_PER_BYTE)[source]
The value of the memory freed by recomputing (rather than storing) one block’s activation.
- estimate_recompute_cost(recompute_flops, flop_cost_per_unit=_DEFAULT_FLOP_COST_PER_UNIT)[source]
The cost of the extra compute spent recomputing one block’s activation during backward.
- class CompressedAdam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, compression_method='auto', **compress_kwargs)[source]
Bases:
OptimizerAdam whose per-parameter
m/vmoment buffers are held COMPRESSED (CompressedOptimizerState) between steps, instead of as two dense fp32 buffers – the “optimizer-state compression” half of F6.Mirrors the well-known 8-bit-Adam pattern (bitsandbytes): decompress -> take the exact Adam update in the parameter’s own dtype -> recompress. The per-step update math is byte-for-byte standard Adam; only the AT-REST storage between steps differs, so
CompressedAdamwithcompression_method="dense"is Adam with no approximation at all (a useful sanity check, exercised by the loss-parity test).Honest cost note: this reference implementation recompresses BOTH moment buffers every step, including (for
compression_method="g4"/"auto") re-fitting G4’s distribution and KS test every step – far more compute than a real deployment would spend (a production system would compress at a coarser cadence, e.g. only on optimizer-state checkpoint/ offload, not every step). Nothing here changes the per-step Adam math itself; only the wall-clock cost of this particular reference cadence differs from a production one.- Parameters: