The same weights, two execution paths, two answers
By Ajish Pradeep · · 5 min read
Claim. A model's evaluation can run a different kernel from its training on the same weights. In my test PyTorch's transformer fast path matched the standard path to seven decimal places on real tokens and returned zeros on padded ones where the standard path did not, so a pooling step that ignored the mask moved logits by up to 7.5.
A trained model is usually thought of as its weights: run them, and you get its function. In practice the function also depends on which code runs them. Frameworks ship several implementations of the same operation, fused kernels for inference, nested tensors that skip padding, different precision paths. They pick one automatically, depending on the mode, the input and the hardware. I wanted to see whether two such paths can disagree on the same weights, and what a model can do to make that disagreement matter.
What the two paths do with padding
PyTorch's transformer encoder (Paszke et al., 20192), built on the architecture of Vaswani et al. (2017)1, has an inference fast path that it uses in evaluation mode when certain conditions hold, and a standard path that it always uses in training mode. With a padding mask, the fast path packs the real tokens into a nested tensor and never computes the padded positions; on the way out it fills them with zeros. The standard path computes every position and lets the mask keep padding out of the attention. Its outputs at padded positions are whatever the layers produce from them: not zeros, but not used by anything that respects the mask either.
I checked a two-layer encoder in evaluation mode with the fast path on and off, on CPU and on GPU, with a batch where the last five of 20 positions were padding:
| Largest difference on real tokens | Largest output on padded positions | |
|---|---|---|
| Fast path | 0 | |
| Standard path | 3.9 on CPU, 4.1 on GPU |
On the real tokens the two paths agree to rounding error. On the padded positions they disagree completely. Neither is wrong; the attention never reads padded positions, so their values are a free choice of the implementation.
A model that depends on the free choice
The disagreement only matters to code that reads the padded positions anyway. The common way to do that by accident is pooling: averaging the encoder's outputs over the whole sequence to get one vector for classification, without excluding padding. I trained a small classifier both ways, with unmasked and with masked mean pooling, on sequences of random length up to 24, three seeds each. Training always runs the standard path. I evaluated each trained model in evaluation mode with the fast path on and off.
| Pooling | Accuracy, fast path | Accuracy, standard path | Largest logit difference |
|---|---|---|---|
| Over every position | 99.95% | 99.96% | 7.5 |
| Over real tokens only | 100% | 100% | 0 |
With unmasked pooling, the same weights produced logits up to 7.5 apart depending on which path ran. The model was trained on the standard path's padded values and evaluated, by default, on the fast path's zeros. On this easy task the top class rarely changed, so accuracy hid it. Calibration, thresholds and anything else that reads the scores themselves would not. With masked pooling the two paths gave identical outputs, to the last bit.
The contract, and how to hold a model to it
The fix for this case is simply correct masking. The broader lesson is that a model has a numerical contract with its framework: the outputs it relies on must be the ones every execution path computes the same way. Values that are only defined by convention, such as padded positions, unused tokens or the result of an operation on an empty set, are outside that contract and may differ from path to path.
A cheap test catches the whole class: run the same weights, in evaluation mode, on the same batch with each optional fast path enabled and disabled, and in each precision the model will be served in, and compare the outputs. Any difference larger than rounding means the model reads something it shouldn't. It belongs next to the model's other tests, because the framework's choice of path can change with a version upgrade.
What else differs between a saved model and a restored one is collected in a checkpoint is more than its state dict, and how the framework's asynchronous execution decides where training time goes is in every .item() call makes the GPU wait for Python.
What I still don't know
I checked one framework version and one module. The conditions under which fast paths engage, and what they put in positions they skip, change between versions; a model that passes this test today can fail it after an upgrade, which is the reason to keep the test. I also haven't looked at export: a model exported to another runtime runs yet another implementation, and comparing its outputs against the training-mode path is the same test with higher stakes.
References
- Attention Is All You Need — Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser and Illia Polosukhin, NeurIPS 2017
- PyTorch: An Imperative Style, High-Performance Deep Learning Library — Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai and Soumith Chintala, NeurIPS 2019