ESMC-AR 846M PLE wide16 bidirectional NCBI mix (10B tokens)

This is a fresh, completed 10B-token ESMC-AR variant with 16 wide Transformer blocks, PLE, gated attention and no gated residual. The FP32 safetensors export contains 845,824,000 parameters and is self-contained for transformers with bundled custom code and the ESMC tokenizer.

• Training task: bidirectional causal next-token training (original:reverse sampled approximately 1:1 by output tokens; no FIM)
• Training token count: 10,000,000,000 exactly (checkpoint step 38,147)
• Validation loss at 9.96B tokens: 2.191812
• Combined corpus: 10,609,635,466 available training prediction tokens from QingWY/protein-pretraining-data and ESMC-tokenized NCBI data
• Validation corpus: 159,600,594 prediction tokens
• Training source commit: 804d0a85c5dfd40ea7229881775a683844e5abfd
• Launch ID: 20260923T141848Z-804d0a85c5df-29b4bd3c
• W&B run: https://wandb.ai/huyiwen/bio-esmc-ar/runs/aec97cc1a4f0
• Context length: 2,048; vocabulary: ESMC token IDs; no automatic BOS/EOS insertion
• FP32 export verified against the original training checkpoint; see verification.json

Data provenance

training_config.json is the exact run config. data_manifest.json records the combined manifest and both source revisions/selections. The packed dataset merged:

1. QingWY/protein-pretraining-data, revision 23949d7e372cf5d0f002bb35dbb6f56d0fd63cb9
2. ittsleep/esmctoken-fromNCBI, revision 11b3ba9bda01fc2b453586e42b7baf9c02b04534

NCBI parquet rows already used ESMC amino-acid IDs 4-23. Every downloaded shard was checked against its Hugging Face LFS SHA256. Rows remained isolated attention segments and no BOS/EOS was synthesized.

Usage

ESMC-AR Hugging Face inference

This directory is self-contained: config.json, configuration_esmc_ar.py, modeling_esmc_ar.py, feature/PLE Python dependencies, safetensors weight shards, tokenizer files, generation_config.json and export_info.json. Tested with PyTorch 2.5.1 and Transformers 4.51.0. Install torch, transformers, safetensors, tokenizers; FlashAttention is optional and the training repository is not required. Attention uses PyTorch SDPA.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

path = '/path/to/export'
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(
    path, trust_remote_code=True, torch_dtype=torch.bfloat16,
).to('cuda').eval()
batch = tokenizer(['MALWMRLLPLL', 'MKWVTFISLLFLFSSAYS'],
                  padding=True, return_tensors='pt').to(model.device)
with torch.inference_mode():
    ids = model.generate(**batch, max_new_tokens=64, do_sample=True,
                         temperature=0.8, top_p=0.95)
print(tokenizer.batch_decode(ids, skip_special_tokens=True))

trust_remote_code=True loads the bundled Python inference files. No network access is needed when the directory is local. For CPU, use torch_dtype=torch.float32 and omit .to('cuda'). FP32 weights under BF16 autocast most closely match the training numerics; directly casting all weights/residuals to BF16 trades a small numerical difference for inference speed and memory.

Supports left-padded batch generation, sampling/greedy/beam search, hidden states/attention outputs and standard shifted causal-LM labels. Baseline models also support inputs_embeds and a tuple KV cache with 2 KV heads. PLE exports require input_ids and use_cache=False; short-convolution exports also disable caching. For those variants, generation recomputes the full context. No packed training interface is exposed here: each batch row is one causal sequence.

Tokenizer IDs are ESMC IDs, not raw ProGen3 IDs. Automatic BOS/EOS insertion is disabled to match training. Supply a nonempty amino-acid prefix. The dataset consists of protein fragments and does not train biological end-of-protein termination. Generation therefore defaults to a bounded length and suppresses nonstandard amino acids, special tokens and unused vocabulary slots. To change this explicitly, pass suppress_tokens=None; this does not make previously unseen tokens trained.

RoPE was trained on fragments up to 2048 positions. Longer inference is mechanically possible, but its quality has not been evaluated. export_info.json records the source step/token count; an intermediate export is not a completed 5B-token model.

From the training repository, export any completed checkpoint with:

.venv/bin/python -m scripts.export_hf \
  --checkpoint /media/huyiwen/bio-esmc-ar-runs/stage1/checkpoint-00001250 \
  --output-dir /media/huyiwen/bio-esmc-ar-runs/hf-checkpoint-00001250