# Tested shape for the GPU port. Pin these; transformers in particular is load-bearing
# because the script relies on cu_seq_lens_q/k reaching flash_attn_varlen_func
# (transformers.modeling_flash_attention_utils, "Case 2" padding-free path).
--extra-index-url https://download.pytorch.org/whl/cu128
torch==2.8.0
transformers==4.57.5
tokenizers>=0.22
datasets>=3.1.0,<4.0.0   # v4 dropped dataset-script support; also what the marin eval branch pins
accelerate>=1.0.0
safetensors>=0.4.5
numpy>=1.26
jinja2>=3.1.0            # only needed by prepare_sft_data.py --verify
sentencepiece
huggingface_hub>=0.26
# FlashAttention-2 is REQUIRED for the 32K varlen path. Build takes ~20 min, or use the
# prebuilt wheel matching your torch/cuda/python/abi from the flash-attention releases page.
flash-attn==2.8.3 --no-build-isolation
# optional
wandb>=0.18