Add FP16 and INT8 ONNX/OpenVINO variants at all MRL widths

76c3670 verified about 2 months ago

5.63 kB

# Training scripts for Bekko Static Embedding Multilingual v1

This directory contains the complete public-data training path for
`hotchpotch/bekko-static-embedding-v1`. The training entry points and export utility are
standalone: they do not import the Bekko package, UniR, or another private
project. The tokenizer is deliberately a separate deterministic stage, so a
complete reproduction consists of the two training scripts, the FP16 export
utility, the YAML recipe, and the data-selection JSON distributed alongside them.

## Files

| File | Purpose |
| --- | --- |
| `train_tokenizer.py` | Builds the uncased 200,000-token WordPiece tokenizer, preserving a 30,000-token English core. |
| `export_fp16.py` | Creates the formal FP16 delivery from the FP32 master without changing the tokenizer or model configs. |
| `train_model.py` | Validates and materializes the selected public Bekko rows through `shared_id`, trains the MNRL + MRL static model, performs intermediate evaluation, exports the model, and runs final HAKARI evaluation. |

In the Hugging Face model repository, this README is published in
`training_scripts/` together with `mnrl-full.yaml` (the formal recipe) and
`bekko-unsupervised-multilingual-40-direct.json` (the pinned public-data
selection). In the source repository, those two assets live in `configs/`.

## Requirements

Use Python 3.11 or later, `uv`, and a CUDA-capable PyTorch installation for
full training. The exact release environment used Python 3.13, PyTorch 2.8.0,
SentenceTransformers 5.6.0, and Transformers 5.14.1. Install at least:

```bash
pip install \
  'torch==2.8.0' \
  'sentence-transformers==5.6.0' \
  'transformers==5.14.1' \
  'datasets>=4.4.2' 'tokenizers>=0.22.2' 'pyyaml>=6.0.3' \
  'numpy' 'simplemma>=1.1.2' 'nltk>=3.9.2' 'xxhash>=3.8.1' \
  'accelerate>=1.1.0' 'wandb>=0.21.0'
```

Install a PyTorch build appropriate for the available CUDA runtime; the command
above is only the version constraint. `flash-attn` is optional and is not
required by these scripts. If Weights & Biases is unavailable, set
`tracking.enabled: false` in the YAML.

Final and intermediate evaluation require a local clone of
[HAKARI-Bench](https://github.com/hakari-bench/hakari-bench). Install that
project according to its own instructions, then set `evaluation.hakari_repo` in
the YAML to its checkout. No private data, private dataset cache, or UniR
checkout is required.

## 1. Train the tokenizer

Run from the directory containing these files, or pass absolute paths:

```bash
uv run train_tokenizer.py \
  --corpus-dir output/tokenizer-corpus \
  --output-dir output/bekko-multilingual-wp200k-en30k-uncased
```

The script downloads deterministic prefixes from the pinned public
`hotchpotch/bekko-embedding-v1-unsupervised` revision. It learns the English
30k core first and then extends it to 200k tokens using non-bitext multilingual
text. Keep the script defaults unchanged for an exact v1 tokenizer.

## 2. Configure the training recipe

Before starting, edit only local paths in `mnrl-full.yaml`:

- `model.tokenizer_model`: the output directory from step 1;
- `run.output_dir` and `data.stage_cache_root`: persistent local storage;
- `evaluation.hakari_repo`: the HAKARI-Bench checkout;
- `evaluation.full_device`: the GPU to use.

Do not change the data spec, sampler, MNRL settings, MRL widths, optimizer, or
seed for an exact v1 reproduction. `data.stage_cache_root` is a resumable local
cache. Its rows come from the public dataset and are keyed by its public
`shared_id` field.

## 3. Train, monitor, export, and evaluate

```bash
python train_model.py --config mnrl-full.yaml
```

The v1 recipe trains one full epoch with MNRL temperature 0.03, learning rate
0.25, 10% warmup, global batch size 16,384, and seed 42. It schedules exactly
100 lightweight intermediate reranking evaluations: five queries per
NanoMIRACL target and three queries per MNanoBEIR target.

After training, the script saves the final SentenceTransformers model, changes
only the exported tokenizer maximum length from the training limit of 512 to
32,768, then starts HAKARI's genuine full `--all` evaluation. Outputs are under
`run.output_dir`, including `final/` and `full-eval/hakari-results/`.

## 4. Export the formal FP16 delivery

The trainer retains the FP32 master in `final/`. Since 2026-09-14, publish a
separate FP16 copy:

```bash
python export_fp16.py \
  --source output/bekko-static-embedding-multilingual-v1/final \
  --output output/bekko-static-embedding-multilingual-v1/final-fp16
```

Adjust these paths to `run.output_dir`. Output must not already exist.
The utility casts only `embedding.weight`; model configs and tokenizer remain
unchanged. The FP16 weight file is 409,600,096 bytes instead of 819,200,096.
The trainer's automatic full evaluation runs on the FP32 master. A separate
matched HAKARI comparison evaluated FP16 storage using `--dtype fp32`: native
micro nDCG@10 × 100 was 45.88353 versus 45.88177 for FP32 storage. This tests
stored-value rounding, not FP16 inference kernels.

## Release checks

The formal v1 export stores FP16 weights, has 1,024 dimensions, and supports
MRL widths 32, 64, 128, 256, 512, and 1,024. Its primary SHA-256 values are:

```text
model.safetensors  586e165a868f306822174e8110020fa8759577cc62a537a702941747732858f8
tokenizer.json     308823060352e35bf9fff6b42ef343f31cc68fa800e284568cc44e9cc0f70a73
```

The 32,768-token export setting is an input ceiling, not a promise that a
32,768-token batch fits in GPU memory. Compare functional invariants and the
full evaluation outputs as well as hashes when reproducing on other hardware.

	# Training scripts for Bekko Static Embedding Multilingual v1

	This directory contains the complete public-data training path for
	`hotchpotch/bekko-static-embedding-v1`. The training entry points and export utility are
	standalone: they do not import the Bekko package, UniR, or another private
	project. The tokenizer is deliberately a separate deterministic stage, so a
	complete reproduction consists of the two training scripts, the FP16 export
	utility, the YAML recipe, and the data-selection JSON distributed alongside them.

	## Files

	\| File \| Purpose \|
	\| --- \| --- \|
	\| `train_tokenizer.py` \| Builds the uncased 200,000-token WordPiece tokenizer, preserving a 30,000-token English core. \|
	\| `export_fp16.py` \| Creates the formal FP16 delivery from the FP32 master without changing the tokenizer or model configs. \|
	\| `train_model.py` \| Validates and materializes the selected public Bekko rows through `shared_id`, trains the MNRL + MRL static model, performs intermediate evaluation, exports the model, and runs final HAKARI evaluation. \|

	In the Hugging Face model repository, this README is published in
	`training_scripts/` together with `mnrl-full.yaml` (the formal recipe) and
	`bekko-unsupervised-multilingual-40-direct.json` (the pinned public-data
	selection). In the source repository, those two assets live in `configs/`.

	## Requirements

	Use Python 3.11 or later, `uv`, and a CUDA-capable PyTorch installation for
	full training. The exact release environment used Python 3.13, PyTorch 2.8.0,
	SentenceTransformers 5.6.0, and Transformers 5.14.1. Install at least:

	```bash
	pip install \
	'torch==2.8.0' \
	'sentence-transformers==5.6.0' \
	'transformers==5.14.1' \
	'datasets>=4.4.2' 'tokenizers>=0.22.2' 'pyyaml>=6.0.3' \
	'numpy' 'simplemma>=1.1.2' 'nltk>=3.9.2' 'xxhash>=3.8.1' \
	'accelerate>=1.1.0' 'wandb>=0.21.0'
	```

	Install a PyTorch build appropriate for the available CUDA runtime; the command
	above is only the version constraint. `flash-attn` is optional and is not
	required by these scripts. If Weights & Biases is unavailable, set
	`tracking.enabled: false` in the YAML.

	Final and intermediate evaluation require a local clone of
	[HAKARI-Bench](https://github.com/hakari-bench/hakari-bench). Install that
	project according to its own instructions, then set `evaluation.hakari_repo` in
	the YAML to its checkout. No private data, private dataset cache, or UniR
	checkout is required.

	## 1. Train the tokenizer

	Run from the directory containing these files, or pass absolute paths:

	```bash
	uv run train_tokenizer.py \
	--corpus-dir output/tokenizer-corpus \
	--output-dir output/bekko-multilingual-wp200k-en30k-uncased
	```

	The script downloads deterministic prefixes from the pinned public
	`hotchpotch/bekko-embedding-v1-unsupervised` revision. It learns the English
	30k core first and then extends it to 200k tokens using non-bitext multilingual
	text. Keep the script defaults unchanged for an exact v1 tokenizer.

	## 2. Configure the training recipe

	Before starting, edit only local paths in `mnrl-full.yaml`:

	- `model.tokenizer_model`: the output directory from step 1;
	- `run.output_dir` and `data.stage_cache_root`: persistent local storage;
	- `evaluation.hakari_repo`: the HAKARI-Bench checkout;
	- `evaluation.full_device`: the GPU to use.

	Do not change the data spec, sampler, MNRL settings, MRL widths, optimizer, or
	seed for an exact v1 reproduction. `data.stage_cache_root` is a resumable local
	cache. Its rows come from the public dataset and are keyed by its public
	`shared_id` field.

	## 3. Train, monitor, export, and evaluate

	```bash
	python train_model.py --config mnrl-full.yaml
	```

	The v1 recipe trains one full epoch with MNRL temperature 0.03, learning rate
	0.25, 10% warmup, global batch size 16,384, and seed 42. It schedules exactly
	100 lightweight intermediate reranking evaluations: five queries per
	NanoMIRACL target and three queries per MNanoBEIR target.

	After training, the script saves the final SentenceTransformers model, changes
	only the exported tokenizer maximum length from the training limit of 512 to
	32,768, then starts HAKARI's genuine full `--all` evaluation. Outputs are under
	`run.output_dir`, including `final/` and `full-eval/hakari-results/`.

	## 4. Export the formal FP16 delivery

	The trainer retains the FP32 master in `final/`. Since 2026-09-14, publish a
	separate FP16 copy:

	```bash
	python export_fp16.py \
	--source output/bekko-static-embedding-multilingual-v1/final \
	--output output/bekko-static-embedding-multilingual-v1/final-fp16
	```

	Adjust these paths to `run.output_dir`. Output must not already exist.
	The utility casts only `embedding.weight`; model configs and tokenizer remain
	unchanged. The FP16 weight file is 409,600,096 bytes instead of 819,200,096.
	The trainer's automatic full evaluation runs on the FP32 master. A separate
	matched HAKARI comparison evaluated FP16 storage using `--dtype fp32`: native
	micro nDCG@10 × 100 was 45.88353 versus 45.88177 for FP32 storage. This tests
	stored-value rounding, not FP16 inference kernels.

	## Release checks

	The formal v1 export stores FP16 weights, has 1,024 dimensions, and supports
	MRL widths 32, 64, 128, 256, 512, and 1,024. Its primary SHA-256 values are:

	```text
	model.safetensors 586e165a868f306822174e8110020fa8759577cc62a537a702941747732858f8
	tokenizer.json 308823060352e35bf9fff6b42ef343f31cc68fa800e284568cc44e9cc0f70a73
	```

	The 32,768-token export setting is an input ceiling, not a promise that a
	32,768-token batch fits in GPU memory. Compare functional invariants and the
	full evaluation outputs as well as hashes when reproducing on other hardware.