Pthahnix/distributed-training-esm2

Distributed Training Corpus for ESM2

A protein sequence corpus for learning distributed training techniques (DDP / TP / PP / FSDP), paired with ESM2 masked language modeling (MLM) continued pretraining.

The dataset is deliberately kept simple -- only two fields -- so that attention stays on the parallelism mechanics rather than on data wrangling. Full traceability is preserved nonetheless: every sequence can be mapped back to its exact row in the source dataset.

Fields

provenance format

train-00000-of-00004.parquet-1633-heavy
|__________________________| |____| |___|
    source filename (with ext)  row   chain

• —Source filename: the parquet filename inside the upstream OpenMed/agab-db snapshot. The .parquet extension is kept as an unambiguous delimiter -- split on ".parquet-" when parsing.
• —Row index: 0-indexed, referring to the physical storage order within that file. Row indices are assigned before any cleaning, filtering or sorting, so they always point at the original row.
• —Chain type: heavy / light / antigen, indicating which column of the original row the sequence came from (heavy_sequence / light_sequence / antigen_sequence).

The source data is a wide table (one antibody-antigen pair per row, carrying three sequences); this dataset explodes those three columns into independent samples. A single source row therefore yields up to 3 records, distinguished by the chain suffix, and provenance is unique across the whole dataset.

Reverse lookup example:

import pyarrow.parquet as pq

prov = "train-00000-of-00004.parquet-1633-heavy"
fname, rest = prov.split(".parquet-", 1)
row, chain = rest.rsplit("-", 1)

raw = pq.read_table(f"{RAW_SNAPSHOT_DIR}/{fname}.parquet")
original_row = raw.slice(int(row), 1)   # all 20 original columns: affinity / CDR / target / ...

Statistics

60,000 rows per shard (the last shard of each split holds the remainder). Chain composition: heavy 862,969 / light 18,873 / antigen 11,490 -- the heavy-chain dominance reflects the composition of the upstream data itself.

Build process

Upstream data: `OpenMed/agab-db`, snapshot commit 345ace3cf7a93eb967ae34b9aa8c5cc27fc5d8c9 (1,227,083 antibody-antigen pairs, 20 columns).

Steps:

1. 1.Index rows: read each parquet file and assign 0-indexed row numbers in physical storage order (before any filtering takes place).
2. 2.Explode: unpivot heavy_sequence / light_sequence / antigen_sequence into independent samples and build provenance. No chain concatenation, no separator tokens -- the ESM2 pretraining corpus itself consists of independent protein sequences.
3. 3.Clean: drop nulls; uppercase; keep only the alphabet [ACDEFGHIKLMNPQRSTVWYXBUZO] (20 standard amino acids plus common ambiguity codes); drop sequences longer than 1022 (ESM2 max_position_embeddings=1026 minus 2 special tokens).
4. 4.Deduplicate: within each split, deduplicate by sequence content. When a sequence occurs in several source rows, the lexicographically smallest provenance is kept, which makes the result reproducible.
5. 5.Leak protection: remove from train any sequence whose content also appears in validation / test (antibody framework regions are highly conserved, so overlap is substantial; left untreated the validation set would certainly have been seen during training). After removal, the train-validation and train-test intersections are both verified to be empty.
6. 6.Shuffle: shuffle each split globally with torch.randperm and seed 42, so that any single shard is a representative mixture of the three chain types (without shuffling, the leading shards would be almost entirely heavy chains).
7. 7.Shard: slice out exactly 60,000 rows per file.

The train / validation / test division follows the three files already present in the upstream snapshot (upstream split 80/10/10, though its dataset card wires only train in configs, which is why the HF page shows a single split).

License

Inherits the terms of the upstream OpenMed/agab-db: non-commercial research use only. The original data is provided by NaturalAntibody; contact them directly for commercial use.

Cite the original dataset:

@dataset{agab_db,
  title={AgAb DB: Antigen Specific Antibody Database},
  author={NaturalAntibody},
  year={2024},
  url={https://naturalantibody.com/agab/}
}

Usage

All examples below were verified on a 2x A100 node.

Daft streaming (recommended for distributed training)

Read straight from the Hub without downloading everything first, which suits handing shards out to different ranks:

import daft

df = daft.read_parquet(
    "hf://datasets/Pthahnix/distributed-training-esm2/data/train-*.parquet"
)
print(df.schema())
print(df.count_rows())   # 674397
df.limit(3).show()

Shard-level reads -- under DDP each rank can pull only its own shards instead of the full dataset:

# rank 0 takes the first 6 shards, rank 1 the last 6
shards = [f"hf://datasets/Pthahnix/distributed-training-esm2/data/train-{i:05d}-of-00012.parquet"
          for i in range(6)]
df = daft.read_parquet(shards)

datasets streaming

from datasets import load_dataset

ds = load_dataset("Pthahnix/distributed-training-esm2", split="train", streaming=True)
for row in ds:
    print(row["provenance"], row["sequence"][:45])
    break
# train-00000-of-00004.parquet-1633-heavy EVQLVESGGGLVQPGGSLRLSCAASGFNIKDTYIHWVRQAPGKGL

datasets full load

from datasets import load_dataset

ds = load_dataset("Pthahnix/distributed-training-esm2")
print(ds)
# DatasetDict({
#     train:      Dataset({features: ['sequence', 'provenance'], num_rows: 674397})
#     validation: Dataset({features: ['sequence', 'provenance'], num_rows: 109373})
#     test:       Dataset({features: ['sequence', 'provenance'], num_rows: 109562})
# })

ESM2 MLM continued pretraining

sequence can be fed to the ESM2 tokenizer directly. The 1022 length cap is already aligned with max_position_embeddings=1026 (leaving room for the <cls> and <eos> special tokens), so no further truncation is needed:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
batch = tokenizer(
    ["EVQLVESGGGLVQPGGSLRLSCAASGFNIKDTYIH"],
    padding=True, return_tensors="pt",
)

The masking policy (15% of positions selected, of which 80% become <mask>, 10% a random amino acid, 10% are left unchanged) belongs to the collator on the training side. This dataset is not pre-tokenized, keeping it model-agnostic.