| # @package _global_ | |
| # All parameters below will be merged with parameters from default configurations set above | |
| # this allows you to overwrite only specified parameters | |
| # Options for training: | |
| # - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme] | |
| # - /model/architecture: | |
| # [foldingdit_100M, foldingdit_360M, foldingdit_700M, | |
| # foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B] | |
| # - /trainer: [default, fsdp] | |
| defaults: | |
| - override /data: pdb | |
| - override /model: simplefold | |
| - override /model/architecture: foldingdit_100M | |
| - override /model/sampler: euler_maruyama | |
| - override /model/processor: protein_processor | |
| - override /callbacks: default | |
| - override /trainer: fsdp # FSDP training | |
| - override /logger: tensorboard | |
| # load_ckpt_path: [PATH_TO_YOUR_CKPT] # e.g. uncomment to load a checkpoint | |
| seed: 12345 | |
| trainer: | |
| max_steps: 300000 | |
| val_check_interval: 10000 # this controls BOTH checkpoint steps and validation | |
| accumulate_grad_batches: 1 # always 1 for FSDP | |
| data: | |
| batch_size: 8 # should be the same as the number of GPUs | |
| num_workers: 16 | |
| model: | |
| _target_: models.simplefold.simplefold.SimpleFold | |
| ema_decay: 0.999 | |
| clip_grad_norm_val: 2.0 | |
| esm_model: "esm2_3B" | |
| scheduler: | |
| _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup | |
| _partial_: true | |
| min_lr: 1e-6 | |
| max_lr: ${model.optimizer.lr} | |
| warmup_steps: 5000 | |
| architecture: | |
| esm_model: ${model.esm_model} | |
| processor: | |
| scale: 16.0 | |
| ref_scale: 5.0 | |
| multiplicity: 16 # number of copies per GPU | |