Add option to freeze weights

[vue1999]

This PR adds frozen transfer learning functionality for the fine-tuning of foundation models.
This implementation is based on the original weight-freezing work by @7radians, who also authored the paper, and the partial reimplementation by @SunZichen-2004.

If you use this functionality in your work, please cite the corresponding paper: https://arxiv.org/abs/2502.15582.

[ilyes319]

@vue1999 Thanks Eszter!!

[ilyes319]

@vue1999 Did test that on multi-GPU? i am bit worried about preserve_grad_state causing synchronisation problems.

[vue1999]

I'm not sure, but I will check that!

[7radians]

@ilyes319 @vue1999 @SunZichen-2004 thank you for working on this! The preserve_grad_state is inherited from the og code, there were no problems with multi-gpu previously, but it’s good to check

[vue1999]

@ilyes319 Yes, it ran fine on multi-GPU. (Also, I think required_grad is a local variable for each process and is not part of the distributed state).