| ### Training | |
| Some settings follow those in the [AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w) paper, The table below shows the training settings for different fine-tuning stages: | |
| | Arguments | Initial training | Fine tuning 1 | Fine tuning 2 | Fine tuning 3 | | |
| |-----------------------------------------|--------|---------|-------|-----| | |
| | `train_crop_size` | 384 | 640 | 768 | 768 | | |
| | `diffusion_batch_size` | 48 | 32 | 32 | 32 | | |
| | `loss.weight.alpha_pae` | 0 | 0 | 0 | 1.0 | | |
| | `loss.weight.alpha_bond` | 0 | 1.0 | 1.0 | 0 | | |
| | `loss.weight.smooth_lddt` | 1.0 | 0 | 0 | 0 | | |
| | `loss.weight.alpha_confidence` | 1e-4 | 1e-4 | 1e-4 | 1e-4| | |
| | `loss.weight.alpha_diffusion` | 4.0 | 4.0 | 4.0 | 0 | | |
| | `loss.weight.alpha_distogram` | 0.03 | 0.03 | 0.03 | 0 | | |
| | `train_confidence_only` | False | False | False | True| | |
| | full BF16-mixed speed(A100, s/step) | ~12 | ~30 | ~44 | ~13 | | |
| | full BF16-mixed peak memory (G) | ~34 | ~35 | ~48 | ~24 | | |
| We recommend carrying out the training on A100-80G or H20/H100 GPUs. If utilizing full BF16-Mixed precision training, the initial training stage can also be performed on A800-40G GPUs. GPUs with smaller memory, such as A30, you'll need to reduce the model size, such as decreasing `model.pairformer.nblocks` and `diffusion_batch_size`. | |
| ### Inference | |
| The model will be infered in BF16 Mixed precision, by **default**, the `SampleDiffusion`,`ConfidenceHead` part will still be infered in FP32 precision. | |
| Below are reference examples of cuda memory usage (G). | |
| | Ntoken | Natom | Default | Full BF16 Mixed | | |
| |--------|-------|-------|------------------| | |
| | 500 | 10000 | 5.6 | 5.1 | | |
| | 1500 | 30000 | 24.8 | 19.2 | | |
| | 2500 | 25000 | 52.2 | 34.8 | | |
| | 3500 | 35000 | 67.6 | 38.2 | | |
| | 4500 | 45000 | 77.0 | 59.2 | | |
| | 5000 | 50000 | OOM | 72.8 | | |
| The script in [runner/inference.py](../runner/inference.py) will automatically change the default precision to compute `SampleDiffusion`,`ConfidenceHead` to avoid OOM as follows: | |
| ```python | |
| def update_inference_configs(configs: Any, N_token: int): | |
| # Setting the default inference configs for different N_token and N_atom | |
| # when N_token is larger than 3000, the default config might OOM even on a | |
| # A100 80G GPUS, | |
| if N_token > 3840: | |
| configs.skip_amp.confidence_head = False | |
| configs.skip_amp.sample_diffusion = False | |
| elif N_token > 2560: | |
| configs.skip_amp.confidence_head = False | |
| configs.skip_amp.sample_diffusion = True | |
| else: | |
| configs.skip_amp.confidence_head = True | |
| configs.skip_amp.sample_diffusion = True | |
| return configs | |
| ``` |