fix: store residual_scale as a scalar and drop the unused conditioning

4e9df37 verified about 1 month ago

5.98 kB

---
library_name: transformers
tags:
  - genomics
  - dna
  - gene-annotation
  - ntv3
  - long-range
license: other
language:
  - code
---
> [!WARNING]
> **TEMPORARY FILE — DELETE BEFORE RELEASE.**
> This README is a working record of the JAX → PyTorch port, written for the NTv3
> release team while the repository is being assembled. It is not the model card.
> Replace it in full with the public model card before this repository is made
> public.

# NTv3 Annotate — insecta (staging)

PyTorch conversion of `new_ntv3_700M_annotate_insecta_1mb`, the released
gene-annotation model for this clade. 703,424,500 parameters across 508
tensors (474 `nnx.Param` plus 34 `nnx.BatchStat`). Modelling code is not in
this repository: `config.json` resolves it from
[`InstaDeepAI/NTv3_pretrained`](https://huggingface.co/InstaDeepAI/NTv3_pretrained)
through `auto_map`, so loading requires `trust_remote_code=True`.

Source checkpoint:
`s3://genomix-research-3e09d/checkpoints/ntv3_annotate/new_ntv3_700M_annotate_insecta_1mb`
(MLflow run `84b02c4b4d0f4346bd115a78a0f6703a`, step 762960, ≈800 B tokens, trained on 93 insect assemblies).

## Verification against the JAX reference

Measured on the 1024 bp fixture the checkpoint bundle ships in its own `test_inputs`
item (`sequence = "ACGT" * 256`), comparing the annotation head's `bed_tracks_logits`
of shape `(1, 1024, 1, 15)`.

| comparison | max abs | mean abs | decoded-class agreement |
| --- | --- | --- | --- |
| **torch fp32 vs JAX fp32** | 4.244e-05 | 1.444e-05 | **100%** |
| torch bf16 autocast vs JAX fp32 | 5.235e-02 | 1.059e-02 | **100%** |

**Conclusion: the port is correct.** The fp32 agreement is accumulation noise across a
703M-parameter network. The bf16 row is dominated by the intrinsic cost of bf16 rather
than by the conversion: on the mammals model, where a true JAX bf16 reference was also
generated, JAX bf16 differs from JAX fp32 by 6.28e-02 on its own, and torch bf16
differs from JAX bf16 by 6.82e-02 — the same magnitude, so the port adds no error of
its own. Most importantly the decoded class is identical at every one of the 1024
positions in both precisions, and since a Viterbi decode consumes the argmax, **neither
setting can change the emitted GFF3**.

The converter refuses to write weights unless decoded-class agreement is 100% in both
precisions, so this file could not have been produced from a failing conversion.

### Precision guidance

Weights ship in **fp32**. bf16 is a *compute* dtype, matching the JAX
`build_bf16_compute_model` path, which lowers five trunk stages (`stem`,
`down_convolution`, `up_convolution`, `transformer_qkvo`, `transformer_ffn`) and leaves
parameters, norms, embeddings, the LM head and the annotation head in fp32:

```python
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    out = model(tokens, conditions_tokens=conds)
```

Do **not** use `torch_dtype="bfloat16"` or `model.to(torch.bfloat16)`. A whole-module
cast raises `expected m1 and m2 to have the same dtype`: `RMSBatchNorm.forward` promotes
its input to ≥fp32 and never casts back, `StandardizedConv1D` then follows the input
dtype, and the first bare `nn.Linear` (`GQABlock.q_layer`) still holds bf16 weights.

## How the weights were converted

Flax variables map one-to-one onto the torch `state_dict`, loaded with `strict=True`.
Leaf renames are `kernel`/`embedding`/`scale` → `weight`, under a `model.` prefix.
Layout rules are keyed on the **owning torch module, never on the tensor shape**: seven
kernels (`decoder_blocks.{0..6}.pointwise_unet_skip.proj.kernel`) are square because
each maps `C → C`, so a shape-based rule cannot tell a transpose from a copy. Copying
them produces a model that loads cleanly under `strict=True`, carries the correct
parameter count, and still agrees 100% on the annotation argmax while the trunk is
silently wrong from the first decoder block. Census check: the model has 108
`Linear.weight` tensors; a converter that transposes 101 has this bug.

## Known gaps in this staging repository

- **Tokenizer.** `vocab.json` here is the 16-token soft-masked vocabulary these models
  were trained on (`A T C G N` plus lowercase `a t c g n`), which differs from the
  11-token vocabulary of the initial NTv3 release. `tokenizer_config.json` resolves the
  tokenizer class from `InstaDeepAI/NTv3_pretrained` through `auto_map`, so
  `AutoTokenizer.from_pretrained(..., trust_remote_code=True)` returns a 16-token
  tokenizer reading this repository's own `vocab.json`.
- **No decoding code.** `config.json` carries the full HMM specification under
  `annotation` — 15 states, the transition grammar, motif rules, the U12 lane, the
  served crop of 0.75 and the 1,048,576 bp window — but the Viterbi decoder and GFF3
  writer are not published here yet. Logits alone are not a gene annotation.
- **Inputs are soft-masked.** Lowercase repeat-masked bases carry their own token ids;
  an uppercase-only assembly loses the repeat signal without raising.
- **Window geometry is not recoverable from the weights.** `ntv3.keep_target_center_fraction`
  is set to the validated serving crop of **0.75**, so a default forward reproduces
  `jax_model/roundtrip_fixture.npz` exactly. Running at a different window produces a
  plausible but wrong annotation rather than an error. Note that
  `annotation.window.trained_keep_target_center_fraction` records 1.0 while
  `jax_model/hyperparameters.json` records 0.75 for the same field; that disagreement is
  unresolved and does not affect serving, since both permit 0.75.
- **The model is unconditioned.** It carries no condition table at all, so `forward` takes an **empty** `conditions_tokens` list — unlike the mammals and plants repositories, which require a length-1 list. Any shared wrapper has to absorb that asymmetry.
- **transformers 5.** `residual_scale` is stored as a scalar, as the model declares it; it was a one-element list before, which transformers 5 refused to load.

	---
	library_name: transformers
	tags:
	- genomics
	- dna
	- gene-annotation
	- ntv3
	- long-range
	license: other
	language:
	- code
	---

	> [!WARNING]
	> TEMPORARY FILE — DELETE BEFORE RELEASE.
	> This README is a working record of the JAX → PyTorch port, written for the NTv3
	> release team while the repository is being assembled. It is not the model card.
	> Replace it in full with the public model card before this repository is made
	> public.

	# NTv3 Annotate — insecta (staging)

	PyTorch conversion of `new_ntv3_700M_annotate_insecta_1mb`, the released
	gene-annotation model for this clade. 703,424,500 parameters across 508
	tensors (474 `nnx.Param` plus 34 `nnx.BatchStat`). Modelling code is not in
	this repository: `config.json` resolves it from
	[`InstaDeepAI/NTv3_pretrained`](https://huggingface.co/InstaDeepAI/NTv3_pretrained)
	through `auto_map`, so loading requires `trust_remote_code=True`.

	Source checkpoint:
	`s3://genomix-research-3e09d/checkpoints/ntv3_annotate/new_ntv3_700M_annotate_insecta_1mb`
	(MLflow run `84b02c4b4d0f4346bd115a78a0f6703a`, step 762960, ≈800 B tokens, trained on 93 insect assemblies).

	## Verification against the JAX reference

	Measured on the 1024 bp fixture the checkpoint bundle ships in its own `test_inputs`
	item (`sequence = "ACGT" * 256`), comparing the annotation head's `bed_tracks_logits`
	of shape `(1, 1024, 1, 15)`.

	\| comparison \| max abs \| mean abs \| decoded-class agreement \|
	\| --- \| --- \| --- \| --- \|
	\| torch fp32 vs JAX fp32 \| 4.244e-05 \| 1.444e-05 \| 100% \|
	\| torch bf16 autocast vs JAX fp32 \| 5.235e-02 \| 1.059e-02 \| 100% \|

	Conclusion: the port is correct. The fp32 agreement is accumulation noise across a
	703M-parameter network. The bf16 row is dominated by the intrinsic cost of bf16 rather
	than by the conversion: on the mammals model, where a true JAX bf16 reference was also
	generated, JAX bf16 differs from JAX fp32 by 6.28e-02 on its own, and torch bf16
	differs from JAX bf16 by 6.82e-02 — the same magnitude, so the port adds no error of
	its own. Most importantly the decoded class is identical at every one of the 1024
	positions in both precisions, and since a Viterbi decode consumes the argmax, **neither
	setting can change the emitted GFF3**.

	The converter refuses to write weights unless decoded-class agreement is 100% in both
	precisions, so this file could not have been produced from a failing conversion.

	### Precision guidance

	Weights ship in fp32. bf16 is a compute dtype, matching the JAX
	`build_bf16_compute_model` path, which lowers five trunk stages (`stem`,
	`down_convolution`, `up_convolution`, `transformer_qkvo`, `transformer_ffn`) and leaves
	parameters, norms, embeddings, the LM head and the annotation head in fp32:

	```python
	with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
	out = model(tokens, conditions_tokens=conds)
	```

	Do not use `torch_dtype="bfloat16"` or `model.to(torch.bfloat16)`. A whole-module
	cast raises `expected m1 and m2 to have the same dtype`: `RMSBatchNorm.forward` promotes
	its input to ≥fp32 and never casts back, `StandardizedConv1D` then follows the input
	dtype, and the first bare `nn.Linear` (`GQABlock.q_layer`) still holds bf16 weights.

	## How the weights were converted

	Flax variables map one-to-one onto the torch `state_dict`, loaded with `strict=True`.
	Leaf renames are `kernel`/`embedding`/`scale` → `weight`, under a `model.` prefix.
	Layout rules are keyed on the owning torch module, never on the tensor shape: seven
	kernels (`decoder_blocks.{0..6}.pointwise_unet_skip.proj.kernel`) are square because
	each maps `C → C`, so a shape-based rule cannot tell a transpose from a copy. Copying
	them produces a model that loads cleanly under `strict=True`, carries the correct
	parameter count, and still agrees 100% on the annotation argmax while the trunk is
	silently wrong from the first decoder block. Census check: the model has 108
	`Linear.weight` tensors; a converter that transposes 101 has this bug.

	## Known gaps in this staging repository

	- Tokenizer. `vocab.json` here is the 16-token soft-masked vocabulary these models
	were trained on (`A T C G N` plus lowercase `a t c g n`), which differs from the
	11-token vocabulary of the initial NTv3 release. `tokenizer_config.json` resolves the
	tokenizer class from `InstaDeepAI/NTv3_pretrained` through `auto_map`, so
	`AutoTokenizer.from_pretrained(..., trust_remote_code=True)` returns a 16-token
	tokenizer reading this repository's own `vocab.json`.
	- No decoding code. `config.json` carries the full HMM specification under
	`annotation` — 15 states, the transition grammar, motif rules, the U12 lane, the
	served crop of 0.75 and the 1,048,576 bp window — but the Viterbi decoder and GFF3
	writer are not published here yet. Logits alone are not a gene annotation.
	- Inputs are soft-masked. Lowercase repeat-masked bases carry their own token ids;
	an uppercase-only assembly loses the repeat signal without raising.
	- Window geometry is not recoverable from the weights. `ntv3.keep_target_center_fraction`
	is set to the validated serving crop of 0.75, so a default forward reproduces
	`jax_model/roundtrip_fixture.npz` exactly. Running at a different window produces a
	plausible but wrong annotation rather than an error. Note that
	`annotation.window.trained_keep_target_center_fraction` records 1.0 while
	`jax_model/hyperparameters.json` records 0.75 for the same field; that disagreement is
	unresolved and does not affect serving, since both permit 0.75.
	- The model is unconditioned. It carries no condition table at all, so `forward` takes an empty `conditions_tokens` list — unlike the mammals and plants repositories, which require a length-1 list. Any shared wrapper has to absorb that asymmetry.
	- transformers 5. `residual_scale` is stored as a scalar, as the model declares it; it was a one-element list before, which transformers 5 refused to load.