[Fix] Restore online MXFP8 quantization for linear layers (#32953)
This commit is contained in:
@@ -617,10 +617,6 @@ class Fp8LinearMethod(LinearMethodBase):
|
|||||||
layer.register_parameter("input_scale", scale)
|
layer.register_parameter("input_scale", scale)
|
||||||
else:
|
else:
|
||||||
layer.register_parameter("input_scale", None)
|
layer.register_parameter("input_scale", None)
|
||||||
elif use_mxfp8:
|
|
||||||
raise ValueError(
|
|
||||||
"MXFP8 requires fp8-serialized checkpoint for linear layers."
|
|
||||||
)
|
|
||||||
|
|
||||||
def create_weights(
|
def create_weights(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -261,15 +261,17 @@ class NemotronHMoE(nn.Module):
|
|||||||
self.fc1_latent_proj = None
|
self.fc1_latent_proj = None
|
||||||
self.fc2_latent_proj = None
|
self.fc2_latent_proj = None
|
||||||
|
|
||||||
self.use_min_latency_fc1_gemm = (
|
self._use_min_latency_fc1_gemm: bool | None = None
|
||||||
self.use_latent_moe
|
|
||||||
and self.fc1_latent_proj is not None
|
|
||||||
and _is_cuda
|
|
||||||
and fused_a_gemm_weight_eligible(self.fc1_latent_proj)
|
|
||||||
)
|
|
||||||
|
|
||||||
def _apply_fc1_latent_proj(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
def _apply_fc1_latent_proj(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||||
if self.use_min_latency_fc1_gemm:
|
if self._use_min_latency_fc1_gemm is None:
|
||||||
|
self._use_min_latency_fc1_gemm = (
|
||||||
|
self.use_latent_moe
|
||||||
|
and self.fc1_latent_proj is not None
|
||||||
|
and _is_cuda
|
||||||
|
and fused_a_gemm_weight_eligible(self.fc1_latent_proj)
|
||||||
|
)
|
||||||
|
if self._use_min_latency_fc1_gemm:
|
||||||
return linear_with_fused_a_gemm(self.fc1_latent_proj, hidden_states)
|
return linear_with_fused_a_gemm(self.fc1_latent_proj, hidden_states)
|
||||||
return self.fc1_latent_proj(hidden_states)[0]
|
return self.fc1_latent_proj(hidden_states)[0]
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user