[Bug Fixed] scale=INF when casting a tensor to scaling FP32/BF16 tensors

[wkcn]

Description
The following script triggers the bug.

import torch
from msamp.common.dtype import Dtypes
torch.tensor([1.0 / 512], device='cuda').cast(Dtypes.kfloat32)

It will return ScalingTensor(tensor([nan], device='cuda:0'), meta=ScalingMeta(qtype=QType(name='kFloat32', value=2), scale=inf, scale_inv=0, amax=0.00195312, window_size=1).

The reason is that the scaling factor is stored as a FP32 scalar and it is equal to the maximum representation value / the maximum absolute value.

When the maximum absolute value is less than 1.0, the scaling factor will be larger than the maximum representation value of FP32, so the scaling factor is INF and its inverse scale_inv is 0.

To address the issue, the scaling factor is set to 1 when the data type is FP32 or BF16. After fixing the bug, it outputs

ScalingTensor(tensor([0.0020], device='cuda:0'), meta=ScalingMeta(qtype=QType(name='kFloat32', value=2), scale=1, scale_inv=1, amax=0.00195312, window_size=1)

[tocean]

Looks like the scaling factor will be always 1 for bf16 and float32. Do we still need to compute amax in TypeCast::cast_to_fp16 for bf16 and fp32? Compute amax may be time consuming.

[wkcn]

Looks like the scaling factor will be always 1 for bf16 and float32. Do we still compute amax in TypeCast::cast_to_fp16 for bf16 and fp32? Compute amax may be time consuming.

Good catch! I will remove the computation of amax for BF16 and FP32.

[wkcn]

@tocean I have tried to remove the amax calculation in ScalingBF16/FP32, but it affects the other calculation logic. I would prefer to optimize this in other late PRs.

[tocean]

@tocean I have tried to remove the amax calculation in ScalingBF16/FP32, but it affects the other calculation logic. I would prefer to optimize this in other late PRs.

Sure.

[tocean]

LGTM