When Activation Checkpointing Calls a Stateful Quantizer Twice

## A delayed-scaling failure in FP8 and packed FP4—and the correctness rules that fixed it

Low-precision training bugs do not always look like low-precision bugs.

They may not crash the kernel. They may not produce an immediate `NaN`. They may survive single-forward unit tests, pass short smoke runs, and appear only as occasional gradient spikes deep into training.

That is what happened in a native CUDA trainer I was developing.

The trainer combined two techniques that are individually reasonable:

1. **Delayed scaling**, where a quantizer uses a previously computed scale while recording the current tensor’s maximum magnitude for a future scale update.

2. **Activation checkpointing**, where a layer’s forward computation is run again during backward to reconstruct activations that were not saved.

The failure appeared when a stateful quantizer lived inside the checkpointed layer body.

The real forward pass called the quantizer and advanced its persistent scaling state. Backward recomputation then called the same quantizer again for the same logical training step. That second call observed state that had already moved forward—and advanced it again.

The recompute pass was therefore not recomputing the original forward pass.

It was calculating a different tensor.

This article describes the bug class, how it was isolated, why the first attempted FP4 fix made the system worse, the public-safe correction pattern, and the tests that should exist anywhere delayed quantization state and activation recomputation meet.

> **Scope note:** This is a report about a custom native trainer. It is not a claim that NVIDIA Transformer Engine, PyTorch, or another library contains the same implementation defect.

However, the hazard class—a stateful quantizer mutated twice per logical step when activation checkpointing re-enters the layer body—is independent of any particular implementation. The correction pattern and testing checklist apply anywhere delayed quantization state meets recomputation.

## Delayed Scaling Is Stateful by Design

A low-precision format has a limited representable range. Before a higher-precision tensor can be quantized, its values must be mapped into that range using a scale.

One possible approach is current or just-in-time scaling:

```text

Read tensor

→ compute current amax

→ derive current scale

→ read or transform tensor again

→ quantize

```

That can be expensive because deriving an exact current scale may require additional tensor reads or synchronization.

Delayed scaling avoids that cost by using information from earlier calls:

```text

Call t:

use scale S\[t\]

quantize tensor X\[t\]

record amax A\[t\]

derive scale S\[t+1\] for a future call

```

NVIDIA’s Transformer Engine documentation describes delayed scaling as deriving scaling factors from historical `amax` values rather than computing them from the current tensor, reducing quantization from two tensor reads to one.[^1]

That means the quantizer is not a pure function of its tensor input.

It depends on persistent state:

```text

quantized_output = Q(input, scale_state)

```

and it mutates persistent state:

```text

scale_state_next = U(scale_state, current_amax)

```

This statefulness is intentional.

The correctness problem begins when call count, call order, ownership scope, or logical-step boundaries differ from what the scale update assumes.

## Activation Checkpointing Re-enters the Forward Path

Activation checkpointing saves memory by not retaining every intermediate activation from the original forward pass.

During backward, the checkpointed region is invoked again to reconstruct the activations required for gradient calculation. PyTorch documents that checkpointing reruns a forward-pass segment during backward and warns that, when the backward invocation differs from the original forward invocation because of persistent or global state, the result may include silently incorrect gradients.[^2]

For a pure function, recomputation is straightforward:

```text

Y = F(X, W)

```

If `X` and `W` are unchanged, calling `F` again should reproduce the same logical result, subject to explicitly managed randomness and numeric tolerance.

But a delayed-scaling quantizer changes the effective function:

```text

Y, state_next = F(X, W, state_current)

```

If the original forward mutates `state_current`, backward recomputation no longer receives the same state:

```text

True forward:

F(X, W, state\[t\])

    → Y_forward, state\[t+1\]

Recompute:

F(X, W, state\[t+1\])

    → Y_recompute, state\[t+2\]

```

The recompute pass is now operating one state transition ahead.

Worse, it advances the state a second time for one logical training step.

## The Core Failure in One Diagram

The intended behavior is:

```text

Logical step t

True forward:

snapshot S\[t\]

    → quantize with S\[t\]

    → record A\[t\]

    → advance live state to S\[t+1\]

Backward recompute:

reuse snapshot S\[t\]

    → reproduce the forward activation

    → do not record A\[t\] again

    → do not advance live state

```

The defective behavior was:

```text

Logical step t

True forward:

quantize with S\[t\]

    → record A\[t\]

    → advance state to S\[t+1\]

    → consume live descale from S\[t+1\]

Backward recompute:

quantize with S\[t+1\]

    → record A\[t\] again

    → advance state to S\[t+2\]

    → consume live descale from S\[t+2\]

```

This created two separate correctness failures:

1. **Same-call metadata mismatch:** quantized data produced with one scale was consumed using a descale value that had already been updated for a future call.

2. **Forward/recompute mismatch:** the recompute call used scaling state that the original forward did not use and advanced the history a second time.

Activation checkpointing amplified the first bug, but the first bug was already incorrect even without checkpointing.

## Why the Bug Was Hard to See

The defect did not initially look like a deterministic state-machine error.

It looked like unstable training.

Across attention backends, the trainer produced occasional extreme gradient-norm spikes. The spikes also appeared in a higher-precision attention control, which suggested that the attention kernel itself was not the common source.

The shared component was the low-precision activation path surrounding several projections.

A sequence of controls narrowed the cause.

### Control 1: Change the Matrix-Math Mode

A reduction-order or TensorFloat-32 explanation was tested by switching to a more conservative accumulation mode.

The failure reproduced at the same point and magnitude.

That made the matrix-math mode an unlikely cause.

### Control 2: Disable Low-Precision Activation Quantization

With the delayed-scaling activation path fully disabled, the original forward and backward recomputation produced matching activations across every instrumented micro-step in the test window.

That isolated the defect away from normalization, general checkpoint nondeterminism, and the attention implementation.

### Control 3: Compare Matched Points in Forward and Recompute

Instrumentation recorded intermediate magnitudes during the true forward and during recomputation for the same layer and micro-step.

At one projection, the input matched while the output did not.

That moved the investigation from “something upstream is unstable” to a much narrower conclusion:

> **The projection was consuming different quantization metadata in forward and recompute even when its input tensor matched.**

### Focused Verification

| Measurement | Before correction | After correction |

|—|—|—|

| Forward/recompute activation comparison | Intermittent mismatches | Bit-identical at all 10 instrumented steps |

| Abnormal gradient-spike behavior in the focused window | Approximately `9.46e6` to `1.66e7` across matched reruns | Approximately `1.05e6`, consistent |

The remaining transient was consistent with ordinary early-training behavior in that engine rather than the previous systemic forward/recompute divergence.

## Bug Class One: Advancing State Before the Consumer Is Finished

The public-safe shape of the defective code was approximately:

```cpp

QuantizedTensor quantize_activation(

Tensor input,

QuantState& state

) {

QuantizedTensor q = quantize_and_record_amax(

    input,

    state.scale,

    state.amax

);

update_delayed_scale(

    state.amax,

    state.history,

    state.scale,

    state.descale

);

return q;

}

// Later:

auto q = quantize_activation(input, state);

auto output = low_precision_gemm(q, state.descale);

```

The error is subtle but unconditional.

The quantized tensor `q` was produced using the scale that existed at the start of the call.

The update then replaced `state.descale`.

The GEMM consumed the new descale—not the descale paired with `q`.

The correct metadata had already been lost from the call site.

### Why Recomputation Made It Worse

The entire quantization function was reached again during activation recomputation.

That meant the recompute path:

- started with already-advanced state;

- recorded a second `amax` for the same logical step;

- advanced the history again;

- used a different scale than the forward pass;

- and again consumed metadata after it had moved.

If the scale is derived from a history window, the second call also inserts a duplicate logical observation into that history.

The bug therefore changes both the current tensor reconstruction and the future evolution of the quantizer.

## The Correctness Fix: Snapshot Before Advance

The correction was not simply “do not update during recompute.”

Every consumer also had to use the metadata that was actually paired with the quantized data.

The safe pattern is:

```cpp

QuantizedTensor quantize_activation(

Tensor input,

QuantState& live,

QuantSnapshot& snapshot,

bool is_recompute

) {

if (is_recompute) {

    return quantize_without_recording(

        input,

        snapshot.scale

    );

}

snapshot.scale = live.scale;

snapshot.descale = live.descale;

QuantizedTensor q = quantize_and_record_amax(

    input,

    snapshot.scale,

    live.amax

);

advance_delayed_state(live);

return q;

}

// Both forward and recompute consume the snapshot:

auto q = quantize_activation(input, live, snapshot, is_recompute);

auto output = low_precision_gemm(q, snapshot.descale);

```

The rules are:

1. **Freeze the state used by the true forward before mutating live state.**

2. **Pair the quantized tensor with its frozen metadata.**

3. **Let the true forward advance the live history exactly once.**

4. **Make recomputation read-only with respect to delayed state.**

5. **Require every downstream consumer to read the snapshot, not the live state.**

The distinction between live state and call-paired state must remain explicit.

A variable named “current descale” is dangerous when “current” can mean either:

- the descale used to encode the tensor currently being consumed; or

- the descale prepared for the next quantization call.

Those are different values.

## The Stronger Invariant

The fix can be stated as a general invariant:

> **For a given logical operation, every encoded value and every consumer of that value must use one immutable metadata snapshot, even if the operation is re-entered for recomputation.**

The live quantizer state belongs to the transition between logical calls.

The snapshot belongs to the data produced by one logical call.

Conflating those lifetimes is what created the defect.

## Bug Class Two: The State Had the Wrong Owner

A similar failure appeared in a custom packed-FP4 attention path, but the first attempted correction made the result dramatically worse.

The initial assumption was that the FP4 path had the same timing problem and therefore needed the same snapshot pattern.

It did have a timing problem—but that was not its first structural problem.

The scale, descale, and `amax` values for attention operands were stored as one shared set for the entire model rather than one set per layer.

That meant the state owner was wrong.

### Why Forward and Backward Order Exposed It

A transformer forward pass visits layers in ascending order:

```text

Layer 0 → Layer 1 → Layer 2 → … → Layer L-1

```

Backward recomputation visits them in reverse:

```text

Layer L-1 → … → Layer 2 → Layer 1 → Layer 0

```

With one shared scalar, the state observed by a layer during recomputation could be the state left behind by a different layer.

A snapshot cannot repair that.

It can faithfully preserve the wrong owner’s value.

The first naive snapshot attempt converted an intermittent defect into a more consistent one, producing a non-finite final gradient norm and causing 18 of 20 optimizer steps to be skipped.

The deeper investigation showed that the scale state needed independent ownership per layer before any timing protocol could be correct.

### The Actual Correction

The persistent state was changed conceptually from:

```cpp

float q_scale;

float q_descale;

float q_amax;

```

to:

```cpp

float q_scale[num_layers];

float q_descale[num_layers];

float q_amax[num_layers];

```

with equivalent per-layer state for the other attention operand.

Large temporary buffers could remain shared because their lifetime ended before the next layer reused them on the same stream.

Persistent state could not remain shared because it carried information across calls.

After per-layer isolation, the skipped-step count in the focused test returned from 18 of 20 to zero, confirming that the ownership defect had been removed.

Residual packed-FP4 instability still required separate investigation. Fixing state scope did not automatically make every coarse quantization recipe numerically suitable.

## Timing Bugs and Scope Bugs Are Different

This distinction is one of the most important lessons from the investigation.

### Timing Defect

The right state belongs to the right layer or worker, but it is read after being advanced for a future call.

Typical correction:

```text

snapshot before mutation

→ recompute reads snapshot

→ recompute does not mutate

```

### Scope Defect

The state is shared across entities that require independent histories.

Typical correction:

```text

allocate state per layer

or per worker

or per stream

or per quantization role

```

Adding a snapshot to wrongly scoped state may preserve corruption more reliably.

The order of operations should be:

1. Identify who owns the state.

2. Give every independent owner independent storage.

3. Define when the state may advance.

4. Snapshot the state used by a logical operation.

5. Make re-entry read-only unless it represents a new logical operation.

## Why FP4 Exposed the Problem More Violently

The FP8 activation defect produced a relatively small multiplicative mismatch per element, but that error propagated across projections, layers, and gradient computation.

The packed-FP4 experiments were less forgiving.

A custom mean-centered variant added another persistent quantity: a running center or offset.

If forward and recompute used different centers, the difference was not merely a small scaling ratio. It could be an additive shift comparable to the magnitude of the data itself.

That produced a useful general lesson:

> **Two state variables can have the same software-lifecycle defect while having very different numerical severity.**

A stale multiplicative scale and a stale additive offset are both examples of old state being used during recompute.

Their impact is not equivalent.

The state transition must be tested in the numerical context where it is used, not approved merely because a similar snapshot pattern worked elsewhere.

Later controlled reruns changed the empirical boundary of some early FP4 conclusions. The durable conclusion was not that all centered FP4 is unusable.

It was that delayed-state/recompute correctness must be solved before the numerical quality of a quantization recipe can be evaluated fairly.

## A Quantizer Should Expose Two Different Operations

One architectural improvement is to stop exposing one function that both computes output and mutates future state.

Instead, separate the API:

```cpp

QuantSnapshot begin_logical_quantization(

const QuantState& live

);

QuantizedTensor quantize_with_snapshot(

Tensor input,

const QuantSnapshot& snapshot

);

QuantObservation observe_for_future_state(

Tensor input

);

void commit_quant_observation(

QuantState& live,

const QuantObservation& observation

);

```

The true forward performs:

```text

snapshot

→ quantize

→ record observation

→ commit one state transition

```

Recompute performs:

```text

reuse snapshot

→ quantize

→ no observation commit

```

This design makes accidental mutation harder.

It also allows tests to assert that a logical training step produces exactly one state commit regardless of how many times the computational body is re-entered.

## State Needs an Epoch

A further safeguard is to associate delayed state with an explicit logical epoch.

For example:

```cpp

struct QuantSnapshot {

uint64_t logical_step;

uint32_t layer_id;

uint32_t role_id;

float scale;

float descale;

};

```

Consumers can then assert:

```cpp

assert(snapshot.logical_step == current_logical_step);

assert(snapshot.layer_id == executing_layer);

assert(snapshot.role_id == expected_quantization_role);

```

The exact structure is implementation-dependent, but the principle is general:

> **Persistent numeric state should carry enough identity to prove that it belongs to the operation consuming it.**

Without that identity, a stale or cross-layer value can remain numerically plausible while being semantically wrong.

## Concurrency Adds a Third Failure Mode

Activation recomputation is sequential re-entry.

Parallel sub-batches or multiple CUDA streams create concurrent re-entry.

If two workers mutate one delayed-scaling state object in parallel, snapshotting alone cannot make the transition correct.

Each worker may need independent state, or the system must serialize the state transition.

The native trainer therefore treated unsupported combinations conservatively: a low-precision profile with shared delayed state could fail closed to a single worker until per-worker state existed.

This is the scope rule one dimension deeper:

```text

per model

is weaker than

per layer

is weaker than

per layer × worker

```

The right shape depends on which entities can enter the quantizer independently.

## How to Test This Bug Class

A single-forward kernel test is insufficient.

The test must model the lifecycle that creates the failure.

### 1. Compare True Forward With Recompute

For the same logical step, capture selected intermediate tensors or stable summaries at matched boundaries.

Useful checks include:

- exact hashes when deterministic;

- bitwise equality when expected;

- maximum absolute difference;

- maximum relative difference;

- matching finite/non-finite masks;

- matching `amax` and scale snapshots.

Do not compare only the final loss.

A late scalar can hide where the divergence began.

### 2. Count State Transitions

Assert:

```text

one logical forward step

→ one amax-history insertion

→ one scale-state advance

```

Recompute should not increment that count.

### 3. Test Reverse Traversal

Run a multi-layer test where forward executes ascending layers and recompute executes descending layers.

This reveals shared-state ownership mistakes that one-layer tests cannot expose.

### 4. Disable Low Precision as a Control

A BF16 or FP32 control can distinguish:

- generic recompute nondeterminism;

- kernel errors unrelated to quantization;

- and stateful low-precision path defects.

### 5. Test the Call Graph, Not Only the Kernel

The quantization kernel may be locally correct while its caller consumes metadata after mutation.

Tests must include:

```text

quantize

→ update state

→ GEMM or attention consumer

```

The bug existed in the relationship between operations.

### 6. Preserve Native Telemetry

Do not suppress the only output stream that reports what the native kernel actually did.

Preserve native per-step metrics separately from smoke-test summaries.

At minimum, retain:

- loss by step;

- learning rate by step;

- gradient norm by step;

- skipped-step counts;

- clipping events;

- dominant gradient source;

- and finite/non-finite status.

### 7. Measure Training-Effective Work

High token throughput does not prove that training is healthy.

A run can keep the GPU busy while the optimizer skips nearly every update because gradients are non-finite.

Report at least:

- optimizer steps attempted;

- optimizer steps applied;

- skipped-step count;

- gradient norm;

- dominant gradient source;

- loss trajectory;

- and throughput.

Throughput without applied updates is not training progress.

## A Public Checklist for Stateful Quantization

Before placing a stateful quantizer inside a recomputed or concurrent region, ask:

### Ownership

- Is the state per tensor role?

- Per layer?

- Per microbatch?

- Per worker or CUDA stream?

- Could reverse traversal expose another owner’s state?

### Timing

- Which call is allowed to update `amax` history?

- Does the consumer use metadata from before or after the update?

- Can one logical step advance the state more than once?

### Recompute

- Does recompute use the exact state used by the true forward?

- Is recompute read-only?

- Can the framework distinguish original forward from recomputation?

### Pairing

- Is every quantized tensor permanently paired with the scale and descale that encoded it?

- Can downstream code accidentally read a newer live value?

### Testing

- Is there a full forward-plus-backward recomputation test?

- Is there a multi-layer reverse-order test?

- Is there a no-low-precision control?

- Are state-transition counts asserted?

- Are per-step native metrics preserved?

### Deployment

- Does the system fail closed for unsupported concurrent entry?

- Can the corrected path be verified outside the live training queue?

- Is the production binary traceable to the verified build?

## What This Finding Does Not Establish

This investigation does **not** establish that:

- delayed scaling is inherently unsafe;

- activation checkpointing is inherently unsafe;

- NVIDIA Transformer Engine contains this bug;

- PyTorch checkpointing contains this bug;

- all FP8 training paths are affected;

- all FP4 recipes are unstable;

- or a snapshot alone is sufficient for every stateful quantizer.

The defect was in a custom integration where delayed state was mutated inside a re-entered layer body and consumed with incorrect lifetime and ownership semantics.

Transformer Engine provides managed low-precision recipes and metadata handling. PyTorch documents that checkpoint recomputation can become incorrect when a recomputed invocation differs because of persistent state.[^1][^2]

The broader lesson is to respect those state boundaries when building custom kernels or bypassing managed framework behavior.

## The Deeper Engineering Lesson

Low-precision formats receive most of the attention, but this failure was not fundamentally about the number of exponent or mantissa bits.

It was about time.

A tensor was encoded at one moment using one state.

The metadata used to interpret it came from a later moment.

Backward recomputation then re-entered the state machine and moved time forward again.

The numerical error was the visible symptom.

The broken temporal contract was the cause.

The durable rule is:

> **When computation is replayed, persistent state must not be replayed as though a new logical event occurred.**

And before that rule can work:

> **The state must belong to the correct layer, role, worker, and stream.**

Delayed scaling and activation checkpointing can coexist.

But the quantizer must distinguish:

- live state from frozen call metadata;

- original execution from recomputation;

- one layer’s history from another’s;

- and one logical training step from repeated physical execution.

Once those distinctions are explicit, the bug becomes testable.

Before they are explicit, a trainer may continue producing tokens, losses, and checkpoints while silently calculating gradients from activations that never existed in the original forward pass.

That is the kind of failure a smoke test will miss—and a production training system cannot afford to ignore.

## Disclosure

This article intentionally omits proprietary implementation names, repository paths, model-family identifiers, full kernel code, private configuration values, and unreleased quantization-recipe details.

The reported mechanism and verification outcomes came from an isolated native CUDA/C++ investigation conducted without modifying the active production training queue.

## References

[^1]: NVIDIA, [FP8 Delayed Scaling — Transformer Engine documentation]( FP8 Delayed Scaling — Transformer Engine 2.16.0 documentation ). Delayed scaling estimates scaling factors from historical `amax` values and reduces quantization from two tensor reads to one.

[^2]: PyTorch, [`torch.utils.checkpoint` documentation]( Redirecting… ). Checkpointing reruns forward-pass segments during backward and warns that differences caused by persistent or global state may lead to silently incorrect gradients.