[Bug] prepare_hidden_states.py needs sglang==0.5.3 or lower

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
• 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 5. Please use English, otherwise it will be closed.

Describe the bug

when prepare_hidden_states for offline training, we meet an error: load_model() missing 1 required positional argument: 'tp_rank'
requirements.txt requeres sglang[all]==0.5.4, but in 0.5.4, func load_model needs 4 argument. Change sglang version to 0.5.3 will solve this problem.

Reproduction

torchrun --nproc_per_node=1
scripts/prepare_hidden_states.py
--model-path /mnt/home/xhd/models/Llama-3.1-8B
--enable-aux-hidden-states
--data-path /mnt/home/xhd/SpecForge-main/longbench-E/longbenche_4-8k.jsonl
--chat-template llama3
--max-length 8192
--tp-size 1
--batch-size 2
--mem-frac=0.8

Environment

specforge==0.1.0
sglang==0.5.4

Start with requirements.txt and scripts/prepare_hidden_states.py, then run the provided torchrun reproduction in the stated environment. Confirm the dependency constraint is compatible with the script and that offline preparation no longer raises the missing tp_rank argument error.