| model: | |
| name: "ELECTRAProteinModel" | |
| max_residues: 512 | |
| embed_dim: 2560 | |
| encoder_depth: 33 | |
| encoder_heads: 40 | |
| decoder_dim: 2560 | |
| decoder_depth: 3 | |
| decoder_heads: 40 | |
| generator_depth: 8 | |
| generator_decoder_depth: 2 | |
| disc_loss_weight: 50.0 # corrective CE same scale as generator CE | |
| generator_lr_multiplier: 1.0 | |
| generator_temperature: 2.0 | |
| relax_temperature_scaling: null # null = 1/sqrt(head_dim), relaxes attention temperature | |
| use_qk_norm: false | |
| # Training configuration | |
| training: | |
| learning_rate: 1.0e-3 | |
| weight_decay: 0.01 | |
| adam_beta1: 0.9 | |
| adam_beta2: 0.95 | |
| warmup_steps: 5000 | |
| batch_size: 32 | |
| num_workers: 8 | |
| max_steps: 500000 | |
| total_training_steps: 500000 | |
| accumulate_grad_batches: 1 | |
| val_check_interval: 100 | |
| test_interval: 100 | |
| gradient_clip_val: 1.0 | |
| # Data configuration | |
| data: | |
| name: "ELECTRADataModule" | |
| lmdb_dir: "/path/to/triprorep-pretrain" | |
| struct_format: "fullatom" | |
| max_length: 512 | |
| bb_vocab_size: 512 | |
| fa_vocab_size: 512 | |
| mask_prob: 0.6 | |
| random_cropping: true | |
| pack_sequences: false | |
| # Lightning configuration | |
| lightning: | |
| gpus: 8 | |
| nodes: 4 | |
| precision: "bf16-mixed" | |
| compile: true | |
| checkpoint_dir: "./checkpoints" | |
| log_dir: "./logs" | |
| wandb_dir: "./wandb" | |
| wandb_project: "k-fold-structure" | |
| wandb_name: "triprorep_3B" | |
| save_every_n_steps: 10000 | |