proteinFM

A protein foundation model framework integrating and extending UniMoMo for protein-ligand design with Mixture-of-Transformers (MoT) architecture.

Quick Links

• Quick Start: See docs/guides/QUICK_START.md
• UniMoMo Integration: See docs/unimomo/ for MoT architecture and setup
• Troubleshooting: See docs/fixes/ for common issues

---

Table of Contents

1. Training
   • Quick Start
   • Hyperparameter Sweep Script
   • Training Modes
   • Understanding Model Control
   • Debug Options
2. Inference
   • MoT Generation
   • CoT Evaluation
3. Evaluation
   • Structural Metrics
4. Configuration Reference
5. Extended JSONL Format
6. Environment Setup

---

Training

Quick Start Training

Single GPU:

python src/lmms_engine/models/bagel_protein/train/train_bagel_protein.py \
    --config src/lmms_engine/models/bagel_protein/config/train_bagel_protein_antibody_dual_leakage.yaml

Multi-GPU (8 GPUs):

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    --master_addr=127.0.0.1 --master_port=12355 \
    src/lmms_engine/models/bagel_protein/train/train_bagel_protein.py \
    --config src/lmms_engine/models/bagel_protein/config/train_bagel_protein_antibody_dual_leakage.yaml

Hyperparameter Sweep Script

The recommended way to train is using the hyperparameter sweep script, which automatically generates output directories and saves modified configs:

./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh [OPTIONS]

Common Examples:

# Basic training with default parameters (8 GPUs)
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh

# Custom learning rates
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh \
    --diffusion_lr 5.0e-5 \
    --understanding_lr 1.0e-5

# Train from antibody-scratch checkpoint
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh --abscratch

# Freeze understanding model (only train diffusion)
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh --no-train-understanding

# Resume from checkpoint
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh \
    --resume ./output/bagel_protein_dual_leakage_diff4_bs8_ga8_dlr1e-4_ulr1e-4_ep300_pretrained

# Enable Weights & Biases logging
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh \
    --wandb_project bagel_protein \
    --wandb_run_name experiment_1

Full Options:

Option	Description	Default
--num_iter_per_und INT	Diffusion iterations per understanding step	4
--diffusion_lr FLOAT	Diffusion model learning rate	1.0e-4
--understanding_lr FLOAT	Understanding model learning rate	1.0e-4
--gradient_accumulation INT	Gradient accumulation steps	8
--batch_size INT	Per-device train batch size	8
--num_train_epochs INT	Number of training epochs	300
--abscratch	Use antibody-scratch checkpoint	false
--no-train-understanding	Freeze understanding model	false
--training_mode MODE	joint, generation_only, understanding_only	from config
--answer_only_qkv	DEBUG: Use answer-only for QKV (no thinking)	false
--num_gpus INT	Number of GPUs	8
--master_port INT	Port for distributed training	12355
--wandb_project NAME	Wandb project (enables logging)	disabled
--wandb_run_name NAME	Wandb run name	auto-generated
--resume PATH	Resume from output directory	-

Output Directory Naming:

bagel_protein_dual_leakage_diff{N}_bs{B}_ga{G}_dlr{DLR}_ulr{ULR}_ep{E}[_frozenund][_answeronly]_{ckpt}

Example: bagel_protein_dual_leakage_diff4_bs8_ga8_dlr1e-4_ulr1e-4_ep300_frozenund_pretrained

Training Modes

Mode	Description	Use Case
joint	Train both understanding + diffusion	Default, best performance
generation_only	Only train diffusion model	Debug diffusion without SFT
understanding_only	Only train understanding model	Debug Qwen fine-tuning

Understanding Model Control

Control whether the understanding model (Qwen3) is trainable:

Flag	Effect
--train-understanding	Qwen3 parameters receive gradients (trainable)
--no-train-understanding	Qwen3 frozen, only diffusion trains
Neither	Uses default (frozen)

Example: Frozen understanding + abscratch:

./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh \
    --no-train-understanding \
    --abscratch

Debug Options

Answer-Only QKV (--answer_only_qkv)

Test if the diffusion model can learn to use ground truth text representations:

# DEBUG: Train diffusion using only answer text for QKV conditioning
# (excludes thinking and foldability sections)
./src/lmms_engine/models/bagel_protein/train/train_hyperparameter_sweep.sh \
    --answer_only_qkv \
    --abscratch

This option:

• QKV conditioning: Uses only the answer field from JSONL
• SFT loss: Still uses full question + thinking + answer
• Purpose: Verify if diffusion can learn from GT answer text representation

---

Inference

MoT Generation (Chain-of-Thought + Diffusion)

Generate protein structures using the BAGEL Protein MoT model:

1. Generate Chain-of-Thought (CoT) from understanding expert (Qwen3)
2. Extract QKV from thinking portion only (before </think>)
3. Run diffusion conditioned on thinking tokens

Single GPU:

python src/lmms_engine/models/bagel_protein/inference/mot_generate.py \
    --config src/lmms_engine/models/bagel_protein/config/test_ab_mot.yaml \
    --checkpoint /path/to/checkpoint.pt \
    --gpu 0 \
    --save_dir ./results/experiment_name

Multi-GPU (8 GPUs):

torchrun --nproc_per_node=8 src/lmms_engine/models/bagel_protein/inference/mot_generate.py \
    --config src/lmms_engine/models/bagel_protein/config/test_ab_mot.yaml \
    --checkpoint /path/to/checkpoint.pt \
    --save_dir ./results/experiment_name

Example (HCDR3 antibody generation):

torchrun --nproc_per_node=8 src/lmms_engine/models/bagel_protein/inference/mot_generate.py \
    --config src/lmms_engine/models/bagel_protein/config/test_ab_mot.yaml \
    --checkpoint ./output/bagel_protein_dual_leakage_diff4_bs8_ga8_dlr1e-4_ulr1e-4_ep300_pretrained/best_model.pt \
    --save_dir ./results/ab_mot_epoch13

Configuration (test_ab_mot.yaml):

Parameter	Description	Default
n_cot_samples	Number of CoT generations per sample	10
n_diff_samples	Number of diffusion samples per CoT	10
dataloader.batch_size	Batch size for inference	8
cdr_type	CDR types to generate (e.g., [HCDR3])	[HCDR3]

Total samples per protein = n_cot_samples × n_diff_samples (default: 100)

Command-Line Options:

Option	Description	Default
--max_new_tokens	Maximum tokens to generate for CoT	4096
--temperature	Temperature for CoT generation	0.7
--seed	Random seed	42
--use_gt_answer	DEBUG: Use GT answer text for QKV	false

Debug Mode (--use_gt_answer):

Tests if diffusion can correctly use ground truth answer text representations:

torchrun --nproc_per_node=8 src/lmms_engine/models/bagel_protein/inference/mot_generate.py \
    --config src/lmms_engine/models/bagel_protein/config/test_ab_mot.yaml \
    --checkpoint /path/to/checkpoint.pt \
    --save_dir ./results/ab_mot_gt_answer \
    --use_gt_answer

When --use_gt_answer is enabled:

• Extracts QKV from ground truth answer field instead of generating CoT
• Tests if diffusion can correctly map GT answer text → structure
• Useful for debugging whether the diffusion model can use text representations

Output Structure:

results/experiment_name/
├── results_mot.jsonl       # Main results for metric evaluation
├── cot_texts.jsonl         # Generated CoT texts for analysis
├── references/             # Ground truth structures (PDB)
│   └── {pdb_id}_ref.pdb
└── candidates/             # Generated structures
    └── {pdb_id}/
        └── {CDR}/
            ├── 0.pdb       # Sample 0
            ├── 1.pdb       # Sample 1
            └── ...         # Up to 99.pdb

CoT Quality Evaluation

Evaluate Chain-of-Thought generation quality across checkpoints:

python src/lmms_engine/models/bagel_protein/inference/evaluate_cot_checkpoints.py \
    --config src/lmms_engine/models/bagel_protein/config/test_ab_mot.yaml \
    --checkpoints /path/to/checkpoint1.pt /path/to/checkpoint2.pt \
    --output_dir ./cot_eval_results \
    --num_samples 50

---

Evaluation

Structural Metrics

Evaluate generated structures using UniMoMo's metrics pipeline:

PYTHONPATH=./external/UniMoMo:$PYTHONPATH python -m scripts.metrics.peptide \
    --results ./results/experiment_name/results_mot.jsonl \
    --antibody \
    --log_suffix HCDR3 \
    --num_workers 96

Example:

PYTHONPATH=./external/UniMoMo:$PYTHONPATH python -m scripts.metrics.peptide \
    --results ./results/ab_mot_epoch13/results_mot.jsonl \
    --antibody \
    --log_suffix HCDR3 \
    --num_workers 96

Metrics Computed:

Metric	Description	Better
RMSD	Root Mean Square Deviation (Å)	Lower
TMscore	Template Modeling score (0-1)	Higher
DockQ	Docking quality score	Higher
AAR	Amino Acid Recovery (%)	Higher
Energy	PyRosetta energy score	Lower

Output:

• Results written to stdout with summary statistics
• Detailed per-sample metrics in the results JSONL

Requirements:

• NumPy < 2.0 (for DockQ compatibility)
• PyRosetta (for energy calculations)
• Ray (for parallel processing)

# Install if needed
uv pip install numpy==1.26.4
uv pip install pyrosetta-installer
python -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()"

---

Configuration Reference

Training Configuration (train_bagel_protein_antibody_dual_leakage.yaml)

Model Config:

model_config:
  training_mode: joint                    # joint, generation_only, understanding_only
  understanding_model_id: thinking-bio-lab/qwen3_4b_it_full_sft_v35_n3416
  extract_layer: 6                        # Qwen layer for QKV extraction
  num_diffusion_timesteps: 100
  generation_latent_dim: 32
  num_iter_per_und: 4                     # Diffusion iters per understanding step
  use_gradient_checkpointing: true        # Save memory (~30% slower)
  understanding_dtype: "bfloat16"         # Qwen dtype (bfloat16 saves memory)
  use_chat_template: true                 # Use Qwen chat template

Dataset Config (Dual Leakage Mode):

dataset_config:
  datasets:
    - class: AntibodyDataset
      mmap_dir: ./datasets/antibody/SAbDab_new/processed
      prompt_jsonl: ./datasets/prompt/train.jsonl
      use_extended_format: true
      prevent_leakage_qkv_only: true      # QKV: no answer, SFT: with answer
      leakage_marker: "**Foldability:**"  # Truncate at this marker for QKV
      use_answer_only_qkv: false          # DEBUG: true = answer-only QKV
      strict_prompt: true

Trainer Args:

trainer_args:
  per_device_train_batch_size: 8
  gradient_accumulation_steps: 8
  learning_rate: 1.0e-4                   # Diffusion LR
  understanding_learning_rate: 1.0e-4     # Qwen LR
  num_train_epochs: 300
  lr_scheduler_type: reduce_on_plateau

Inference Configuration (test_ab_mot.yaml)

# Sampling parameters
n_cot_samples: 10      # CoT generations per sample
n_diff_samples: 10     # Diffusion samples per CoT

# Dataset
eval_dataset_config:
  datasets:
    - class: AntibodyDataset
      specify_index: ./datasets/antibody/SAbDab_new/processed/test_index.txt
      cdr_type:
        - HCDR3

---

Extended JSONL Format for SFT Training

The system supports two JSONL formats for training data:

Format Comparison

Legacy Format (Simple):

{"id": "7abc", "prompt": "Design a protein that binds to DNA"}

Extended Format (SFT-Ready):

{"complex_id": "7abc", "question": "Design...", "thinking": "We need helices... **Foldability:** ANSWER", "answer": "ANSWER"}

Dual Leakage Prevention Mode

The recommended configuration prevents data leakage while enabling proper SFT:

prevent_leakage_qkv_only: true  # QKV: truncated thinking, SFT: full response
leakage_marker: "**Foldability:**"

What happens:

• QKV extraction: Uses question + thinking (truncated at marker) - NO answer leakage
• SFT loss: Uses question + thinking + answer - FULL supervision

Data Flow

JSONL: {"question": "Design...", "thinking": "Analysis... **Foldability:** SEQ", "answer": "SEQ"}
                    ↓
┌──────────────────────────────────────────────────────────┐
│  QKV Version (for diffusion conditioning):               │
│  "Design..." + "Analysis..."                             │
│  (truncated at marker, no answer leakage)                │
├──────────────────────────────────────────────────────────┤
│  SFT Version (for language model loss):                  │
│  "Design..." + "Analysis... **Foldability:** SEQ" + SEQ  │
│  (full response for supervision)                         │
└──────────────────────────────────────────────────────────┘

---

Environment Setup

Quick Install (uv)

# Clone with submodules
git clone --recursive https://github.com/ThinkProteo/proteinFM.git
cd proteinFM

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh
source "$HOME/.cargo/env"

# Create venv and install
uv venv && source .venv/bin/activate
uv pip install -e .

Key Dependencies

# PyTorch (CUDA 12.1)
uv pip install torch --find-links https://download.pytorch.org/whl/cu121

# PyTorch Geometric
uv pip install torch_scatter torch_sparse torch_cluster \
    -f https://data.pyg.org/whl/torch-2.4.0+cu121.html --no-build-isolation

# Flash Attention (optional, for speed)
uv pip install flash-attn --no-build-isolation

# For metrics evaluation
uv pip install numpy==1.26.4  # DockQ requires NumPy < 2
uv pip install pyrosetta-installer
python -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()"

Environment Variables

export UNIMOMO_CODE_ROOT=/path/to/UniMoMo
export UNIMOMO_DATA_ROOT=/path/to/data
export UNIMOMO_CKPT=/path/to/checkpoint.ckpt

---

Architecture Overview

proteinFM/
├── external/UniMoMo/           # Git submodule (unmodified)
├── src/lmms_engine/models/
│   ├── bagel_protein/          # Main BAGEL Protein implementation
│   │   ├── config/             # YAML configurations
│   │   ├── model/              # Model implementations
│   │   ├── train/              # Training scripts
│   │   └── inference/          # Inference scripts
│   └── unimomo/                # UniMoMo adapters & MoT
├── datasets/                   # Data directory
└── output/                     # Training outputs

Key Components

Component	Description
BagelProteinSimpleModel	Main model combining Qwen3 + LDM
EPTMoT	Mixture-of-Transformers encoder
train_hyperparameter_sweep.sh	Training launcher with hyperparameter control
mot_generate.py	MoT inference (CoT + diffusion)
BagelChatCollate	Collate function with Qwen chat template

---

Troubleshooting

Q: NCCL timeout during evaluation

• Set eval_generation_mode: false in config for multi-GPU training

Q: NumPy/DockQ incompatibility

• Install numpy<2: uv pip install numpy==1.26.4
• Reinstall DockQ: uv pip install --force-reinstall DockQ

Q: PyRosetta import error

• Use pyrosetta-installer: python -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()"

Q: OOM during training

• Enable gradient checkpointing: use_gradient_checkpointing: true
• Use bfloat16 for understanding model: understanding_dtype: "bfloat16"
• Reduce batch size or increase gradient accumulation

Q: SFT loss is always zero

• Ensure training_mode: joint (not generation_only)
• Verify use_extended_format: true in dataset config

---

License

This project builds upon UniMoMo. Please see the respective licenses.