Skip to content

feat(fx-importer): support for importing fp8 model parameters #4020

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(fx-importer): support for importing fp8 model parameters #4020

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions python/torch_mlir/extras/fx_importer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,6 @@

TORCH_DTYPE_TO_MLIR_TYPE: Dict[torch.dtype, Callable[[], IrType]] = {
torch.float16: lambda: F16Type.get(),
torch.bfloat16: lambda: BF16Type.get(),
torch.float32: lambda: F32Type.get(),
torch.float64: lambda: F64Type.get(),
torch.uint8: lambda: IntegerType.get_unsigned(8),
Expand All @@ -191,16 +190,20 @@
torch.complex64: lambda: ComplexType.get(F32Type.get()),
torch.complex128: lambda: ComplexType.get(F64Type.get()),
}
# Type entries added only in torch with higher version
# Type entries added only in torch with higher version. bfloat16 is present from before
# but handling it here keeps the numpy and ml_dtypes business clean.
OPTIONAL_TORCH_DTYPE_TO_MLIR_TYPE = {
"bfloat16": lambda: BF16Type.get(),
"float8_e5m2": lambda: Float8E5M2Type.get(),
"float8_e4m3fn": lambda: Float8E4M3FNType.get(),
"float8_e5m2fnuz": lambda: Float8E5M2FNUZType.get(),
"float8_e4m3fnuz": lambda: Float8E4M3FNUZType.get(),
}
OPTIONAL_TORCH_DTYPES: List[TorchDtype] = list()
for dtype_str, mlir_type in OPTIONAL_TORCH_DTYPE_TO_MLIR_TYPE.items():
if hasattr(torch, dtype_str):
TORCH_DTYPE_TO_MLIR_TYPE[getattr(torch, dtype_str)] = mlir_type
OPTIONAL_TORCH_DTYPES.append(getattr(torch, dtype_str))

TORCH_DTYPE_TO_NPY_TYPE = {
# torch.qint8: None, # no equivalent np datatype
Expand All @@ -218,8 +221,15 @@
torch.complex64: np.complex64,
torch.complex128: np.complex128,
}

if ml_dtypes is not None:
TORCH_DTYPE_TO_NPY_TYPE[torch.bfloat16] = ml_dtypes.bfloat16
# Type entries added only in torch with higher version. ml_dtypes follows the same
# naming but we should check both regardless.
for dtype_str in OPTIONAL_TORCH_DTYPE_TO_MLIR_TYPE.keys():
if hasattr(torch, dtype_str) and hasattr(ml_dtypes, dtype_str):
TORCH_DTYPE_TO_NPY_TYPE[getattr(torch, dtype_str)] = getattr(
ml_dtypes, dtype_str
)

TORCH_DTYPE_TO_INT = {
torch.uint8: 0,
Expand Down Expand Up @@ -2070,10 +2080,10 @@ def _make_vtensor_literal_op(
) -> Operation:
mapping = py_attr_tracker.track(tensor)
if mapping.is_empty:
# check support for bfloat16
# check support for bfloat16 and optional types.
assert not (
tensor.dtype == torch.bfloat16 and ml_dtypes is None
), f"torch.bfloat16 requires the ml_dtypes package, please run:\n\npip install ml_dtypes\n"
tensor.dtype in OPTIONAL_TORCH_DTYPES and ml_dtypes is None
), f"{tensor.dtype} requires the ml_dtypes package, please run:\n\npip install ml_dtypes\n"
# Resolve the attribute.
npy_dtype = TORCH_DTYPE_TO_NPY_TYPE.get(tensor.dtype)
assert (
Expand Down
Loading