[Question] How to apply MS-AMP to only part of the model?

[veritas9872]

What's the issue, what's expected?: I would like to apply MS-AMP to only parts of the model that are less sensitive to reduced precision.

Additional information:

Some parts of models are more sensitive to reduced precision than others but the current API makes it difficult to apply MS-AMP only to desired parts of the model. Is there an easy way of doing this?

[tocean]

Hi @veritas9872, thanks for your attention to MS-AMP. You can set use_fp32_linear to toch.nn.Linear if you don't want to use FP8 for this layer.
https://github.com/Azure/MS-AMP/blob/9ac98df5371f3d4174d8f103a1932b3a41a4b8a3/msamp/nn/linear.py#L153

[veritas9872]

Thanks! I will try it out!

[tocean]

Close this issue since there are no activities for a long time