python -m math_jlens.explore \
    --model ../Lightning-OPD/checkpoints/qwen3-4b-base-sft-qwen3-8b/ \
    --data data/dapo-math-17k/dapo-math-17k.jsonl  \
    --lens outputs/multi-gpu-20/lens-bf16.pt \
    --fit-count 20  \
    --seed 17 \
    --sample-index 0  \
    --device cuda:4