Buckets:
| """Numerical-state guards shared by training modules.""" | |
| from __future__ import annotations | |
| from collections.abc import Iterable | |
| import torch | |
| def finite_tensors(tensors: Iterable[torch.Tensor], *, chunk_size: int = 64) -> bool: | |
| if chunk_size <= 0: | |
| raise ValueError("chunk_size must be positive") | |
| grouped: dict[tuple[torch.device, torch.dtype], list[torch.Tensor]] = {} | |
| for tensor in tensors: | |
| if tensor.is_floating_point() or tensor.is_complex(): | |
| grouped.setdefault((tensor.device, tensor.dtype), []).append( | |
| tensor.detach() | |
| ) | |
| for values in grouped.values(): | |
| for start in range(0, len(values), chunk_size): | |
| norms = torch._foreach_norm(values[start : start + chunk_size], float("inf")) | |
| if not bool(torch.isfinite(torch.stack(norms)).all()): | |
| return False | |
| return True | |
| def finite_parameters(parameters: Iterable[torch.nn.Parameter]) -> bool: | |
| return finite_tensors(parameter for parameter in parameters) | |
| def finite_optimizer_state(optimizer) -> bool: | |
| return finite_tensors( | |
| value | |
| for state in optimizer.state.values() | |
| for value in state.values() | |
| if isinstance(value, torch.Tensor) | |
| ) | |
| def assert_finite_training_state(parameters, optimizer) -> None: | |
| if not finite_parameters(parameters) or not finite_optimizer_state(optimizer): | |
| raise FloatingPointError( | |
| "student parameters or optimizer state became non-finite after step" | |
| ) | |
| __all__ = [ | |
| "assert_finite_training_state", | |
| "finite_optimizer_state", | |
| "finite_parameters", | |
| "finite_tensors", | |
| ] | |
Xet Storage Details
- Size:
- 1.66 kB
- Xet hash:
- e4fc57db965830e3b1919acb415237e1b8f2678622ac9b0b7247d336682ddd11
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.