Pull requests / #2247

#2247 --session-state-codec bf16: a saved conversation's running state at half the bytes

open · @sumomomomomo · 0 Kommentare · Auf GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentationLinux

Beschreibung

## What changes

Adds an opt-in `--session-state-codec bf16` engine argument to store saved conversation state as BF16 instead of FP32. It applies to parked conversation-cache images and session files. The default remains `fp32`; pass the option through the server config’s `args`.

The codec covers the GDN recurrence and convolution history, PLE history, and QSA indexer tails and dead keys. Integer `block_pos` values are copied unchanged.

For the documented model geometry, one running-state payload shrinks from **118,063,152 B to 59,031,600 B**. Parked images retain the checkpoint chain, so their savings grow with the number of checkpoints. Session files retain the live state and the selected resume checkpoint.

## What stays unchanged

- Kernels continue to compute with FP32 state. Stored BF16 values are widened to FP32 before use.
- K/V storage, checkpoint frequency, model/config fingerprints, and inference kernels are unchanged.
- Conversion runs on the host in slices using at most 4 MiB of FP32 scratch.
- This PR does not add denser checkpoints, disk-backed parking, or FP8 storage.

BF16 storage is lossy. Encoding uses round-to-nearest-even; decoding widens the stored value exactly. Very small values can round to zero, and values near the FP32 maximum can round to infinity.

## File compatibility

Each checkpoint records its codec. Files containing coded state use format version 2; all-FP32 saves continue to use version 1.

- A BF16 file can restore on a supported engine configured for FP32, and a version-1 file can restore on one configured for BF16.
- Unknown codec codes are rejected.
- Older builds reject version 2.
- Version-1 serialization is preserved. A 63,459-token FP32 save was byte-identical to the parent build, `a0116b27`, using the same input and settings.

The SYCL port does **not** implement the BF16 flag. Its build dependencies and SAVE caller are updated for the shared file API; its own saves remain FP32.

## Measurements

Measured on Linux with CUDA 12.8, a i9 13950HX QS, Tesla V100 32 GB, 96 GB RAM, IQ3_S, INT8 K/V, and a 262,144-token context. Configurations differed only in the state codec. Sizes below use decimal MB.

| Tokens | FP32 file | BF16 file | Reduction | Save ms, FP32 → BF16 | Restore ms, FP32 → BF16 |
| ---: | ---: | ---: | ---: | ---: | ---: |
| 63,459 | 1,205.3 MB | 1,087.2 MB | 9.8% | 546.4 → 503.2 | 759.5 → 714.3 |
| 32,650 | 734.8 MB | 616.7 MB | 16.1% | 444.1 → 315.5 | 468.5 → 400.3 |
| 37 | 236.7 MB | approximately 118.7 MB | approximately 50% | — | — |

Total savings depend on the proportion occupied by running state versus K/V. These are individual measurements, not a repeated performance benchmark.

On the measured 63K state, the largest relative rounding error was **0.3891%**. No sampled value changed sign, became zero, or became non-finite.

Needle recall after saving and restoring a 32K conversation succeeded at depths 10%, 50%, and 90% with both codecs. These checks provide limited functional validation, not a general quality guarantee.

## Validation and remaining work

The codec, session-file, SAVE-reclaim, snapshot, validation, conversation-cache, and memory tests pass. Coverage includes default serialization, codec rejection, BF16 round trips, cross-codec SAVE memory accounting, parking estimates, and allocation failure during checkpoint widening.

`ple_parity` fails on this checkout because its required model artifact is missing; the same failure occurs before this change.

Remaining validation:

- HIP and SYCL builds were not run; this machine has neither ROCm nor oneAPI.
- Multi-checkpoint parking followed by resume from an older checkpoint has not been tested with the model.
- Parked-image savings and 4K/128K cases have not been measured.
- No logits/KL comparison or error-versus-continuation-length study has been performed.
- `tools/ab_engine.py` has not been run.

Full methodology and measurements are in `docs/STATE_SNAPSHOT_BF16.md`.

Mehr auf der Site

Links zu Install, Modellen, Releases.