tribunal-kit 9.2.3 → 9.2.4

package/.agents/skills/train-sentence-transformers/references/hardware_guide.md

@@ -0,0 +1,105 @@
+# Hardware Guide
+
+Training embedding models is memory-bound more often than compute-bound.
+
+## If you hit OOM
+
+Try in this order:
+
+1. **Reduce `per_device_train_batch_size`**. Raise `gradient_accumulation_steps` to keep the effective batch size for regression losses. (For MNRL, effective batch via grad-accum is **not** equivalent. See point 3.)
+2. **Enable `gradient_checkpointing=True`**. ~30% slower, ~40% less activation memory. Incompatible with `Cached*` losses.
+3. **Switch to a `Cached*` loss**:
+   - `CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)`: forwards in mini-batches, accumulates the contrastive loss over the full batch. Can simulate batch sizes of 1024+ on a 24GB GPU.
+   - `CachedSpladeLoss(model, loss=..., mini_batch_size=16)`: same trick for sparse.
+   - `CachedGISTEmbedLoss(model, guide_model, mini_batch_size=32)`: GIST variant.
+4. **Enable PEFT / LoRA** for decoder models >1B. `LoraConfig(r=64, lora_alpha=128, task_type="FEATURE_EXTRACTION")`. See `../scripts/train_sentence_transformer_with_lora_example.py` (docstring covers when to use, hyperparams, QLoRA, sharing).
+5. **Move to multi-GPU**. See below.
+6. **Shorten sequences**. If truncating to 128 is already sufficient for your task, set `max_seq_length` on the transformer module.
+
+## Multi-GPU
+
+`sentence-transformers` uses `accelerate` under the hood. Distributed training works without code changes.
+
+### Data parallel (DDP)
+
+Launch:
+
+```bash
+accelerate launch train.py
+
+# or explicitly:
+accelerate launch --multi_gpu --num_processes=4 train.py
+```
+
+`per_device_train_batch_size` stays per-GPU. Effective batch size scales linearly. MNRL's in-batch negatives remain **per-device**, not global, unless you pass `gather_across_devices=True` to losses that support it (`MultipleNegativesRankingLoss`, `CachedMultipleNegativesRankingLoss`, the symmetric variants, `GISTEmbedLoss`, `CachedGISTEmbedLoss`, `SparseMultipleNegativesRankingLoss`).
+
+### FSDP / DeepSpeed
+
+For models >3B, use `accelerate config` to enable FSDP or DeepSpeed ZeRO. Both are supported: `sentence-transformers` doesn't require any code changes, only the launch config.
+
+```bash
+accelerate config                   # interactive; choose FSDP or DeepSpeed
+accelerate launch train.py
+```
+
+With FSDP full-shard: a 7B model trains on 4×24GB GPUs that would OOM on any single one of them.
+
+**FSDP caveats** (from the [distributed training docs](https://sbert.net/docs/sentence_transformer/training/distributed.html)):
+- **Evaluators don't run under FSDP** as of writing: the eval hooks call `model.encode()` which FSDP-wrapped modules can't service mid-training. Plan to evaluate post-training on a single-GPU load of the final checkpoint instead, or train with DDP if you need mid-training evaluation.
+- **Layer wrapping must be specified**, e.g. `fsdp_config={"transformer_layer_cls_to_wrap": "BertLayer"}` (substitute the right layer class for your model: `BertLayer`, `LlamaDecoderLayer`, `Qwen2DecoderLayer`, etc.). Without this FSDP sharding can silently misbehave.
+- **Slower than DDP** for models that fit on a single GPU. Only reach for FSDP when you actually need the memory savings.
+
+DeepSpeed ZeRO-2/3 is an alternative with its own config. It works identically at the `accelerate config` level.
+
+## Effective batch size for contrastive losses
+
+For `MultipleNegativesRankingLoss` and its variants, **batch size is a quality knob**, not just a speed knob. Bigger batches = more in-batch negatives = richer gradients.
+
+The effective pool of in-batch negatives per anchor:
+
+| Setup | In-batch negatives per anchor |
+|---|---|
+| Single GPU, batch 64 | 63 |
+| 4× DDP, per-device batch 64 | 63 local only by default. 255 with `MultipleNegativesRankingLoss(model, gather_across_devices=True)` |
+| Single GPU, CachedMNRL, mini_batch 32, batch 256 | 255 |
+| 4× DDP, CachedMNRL, per-device 256 | 255 local. 1023 with `gather_across_devices=True` |
+
+For large corpora (retrieval), push toward 512+ effective negatives. For small, clean datasets (STS), 64 is plenty.
+
+## Precision choice by GPU
+
+| GPU generation | Recommended |
+|---|---|
+| T4, V100, GTX 1xxx, RTX 2xxx | `fp16=True` |
+| RTX 3xxx, A10G, A100, L4 | `bf16=True` |
+| RTX 4xxx, H100, B200 | `bf16=True` (or fp8 on H100 via specific kernels, not default) |
+| Apple M-series / ROCm | MPS/ROCm support is variable. `fp16` or `fp32` most reliable |
+
+bf16 is more numerically stable and almost always preferred when available.
+
+## Hugging Face Jobs flavor guide
+
+Hugging Face Jobs requires a Pro/Team/Enterprise plan. Pricing is approximate and subject to change. See the [Jobs pricing page](https://huggingface.co/docs/huggingface_hub/guides/jobs).
+
+| Flavor | Memory | Typical use | Est. $/hr |
+|---|---|---|---|
+| `cpu-basic` | ~2 GB | Dataset prep, validation, hard-neg mining (small) | <$0.10 |
+| `cpu-upgrade` | ~4 GB | Same, slightly bigger | $0.10 |
+| `t4-small` | 16 GB | Demos, MiniLM/DistilBERT with small batches | ~$0.75 |
+| `t4-medium` | 16 GB | MiniLM / DistilBERT with larger batch | ~$1.50 |
+| `l4x1` | 24 GB | BERT-base, MPNet, ModernBERT-base | ~$2.50 |
+| `a10g-small` | 24 GB | BERT-base to BERT-large | ~$3.50 |
+| `a10g-large` | 48 GB | ModernBERT-large, Qwen3-0.6B | ~$5.00 |
+| `a10g-largex2` | 96 GB (2× 48GB) | Mid-size multi-GPU | ~$10 |
+| `a100-large` | 80 GB | Large models or big contrastive batches | ~$10-12 |
+| `h100` | 80 GB | Biggest single-GPU | ~$12 |
+| `h100x8` | 640 GB | LLM-scale distributed | ~$96 |
+
+Defaults by base model:
+- MiniLM / DistilBERT -> `t4-small`
+- BERT-base / MPNet / ModernBERT-base -> `a10g-small` or `l4x1`
+- BERT-large / ModernBERT-large -> `a10g-large`
+- Qwen3-0.6B decoder base -> `a10g-large`
+- 1B+ decoder bases with LoRA -> `a10g-large` or `a100-large`
+
+Always start one flavor **smaller** than you think you need: OOM on Jobs is cheap ($0.50-$5 for a failed run). Underprovisioned is better than overprovisioned for the first attempt. When budgeting `timeout`, add **20-30% buffer** for model loading, checkpoint saving, and Hub push.

package/.agents/skills/train-sentence-transformers/references/hf_jobs_execution.md

@@ -0,0 +1,173 @@
+# Hugging Face Jobs Execution
+
+Run training on Hugging Face's managed GPUs without provisioning any local infrastructure. The same training script runs locally and on Jobs. This reference covers only the Jobs-specific concerns.
+
+## Prerequisites
+
+- Hugging Face account with a **Pro, Team, or Enterprise** plan. Jobs are paid.
+- `HF_TOKEN` with **write** permission. Log in once locally with `hf auth login` (the modern command from the `hf` CLI, replacing the deprecated `huggingface-cli login`).
+- Access to the `hf_jobs()` MCP tool, or the `hf` CLI (`curl -LsSf https://hf.co/cli/install.sh | bash -s`).
+
+## The three submission paths
+
+### 1. Inline script via MCP (recommended in Claude Code)
+
+Pass the full training script as `script`. Dependencies come from the PEP 723 header.
+
+```python
+hf_jobs("uv", {
+    "script": """
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["sentence-transformers[train]>=5.0", "trackio"]
+# ///
+
+# <full training script content>
+""",
+    "flavor": "a10g-large",
+    "timeout": "3h",
+    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
+})
+```
+
+### 2. Script-from-URL via MCP
+
+Upload the script to the Hub (as a model or dataset repo file) or a Gist, then reference by URL:
+
+```python
+hf_jobs("uv", {
+    "script": "https://huggingface.co/USERNAME/scripts/resolve/main/train_bi_encoder.py",
+    "flavor": "a10g-large",
+    "timeout": "3h",
+    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
+})
+```
+
+Local file paths (`./train.py`, `/path/to/train.py`) **do not work**. Jobs run in isolated containers without access to your filesystem.
+
+### 3. CLI
+
+```bash
+hf jobs uv run \
+    --flavor a10g-large \
+    --timeout 3h \
+    --secrets HF_TOKEN \
+    "https://huggingface.co/USERNAME/scripts/resolve/main/train.py"
+```
+
+Syntax gotchas:
+- Command order is `hf jobs uv run`, **not** `hf jobs run uv`.
+- Flags (`--flavor`, `--timeout`, `--secrets`) go **before** the script URL.
+- `--secrets` (plural), not `--secret`.
+
+## Required script modifications for Jobs
+
+Add these to your `TrainingArguments`:
+
+```python
+args = SentenceTransformerTrainingArguments(
+    ...,
+    push_to_hub=True,
+    hub_model_id="your-username/my-model",
+    hub_strategy="every_save",        # push each checkpoint; timeout-safe
+    save_strategy="steps",
+    save_steps=0.1,                   # 10 saves/pushes per epoch; scales with dataset size
+)
+```
+
+Why each matters:
+
+| Argument | Why |
+|---|---|
+| `push_to_hub=True` | The Jobs container is destroyed after the job finishes. Without Hub push, all weights are lost. |
+| `hub_model_id` | Required to identify the destination repo. |
+| `hub_strategy="every_save"` | Default, but worth being deliberate about on Jobs: each checkpoint is pushed as it's written, so a timeout leaves all completed checkpoints on the Hub. `"end"` only pushes once `trainer.train()` returns, so a timeout loses everything. |
+| `save_strategy="steps"` + `save_steps=0.1` | Checkpoints must actually be saved for `hub_strategy="every_save"` to push them. Fractional `0.1` = save every 10% of training, auto-scales with dataset size. |
+
+## Secrets
+
+Secrets are environment variables injected into the Jobs container. They never appear in logs and are not part of the script.
+
+| Secret | Required when |
+|---|---|
+| `HF_TOKEN` | Always, for Hub push. Also covers Trackio auth. |
+| `WANDB_API_KEY` | Using `report_to="wandb"`. |
+| `MLFLOW_TRACKING_URI`, `MLFLOW_TRACKING_TOKEN` | Using MLflow with a remote server. |
+
+The `$HF_TOKEN` syntax in the job config references the value from your local environment at submission time. The literal string `$HF_TOKEN` is replaced with your token's value. Never hardcode tokens in the script itself.
+
+Trackio (the default tracker in this skill) uses `HF_TOKEN` for auth, so no extra secrets are needed. Only switch to the W&B / MLflow rows above if you're using those trackers.
+
+## Timeout
+
+Default is **30 minutes**, which is too short for almost any real training. Set explicitly:
+
+```python
+"timeout": "2h"       # 2 hours
+"timeout": "90m"      # 90 minutes
+"timeout": "1.5h"     # 90 minutes
+"timeout": 7200       # seconds, as integer
+```
+
+Rule: **estimated training time × 1.3**. The extra buffer covers model loading, dataset caching, checkpoint saving, and Hub push.
+
+On timeout, the container is killed immediately. Only data on the Hub (`hub_strategy="every_save"` saves you here) or in persistent volumes survives.
+
+## Dataset caching
+
+Hugging Face datasets are cached at `~/.cache/huggingface/datasets` by default. That's **inside the container**, which is destroyed after the job. Each Jobs run re-downloads the dataset.
+
+For large datasets (>5 GB), this matters. Options:
+
+- **Persistent `/data` volume** (Jobs feature, check current documentation): set `HF_DATASETS_CACHE=/data/datasets` so caches persist across jobs.
+- **Pre-cache locally, push to Hub**: if the dataset is on Hub already, nothing to do. If it's local-only, `dataset.push_to_hub(...)` once so subsequent jobs load from Hub.
+
+## Monitoring a running job
+
+```bash
+hf jobs ps [--all]                        # running (or all) jobs
+hf jobs inspect <job-id>                  # full config + status
+hf jobs logs <job-id> [--follow|--tail N] # tail or stream
+hf jobs cancel <job-id>
+hf jobs hardware                          # list flavors + hourly rates
+```
+
+`hf jobs logs <id> --follow` under `Bash run_in_background` pairs nicely with a `Monitor` watching for the `VERDICT:` line emitted by your training script's verdict block.
+
+MCP equivalents (signatures may vary by server version, so check the actual
+tool listing): `hf_jobs("ps")`, `hf_jobs("logs", {"job_id": ...})`,
+`hf_jobs("cancel", {"job_id": ...})`.
+
+For recurring runs, `hf jobs scheduled uv run "<cron>" <script> ...`
+schedules. `hf jobs scheduled ps/suspend/delete` manages.
+
+## Common failures
+
+### "Model not found on Hub" after a successful-looking run
+
+The run succeeded but `push_to_hub` was not enabled. The container is gone. The weights are gone.
+
+Fix: always set `push_to_hub=True` + `hub_model_id=...` + `secrets={"HF_TOKEN": "$HF_TOKEN"}`.
+
+### Tracker not connecting
+
+- **Trackio:** `HF_TOKEN` missing or lacks write permission. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN"}` and make sure the token has write access.
+- **W&B:** `WANDB_API_KEY` missing. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN", "WANDB_API_KEY": "$WANDB_API_KEY"}`.
+
+### OOM on first step
+
+Flavor too small. Move up one tier (see `hardware_guide.md`).
+
+### Training starts but eval hangs forever
+
+`eval_strategy="steps"` with no `eval_dataset`. Always provide an eval dataset, or set `eval_strategy="no"`.
+
+### Dataset download times out
+
+Large dataset or slow cold-cache. Increase `timeout` or pre-cache to a persistent volume.
+
+### `CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True` crash
+
+The cached losses are incompatible with gradient checkpointing. Disable `gradient_checkpointing`.
+
+After submission, the MCP returns a job ID. Monitor with `hf_jobs("logs", {"job_id": ...})` when you want an update. Don't poll in a tight loop. End-to-end submission templates live in `scripts/train_sentence_transformer_example.py` / `scripts/train_cross_encoder_example.py` / `scripts/train_sparse_encoder_example.py`. Wrap the script contents in the inline pattern from §1 above.

package/.agents/skills/train-sentence-transformers/references/losses_cross_encoder.md

@@ -0,0 +1,137 @@
+# Cross-Encoder Losses (reranker)
+
+All losses live in `sentence_transformers.cross_encoder.losses`.
+
+Cross-encoder losses fall into three families: **pointwise** (score each pair independently), **pairwise** (score pairs of pairs), and **listwise** (score a whole ranked list at once). Plus contrastive/distillation variants for specific setups.
+
+## Top-line decision table
+
+| You have | Use | Family |
+|---|---|---|
+| `(query, passage, label)` with `label ∈ {0, 1}` or `[0, 1]` | `BinaryCrossEntropyLoss` | Pointwise |
+| `(query, passage, class_id)` multi-class | `CrossEntropyLoss` | Pointwise |
+| `(query, passage)` with implicit positives, want contrastive | `CachedMultipleNegativesRankingLoss` | Contrastive |
+| `(query, passages, labels)` (one row per query, with parallel lists of candidate passages and their relevance scores) | `LambdaLoss` | Listwise |
+| Same listwise shape, want well-tested simpler loss | `ListNetLoss` or `ListMLELoss` | Listwise |
+| `(query, positive, negative)` pairwise | `RankNetLoss` | Pairwise |
+| `(query, passage, teacher_score)` distillation from stronger reranker | `MSELoss` or `MarginMSELoss` | Distillation |
+
+## Pointwise losses
+
+### `BinaryCrossEntropyLoss`
+
+**The default cross-encoder loss.** For `(query, passage, label)` where label is 0/1 or a continuous [0, 1] relevance score.
+
+```python
+loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(5.0))
+```
+
+- `pos_weight`: upweight positives when your dataset is imbalanced (e.g. 1 positive + N hard negatives). A good default is `pos_weight = num_hard_negatives`.
+- Handles both binary labels and continuous relevance scores. If you have graded relevance (0.0, 0.5, 1.0), BCE still works.
+
+### `CrossEntropyLoss`
+
+Multi-class classification over passages. Use with `CrossEncoder("...", num_labels=N)` where N is the number of classes (e.g. 3 for NLI: entailment / neutral / contradiction).
+
+## Contrastive losses (for reranker training)
+
+### `MultipleNegativesRankingLoss` (cross-encoder)
+
+Contrastive training mirroring bi-encoder MNRL. Applies in-batch negatives: for each `(query, positive)`, every other `positive` in the batch is used as a negative.
+
+- **Data**: `(query, positive)` or `(query, positive, negative_0, negative_1, ...)`.
+- For CrossEncoder training, **prefer `BinaryCrossEntropyLoss` with mined hard negatives** as the default. MNRL is the fallback when you only have `(query, positive)` pairs and want to avoid a separate mining pass.
+
+### `CachedMultipleNegativesRankingLoss`
+
+Cached variant (GradCache). Decouples per-device batch size from effective in-batch negatives, same as the bi-encoder cached version.
+
+```python
+loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)
+```
+
+- Incompatible with `gradient_checkpointing=True`.
+- Key choice for training rerankers with effective batch 256+ on a single GPU.
+
+## Distillation losses
+
+> **`activation_fn=nn.Identity()` is mandatory** for all distillation, listwise, and pairwise losses below. Only `BinaryCrossEntropyLoss` and `CrossEntropyLoss` tolerate the default `Sigmoid`. The loss sees raw logits during training, but the model's `activation_fn` is applied at evaluation via `predict()`. The default `Sigmoid` (with `num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently collapsing eval ranking (training loss looks healthy while nDCG crashes from ~0.59 to ~0.14). Construct as `CrossEncoder("...", num_labels=1, activation_fn=torch.nn.Identity())`. See `troubleshooting.md` ("CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training") for the failure-mode walkthrough.
+
+### `MSELoss` (cross-encoder)
+
+Regress the student's output score to the teacher's score. Construct the model with `activation_fn=nn.Identity()` (see callout above).
+
+- **Data**: `(query, passage, teacher_score)`.
+- The teacher is typically a larger/stronger cross-encoder. Precompute its scores once, store as the label.
+
+### `MarginMSELoss` (cross-encoder)
+
+Regress the **difference** between positive and negative scores against the teacher's margin. Typically gives better distillation than plain MSE.
+
+- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
+- Popular recipe for MS MARCO-style distillation.
+
+## Listwise losses
+
+All listwise losses expect the dataset to have per-query **lists** of candidate passages and their scores, typically via a collator that groups rows by query.
+
+> Listwise and pairwise losses below also require `activation_fn=nn.Identity()` (see the Distillation-section callout above).
+
+### `LambdaLoss`
+
+The state-of-the-art listwise ranking loss. Optimizes a surrogate of nDCG via weighted pairwise comparisons.
+
+**Data shape**: `(query, [doc1...docN], [score1...scoreN])` per row. One query, a list of candidate documents, and a parallel list of relevance scores. Use `mine_hard_negatives(..., output_format="labeled-list", ...)` to build this from `(query, positive)` pairs. For ID-only datasets with separate text splits (e.g. `lightonai/ms-marco-en-bge`), `sentence_transformers.util.resolve_ids(..., output_format="lists")` produces this nested shape directly.
+
+```python
+import torch.nn as nn
+from sentence_transformers.cross_encoder.losses import LambdaLoss, NDCGLoss2PPScheme
+
+model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
+loss = LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())
+```
+
+- Strong default when you have multiple candidates per query and graded relevance.
+- `weighting_scheme`: `NDCGLoss2PPScheme` (default), `NDCGLoss2Scheme`, `LambdaRankScheme`. The `NDCGLoss2PPScheme` default was shown to reach the strongest performance in the original LambdaLoss paper.
+
+#### LambdaLoss-specific operational notes
+
+- **OOM recovery order**: drop `mini_batch_size` first (chunking inside the loss preserves the K-list semantic), then `per_device_train_batch_size` paired with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list length) only as a last resort. Lowering K changes the experiment. The others don't.
+- **Tiny loss at large K is expected.** With `NDCGLoss2PPScheme`, the loss normalizes by the discount-weighted pair count. At K=128 the loss can scale to ~`1e-4` numerically. That's not "training broken". Read `eval_NanoBEIR_R100_mean_ndcg@10` (or your equivalent) instead of the training loss to judge progress.
+- For very large K (>=128), the O(K²) weight buffer per query is materialized outside of the forward chunking, so even small `mini_batch_size` may not be enough. Consider scoring strategy: top-K hard negatives often beat random K.
+
+### `ListNetLoss`
+
+Treats the ranking as a probability distribution (via softmax) and minimizes cross-entropy against the teacher distribution.
+
+### `ListMLELoss`
+
+Maximum-likelihood estimation over permutations. Simpler than LambdaLoss, decent default.
+
+### `PListMLELoss`
+
+Position-aware ListMLE. Weights higher-ranked items more heavily. Often outperforms plain ListMLE on top-k metrics.
+
+### `RankNetLoss`
+
+Pairwise classification: for each pair of candidates, predict which is ranked higher via cross-entropy.
+
+- Simpler and faster than LambdaLoss.
+- Scales quadratically with list length. Not great for long candidate lists (>20).
+
+### `ADRMSELoss`
+
+Alternative listwise formulation (Approx Discounted Rank MSE) from the Rank-DistiLLM paper. Same data shape as `LambdaLoss`. In practice LambdaLoss is the stronger default. `RankNetLoss` was reported as marginally more effective (~0.002 nDCG@10) than ADRMSELoss on the paper's LLM-distillation setup, and LambdaLoss generally beats both.
+
+Hard-negative mining is essential for any contrastive reranker (random negatives teach nothing). See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.
+
+## Gotchas
+
+- **`BinaryCrossEntropyLoss` without `pos_weight`** when you have 5+ hard negatives per positive: the loss under-weights the positive signal. Set `pos_weight=num_hard_negatives`.
+- **`CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
+- **Listwise losses with wildly different list lengths per query**: some losses don't handle ragged lists well. Pad or truncate to a fixed length.
+- **`MarginMSELoss` without precomputed teacher score diffs**: this loss needs the `score_diff` label column populated from a teacher pass. It does not compute the teacher internally.
+- **Training a cross-encoder with `num_labels=1` then using `CrossEntropyLoss`**: mismatch. BCE is for num_labels=1, CE is for num_labels>=2.
+- **Default `Sigmoid` activation under any non-BCE loss**: silently destroys eval ranking. Pass `activation_fn=torch.nn.Identity()` to `CrossEncoder(...)` for distillation, listwise, and pairwise losses (everything except BCE/CE). See the callouts in the Distillation and Listwise sections.
+- **Custom CE head writing to the wrong feature key**: a custom scoring head must populate `features["scores"]` (not `features["sentence_embedding"]`). Otherwise `CrossEncoder.predict()` raises `KeyError: 'scores'` at inference time, even though training succeeds.
+- **Loading a CE with a custom-class head fails from a different script**: if you defined a `class ClassifierHead(nn.Module)` inline in `train.py` and saved the model, `modules.json` records `__main__.ClassifierHead`. Loading via `CrossEncoder("path")` from any other entry point raises `ImportError: Module '__main__' does not define a 'ClassifierHead' attribute`. Either move the class into an importable file (`my_pkg/heads.py`), build the same shape from stock ST modules (`Dense + LayerNorm + Dense`), or document that the model is only loadable from the same script.

package/.agents/skills/train-sentence-transformers/references/losses_multi_vector_encoder.md

@@ -0,0 +1,100 @@
+# Multi-Vector-Encoder Losses (ColBERT / late-interaction)
+
+All losses live in `sentence_transformers.multi_vector_encoder.losses`.
+
+Multi-vector models emit **one embedding per token** and score `(query, document)` via MaxSim (for each query token, take the max similarity to any document token, then sum across query tokens). This is fundamentally different from single-vector cosine, which changes what "temperature" means: MaxSim is an unbounded sum over query-token similarities, so for it a scale near 1.0 makes sense. **The default `scale=1.0` (temperature=1.0) is correct for unnormalized MaxSim: do not copy `scale=20.0` from bi-encoder MNRL.** With MeanMaxSim scoring (`similarity_fct=mean_colbert_scores`, and set `model.similarity_fn_name = "meanmaxsim"` so evaluation matches), each score is divided by its query's token count, so a scale of roughly the average length is a reasonable start, although some works use larger scales.
+
+## Top-line decision table
+
+| You have | Use |
+|---|---|
+| `(anchor, positive)` or `(anchor, positive, negative)` triplets | `MultiVectorMultipleNegativesRankingLoss` |
+| Same, want effective batch size of 128+ | `CachedMultiVectorMultipleNegativesRankingLoss` |
+| Cross-encoder teacher scores, `(query, positive, negative, score_diff)` | `MultiVectorMarginMSELoss` |
+| Listwise distillation `(query, [doc_1..doc_N], scores)` | `MultiVectorDistillKLDivLoss` |
+
+Hard-negative mining is essential for competitive results, because random in-batch negatives leave a lot on the table for late-interaction models. See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.
+
+## Contrastive losses
+
+### `MultiVectorMultipleNegativesRankingLoss`
+
+The default late-interaction contrastive loss. In-batch positives plus explicit hard negatives, scored with MaxSim (or XTR).
+
+```python
+from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss
+
+loss = MultiVectorMultipleNegativesRankingLoss(model=model)  # scale=1.0 (temperature=1.0) is the correct default
+```
+
+- **Data**: `(anchor, positive)` or `(anchor, positive, negative_1, ..., negative_n)`. The collator stamps each column's task (column 0 becomes query, others become document). Pass `task=...` on a column to override.
+- Set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args (same reason as bi-encoder MNRL).
+- `similarity_fct=colbert_scores` by default. Pass `XTRScores(top_k=...)` for the sparser XTR-style scoring (from `sentence_transformers.multi_vector_encoder.scoring`). XTR applies to training only: evaluation always scores with MaxSim (see Gotchas).
+- `scale=1.0` (temperature=1.0) matches PyLate and is correct for unnormalized MaxSim. **Do NOT copy `scale=20.0` from bi-encoder MNRL.** With MeanMaxSim (`similarity_fct=mean_colbert_scores`), raise the scale to roughly the average query length instead, which is the same objective at the same strength.
+
+### `CachedMultiVectorMultipleNegativesRankingLoss`
+
+GradCache variant: chunked embedding forward, cached gradients, and a second re-embedding pass. Decouples per-device batch size from effective in-batch negatives.
+
+```python
+loss = CachedMultiVectorMultipleNegativesRankingLoss(
+    model=model,
+    mini_batch_size=8,           # or use mini_batch_num_tokens=... for token-budgeted packing
+    score_mini_batch_size=4,     # optional, smaller trims the transient (Q, Q*N, q_tok, d_tok) buffer
+)
+```
+
+- **Incompatible with `gradient_checkpointing=True`** (same as every `Cached*` loss).
+- `mini_batch_num_tokens=N` packs sequences until N real tokens per mini-batch instead of a fixed `mini_batch_size`. Big win on variable-length data with flash-attention / input flattening.
+- `score_mini_batch_size` chunks the SCORING phase (which builds `(Q, Q*N, q_tokens, d_tokens)` intermediates) independently from the embedding phase. Drop it first when hitting OOM in the loss stage.
+- `chunk_elements` is the second lever on that same phase, and it cuts the other axis: `score_mini_batch_size` chunks queries, `chunk_elements` chunks documents inside each MaxSim call. Pass it through the scorer, e.g. `similarity_fct=partial(colbert_scores, chunk_elements=1_000_000)`. Scores and gradients are unchanged (it is a pure memory knob), and it composes with `score_mini_batch_size`.
+- `gather_across_devices=True` gathers document embeddings across DDP ranks. Enables cross-rank in-batch negatives.
+
+## Distillation losses
+
+Distillation is where multi-vector models learn most efficiently: cross-encoder teachers (e.g. `gte-modernbert-base`) score `(query, doc)` pairs offline, and the student MaxSim model regresses to that signal.
+
+For the standard MS MARCO KD format (LightOn's `ms-marco-en-bge` and similar: `(query_id, document_ids, scores)` with separate `queries` / `documents` splits), you can use `sentence_transformers.util.resolve_ids` to join the IDs against the text splits at load time. It's a factory that returns a batched transform: pass a mapping from each input ID column to its lookup dataset, e.g. `dataset.set_transform(resolve_ids({"query_id": queries, "document_ids": documents}, max_list_length=32))` (lazy, no caching) or via `dataset.map(..., batched=True, remove_columns=["query_id", "document_ids"])` (eager, cached: `map` keeps input columns unless told to drop them, so pass the ID columns to `remove_columns`), yielding `(query, document_1, ..., document_32, scores)` rows ready for `MultiVectorDistillKLDivLoss`. The same mapping shape covers triplets with IDs (`{"query_id": queries, "positive_id": documents, "negative_id": documents}`) and ID-only contrastive datasets without a `scores` column.
+
+### `MultiVectorMarginMSELoss`
+
+Regress the **margin** between positive and negative MaxSim scores against the teacher's margin.
+
+```python
+loss = MultiVectorMarginMSELoss(model=model)
+```
+
+- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`. Raw teacher scores `[score(q, pos), score(q, neg_1), ...]` also work, converted to margins internally.
+- Popular recipe from PyLate / colpali-engine.
+- Teacher scores are precomputed once, stored as the label column. The loss does not run the teacher inline.
+- Same `similarity_fct` convention as the other losses, but pairwise-shaped: defaults to `colbert_scores_pairwise`, or pass `xtr_scores_pairwise` for XTR.
+
+### `MultiVectorDistillKLDivLoss`
+
+Listwise KL-div: student's softmax distribution over N candidates should match the teacher's.
+
+```python
+loss = MultiVectorDistillKLDivLoss(model=model)
+```
+
+- **Data**: `(query, document_1, ..., document_N, scores)` where `scores` is a list of N teacher scores per row. One column per candidate document (the standard positional multi-column convention), and the label column must use a recognized name (`label`, `labels`, `score`, `scores`). `resolve_ids` expands a stored ID list into the numbered columns.
+- Stronger training signal than `MarginMSE` when you have full `N`-way teacher scores (not just positive/negative margins).
+- `similarity_fct` defaults to `colbert_kd_scores`, the listwise KD variant, not the `colbert_scores` used by the MNRL family. `XTRKDScores` is the XTR counterpart.
+- `temperature` softens both distributions before the softmax, and can be overridden per side with `student_temperature` / `teacher_temperature`. Set `teacher_temperature` from the spread of your teacher's scores rather than from a value copied out of a paper: a float32 softmax underflows to exact zeros once a row's spread divided by the temperature exceeds about 100, and every candidate that underflows drops out of the KL, which is the ranking information you are distilling. Sharpening `student_temperature` far below the student's own score spread collapses that side too, and once both are one-hot the loss and its gradient are exactly zero. **Start at 1.0 / 1.0 and only move a side you have a measured reason to move.**
+- The loss is scaled by the student temperature squared, which keeps gradient magnitudes comparable only for a shared temperature at or above 1.0 (the regime Hinton et al. cover). Below that the factor shrinks the loss instead, by 1e-6 at `student_temperature=0.001`: when combining KD with a contrastive loss, scale the KD weight back up accordingly.
+- **OOM**: drop `per_device_train_batch_size` first and raise `gradient_accumulation_steps` to hold the effective batch. Only reduce `N` (candidate-list length) as a last resort, since lowering N changes the experiment.
+
+## Data-shape gotchas
+
+- **Column 0 is always the query.** Losses call `self.model(sf, task="query" if idx==0 else "document")`. Use the collator's `router_mapping` to override if you need a non-standard column layout.
+- **Cross-column varlen `T`** is handled: the batch's positive and negative columns can have different token counts (Qwen2-VL family etc.). `stack_padded_token_embeddings` pads to the per-batch max before MaxSim.
+- **Mask flow**: MVE reads the SCORING mask from the model OUTPUT dict (`MultiVectorMask` rewrites the input mask). Custom loss code that reads `sentence_features[i]["attention_mask"]` directly instead of `outputs[i]["attention_mask"]` will silently score against skiplisted tokens.
+
+## Gotchas
+
+- **`scale=20.0` copied from bi-encoder MNRL**: misscaled for the unnormalized MaxSim sum. Keep `scale=1.0` there. Only length-normalized MeanMaxSim wants a larger scale, roughly the average query length.
+- **Missing `Normalize` at the token level in the pipeline**: `colbert_scores` assumes L2-normalized token embeddings. If your custom pipeline drops the token-level `Normalize`, either add one or pass `normalize_embeddings=True` semantics via a wrapper.
+- **`CachedMultiVectorMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
+- **`MultiVectorMarginMSELoss` without precomputed `score_diff`**: label column must be populated from a teacher pass ahead of training. The loss does not compute the teacher inline.
+- **Expecting XTR scoring at eval**: `XTRScores` is a train-only `similarity_fct`. XTR takes its top-k across the whole candidate set, so a `(query, document)` pair has no standalone score, which means it cannot be the model's `similarity_fn_name` and the evaluators reject it outright. Evaluation and inference score with MaxSim, also for XTR-trained models. That is by design, not a mismatch to fix. To score a fixed candidate set ad hoc, call `xtr_scores` directly.
+- **Distillation with a weak teacher**: multi-vector students easily match a small-model teacher and then plateau. Use a strong cross-encoder (e.g. `gte-modernbert-base`, `mxbai-rerank-large-v2`) for the teacher pass.