DeepSpeedPlugin: non-float16 metrics

[leezu]

PL DeepSpeedPlugin passes the whole LightningModule to the DeepSpeedEngine. DeepSpeedEngine casts the model provided to float16 during initialization. This includes casting the floating point states of torchmetrics associated with the module.

It would be helpful to allow using non-float16 states for metrics together with DeepSpeedPlugin.

cc @SeanNaren

[Borda]

we have introduced the float16 metrics in v0.3.0
mind check https://github.com/PyTorchLightning/metrics/releases/tag/v0.3.0

[leezu]

Thank you @Borda. I have a custom metric to represent a per-token loss where the number of tokens differs across workers. This metric goes to inf early in the training due to overflow if float16 states are used. It's possible to workaround with manual scaling. I'm just raising this discussion to point out that the DeepSpeedPlugin behaves differently to the other plugins

[Borda]

So you say that the inputs would be float16 but internal cumulative states float64 to overcome potential overflow, right?

[SeanNaren]

Thanks for making this issue @leezu, this is also a known issue for certain aggressive FP16 modes like APEX as well, since the entire model is kept in half precision. Might need some workaround to ensure metrics are kept in their default precision.

[leezu]

So you say that the inputs would be float16 but internal cumulative states float64 to overcome potential overflow, right?

Yes. The inputs might even be float64 as well, given the user controls the inputs when calling the metric.

[leezu]

@Borda @SeanNaren torchmetrics 0.3 avoids the down-cast, but it's broken in version 0.4 and 0.5: Lightning-AI/torchmetrics#484