---

name: training description: "Build Tevatron dense, sparse, distillation, GradCache, LoRA, DeepSpeed, and JAX training commands safely." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

Training

Use this sub-skill when the task is to assemble, review, or troubleshoot Tevatron retriever training commands. It covers dense PyTorch training, teacher distillation, GradCache, LoRA, DeepSpeed launcher/config choices, SPLADE and UniCOIL sparse patterns, RepLLaMA-style LLM retriever training, and JAX/Tevax TPU or GPU routes.

Do not use this sub-skill for embedding checkpoints, FAISS/Pyserini search, reranker training, raw data conversion, or multimodal/Qwen training details; route those to the sibling workflow that owns encoding/retrieval, reranking, data preparation, or multimodal LLM usage.

Start Here

1. Read references/training-workflows.md for command templates and validation checkpoints.
2. Read references/arguments-reference.md when choosing ModelArguments, DataArguments, training arguments, or sparse/JAX-specific names.
3. Use scripts/build_training_command.py --help to generate a command plan without launching training.
4. Use references/troubleshooting.md before recommending dependency installs, checkpoint deletion, or memory fixes.

Route by Request

• Dense retriever training: use tevatron.retriever.driver.train with --do_train, model/data arguments, and PyTorch/Transformers training flags.
• Memory-constrained training: add --grad_cache plus query/passage chunk sizes, or choose DeepSpeed ZeRO-3 for large models.
• LoRA LLM retriever training: add --lora, LLM pooling/prefix options, PEFT dependency notes, and usually DeepSpeed or gradient checkpointing.
• Teacher distillation: use tevatron.retriever.driver.train_distil; verify every positive and negative passage carries a numeric teacher score.
• Sparse retriever training: use the SPLADE/UniCOIL patterns in the workflow reference; SPLADE has current package model support, while UniCOIL requires an adapted driver in this package version.
• JAX/TPU training: use tevatron.retriever.driver.jax_train for the current HF-style driver or tevatron.tevax.experimental.mp.train_lora for the experimental Tevax LoRA route.

Required Checks

Before handing a command to a user, confirm:

• --do_train is present for PyTorch/Transformers-style drivers.
• --output_dir is new, intentionally resumable, or paired with --overwrite_output_dir.
• The data route matches the schema: inline positive_passages/negative_passages, or ID lists with --corpus_name/--corpus_path.
• Optional dependencies are explicit: torch, peft, deepspeed, grad_cache, flash-attn, jax, flax, optax, and magix are not installed by Tevatron's core package.
• Long-running GPU/TPU commands are presented as plans unless the user explicitly asks to run training.