ESMC-AR 846M PLE wide16 (5B tokens)
This is a fresh, completed 5B-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: original next-token prediction only
- Training data: fixed shuffled selection from
QingWY/protein-pretraining-data, revision23949d7e372cf5d0f002bb35dbb6f56d0fd63cb9 - Final training token count: 5,000,000,000
- Last reported original-stream validation loss: 2.115281 (task mixtures make cross-model comparisons non-equivalent)
- Training source commit:
a4da29a3504a301b37f4e9e3071f4187de1f0f3a - W&B run: https://wandb.ai/huyiwen/bio-esmc-ar/runs/d5e0765abba7
- Context length: 2,048; vocabulary: ESMC token IDs; no automatic BOS/EOS insertion
- FP32 export verified against the original training checkpoint; see
verification.json
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