anthropics / anthropics/jacobian-lens
JacobianLens.save() silently overflows large finite Jacobians to inf with default float16
- Lenguaje dominante
- Python
- Estrellas
- 1.9k
- Forks
- 282
- Métricas de merge de PR
- Sin PR fusionados en 30 d
Descripción
## Summary
`JacobianLens.save()` defaults to `torch.float16` and casts every Jacobian with `J.to(dtype)` without checking whether the values are representable. Any finite value with magnitude above 65,504 becomes `inf`; after loading, ordinary matrix products can then become `NaN`.
This is reproducible independently of the fitting procedure:
```python
import tempfile
import torch
from jlens import JacobianLens
path = tempfile.mktemp(suffix=".pt")
J = torch.full((2, 2), 1e16, dtype=torch.float32)
lens = JacobianLens({0: J}, n_prompts=1, d_model=2)
lens.save(path) # default dtype=torch.float16
loaded = JacobianLens.load(path)
print(torch.isfinite(lens.jacobians[0]).all()) # tensor(True)
print(torch.isinf(loaded.jacobians[0]).all()) # tensor(True)
print(loaded.transport(0, torch.ones(2))) # tensor([inf, inf])
```
A mixed-sign matrix or downstream subtraction readily turns those infinities into `NaN`.
## Real checkpoint where this occurred
I hit this while fitting a lens for a LoRA-finetuned Qwen3-1.7B model:
- In the fp32 fit checkpoint, layers 0-16 were finite, with per-layer maximum absolute entries between approximately `1.17e16` and `5.43e16`.
- After the default `save()`, all `4,194,304` entries in each of layers 0-16 were `inf` in the fp16 lens file.
- Layers 17-26 remained finite because their values were O(1).
- Rebuilding the lens with `dtype=torch.float32` restored finite downstream computations.
A separate fit in the same project had O(1) Jacobians at every layer and saved correctly. Therefore the huge Jacobians appear to be a fit-specific numerical-conditioning problem, while the silent fp16 overflow is a separate serialization problem that can corrupt any out-of-range lens.
The pathological fit also had extremely negative shallow-layer R-squared values (around `-1e25`), so I do not claim that `1e16` is a valid expected Jacobian scale. The issue is that `save()` silently turns finite tensors into a non-finite checkpoint instead of preserving them or reporting that the requested storage dtype is unsafe.
## Suggested behavior
Any of the following would prevent silent corruption:
1. Keep fp32 as the default storage dtype.
2. Before casting, check the target dtype range and raise a clear error or warning if any entry is out of range.
3. Automatically fall back to bf16/fp32 when fp16 cannot represent the data, and report the fallback.
4. At minimum, check `torch.isfinite()` after conversion and refuse to save a corrupted checkpoint.
Tested against commit `581d398613e5602a5af361e1c34d3a92ea82ba8e` with Python 3.12.12 and PyTorch 2.11.0+cu128.
Guía de contribución
No hay ninguna guía de contribución indexada para este repositorio
Evaluación
Este issue todavía no se ha evaluado.