[Fix] import MXLinear from Primus Turbo (#272)

Signed-off-by: Gene Der Su <[email protected]>

Author: GeneDer

.github/workflows/ci.yaml

@@ -10,7 +10,7 @@ on:
   pull_request:
 
 env:
-  PRIMUS_TURBO_COMMIT: 0385cdd615cb4eff28a1cbbf3583fccf95d11fe9 # chore: refactor grouped gemm blockwise python code (#142)
+  PRIMUS_TURBO_COMMIT: 7ecd1edc31a3f1607e92b536b4250f9d98bfe423 # feat: unify fp8 gemm (#148)
 
 jobs:
   code-lint:

primus/backends/torchtitan/components/quantization/mx.py

@@ -7,7 +7,7 @@
 import torch
 import torch.nn as nn
 from primus_turbo.pytorch.core.float8 import Float8QuantConfig, ScalingGranularity
-from primus_turbo.pytorch.modules import MXLinear
+from primus_turbo.pytorch.modules import Float8Linear
 from torchtitan.config.job_config import JobConfig
 from torchtitan.distributed import ParallelDims
 from torchtitan.protocols.model_converter import (
@@ -21,8 +21,8 @@
 
 def replace_turbo_mxlinear_modules(model: nn.Module, config: Float8QuantConfig):
     for name, module in model.named_children():
-        if isinstance(module, torch.nn.Linear) and not isinstance(module, MXLinear):
-            mx_linear = MXLinear.from_float(module, config)
+        if isinstance(module, torch.nn.Linear) and not isinstance(module, Float8Linear):
+            mx_linear = Float8Linear.from_float(module, config)
             setattr(model, name, mx_linear)
         else:
             replace_turbo_mxlinear_modules(module, config)