deepseek-ai / deepseek-ai/DeepSpec
feat: add support for FP8 target cache quantization
- Dominant language
- Python
- Stars
- 7.1k
- Forks
- 667
- PR merge metrics
- No merged PRs in 30d
Description
# Feature Description
Currently, preparing the target cache (saving the intermediate target hidden states and last hidden states) requires a significant disk footprint. For large-scale pretraining or instruction-tuning datasets, target caches can easily scale up to multiple terabytes of storage because activations are stored in 16-bit formats (`bfloat16`).
I would like to propose adding support for **FP8 quantization** (`float8_e4m3fn` and `float8_e5m2`) to reduce the target cache storage footprint by **50%** (1 byte per element instead of 2 bytes).
---
## Proposed Solution
The implementation will introduce:
1. **Configurable CLI parameter**: Add a `--hidden-dtype` argument to `prepare_target_cache.py` to allow users to generate caches in `bfloat16`, `float8_e4m3fn`, or `float8_e5m2`.
2. **Quantized Serialization**: Quantize the hidden state tensors to the specified FP8 format and save them as raw bytes in the cache shards.
3. **Transparent Dequantization**: Read the float8 bytes back as float8 tensors and cast them back to `bfloat16` on the fly inside `CacheDataset.__getitem__`. This prevents any downstream model or trainer code from having to be modified, as they will still receive standard `bfloat16` tensors.
4. **Full Backward Compatibility**: Caches generated in older versions (or without specifying a dtype) will default to `bfloat16` and load normally.
---
## Implementation Details & Impact
I have already implemented a working prototype of this and validated it:
* **Storage savings**: Exactly 50% reduction in shard file sizes on disk.
* **Precision impact**:
* `float8_e4m3fn` introduces a Mean Absolute Error (MAE) of only `~0.0090` (near-lossless).
* `float8_e5m2` introduces an MAE of only `~0.0179`.
* **Testing**: Includes a comprehensive suite of unit tests covering precision degradation, dataloader collation, and end-to-end integration tests.
I'd love to get feedback on this feature and submit a pull request!
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.