Training details

[andrewliu2001]

Hi! Thank you for your excellent work! What GPUs did you train on what how many?

Also, does the diffusion module not need to be trained in fp32 (as previous versions of Boltz) did? Can you give me some general tips on train stability you found important?

Thanks!

[yuyangw]

Hi @andrewliu2001 thanks for your interest in our work. For largest 1.6B and 3B models, we use 128 H100s and use 32 H100s for all other model sizes.

We didn't find FP32 is required for stable training and we used mixed precision with BF16 in training SimpleFold. You may want to check RMS normalization for QK which helps stabilize the training.

[andrewliu2001]

Ok. Also, how important is the rescaling of coordinates by self.scale in the preprocessing stage? I'm trying to figure out why when I combine the model with Boltz training pipeline I get irrecoverable loss spikes.

[yuyangw]

My experience is rescaling coordinates does lead to better training efficiency but may not be the key in training stability. When using the Boltz training pipeline, did you apply all the losses in Boltz?

[andrewliu2001]

Yes I am using the diffusion MSE loss and a custom kernel for SmoothLDDT loss.

[andrewliu2001]

I realized Boltz uses checkpoint_wrapper which is mysteriously converting some activations to fp16. I'll try again with this fixed.

[yuyangw]

Thanks for letting us know. Please keep us posted.