DPO reference log-prob cache is not found when the DeepSpeed stage changes

Description

With dpo.precompute_ref_log_probs: true, TRL computes the reference log-probs once and saves them to a
.npz file next to the dataset cache. A later run loads that file instead of computing them again.

The cache only works when the DeepSpeed stage of the later run matches the stage of the run that wrote
it. TRL names the file after
Hasher.hash((dataset._fingerprint, hash_module(model)))
(dpo_trainer.py L917–L921),
and hash_module hashes the bytes of model.state_dict()
(utils.py L1071–L1079).
Under ZeRO-3 the parameters are already partitioned when DPOTrainer.__init__ runs, so the same
checkpoint gives a different hash under ZeRO-3 than under ZeRO-2. The file name changes with it.

A run with a different stage looks for a file that does not exist, computes the log-probs again, and
writes a second cache file. Nothing is logged. This rules out the obvious workflow: precompute once
under ZeRO-3, then train under ZeRO-2.

Reproduction

1. Precompute under ZeRO-3:

   python scripts/submit.py --config configs/trl/dpo.yaml --tokenize-only \
       dpo.precompute_ref_log_probs=true deepspeed.zero_optimization.stage=3
   

2. Train with the same model and data under ZeRO-2 (the stage in configs/trl/dpo.yaml):

   python scripts/submit.py --config configs/trl/dpo.yaml dpo.precompute_ref_log_probs=true
   

3. The second job shows Computing reference log probs for train dataset again, and the dataset cache
   directory now holds two .npz files.

Expected Behavior

The second run loads the cache that the first run wrote. The cache key depends on the model weights and
the data, not on how DeepSpeed partitions the weights.

Environment (if relevant)

TRL 0.29.0, DeepSpeed 0.18.6, PyTorch 2.9.1, Python 3.12

Start in trl/trainer/dpo_trainer.py around the cache naming logic at lines 917–921, then inspect hash_module in trl/trainer/utils.py at lines 1071–1079. Reproduce the issue with the two commands in the report using DeepSpeed stages 3 and 2. Done means the second run reuses the first .npz cache and does not create a second file or recompute the reference log-probs.