FP8 weight caching shows no speedup (sometimes slowdown) and causes training curve differences under Float8BlockScaling

[Mnb66]

Hi,

We are testing FP8 weight caching in Transformer Engine by passing is_first_microbatch=True/False during gradient accumulation. We are using the Float8BlockScaling recipe and we get two unexpected behaviors:

• Enabling FP8 weight caching does not improve training throughput in our setup. In some runs it even makes training slower.
• Enabling FP8 weight caching changes the training curve at some steps. Under Float8BlockScaling, we expected the runs with and without caching to be numerically identical, or at least fully aligned in practice.
  For Float8BlockScaling, our understanding from the implementation is that quantization is blockwise and does not rely on delayed amax_history state like DelayedScaling.
  Because of that, we expected:
• FP8 weight caching to provide some speedup when the same frozen weights are reused across microbatches.
• Runs with and without caching to remain fully aligned under Float8BlockScaling, since the weights are unchanged within a gradient accumulation cycle.

Questions:

• Is FP8 weight caching expected to provide a measurable speedup under Float8BlockScaling?
  Is there any known case where FP8 weight caching can be neutral or even slower with Float8BlockScaling?
• Under Float8BlockScaling, should runs with and without FP8 weight caching be bitwise identical, or at least numerically aligned step by step?
• If differences are expected, what is the source of the difference for Float8BlockScaling, given that it does not appear to use delayed amax_history state?

The current documentation warning about FP8 weight caching causing non-bitwise-identical outputs seems understandable for delayed-scaling recipes, but we are unsure whether that warning is also intended to apply to Float8BlockScaling.

Thank you for any info or advice in advance!

[ptrendx]

Hi @Mnb66, that is definitely unexpected behavior. You are correct that in the case of non-delayed recipes you should get exactly the same results with and without using this feature, and also you should see some speedup (since this would avoid requantizing weights every iteration). I will investigate, but if you could share some more information of your setup (or even better some script showing this behavior) that would be very helpful in narrowing down the issue.

[Mnb66]

Hi @Mnb66, that is definitely unexpected behavior. You are correct that in the case of non-delayed recipes you should get exactly the same results with and without using this feature, and also you should see some speedup (since this would avoid requantizing weights every iteration). I will investigate, but if you could share some more information of your setup (or even better some script showing this behavior) that would be very helpful in narrowing down the issue.

Hi @ptrendx thank you the rapid reply. I can briefly outline our implementation of FP8 weight caching, using our custom training framework built on TE and PyTorch (excluding Megatron).

First, we collect FP8 layers once before the training loop starts to avoid the overhead of traversing model.modules() during every iteration. Next, we determine the current step and inject a state into the module to indicate whether it is the first microbatch (setting module._is_first_microbatch to True or False). This logic is applied to instances such as LayerNormMLP and MHA. Finally, after the loop ends, we reset the first_microbatch state to None to prevent any interference with evaluation or prediction modes.

Given that the implementation is relatively straightforward, I don’t think there are any issues with my current approach (which has also undergone repeated code reviews using Claude Opus 4.6). However, if you can provide specific assurances regarding the performance and numerical consistency of the native TE fp8 cache implementation, I would be willing to conduct a more thorough investigation into my own code.