[Feature] Optimization of Storage Efficiency for Hidden States in Large-Scale Models

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Feature Request: Optimization of Storage Efficiency for Hidden States in Large-Scale Models

Subject: Addressing prohibitive disk storage requirements for long-context hidden states (100B+ models)

1. Problem Description

The current implementation of prepare_hidden_states.py results in unsustainable disk usage when processing large-scale models with long contexts.

Observed Metrics:

• Target Model: GPT-OSS 120B (Hidden size = 2880).
• Context Length: 65,536 tokens.
• Storage Density: 53 samples currently occupy 62 GB (approx. 1.17 GB per sample).
• Projected Scaling: * 1,000 samples 1.2 TB
• 30,000 samples 35 TB

This magnitude of storage is a major barrier for standard research environments and creates a massive I/O bottleneck during the subsequent training phase.

2. Inquiry for Improvement

Given the current architecture of SpecForge, what are the recommended strategies or planned features to mitigate this storage crisis? Specifically, we would like to inquire about the feasibility and roadmap for the following:

• Dynamic Sequence Length Storage: How can the script be modified to save only the active tokens (unpadded) based on the attention_mask instead of the full max_length buffer?
• Low-Precision Storage Formats: Is there a plan to support storing hidden states in lower precision (e.g., FP8 or INT8 with scaling factors) to reduce the per-sample footprint by 50-75%?
• On-the-fly Compression: How can integration with high-speed compression libraries (like Zstd or LZ4) be implemented within the _save_tensor_async pipeline?
• Sparse Storage Options: For models with extremely high dimensions, are there plans to support sparse saving or dimensionality reduction for the auxiliary states?
• Hybrid GPU Partitioning for Online Extraction: Alternatively, is there support (or a plan to support) a hybrid execution mode to avoid massive disk storage entirely? For example, in a 4-GPU setup, could we partition the resources such that 2 GPUs are dedicated to training the speculative heads while the other 2 GPUs perform online extraction of hidden states from the target model in real-time?

3. Goal

The objective is to find a path that reduces the storage requirement from ~1.2 GB/sample to a manageable level (e.g., <200 MB/sample) without compromising the training integrity of the speculative heads.

We look forward to hearing your insights on how to optimize the storage backend for these high-context, large-model scenarios.

Related resources

No response

Start with prepare_hidden_states.py and the _save_tensor_async pipeline, then trace how attention_mask and hidden states are currently stored. The request presents several possible strategies rather than one scoped change, so a maintainer must first choose and define an approach. Done would require an agreed implementation and validation against the stated goal of reducing storage below 200 MB per sample.