[NPU] [Bugfix] [diffusion] Fix NZ performance bug for diffusion models (#20684)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
@@ -22,6 +22,7 @@ from torch.distributed.fsdp import (
|
||||
)
|
||||
from torch.nn.modules.module import _IncompatibleKeys
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
get_param_names_mapping,
|
||||
hf_to_custom_state_dict,
|
||||
@@ -145,7 +146,7 @@ def maybe_load_fsdp_model(
|
||||
if quant_method is not None and hasattr(
|
||||
quant_method, "process_weights_after_loading"
|
||||
):
|
||||
if _is_npu:
|
||||
if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod):
|
||||
# Activate the NZ format for storing weights,
|
||||
# which is a specific optimization for Ascend NPU
|
||||
torch.npu.config.allow_internal_format = True
|
||||
|
||||
@@ -124,7 +124,12 @@ def npu_format_cast(
|
||||
return tensor
|
||||
|
||||
if tensor.device == torch.device("cpu"):
|
||||
return torch.ops.npu.npu_format_cast(tensor.npu(), acl_format.value).cpu()
|
||||
logger.warning_once(
|
||||
"Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. "
|
||||
"Please disable offloading, otherwise the performance will be "
|
||||
"significantly reduced."
|
||||
)
|
||||
return tensor
|
||||
else:
|
||||
return torch.ops.npu.npu_format_cast(tensor, acl_format.value)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user