L2 norm for fp16 goes to nan, needs to be fixed

[rahul003]

When .norm() method is called on a float16 array it tries to compute the norm in fp16, which causes the squared sum to go out of range and causes nan as output. Large normalizations should be kept in float32, similar to how it's done for softmax and batchnorm. Ideally this should be done by accumulating the squared sum into a float32 buffer.
This problem might affect the L2 norm operator as well. Haven't verified that.

[anirudhacharya]

Yes, I have encountered this bug. Currently norm operator does not have tests for gradient or symbolic_backward and for different precision formats like fp16 and fp32. There is just one check in test_ndarray.py

@sandeep-krishnamurthy also please label this - "Operator"

[leezu]

How about using a numerically stable L2 norm algorithm by scaling the elements? Ie. BLAS nrm2

[rahul003]

That sounds good too

[leezu]

@rahul003 would this fix your issue #11573 ?

[haojin2]

@leezu I can verify the fix through a simple example:

>>> a = mx.nd.array([256, 256, 256, 256, 256], dtype=np.float16)
>>> mx.nd.norm(a)
[ 572.5]
<NDArray 1 @cpu(0)>

@rahul003 Do you think leo's fix would solve the issue?

[rahul003]

Yeah, I think we can close this.

[haojin2]

@leezu Thanks for the fix!