Pretraining Models

Overview

ChemLactica supports pretraining language models on chemical data using the train.py script. The training system is built on top of Hugging Face Transformers with custom features for efficient distributed training.

Quick Start

Prepare Your Data

Organize your training data in JSONL format in directories. You’ll need:

Training data directories (can be multiple)
Validation data directory
Data type labels for each directory

Choose Model Configuration

Select a model configuration from the available options:

125m - 125 million parameter model
1.3b - 1.3 billion parameter model
6.7b - 6.7 billion parameter model
mistral7b - Mistral 7B based model
llama2 - Llama 2 based model

Run Training

Execute the training script with your configuration.

Basic Training Command

python chemlactica/train.py \
  --train_type pretrain \
  --from_pretrained facebook/galactica-125m \
  --model_config 125m \
  --training_data_dirs /path/to/train/data \
  --dir_data_types smiles \
  --valid_data_dir /path/to/valid/data \
  --train_batch_size 16 \
  --eval_steps 500 \
  --save_steps 1000 \
  --max_steps 10000 \
  --checkpoints_root_dir ./checkpoints \
  --experiment_name my_experiment

Command-Line Arguments

Required Arguments

string

required

Type of training to perform. Options: pretrain, sft, isft, dpo

string

required

Path to pretrained model directory or Hugging Face model identifier

string

required

Model configuration name (e.g., 125m, 1.3b, llama2)

string[]

required

List of directories containing training data in JSONL format

string[]

required

Data type labels for each training directory (same order as training_data_dirs)

string

required

Directory containing validation data

integer

required

Training batch size per GPU

integer

required

Number of training steps between evaluations

integer

required

Number of steps between checkpoint saves

string

required

Root directory for saving model checkpoints

Optional Arguments

float

Learning rate (defaults to config value if not specified)

integer

Number of warmup steps for learning rate scheduler

integer

default:"-1"

Maximum number of training steps (overrides num_train_epochs)

integer

Number of training epochs

integer

Number of steps for LR scheduler (defaults to max_steps)

integer

Validation batch size (defaults to train_batch_size)

integer

default:"4"

Buffer size for dataset shuffling

string

default:"none"

Name for the experiment (used in tracking)

integer

default:"0"

Number of dataloader worker processes

integer

default:"1"

Number of gradient accumulation steps

boolean

default:"false"

Enable gradient checkpointing to save memory

boolean

default:"false"

Use Flash Attention 2 for faster training

Tracking and Debugging

boolean

default:"true"

Enable experiment tracking with Aim

string

Directory for saving tracking data

boolean

default:"false"

Enable PyTorch profiler for performance analysis

string

Directory for profiling output

boolean

default:"false"

Enable reproducibility checks (for testing only)

Advanced Usage

Multi-GPU Training

ChemLactica uses Accelerate for distributed training. Launch with:

accelerate launch --config_file config/accelerate_config.yaml \
  chemlactica/train.py \
  --train_type pretrain \
  --from_pretrained facebook/galactica-125m \
  --model_config 125m \
  # ... other arguments

Resume from Checkpoint

To resume training from a checkpoint, pass the checkpoint directory path to --from_pretrained:

python chemlactica/train.py \
  --train_type pretrain \
  --from_pretrained ./checkpoints/facebook/galactica-125m/hash123/checkpoint-5000 \
  --model_config 125m \
  # ... other arguments

Multiple Data Sources

You can train on multiple data directories with different types:

python chemlactica/train.py \
  --training_data_dirs /path/to/smiles /path/to/reactions /path/to/properties \
  --dir_data_types smiles reactions properties \
  # ... other arguments

Dynamic Gradient Accumulation

Enable automatic gradient accumulation scheduling:

# In your config YAML
train_config:
  grad_accumulation_scheduler: true
  dynamic_grad_accumulation: true
  grad_accumulation_max: 256
  grad_accumulation_delta_steps: 100
  grad_accumulation_delta_percentage: 0.02
  grad_accumulation_patience: 4000

Training Features

Custom Callbacks

ChemLactica includes several custom callbacks:

WPSCounterCallback: Tracks words per second
CustomProgressCallback: Enhanced progress reporting with FLOPS tracking
EarlyStoppingCallback: Stops training at specified steps
JsonlDatasetResumeCallback: Handles dataset resumption for streaming data
ReproducibilityCallback: Validates training reproducibility
GradientAccumulationScheduler: Dynamically adjusts gradient accumulation

Memory Optimization

python chemlactica/train.py \
  --gradient_checkpointing \
  # ... other arguments

python chemlactica/train.py \
  --flash_attn \
  # ... other arguments

python chemlactica/train.py \
  --gradient_checkpointing \
  --flash_attn \
  --gradient_accumulation_steps 4 \
  # ... other arguments

Checkpoint Management

Checkpoints are saved with the following structure:

checkpoints_root_dir/
└── organization/
    └── model_name/
        └── experiment_hash/
            ├── checkpoint-1000/
            ├── checkpoint-2000/
            └── last/

The save_total_limit configuration controls how many checkpoints are kept:

# From train.py:259
save_total_limit=train_config.save_total_limit  # Default: 4

Monitoring Training

Aim Tracking

When --track is enabled, training metrics are logged to Aim:

python chemlactica/train.py \
  --track \
  --track_dir ./aim_logs \
  --experiment_name chemistry_pretrain \
  # ... other arguments

View metrics with:

aim up --repo ./aim_logs

Key Metrics

Training loss
Validation loss
Learning rate
Words per second (WPS)
Gradient norm
FLOPS utilization

Example Configurations

python chemlactica/train.py \
  --train_type pretrain \
  --from_pretrained facebook/galactica-125m \
  --model_config 125m \
  --training_data_dirs ./data/train \
  --dir_data_types smiles \
  --valid_data_dir ./data/valid \
  --learning_rate 1.4e-3 \
  --warmup_steps 500 \
  --train_batch_size 32 \
  --eval_steps 500 \
  --save_steps 1000 \
  --max_steps 50000 \
  --checkpoints_root_dir ./checkpoints \
  --experiment_name galactica_125m_pretrain

python chemlactica/train.py \
  --train_type pretrain \
  --from_pretrained facebook/galactica-1.3b \
  --model_config 1.3b \
  --training_data_dirs ./data/train \
  --dir_data_types smiles \
  --valid_data_dir ./data/valid \
  --learning_rate 1.4e-3 \
  --warmup_steps 500 \
  --train_batch_size 4 \
  --gradient_accumulation_steps 8 \
  --gradient_checkpointing \
  --flash_attn \
  --eval_steps 500 \
  --save_steps 1000 \
  --max_steps 100000 \
  --checkpoints_root_dir ./checkpoints \
  --experiment_name galactica_1.3b_pretrain

Troubleshooting

Out of Memory

If you encounter OOM errors:

Reduce train_batch_size
Increase gradient_accumulation_steps
Enable --gradient_checkpointing
Use --flash_attn if available
Reduce save_total_limit to save memory

Slow Training

Enable --flash_attn for 2-3x speedup
Increase dataloader_num_workers
Adjust shuffle_buffer_size
Use mixed precision training (bf16 is enabled by default)

Checkpoint Loading Issues

Make sure the checkpoint path includes the full path to the checkpoint directory, not just the parent directory.

Getting Started

Core Concepts

Training

Molecular Optimization

Generation

Guides

Overview

Quick Start

Basic Training Command

Command-Line Arguments

Required Arguments

Optional Arguments

Tracking and Debugging

Advanced Usage

Multi-GPU Training

Resume from Checkpoint

Multiple Data Sources

Dynamic Gradient Accumulation

Training Features

Custom Callbacks

Memory Optimization

Checkpoint Management

Monitoring Training

Aim Tracking

Key Metrics

Example Configurations

Troubleshooting

Out of Memory

Slow Training

Checkpoint Loading Issues

Next Steps

Fine-tuning

Configuration

​Overview

​Quick Start

​Basic Training Command

​Command-Line Arguments

​Required Arguments

​Optional Arguments

​Tracking and Debugging

​Advanced Usage

​Multi-GPU Training

​Resume from Checkpoint

​Multiple Data Sources

​Dynamic Gradient Accumulation

​Training Features

​Custom Callbacks

​Memory Optimization

​Checkpoint Management

​Monitoring Training

​Aim Tracking

​Key Metrics

​Example Configurations

​Troubleshooting

​Out of Memory

​Slow Training

​Checkpoint Loading Issues

​Next Steps

Fine-tuning

Configuration

Overview

Quick Start

Basic Training Command

Command-Line Arguments

Required Arguments

Optional Arguments

Tracking and Debugging

Advanced Usage

Multi-GPU Training

Resume from Checkpoint

Multiple Data Sources

Dynamic Gradient Accumulation

Training Features

Custom Callbacks

Memory Optimization

Checkpoint Management

Monitoring Training

Aim Tracking

Key Metrics

Example Configurations

Troubleshooting

Out of Memory

Slow Training

Checkpoint Loading Issues

Next Steps