[Bugfix] LinearReplacer.replace(linear, Dtypes.kbfloat16) raises error.

[wkcn]

Description
When replacing the weight with ScalingBF16, it will raise the following error.

Traceback (most recent call last):
  File "test_code.py", line 9, in <module>
    model = LinearReplacer.replace(linear, qtype)
  File "/usr/local/lib/python3.8/dist-packages/msamp/nn/linear.py", line 176, in replace
    model = cls._replace(model, weight_qtype)
  File "/usr/local/lib/python3.8/dist-packages/msamp/nn/linear.py", line 154, in _replace
    fp8_net = cls._build_fp8linear(model, weight_qtype)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/msamp/nn/linear.py", line 100, in _build_fp8linear
    raise ValueError(
ValueError: weight dtype is not same, dtype ins fp8 linear is torch.float16,dtype in weight is torch.bfloat16

Reproduce code:

import torch

from msamp.common.dtype import Dtypes
from msamp.nn import LinearReplacer

input = torch.randn((4, 4), device='cuda')
linear = torch.nn.Linear(4, 8).cuda()
qtype = Dtypes.kbfloat16
model = LinearReplacer.replace(linear, qtype)

The reason is that LinearReplacer did not pass the argument weight_qtype in FP8Linear, so weight_qtype is the default value Dtypes.float16, even if calling LinearReplacer.replace(linear, Dtypes.kbfloat16).

Major Revision

• pass the argument weight_qtype
• unittest

[tocean]

LGTM.