anthropics / anthropics/jacobian-lens

JacobianLens.save() silently overflows large finite Jacobians to inf with default float16

Abierto
#6 0 comentarios 0 reacciones 0 asignados Ver en GitHub
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.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.