[Bugfix][log]Unify MindIE SD logging and improve diagnostics-MindIE-SD-AtomGit

Mmindiesd/compilation/aclgraph_backend.py+22-5

	@@ -13,12 +13,13 @@
13	import contextlib	13	import contextlib
14	import dataclasses	14	import dataclasses
15	import importlib	15	import importlib
16	-import logging
17	from typing import Any	16	from typing import Any
18		17
19	import torch	18	import torch
20		19
21	-logger = logging.getLogger(__name__)	20	+from ..utils.logs.logging import logger
		21	+
		22	+DEBUG_LOG_LEVEL = 10
		G G guowenna16月3日 Python 标准 logging 里 DEBUG == 10。
22		23
23	# ---------------------------------------------------------------------------	24	# ---------------------------------------------------------------------------
24	# NPU 可用性检测	25	# NPU 可用性检测
	@@ -130,13 +131,17 @@ def create_aclgraph_backend():
130	entry = entries[input_shape]	131	entry = entries[input_shape]
131		132
132	# D2: input address debug validation	133	# D2: input address debug validation
133	- if logger.isEnabledFor(logging.DEBUG) and entry.input_addresses is not None:	134	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL) and entry.input_addresses is not None:
134	new_addrs = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]	135	new_addrs = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
135	for i, (old_addr, new_addr) in enumerate(zip(entry.input_addresses, new_addrs)):	136	for i, (old_addr, new_addr) in enumerate(zip(entry.input_addresses, new_addrs)):
136	if old_addr != new_addr:	137	if old_addr != new_addr:
137	logger.warning(	138	logger.warning(
138	- "ACLGraph input address mismatch at position %d: "	139	+ "[MindIE-SD/compilation] ACLGraph input address changed. "
		G G guowenna16月3日 ACLGraph 输入地址不一致 warning 改为 MindIE 格式，补充 index、captured/current data_ptr、可能原因和排查建议。
139	- "captured=%d, current=%d",	140	+ "issue=input data_ptr differs from captured graph buffer, index=%d, "
		141	+ "expected_data_ptr=%d, actual_data_ptr=%d. "
		142	+ "possible_cause=caller reused the graph with a different tensor storage. "
		143	+ "Troubleshooting: confirm input tensors are copied into static buffers before replay; "
		144	+ "enable DEBUG graph logs to inspect the capture and replay sequence.",
140	i,	145	i,
141	old_addr,	146	old_addr,
142	new_addr,	147	new_addr,
	@@ -153,6 +158,18 @@ def create_aclgraph_backend():
153	if static_buf.data_ptr() == new_inp.data_ptr():	158	if static_buf.data_ptr() == new_inp.data_ptr():
154	continue	159	continue
155	if static_buf.shape != new_inp.shape or static_buf.dtype != new_inp.dtype:	160	if static_buf.shape != new_inp.shape or static_buf.dtype != new_inp.dtype:
		161	+ logger.error(
		162	+ "[MindIE-SD/compilation] ACLGraph input validation failed. "
		G G guowenna16月3日 shape/dtype 不匹配抛异常前新增 error 日志，打印 expected/actual shape、dtype，方便定位 graph replay 输入不兼容。
		163	+ "issue=input shape or dtype differs from captured graph, index=%d, "
		164	+ "expected_shape=%s, expected_dtype=%s, actual_shape=%s, actual_dtype=%s. "
		165	+ "possible_cause=model graph was replayed with incompatible inputs. "
		166	+ "Troubleshooting: use the same shape and dtype as graph capture or trigger a new graph capture.",
		167	+ i,
		168	+ tuple(static_buf.shape),
		169	+ static_buf.dtype,
		170	+ tuple(new_inp.shape),
		171	+ new_inp.dtype,
		172	+ )
156	raise RuntimeError(	173	raise RuntimeError(
157	f"ACLGraph input mismatch at position {i}: "	174	f"ACLGraph input mismatch at position {i}: "
158	f"captured {tuple(static_buf.shape)}/{static_buf.dtype}, "	175	f"captured {tuple(static_buf.shape)}/{static_buf.dtype}, "

Mmindiesd/compilation/mindie_sd_backend.py+17-13

	@@ -10,12 +10,10 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
13	-import functools
14	-import logging
15	from typing import Any, Callable, Optional, Sequence	13	from typing import Any, Callable, Optional, Sequence
16		14
17	import torch	15	import torch
18	-import torch.fx as fx	16	+from torch import fx
		G G guowenna16月3日 `删除未使用 functools/logging；torch.fx 改为 from torch import fx，修 lint。`
19	from torch._dynamo.backends.common import aot_autograd	17	from torch._dynamo.backends.common import aot_autograd
20	from torch._inductor.freezing import freeze	18	from torch._inductor.freezing import freeze
21	from ._custom_decomposition import select_custom_decomp_table	19	from ._custom_decomposition import select_custom_decomp_table
	@@ -23,6 +21,7 @@ from ._custom_decomposition import select_custom_decomp_table
23	try:	21	try:
24	from torch.fx.passes.graph_transform_observer import GraphTransformObserver	22	from torch.fx.passes.graph_transform_observer import GraphTransformObserver
25	except ImportError:	23	except ImportError:
		24	+
26	class GraphTransformObserver:	25	class GraphTransformObserver:
27	def __init__(self, gm, passname, subsystem=None, log_url=None):	26	def __init__(self, gm, passname, subsystem=None, log_url=None):
28	self.gm = gm	27	self.gm = gm
	@@ -36,6 +35,7 @@ except ImportError:
36	def apply_graph_pass(self, pass_func):	35	def apply_graph_pass(self, pass_func):
37	pass_func(self.gm.graph)	36	pass_func(self.gm.graph)
38		37
		38	+
39	from .compiliation_config import CompilationConfig	39	from .compiliation_config import CompilationConfig
40		40
41	from .aclgraph_backend import npu_graph_available, create_aclgraph_backend	41	from .aclgraph_backend import npu_graph_available, create_aclgraph_backend
	@@ -43,13 +43,15 @@ from .aclgraph_backend import npu_graph_available, create_aclgraph_backend
43	from .passes import activate_pattern_once	43	from .passes import activate_pattern_once
44	from .passes.register_pattern_to_pass import patterns	44	from .passes.register_pattern_to_pass import patterns
45	from .passes.redundant_node_elimination_pass import ReduandantNodeEliminationPass	45	from .passes.redundant_node_elimination_pass import ReduandantNodeEliminationPass
		46	+from ..utils.logs.logging import logger
46		47
47	-logger = logging.getLogger(__name__)	48	+DEBUG_LOG_LEVEL = 10
		G G guowenna16月3日引入 MindIE SD logger 和 DEBUG_LOG_LEVEL。
48		49
49		50
50	def decompose_auto_functionalized(graph: fx.Graph):	51	def decompose_auto_functionalized(graph: fx.Graph):
51	try:	52	try:
52	from torch._inductor.fx_passes.post_grad import decompose_auto_functionalized as original_decompose	53	from torch._inductor.fx_passes.post_grad import decompose_auto_functionalized as original_decompose
		54	+
53	return original_decompose(graph)	55	return original_decompose(graph)
54	except ImportError:	56	except ImportError:
55	for node in list(graph.nodes):	57	for node in list(graph.nodes):
	@@ -83,7 +85,11 @@ class MindieSDBackend:
83	graph = self.compile(graph, example_inputs)	85	graph = self.compile(graph, example_inputs)
84	return create_aclgraph_backend()(graph, example_inputs)	86	return create_aclgraph_backend()(graph, example_inputs)
85	if CompilationConfig.aclgraph_only and npu_graph_available:	87	if CompilationConfig.aclgraph_only and npu_graph_available:
86	- logger.info("Using ACLGraph backend with torch.npu.graph")	88	+ logger.debug(
		89	+ "[MindIE-SD/compilation] ACLGraph backend selected. aclgraph_only=%s, npu_graph_available=%s.",
		90	+ CompilationConfig.aclgraph_only,
		91	+ npu_graph_available,
		92	+ )
87	return create_aclgraph_backend()(graph, example_inputs)	93	return create_aclgraph_backend()(graph, example_inputs)
88	else:	94	else:
89	# Use default backend	95	# Use default backend
	@@ -99,7 +105,7 @@ class MindieSDBackend:
99	log_url=CompilationConfig.graph_log_url,	105	log_url=CompilationConfig.graph_log_url,
100	).apply_gm_pass(ReduandantNodeEliminationPass())	106	).apply_gm_pass(ReduandantNodeEliminationPass())
101	logger.debug("Graph after redundant node elimination pass:")	107	logger.debug("Graph after redundant node elimination pass:")
102	- if logger.isEnabledFor(logging.DEBUG):	108	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL):
103	logger.debug(graph.print_readable(print_output=False))	109	logger.debug(graph.print_readable(print_output=False))
104		110
105	@classmethod	111	@classmethod
	@@ -112,7 +118,7 @@ class MindieSDBackend:
112	log_url=CompilationConfig.graph_log_url,	118	log_url=CompilationConfig.graph_log_url,
113	).apply_gm_pass(patterns)	119	).apply_gm_pass(patterns)
114	logger.debug("Graph after pattern matching:")	120	logger.debug("Graph after pattern matching:")
115	- if logger.isEnabledFor(logging.DEBUG):	121	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL):
116	logger.debug(graph.print_readable(print_output=False))	122	logger.debug(graph.print_readable(print_output=False))
117		123
118	@classmethod	124	@classmethod
	@@ -132,9 +138,7 @@ class MindieSDBackend:
132	) -> tuple[Callable, Optional[Any]]:	138	) -> tuple[Callable, Optional[Any]]:
133	def freezing_compile(compile_inner, aot_autograd_gm, example_inputs):	139	def freezing_compile(compile_inner, aot_autograd_gm, example_inputs):
134	# Freeze the graph first before passing to AOT Autograd.	140	# Freeze the graph first before passing to AOT Autograd.
135	- frozen_gm, preserved_arg_indices = freeze(	141	+ frozen_gm, preserved_arg_indices = freeze(gm, aot_autograd_gm, example_inputs)
136	- gm, aot_autograd_gm, example_inputs
137	- )
138	example_inputs = [example_inputs[ind] for ind in preserved_arg_indices]	142	example_inputs = [example_inputs[ind] for ind in preserved_arg_indices]
139	optimized_function = compile_inner(frozen_gm, example_inputs)	143	optimized_function = compile_inner(frozen_gm, example_inputs)
140		144
	@@ -149,7 +153,7 @@ class MindieSDBackend:
149		153
150	def graph_rewrite_before_freezing(fx_graph, inputs):	154	def graph_rewrite_before_freezing(fx_graph, inputs):
151	logger.debug("Graph before compiling:")	155	logger.debug("Graph before compiling:")
152	- if logger.isEnabledFor(logging.DEBUG):	156	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL):
153	logger.debug(fx_graph.print_readable(print_output=False))	157	logger.debug(fx_graph.print_readable(print_output=False))
154	self.__class__.apply_redundant_node_elimination_pass(fx_graph, inputs)	158	self.__class__.apply_redundant_node_elimination_pass(fx_graph, inputs)
155	self.__class__.apply_pattern_match_passes(fx_graph, inputs)	159	self.__class__.apply_pattern_match_passes(fx_graph, inputs)
	@@ -160,7 +164,7 @@ class MindieSDBackend:
160	# make sure we add freezing passes after constant folding	164	# make sure we add freezing passes after constant folding
161	self.__class__.apply_decompose_auto_functionalized_pass(fx_graph)	165	self.__class__.apply_decompose_auto_functionalized_pass(fx_graph)
162	logger.debug("Graph after compiling:")	166	logger.debug("Graph after compiling:")
163	- if logger.isEnabledFor(logging.DEBUG):	167	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL):
164	logger.debug(fx_graph.print_readable(print_output=False))	168	logger.debug(fx_graph.print_readable(print_output=False))
165	return fx_graph	169	return fx_graph
166		170
	@@ -179,4 +183,4 @@ class MindieSDBackend:
179	return aot_autograd(	183	return aot_autograd(
180	fw_compiler=compile_inner,	184	fw_compiler=compile_inner,
181	decompositions=decompositions,	185	decompositions=decompositions,
182	- )(gm, example_inputs)	186	+ )(gm, example_inputs)

Mmindiesd/compilation/passes/pattern_match_pass.py+37-24

	@@ -10,7 +10,6 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
13	-import logging
14	from typing import Any, Callable, Dict, List, Tuple, Optional, Sequence	13	from typing import Any, Callable, Dict, List, Tuple, Optional, Sequence
15	import re	14	import re
16	import torch	15	import torch
	@@ -19,12 +18,14 @@ from torch._inductor.pattern_matcher import PatternMatcherPass
19		18
20	from .gm_pass_base import GraphModulePass	19	from .gm_pass_base import GraphModulePass
21	from .._custom_decomposition import select_custom_decomp_table	20	from .._custom_decomposition import select_custom_decomp_table
		21	+from ...utils.logs.logging import logger
22		22
23	-logger = logging.getLogger(__name__)	23	+DEBUG_LOG_LEVEL = 10
24		24
25	torch_version = re.match(r"(\d+\.\d+)", torch.__version__).group(1)	25	torch_version = re.match(r"(\d+\.\d+)", torch.__version__).group(1)
26	IS_TORCH_21 = torch_version == "2.1"	26	IS_TORCH_21 = torch_version == "2.1"
27	if IS_TORCH_21:	27	if IS_TORCH_21:
		28	+
28	def mindie_inference_graph(fn, args):	29	def mindie_inference_graph(fn, args):
29	from torch.fx.experimental.proxy_tensor import make_fx	30	from torch.fx.experimental.proxy_tensor import make_fx
30	from torch._subclasses.fake_tensor import FakeTensor	31	from torch._subclasses.fake_tensor import FakeTensor
	@@ -45,13 +46,9 @@ if IS_TORCH_21:
45		46
46	class PatternMatchPass(GraphModulePass):	47	class PatternMatchPass(GraphModulePass):
47	def __init__(self):	48	def __init__(self):
48	- self.pattern_replacements: Dict[	49	+ self.pattern_replacements: Dict[str, Tuple[Callable[..., Any], Callable[..., Any]]] = {}
		G G guowenna16月3日 ruff 格式化，简化多行类型/函数调用，无行为变化。
49	- str, Tuple[Callable[..., Any], Callable[..., Any]]
50	- ] = {}
51	try:	50	try:
52	- self.pattern_pass: PatternMatcherPass = PatternMatcherPass(	51	+ self.pattern_pass: PatternMatcherPass = PatternMatcherPass(pass_name="pattern_match_pass") # nosec B106
53	- pass_name="pattern_match_pass"
54	- )
55	except TypeError:	52	except TypeError:
56	self.pattern_pass: PatternMatcherPass = PatternMatcherPass()	53	self.pattern_pass: PatternMatcherPass = PatternMatcherPass()
57		54
	@@ -62,12 +59,13 @@ class PatternMatchPass(GraphModulePass):
62	if cnt == 0:	59	if cnt == 0:
63	break	60	break
64	matched_cnt += cnt	61	matched_cnt += cnt
65	- if logger.isEnabledFor(logging.DEBUG):	62	+ if logger.isEnabledFor(DEBUG_LOG_LEVEL):
66	logger.debug("PatternMatchPass replace %d patterns.", matched_cnt)	63	logger.debug("PatternMatchPass replace %d patterns.", matched_cnt)
67	pattern_idx = 0	64	pattern_idx = 0
68	logger.debug("Patterns registered for replacement:")	65	logger.debug("Patterns registered for replacement:")
69	try:	66	try:
70	from torch._inductor.pattern_matcher import PatternPrettyPrinter	67	from torch._inductor.pattern_matcher import PatternPrettyPrinter
		68	+
71	for pattern_entry in self.pattern_pass.patterns.values():	69	for pattern_entry in self.pattern_pass.patterns.values():
72	for p in pattern_entry:	70	for p in pattern_entry:
73	p_str = PatternPrettyPrinter.run(p.pattern)	71	p_str = PatternPrettyPrinter.run(p.pattern)
	@@ -77,9 +75,6 @@ class PatternMatchPass(GraphModulePass):
77	logger.debug("PatternPrettyPrinter not available, skipping pattern printing")	75	logger.debug("PatternPrettyPrinter not available, skipping pattern printing")
78	return graph	76	return graph
79		77
		G G guowenna16月3日删除仅调用父类的 uuid，修 pylint 无意义重写。
80	- def uuid(self) -> Any:
81	- return super().uuid()
82	-
83	def register_pattern(	78	def register_pattern(
84	self,	79	self,
85	name: str,	80	name: str,
	@@ -88,6 +83,13 @@ class PatternMatchPass(GraphModulePass):
88	example_inputs: List[torch.Tensor],	83	example_inputs: List[torch.Tensor],
89	):	84	):
90	if name in self.pattern_replacements:	85	if name in self.pattern_replacements:
		86	+ logger.error(
		G G guowenna16月3日 pattern 重复注册前新增 error，打印 pattern name、原因和排查建议。
		87	+ "[MindIE-SD/compilation] Pattern registration failed. "
		88	+ "issue=pattern name already registered, pattern_name=%s, expected=unique pattern name. "
		89	+ "possible_cause=activate_pattern_once or custom registration was called repeatedly with the same name. "
		90	+ "Troubleshooting: check pattern registration order and avoid duplicate names.",
		91	+ name,
		92	+ )
91	raise ValueError(f"Pattern '{name}' is already registered.")	93	raise ValueError(f"Pattern '{name}' is already registered.")
92		94
93	self.pattern_replacements[name] = (pattern, replacement)	95	self.pattern_replacements[name] = (pattern, replacement)
	@@ -97,7 +99,14 @@ class PatternMatchPass(GraphModulePass):
97	if IS_TORCH_21:	99	if IS_TORCH_21:
98	pm.fwd_only = mindie_inference_graph	100	pm.fwd_only = mindie_inference_graph
99	else:	101	else:
100	- logger.warning("fwd_only not available in current torch version")	102	+ logger.warning(
		103	+ "[MindIE-SD/compilation] Pattern replacement preparation failed. "
		G G guowenna16月3日 fwd_only 不可用 warning 改为 MindIE 格式，补充 torch 版本和兼容性排查。
		104	+ "issue=torch._inductor.pattern_matcher.fwd_only is unavailable, torch_version=%s, "
		105	+ "expected=fwd_only API exists or torch version is 2.1 for compatibility patch. "
		106	+ "possible_cause=current torch version does not provide the expected inductor API. "
		107	+ "Troubleshooting: verify torch version compatibility and pattern registration stack.",
		108	+ torch.__version__,
		109	+ )
101		110
102	def fwd_only_with_custom_decomp(	111	def fwd_only_with_custom_decomp(
103	fn: Callable[..., Any],	112	fn: Callable[..., Any],
	@@ -107,16 +116,10 @@ class PatternMatchPass(GraphModulePass):
107	get_decomp_fn: Optional[Callable[..., Any]] = select_custom_decomp_table,	116	get_decomp_fn: Optional[Callable[..., Any]] = select_custom_decomp_table,
108	) -> torch.fx.GraphModule:	117	) -> torch.fx.GraphModule:
109	if IS_TORCH_21:	118	if IS_TORCH_21:
110	- return pm.fwd_only(	119	+ return pm.fwd_only(fn=fn, args=args)
111	- fn=fn,
112	- args=args
113	- )
114	else:	120	else:
115	return pm.fwd_only(	121	return pm.fwd_only(
116	- fn=fn,	122	+ fn=fn, args=args, run_functional_passes=run_functional_passes, get_decomp_fn=get_decomp_fn
117	- args=args,
118	- run_functional_passes=run_functional_passes,
119	- get_decomp_fn=get_decomp_fn
120	)	123	)
121		124
122	try:	125	try:
	@@ -130,9 +133,19 @@ class PatternMatchPass(GraphModulePass):
130	logger.debug("Successfully register pattern: %s", name)	133	logger.debug("Successfully register pattern: %s", name)
131	except RuntimeError as e:	134	except RuntimeError as e:
132	if "Duplicate pattern" in str(e):	135	if "Duplicate pattern" in str(e):
133	- logger.warning(	136	+ logger.debug(
134	- "Pattern '%s' is already registered. Skipping duplicate registration.",	137	+ "[MindIE-SD/compilation] Duplicate pattern registration skipped. "
		138	+ "pattern_name=%s, possible_cause=the same pattern was activated more than once.",
135	name,	139	name,
136	)	140	)
137	else:	141	else:
138	- raise e	142	+ logger.error(
		G G guowenna16月3日其他 RuntimeError 抛出前新增 error，打印 pattern、actual_error、shape/dtype 排查建议。
		143	+ "[MindIE-SD/compilation] Pattern registration failed. "
		144	+ "issue=torch inductor register_replacement raised RuntimeError, pattern_name=%s, "
		145	+ "actual_error=%s. possible_cause=pattern, replacement, or example_inputs are incompatible. "
		146	+ "Troubleshooting: inspect the pattern definition, replacement function schema, example input "
		147	+ "shape/dtype, and torch._inductor stack.",
		148	+ name,
		149	+ e,
		150	+ )
		151	+ raise e

Mmindiesd/compilation/passes/redundant_node_elimination_pass.py+5-20

	@@ -10,16 +10,13 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
13	-import logging
14	import operator	13	import operator
15		14
16	import torch	15	import torch
17	-import torch.fx as fx	16	+from torch import fx
18		17
19	from .gm_pass_base import GraphModulePass	18	from .gm_pass_base import GraphModulePass
20		19
21	-logger = logging.getLogger(__name__)
22	-
23		20
24	def get_node_shape(node: torch.fx.node) -> torch.Size \| None:	21	def get_node_shape(node: torch.fx.node) -> torch.Size \| None:
25	"""Retrieve the shape of the tensor represented by the node, if available."""	22	"""Retrieve the shape of the tensor represented by the node, if available."""
	@@ -57,9 +54,7 @@ class ReduandantNodeEliminationPass(GraphModulePass):
57	memory_format = node.kwargs.get("memory_format")	54	memory_format = node.kwargs.get("memory_format")
58		55
59	# None (default) and torch.preserve_format are both no-ops	56	# None (default) and torch.preserve_format are both no-ops
60	- is_noop_clone = (memory_format is None) or (	57	+ is_noop_clone = (memory_format is None) or (memory_format == torch.preserve_format)
61	- memory_format == torch.preserve_format
62	- )
63		58
64	if is_noop_clone:	59	if is_noop_clone:
65	# The input to the clone is the first argument	60	# The input to the clone is the first argument
	@@ -80,10 +75,7 @@ class ReduandantNodeEliminationPass(GraphModulePass):
80	split_sizes_arg = node.args[1]	75	split_sizes_arg = node.args[1]
81		76
82	# We can only optimize if split_sizes is a constant list/tuple	77	# We can only optimize if split_sizes is a constant list/tuple
83	- if (	78	+ if isinstance(split_sizes_arg, (list, tuple)) and len(split_sizes_arg) == 1:
84	- isinstance(split_sizes_arg, (list, tuple))
85	- and len(split_sizes_arg) == 1
86	- ):
87	# This split operation produces a list of 1 tensor,	79	# This split operation produces a list of 1 tensor,
88	# which is just the original input tensor.	80	# which is just the original input tensor.
89	input_tensor_node = node.args[0]	81	input_tensor_node = node.args[0]
	@@ -98,11 +90,7 @@ class ReduandantNodeEliminationPass(GraphModulePass):
98	# Iterate over a static list of users	90	# Iterate over a static list of users
99	for user_node in list(node.users.keys()):	91	for user_node in list(node.users.keys()):
100	# Check if the user is `getitem(self, 0)`	92	# Check if the user is `getitem(self, 0)`
101	- if (	93	+ if user_node.target == operator.getitem and len(user_node.args) == 2 and user_node.args[1] == 0:
102	- user_node.target == operator.getitem
103	- and len(user_node.args) == 2
104	- and user_node.args[1] == 0
105	- ):
106	users_to_replace.append(user_node)	94	users_to_replace.append(user_node)
107	else:	95	else:
108	# This node is used in a way we don't support	96	# This node is used in a way we don't support
	@@ -121,10 +109,7 @@ class ReduandantNodeEliminationPass(GraphModulePass):
121	graph.erase_node(node)	109	graph.erase_node(node)
122	modified = True	110	modified = True
123		111
124	- elif (	112	+ elif node.target in (torch.ops.aten.view.default, torch.ops.aten.reshape.default):
125	- node.target == torch.ops.aten.view.default
126	- or node.target == torch.ops.aten.reshape.default
127	- ):
128	# remove unused view nodes if the shape is the same as input	113	# remove unused view nodes if the shape is the same as input
129	# we rely on the shape info via shape propagation	114	# we rely on the shape info via shape propagation
130	input_node_shape = get_node_shape(node.args[0])	115	input_node_shape = get_node_shape(node.args[0])

Mmindiesd/eplb/eplb_scheduler.py+133-73

	@@ -13,6 +13,7 @@
13	import os	13	import os
14	import multiprocessing	14	import multiprocessing
15	from multiprocessing.managers import BaseManager	15	from multiprocessing.managers import BaseManager
		16	+from dataclasses import dataclass
		G G guowenna16月3日新增 dataclass，用于把 scheduler 运行上下文收拢到结构体里，避免 run_scheduler 里变量过多、嵌套过深。
16	import threading	17	import threading
17	import queue	18	import queue
18	import time	19	import time
	@@ -31,10 +32,23 @@ instruction_queues = {}
31	class ScheduleManager(BaseManager):	32	class ScheduleManager(BaseManager):
32	pass	33	pass
33		34
		35	+
34	ScheduleManager.register('get_upload_queues', callable=lambda rank: upload_queues[rank])	36	ScheduleManager.register('get_upload_queues', callable=lambda rank: upload_queues[rank])
35	ScheduleManager.register('get_instruction_queues', callable=lambda rank: instruction_queues[rank])	37	ScheduleManager.register('get_instruction_queues', callable=lambda rank: instruction_queues[rank])
36		38
37		39
		40	+@dataclass
		G G guowenna16月3日新增 SchedulerContext，集中保存 scheduler_args/world_size/redundant/buffer/update_count。原因是后续拆函数都需要这些状态，直接传一堆参数会更乱。
		41	+class SchedulerContext: # pylint: disable=too-many-instance-attributes
		42	+ scheduler_args: argparse.Namespace
		43	+ world_size: int
		44	+ redundant: int
		45	+ experts_set: set
		46	+ experts_per_rank: int
		47	+ load_report_buffer: dict
		48	+ local_expert_buffer: dict
		49	+ update_count: int = 0
		50	+
		51	+
38	def get_args():	52	def get_args():
39	parser = argparse.ArgumentParser(description="EPLB scheduler")	53	parser = argparse.ArgumentParser(description="EPLB scheduler")
40	parser.add_argument("--world_size", type=int, default=8)	54	parser.add_argument("--world_size", type=int, default=8)
	@@ -65,21 +79,18 @@ def get_manager_client(addr, auth_key):
65	return manager	79	return manager
66		80
67		81
68	-def run_scheduler(args):	82	+def _init_scheduler_context(scheduler_args):
		G G guowenna16月3日把原 run_scheduler 前半段初始化逻辑拆成 _init_scheduler_context。原因是修复 pylint too-many-nested-blocks，同时让初始化、循环处理、结果下发职责分开。
69	- world_size = args.world_size	83	+ world_size = scheduler_args.world_size
70	- server_addr = (args.host, args.port)	84	+ server_addr = (scheduler_args.host, scheduler_args.port)
71	- redundant = args.redundant	85	+ redundant = scheduler_args.redundant
72	- auth_key = args.auth_key	86	+ auth_key = scheduler_args.auth_key
73	- experts_set = set(range(args.expert_num))	87	+ experts_set = set(range(scheduler_args.expert_num))
74	- experts_per_rank = args.expert_num // world_size	88	+ experts_per_rank = scheduler_args.expert_num // world_size
75		89
76	- num_moe_layers = args.block_num	90	+ num_moe_layers = scheduler_args.block_num
77	load_report_buffer = {idx: {} for idx in range(num_moe_layers)}	91	load_report_buffer = {idx: {} for idx in range(num_moe_layers)}
78	local_expert_buffer = {idx: {} for idx in range(num_moe_layers)}	92	local_expert_buffer = {idx: {} for idx in range(num_moe_layers)}
79		93
80	- count = 0
81	-
82	- global upload_queues, instruction_queues
83	# zmq	94	# zmq
84	for rank in range(world_size):	95	for rank in range(world_size):
85	upload_queues[rank] = queue.Queue()	96	upload_queues[rank] = queue.Queue()
	@@ -87,76 +98,125 @@ def run_scheduler(args):
87		98
88	start_manager_server(server_addr, auth_key)	99	start_manager_server(server_addr, auth_key)
89		100
90	- logger.debug(f"[Scheduler] starting moniter")	101	+ return SchedulerContext(
		102	+ scheduler_args=scheduler_args,
		103	+ world_size=world_size,
		104	+ redundant=redundant,
		105	+ experts_set=experts_set,
		106	+ experts_per_rank=experts_per_rank,
		107	+ load_report_buffer=load_report_buffer,
		108	+ local_expert_buffer=local_expert_buffer,
		109	+ )
		110	+
		111	+
		112	+def _complete_local_expert_list(context, local_expert_list):
		G G guowenna16月3日新增 _complete_local_expert_list，单独处理 EX 模式下冗余专家补齐逻辑。random.sample(... ) # nosec B311 是因为这里不是安全随机场景，避免 bandit 误报。
		113	+ scheduler_args = context.scheduler_args
		114	+ if (
		115	+ scheduler_args.mode == "EX"
		116	+ and context.redundant > 0
		117	+ and len(local_expert_list) != (context.experts_per_rank + context.redundant)
		118	+ ):
		119	+ random_range = list(context.experts_set - set(local_expert_list))
		120	+ redundant_expert = random.sample(random_range, context.redundant) # nosec B311
		121	+ return local_expert_list + redundant_expert
		122	+ return local_expert_list
		123	+
		124	+
		125	+def _emit_layer_update(context, layer_idx, transfer):
		G G guowenna16月3日新增 _emit_layer_update，单独负责 greedy 计算和任务下发。原因是原来嵌套在队列循环里，层级太深且日志不清晰。
		126	+ scheduler_args = context.scheduler_args
		127	+ response = context.load_report_buffer[layer_idx]
		128	+ expert_dict = dict(sorted(context.local_expert_buffer[layer_idx].items()))
		129	+ context.load_report_buffer[layer_idx] = {}
		130	+ context.local_expert_buffer[layer_idx] = {}
		131	+
		132	+ logger.debug(
		G G guowenna16月3日新增 MindIE 风格 debug 日志，记录 layer_idx/world_size/mode，方便定位是哪层、哪种 EPLB 模式触发计算。
		133	+ "[MindIE-SD/eplb] EPLB greedy compute started. layer_idx=%s, world_size=%s, mode=%s.",
		134	+ layer_idx,
		135	+ context.world_size,
		136	+ scheduler_args.mode,
		137	+ )
		138	+ result = eplb_greedy(
		139	+ response=response,
		140	+ algorithm_type=scheduler_args.mode,
		141	+ device_to_expert=expert_dict,
		142	+ world_size=context.world_size,
		143	+ expert_num=scheduler_args.expert_num,
		144	+ max_move=scheduler_args.max_move,
		145	+ redundant=context.redundant,
		146	+ )
		147	+ update, device_indices_list, local_expert_indices_list, local_expert_list, expert_trans_tensor = result
		148	+ if not update:
		149	+ return
		150	+
		151	+ transfer.update_emit_task(
		152	+ device_indices_list,
		153	+ local_expert_indices_list,
		154	+ local_expert_list,
		155	+ expert_trans_tensor,
		156	+ context.world_size,
		157	+ )
		158	+ context.update_count += 1
		159	+ logger.debug(
		160	+ "[MindIE-SD/eplb] Layer layout computed. layer_idx=%s, update_count=%s.",
		G G guowenna16月3日新增布局计算完成日志，打印 layer_idx/update_count，替代原来较泛的 layer_x layout has computed。
		161	+ layer_idx,
		162	+ context.update_count,
		163	+ )
		164	+
		165	+
		166	+def _process_rank_report(context, rank):
		G G guowenna16月3日新增 _process_rank_report，单独处理单个 rank 的上报数据。
		167	+ try:
		G G guowenna16月3日 queue.Empty 直接返回 True，表示该 rank 当前队列为空，减少原来的内层 try/except 嵌套。
		168	+ report = upload_queues[rank].get_nowait()
		169	+ except queue.Empty:
		170	+ return True
		171	+
		172	+ try:
		G G guowenna16月3日异常改成带上下文的 ModelExecError，写清楚问题、rank、world_size、mode、原始异常、可能原因和排查方向。原来的 "[Scheduler] error : {e}" 没有 f-string 生效，也没有根因信息。
		173	+ layer_idx = report['moe_layer_idx']
		174	+ load_data = report['load']
		175	+ local_expert_list = _complete_local_expert_list(context, report['local_expert_list'])
		176	+
		177	+ context.load_report_buffer[layer_idx][rank] = load_data
		178	+ context.local_expert_buffer[layer_idx][rank] = local_expert_list
		179	+ transfer = UpdateTaskTransfer(instruction_queues, layer_idx)
		180	+
		181	+ if len(context.load_report_buffer[layer_idx]) == context.world_size:
		182	+ _emit_layer_update(context, layer_idx, transfer)
		183	+ return False
		184	+ except Exception as e:
		185	+ raise ModelExecError(
		186	+ "[MindIE-SD/eplb] EPLB scheduler failed. "
		187	+ f"issue=failed to process upload queue, rank={rank}, world_size={context.world_size}, "
		188	+ f"mode={context.scheduler_args.mode}, "
		189	+ f"actual_error={e}. possible_cause=invalid load report, greedy algorithm failure, or queue state "
		190	+ "mismatch. Troubleshooting: inspect worker load report fields, EPLB mode/redundant settings, "
		191	+ "and scheduler traceback."
		192	+ ) from e
		193	+
		194	+
		195	+def run_scheduler(scheduler_args):
		G G guowenna16月3日 run_scheduler 只保留主循环，调用 _process_rank_report。原因是解决 pylint 嵌套过深，同时提升可读性。
		196	+ context = _init_scheduler_context(scheduler_args)
		197	+ logger.debug(
		G G guowenna16月3日启动日志改为 MindIE 格式，打印 world_size/host/port/mode。
		198	+ "[MindIE-SD/eplb] Scheduler monitor started. world_size=%s, host=%s, port=%s, mode=%s.",
		199	+ context.world_size,
		200	+ scheduler_args.host,
		201	+ scheduler_args.port,
		202	+ scheduler_args.mode,
		203	+ )
91		204
92	while True:	205	while True:
93	all_queues_empty = True	206	all_queues_empty = True
94	try:	207	try:
95	- for rank in range(world_size):	208	+ for rank in range(context.world_size):
96	- try:	209	+ all_queues_empty = _process_rank_report(context, rank) and all_queues_empty
97	- report = upload_queues[rank].get_nowait()
98	-
99	- layer_idx = report['moe_layer_idx']
100	- load_data = report['load']
101	- local_expert_list = report['local_expert_list']
102	-
103	- if args.mode == "EX" and redundant > 0 and len(local_expert_list) != (experts_per_rank + redundant):
104	- random_range = list(experts_set - set(local_expert_list))
105	- redundant_expert = random.sample(random_range, redundant)
106	- local_expert_list = local_expert_list + redundant_expert
107	-
108	- load_report_buffer[layer_idx][rank] = load_data
109	- local_expert_buffer[layer_idx][rank] = local_expert_list
110	- transfer = UpdateTaskTransfer(instruction_queues, layer_idx)
111	-
112	- all_queues_empty = False
113	-
114	- if len(load_report_buffer[layer_idx]) == world_size:
115	-
116	- response = load_report_buffer[layer_idx]
117	- expert_dict = local_expert_buffer[layer_idx]
118	- expert_dict = dict(sorted(expert_dict.items()))
119	- load_report_buffer[layer_idx] = {}
120	- local_expert_buffer[layer_idx] = {}
121	-
122	- logger.debug(f"[greedy] eplb greedy compute")
123	- result = eplb_greedy(
124	- response=response, algorithm_type=args.mode,
125	- device_to_expert=expert_dict, world_size=world_size,
126	- expert_num=args.expert_num, max_move=args.max_move, redundant=redundant)
127	- (
128	- update,
129	- device_indices_list,
130	- local_expert_indices_list,
131	- local_expert_list,
132	- expert_trans_tensor
133	- ) = result
134	-
135	- if not update:
136	- continue
137	-
138	- transfer.update_emit_task(
139	- device_indices_list,
140	- local_expert_indices_list,
141	- local_expert_list,
142	- expert_trans_tensor,
143	- world_size
144	- )
145	- count += 1
146	- logger.info(f"[Scheduler] layer_{layer_idx} layout has computed.")
147	- except queue.Empty:
148	- pass
149	- except Exception as e:
150	- raise ModelExecError("[Scheduler] error : {e}") from e
151	except (KeyboardInterrupt, SystemExit):	210	except (KeyboardInterrupt, SystemExit):
		G G guowenna16月3日退出信号日志改为 debug，且使用 [MindIE-SD/eplb] 模块前缀。
152	- logger.info("[Scheduler] exit sign!")	211	+ logger.debug("[MindIE-SD/eplb] Scheduler received exit signal.")
153	break	212	break
154	if all_queues_empty:	213	if all_queues_empty:
155	time.sleep(0.1)	214	time.sleep(0.1)
156		215
157	- logger.info(f"Already has update {count} times")	216	+ logger.debug("[MindIE-SD/eplb] Scheduler update count. count=%s.", context.update_count)
		G G guowenna16月3日结束日志改为 debug，并用参数化日志打印 update_count，避免普通运行时 info 噪声。
158	- logger.info("[Scheduler] Scheduler cycle end.")	217	+ logger.debug("[MindIE-SD/eplb] Scheduler cycle ended.")
		218	+
159		219
160	if __name__ == '__main__':	220	if __name__ == '__main__':
161	- args = get_args()	221	+ cli_args = get_args()
		G G guowenna16月3日 args 改名 cli_args，语义更明确，避免和函数参数混淆。
162	- run_scheduler(args)	222	+ run_scheduler(cli_args)

Mmindiesd/eplb/greedy_algorithm.py+199-78

	@@ -10,6 +10,7 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
		13	+# pylint: disable=logging-fstring-interpolation,pointless-string-statement,unbalanced-tuple-unpacking,useless-parent-delegation
13	import random	14	import random
14	from dataclasses import dataclass	15	from dataclasses import dataclass
15	import numpy as np	16	import numpy as np
	@@ -31,7 +32,7 @@ class LoadData:
31	expert_trans_tensor: torch.Tensor = None	32	expert_trans_tensor: torch.Tensor = None
32		33
33		34
34	-class EPLBService():	35	+class EPLBService:
		G G guowenna16月3日 Python3 不需要空括号，修 pylint/ruff 风格问题。
35	"""	36	"""
36	EPLB算法的基类	37	EPLB算法的基类
37		38
	@@ -43,8 +44,18 @@ class EPLBService():
43	cost_local (float): 单个Token的本地计算成本 (C_comp)	44	cost_local (float): 单个Token的本地计算成本 (C_comp)
44	cost_remote (float): 单个Token的远程通信成本 (C_comm)	45	cost_remote (float): 单个Token的远程通信成本 (C_comm)
45	"""	46	"""
46	- def __init__(self, num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,	47	+
47	- max_move_number, load_balance_threshold):	48	+ def __init__(
		G G guowenna16月3日构造函数和 LoadData 初始化改成多行，是 ruff format 自动格式化，便于通过 pre-commit，无行为变化。
		49	+ self,
		50	+ num_devices,
		51	+ num_experts,
		52	+ expert_mems,
		53	+ device_mems,
		54	+ cost_local,
		55	+ cost_remote,
		56	+ max_move_number,
		57	+ load_balance_threshold,
		58	+ ):
48	self.num_devices = num_devices	59	self.num_devices = num_devices
49	self.num_experts = num_experts	60	self.num_experts = num_experts
50	self.expert_mems = expert_mems	61	self.expert_mems = expert_mems
	@@ -88,28 +99,35 @@ class EPLBService():
88	# --- 2. 处理共享专家 ---	99	# --- 2. 处理共享专家 ---
89	self.process_share_expert(placement, shared_expert_id, used_mems)	100	self.process_share_expert(placement, shared_expert_id, used_mems)
90	# --- 3. 阶段1：满足约束的初始分配 ---	101	# --- 3. 阶段1：满足约束的初始分配 ---
91	- initial_load_data = LoadData(placement=placement, shared_expert_id=shared_expert_id,	102	+ initial_load_data = LoadData(
92	- total_traffic=total_traffic, used_mems=used_mems,	103	+ placement=placement,
93	- origin_device_to_expert=origin_device_to_expert,	104	+ shared_expert_id=shared_expert_id,
94	- sorted_experts=sorted_experts, expert_trans_tensor=expert_trans_tensor)	105	+ total_traffic=total_traffic,
		106	+ used_mems=used_mems,
		107	+ origin_device_to_expert=origin_device_to_expert,
		108	+ sorted_experts=sorted_experts,
		109	+ expert_trans_tensor=expert_trans_tensor,
		110	+ )
95	device_to_expert = self.initial_placement(initial_load_data)	111	device_to_expert = self.initial_placement(initial_load_data)
96	# --- 4. 阶段2：迭代复制优化minmax目标 ---	112	# --- 4. 阶段2：迭代复制优化minmax目标 ---
97	- optimize_load_data = LoadData(placement=placement, shared_expert_id=shared_expert_id,	113	+ optimize_load_data = LoadData(
98	- total_traffic=total_traffic, used_mems=used_mems,	114	+ placement=placement,
99	- global_expert_load=global_expert_load, device_to_expert=device_to_expert)	115	+ shared_expert_id=shared_expert_id,
		116	+ total_traffic=total_traffic,
		117	+ used_mems=used_mems,
		118	+ global_expert_load=global_expert_load,
		119	+ device_to_expert=device_to_expert,
		120	+ )
100	self.optimize_min_max(optimize_load_data)	121	self.optimize_min_max(optimize_load_data)
101		122
102	return {	123	return {
103	"final_placement": placement,	124	"final_placement": placement,
104	"final_memory_usage": used_mems,	125	"final_memory_usage": used_mems,
105	"device_to_expert_map": device_to_expert,	126	"device_to_expert_map": device_to_expert,
106	- "expert_trans_tensor": expert_trans_tensor	127	+ "expert_trans_tensor": expert_trans_tensor,
107	}	128	}
108		129
109	- def optimize_min_max(	130	+ def optimize_min_max(self, load_data: LoadData):
110	- self,
111	- load_data: LoadData
112	- ):
113	pass	131	pass
114		132
115	def initial_placement(self, load_data: LoadData):	133	def initial_placement(self, load_data: LoadData):
	@@ -139,9 +157,22 @@ class EPLBService():
139	load_data.used_mems[best_device] += self.expert_mems[expert_id]	157	load_data.used_mems[best_device] += self.expert_mems[expert_id]
140	logger.debug(	158	logger.debug(
141	f" -> Place high-load expert {expert_id} to device {best_device} "	159	f" -> Place high-load expert {expert_id} to device {best_device} "
142	- f"(Current memory: {load_data.used_mems[best_device]:.1f}GB)")	160	+ f"(Current memory: {load_data.used_mems[best_device]:.1f}GB)"
		161	+ )
143	else:	162	else:
144	# 报错：如果每个专家连一个初始位置都找不到，说明无解	163	# 报错：如果每个专家连一个初始位置都找不到，说明无解
		164	+ logger.error(
		G G guowenna16月3日在初始专家放置失败前新增 logger.error。原因是原来只抛 MemoryError，日志里没有 MindIE 格式，也没有参数对比。现在打印 expert_id、期望内存、设备内存、已用内存、可能原因和排查建议。
		165	+ "[MindIE-SD/eplb] Expert initial placement failed. "
		166	+ "issue=no device has enough memory for expert, expert_id=%s, expected_free_memory>=%sGB, "
		167	+ "actual_device_memory=%sGB, actual_used_memory=%sGB. "
		168	+ "possible_cause=EPLB memory configuration cannot fit every expert. "
		169	+ "Troubleshooting: increase expert_per_rank/redundant memory budget, reduce expert size, "
		170	+ "or check world_size and expert_num settings.",
		171	+ expert_id,
		172	+ self.expert_mems[expert_id],
		173	+ self.device_mems,
		174	+ load_data.used_mems,
		175	+ )
145	raise MemoryError(	176	raise MemoryError(
146	f"Error：expert {expert_id} (need {self.expert_mems[expert_id]}GB) "	177	f"Error：expert {expert_id} (need {self.expert_mems[expert_id]}GB) "
147	f"Unable to locate initial position for this expert on any device, check memory configuration."	178	f"Unable to locate initial position for this expert on any device, check memory configuration."
	@@ -166,12 +197,25 @@ class EPLBService():
166		197
167	def process_share_expert(self, placement, shared_expert_id, used_mems):	198	def process_share_expert(self, placement, shared_expert_id, used_mems):
168	if shared_expert_id is not None:	199	if shared_expert_id is not None:
169	- logger.debug(f"--- [Preprocessing] Shared Expert {shared_expert_id} detected, "	200	+ logger.debug(
170	- f"forcing deployment on all devices ---")	201	+ f"--- [Preprocessing] Shared Expert {shared_expert_id} detected, forcing deployment on all devices ---"
		202	+ )
171	shared_mem = self.expert_mems[shared_expert_id]	203	shared_mem = self.expert_mems[shared_expert_id]
172	for i in range(self.num_devices):	204	for i in range(self.num_devices):
173	# 检查内存是否足够	205	# 检查内存是否足够
174	if self.device_mems[i] < shared_mem:	206	if self.device_mems[i] < shared_mem:
		207	+ logger.error(
		G G guowenna16月3日共享专家放置失败前新增 logger.error。原因同上，补齐 device_id/shared_expert_id/expected/actual，方便判断是哪个卡内存不足。
		208	+ "[MindIE-SD/eplb] Shared expert placement failed. "
		209	+ "issue=device memory is insufficient for shared expert, device_id=%s, shared_expert_id=%s, "
		210	+ "expected_device_memory>=%sGB, actual_device_memory=%sGB. "
		211	+ "possible_cause=shared expert memory requirement exceeds one or more device budgets. "
		212	+ "Troubleshooting: increase per-device expert memory budget, reduce shared expert size, "
		213	+ "or adjust EPLB deployment configuration.",
		214	+ i,
		215	+ shared_expert_id,
		216	+ shared_mem,
		217	+ self.device_mems[i],
		218	+ )
175	raise MemoryError(	219	raise MemoryError(
176	f"Device {i} (memory {self.device_mems[i]}GB) cannot accommodate the shared expert."	220	f"Device {i} (memory {self.device_mems[i]}GB) cannot accommodate the shared expert."
177	f"{shared_expert_id} (need {shared_mem}GB)。"	221	f"{shared_expert_id} (need {shared_mem}GB)。"
	@@ -185,16 +229,30 @@ class A2ARedundantExpertService(EPLBService):
185	"""	229	"""
186	面向All-to-all通信方式下的冗余专家动态调度	230	面向All-to-all通信方式下的冗余专家动态调度
187	"""	231	"""
188	- def __init__(self, num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,
189	- max_move_number, load_balance_threshold):
190	- super().__init__(
191	- num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,
192	- max_move_number, load_balance_threshold)
193		232
194	- def optimize_min_max(	233	+ def __init__(
195	- self,	234	+ self,
196	- load_data: LoadData	235	+ num_devices,
197	- ):	236	+ num_experts,
		237	+ expert_mems,
		238	+ device_mems,
		239	+ cost_local,
		240	+ cost_remote,
		241	+ max_move_number,
		242	+ load_balance_threshold,
		243	+ ):
		244	+ super().__init__(
		245	+ num_devices,
		246	+ num_experts,
		247	+ expert_mems,
		248	+ device_mems,
		249	+ cost_local,
		250	+ cost_remote,
		251	+ max_move_number,
		252	+ load_balance_threshold,
		253	+ )
		254	+
		255	+ def optimize_min_max(self, load_data: LoadData):
198	logger.debug("--- [Phase 2] Starting iterative replication optimization.---")	256	logger.debug("--- [Phase 2] Starting iterative replication optimization.---")
199	iteration = 1	257	iteration = 1
200	while True:	258	while True:
	@@ -216,19 +274,19 @@ class A2ARedundantExpertService(EPLBService):
216		274
217	bottleneck_device = np.argmax(current_loads)	275	bottleneck_device = np.argmax(current_loads)
218	max_load = np.max(current_loads)	276	max_load = np.max(current_loads)
219	- logger.debug(f"\nRound {iteration} \| The current system bottleneck (maximum load): {max_load:,.0f} "	277	+ logger.debug(
220	- f"(on device {bottleneck_device})")	278	+ f"\nRound {iteration} \| The current system bottleneck (maximum load): {max_load:,.0f} "
		279	+ f"(on device {bottleneck_device})"
		280	+ )
221		281
222	# 遍历所有可能的“复制操作”	282	# 遍历所有可能的“复制操作”
223	# “复制操作”是指将专家 j 复制到设备 i，前提是 i 上没有 j，且内存足够。	283	# “复制操作”是指将专家 j 复制到设备 i，前提是 i 上没有 j，且内存足够。
224	for expert_id in range(self.num_experts):	284	for expert_id in range(self.num_experts):
225	-
226	# 跳过已被强制部署的共享专家	285	# 跳过已被强制部署的共享专家
227	if load_data.shared_expert_id is not None and expert_id == load_data.shared_expert_id:	286	if load_data.shared_expert_id is not None and expert_id == load_data.shared_expert_id:
228	continue	287	continue
229		288
230	for device_id in range(self.num_devices):	289	for device_id in range(self.num_devices):
231	-
232	# 如果专家已经存在于此设备，或内存不足，则跳过	290	# 如果专家已经存在于此设备，或内存不足，则跳过
233	if load_data.placement[device_id, expert_id] == 1:	291	if load_data.placement[device_id, expert_id] == 1:
234	continue	292	continue
	@@ -259,8 +317,9 @@ class A2ARedundantExpertService(EPLBService):
259	# --- 迭代停止 ---	317	# --- 迭代停止 ---
260	# 当找不到任何一个收益分数更大的机会时 (所有可能的专家复制，带来的内存开销都得不偿失)，或者内存不足而无法操作时，算法停止。	318	# 当找不到任何一个收益分数更大的机会时 (所有可能的专家复制，带来的内存开销都得不偿失)，或者内存不足而无法操作时，算法停止。
261	if best_move is None:	319	if best_move is None:
262	- logger.debug("\n--- [Phase 2] Optimization completed: No more beneficial replication operations "	320	+ logger.debug(
263	- "found.---")	321	+ "\n--- [Phase 2] Optimization completed: No more beneficial replication operations found.---"
		322	+ )
264	break	323	break
265		324
266	# 执行本轮找到的最佳移动	325	# 执行本轮找到的最佳移动
	@@ -269,11 +328,13 @@ class A2ARedundantExpertService(EPLBService):
269	load_data.device_to_expert[dev_to_add].append(exp_to_add)	328	load_data.device_to_expert[dev_to_add].append(exp_to_add)
270	load_data.used_mems[dev_to_add] += self.expert_mems[exp_to_add]	329	load_data.used_mems[dev_to_add] += self.expert_mems[exp_to_add]
271		330
272	- logger.debug(f" -> Best Move: Copy Expert {exp_to_add} to Device {dev_to_add} "	331	+ logger.debug(
273	- f"(Benefit Score: {max_score:,.2f})")	332	+ f" -> Best Move: Copy Expert {exp_to_add} to Device {dev_to_add} (Benefit Score: {max_score:,.2f})"
		333	+ )
274	logger.debug(f" Load benefit: Reduced load by {move_gain:,.0f} for Device {dev_to_add}")	334	logger.debug(f" Load benefit: Reduced load by {move_gain:,.0f} for Device {dev_to_add}")
275	- logger.debug(f" New memory status: "	335	+ logger.debug(
276	- f"{load_data.used_mems[dev_to_add]:.1f}GB / {self.device_mems[dev_to_add]}GB")	336	+ f" New memory status: {load_data.used_mems[dev_to_add]:.1f}GB / {self.device_mems[dev_to_add]}GB"
		337	+ )
277		338
278	iteration += 1	339	iteration += 1
279		340
	@@ -282,10 +343,28 @@ class AGRedundantExpertService(EPLBService):
282	"""	343	"""
283	面向All-Gather通信方式下的冗余专家动态调度	344	面向All-Gather通信方式下的冗余专家动态调度
284	"""	345	"""
285	- def __init__(self, num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,	346	+
286	- max_move_number, load_balance_threshold):	347	+ def __init__(
287	- super().__init__(num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,	348	+ self,
288	- max_move_number, load_balance_threshold)	349	+ num_devices,
		350	+ num_experts,
		351	+ expert_mems,
		352	+ device_mems,
		353	+ cost_local,
		354	+ cost_remote,
		355	+ max_move_number,
		356	+ load_balance_threshold,
		357	+ ):
		358	+ super().__init__(
		359	+ num_devices,
		360	+ num_experts,
		361	+ expert_mems,
		362	+ device_mems,
		363	+ cost_local,
		364	+ cost_remote,
		365	+ max_move_number,
		366	+ load_balance_threshold,
		367	+ )
289		368
290	@staticmethod	369	@staticmethod
291	def get_expert_total_demand(total_traffic):	370	def get_expert_total_demand(total_traffic):
	@@ -313,8 +392,10 @@ class AGRedundantExpertService(EPLBService):
313		392
314	bottleneck_device = np.argmax(current_loads)	393	bottleneck_device = np.argmax(current_loads)
315	max_load = np.max(current_loads)	394	max_load = np.max(current_loads)
316	- logger.debug(f"\nRound {iteration} \| The current system bottleneck (maximum load): {max_load:,.0f} "	395	+ logger.debug(
317	- f"(on device {bottleneck_device})")	396	+ f"\nRound {iteration} \| The current system bottleneck (maximum load): {max_load:,.0f} "
		397	+ f"(on device {bottleneck_device})"
		398	+ )
318	logger.debug(f" Devices load: {[f'{device_load:,.0f}' for device_load in current_loads]}")	399	logger.debug(f" Devices load: {[f'{device_load:,.0f}' for device_load in current_loads]}")
319		400
320	# --- 修改点 4: 重写收益评估逻辑 ---	401	# --- 修改点 4: 重写收益评估逻辑 ---
	@@ -367,8 +448,9 @@ class AGRedundantExpertService(EPLBService):
367	best_move = (device_id, expert_id, gain)	448	best_move = (device_id, expert_id, gain)
368		449
369	if best_move is None:	450	if best_move is None:
370	- logger.debug("\n--- [Phase 2] Optimization completed: No more beneficial replication operations "	451	+ logger.debug(
371	- "found.---")	452	+ "\n--- [Phase 2] Optimization completed: No more beneficial replication operations found.---"
		453	+ )
372	break	454	break
373		455
374	dev_to_add, exp_to_add, move_gain = best_move	456	dev_to_add, exp_to_add, move_gain = best_move
	@@ -376,11 +458,13 @@ class AGRedundantExpertService(EPLBService):
376	load_data.device_to_expert[dev_to_add].append(exp_to_add)	458	load_data.device_to_expert[dev_to_add].append(exp_to_add)
377	load_data.used_mems[dev_to_add] += self.expert_mems[exp_to_add]	459	load_data.used_mems[dev_to_add] += self.expert_mems[exp_to_add]
378		460
379	- logger.debug(f" -> Best Move: Copy Expert {exp_to_add} to Device {dev_to_add} "	461	+ logger.debug(
380	- f"(Benefit Score: {max_score:,.2f})")	462	+ f" -> Best Move: Copy Expert {exp_to_add} to Device {dev_to_add} (Benefit Score: {max_score:,.2f})"
		463	+ )
381	logger.debug(f" Load benefit: System maximum load reduced by {move_gain:,.0f}")	464	logger.debug(f" Load benefit: System maximum load reduced by {move_gain:,.0f}")
382	- logger.debug(f" New memory status: "	465	+ logger.debug(
383	- f"{load_data.used_mems[dev_to_add]:.1f}GB / {self.device_mems[dev_to_add]}GB")	466	+ f" New memory status: {load_data.used_mems[dev_to_add]:.1f}GB / {self.device_mems[dev_to_add]}GB"
		467	+ )
384		468
385	iteration += 1	469	iteration += 1
386		470
	@@ -389,10 +473,28 @@ class ExpertExchangeService(EPLBService):
389	"""	473	"""
390	基于专家交换的动态调度方案	474	基于专家交换的动态调度方案
391	"""	475	"""
392	- def __init__(self, num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,	476	+
393	- max_move_number, load_balance_threshold):	477	+ def __init__(
394	- super().__init__(num_devices, num_experts, expert_mems, device_mems, cost_local, cost_remote,	478	+ self,
395	- max_move_number, load_balance_threshold)	479	+ num_devices,
		480	+ num_experts,
		481	+ expert_mems,
		482	+ device_mems,
		483	+ cost_local,
		484	+ cost_remote,
		485	+ max_move_number,
		486	+ load_balance_threshold,
		487	+ ):
		488	+ super().__init__(
		489	+ num_devices,
		490	+ num_experts,
		491	+ expert_mems,
		492	+ device_mems,
		493	+ cost_local,
		494	+ cost_remote,
		495	+ max_move_number,
		496	+ load_balance_threshold,
		497	+ )
396		498
397	def initial_placement(self, load_data: LoadData):	499	def initial_placement(self, load_data: LoadData):
398	device_to_expert = load_data.origin_device_to_expert.copy()	500	device_to_expert = load_data.origin_device_to_expert.copy()
	@@ -403,8 +505,8 @@ class ExpertExchangeService(EPLBService):
403	device_to_expert = {}	505	device_to_expert = {}
404	for device_id in range(self.num_devices):	506	for device_id in range(self.num_devices):
405	expert_end_index = (device_id + 1) * (self.num_experts // self.num_devices)	507	expert_end_index = (device_id + 1) * (self.num_experts // self.num_devices)
406	- device_loads[device_id] = load_data.total_traffic[device_id, expert_start_index: expert_end_index].sum()	508	+ device_loads[device_id] = load_data.total_traffic[device_id, expert_start_index:expert_end_index].sum()
407	- device_to_expert[device_id] = [i for i in range(expert_start_index, expert_end_index)]	509	+ device_to_expert[device_id] = list(range(expert_start_index, expert_end_index))
		G G guowenna16月3日切片空格格式化，ruff 要求，无行为变化。
408	for expert_id in range(expert_start_index, expert_end_index):	510	for expert_id in range(expert_start_index, expert_end_index):
409	load_data.used_mems[device_id] += self.expert_mems[expert_id]	511	load_data.used_mems[device_id] += self.expert_mems[expert_id]
410	expert_start_index = expert_end_index	512	expert_start_index = expert_end_index
	@@ -426,8 +528,10 @@ class ExpertExchangeService(EPLBService):
426	delta_load = (max_load - min_load) // 2	528	delta_load = (max_load - min_load) // 2
427	logger.debug(f"--- Max-min device load diff {delta_load * 2}")	529	logger.debug(f"--- Max-min device load diff {delta_load * 2}")
428	if delta_load * 2 < self.load_balance_threshold:	530	if delta_load * 2 < self.load_balance_threshold:
429	- logger.debug(f"------------ Max-min device load diff less than {self.load_balance_threshold} "	531	+ logger.debug(
430	- f"End the iteration ---------------------")	532	+ f"------------ Max-min device load diff less than {self.load_balance_threshold} "
		533	+ f"End the iteration ---------------------"
		534	+ )
431	break	535	break
432		536
433	# 通过两个设备上的专家负载，计算专家间两两交换后带来的负载差与设备负载差的差值矩阵，形成一个二维矩阵，此时值最小的i,j，就是我们要交换的两个专家	537	# 通过两个设备上的专家负载，计算专家间两两交换后带来的负载差与设备负载差的差值矩阵，形成一个二维矩阵，此时值最小的i,j，就是我们要交换的两个专家
	@@ -437,7 +541,7 @@ class ExpertExchangeService(EPLBService):
437	max_load_device_traffic = load_data.total_traffic[max_device_index, expert_max_idx]	541	max_load_device_traffic = load_data.total_traffic[max_device_index, expert_max_idx]
438	min_load_device_traffic = load_data.total_traffic[min_device_index, expert_min_idx]	542	min_load_device_traffic = load_data.total_traffic[min_device_index, expert_min_idx]
439	else:	543	else:
440	- raise ParametersInvalid(f"[greedy] expert not in index list")	544	+ raise ParametersInvalid("[greedy] expert not in index list")
441		545
442	trans_traffic = np.abs((max_load_device_traffic[:, np.newaxis] - min_load_device_traffic) - delta_load)	546	trans_traffic = np.abs((max_load_device_traffic[:, np.newaxis] - min_load_device_traffic) - delta_load)
443	# 找到 trans_traffic 中最小值的全局索引	547	# 找到 trans_traffic 中最小值的全局索引
	@@ -454,28 +558,39 @@ class ExpertExchangeService(EPLBService):
454	experid_from_min_to_max = expert_min_idx[experid_from_min_to_max_index]	558	experid_from_min_to_max = expert_min_idx[experid_from_min_to_max_index]
455		559
456	# 检查条件：目标设备上不能已包含对方专家	560	# 检查条件：目标设备上不能已包含对方专家
457	- if (experid_from_min_to_max not in device_to_expert.get(max_device_index, []) and	561	+ if experid_from_min_to_max not in device_to_expert.get(
458	- experid_from_max_to_min not in device_to_expert.get(min_device_index, [])):	562	+ max_device_index, []
		563	+ ) and experid_from_max_to_min not in device_to_expert.get(min_device_index, []):
459	# 找到满足条件的交换对，跳出循环	564	# 找到满足条件的交换对，跳出循环
460	break	565	break
461	else:	566	else:
462	logger.debug("\n--- No more beneficial replication operations found.---")	567	logger.debug("\n--- No more beneficial replication operations found.---")
463	break	568	break
464		569
465	- move_expert_load = (load_data.total_traffic[max_device_index, experid_from_max_to_min]	570	+ move_expert_load = (
466	- - load_data.total_traffic[min_device_index, experid_from_min_to_max])	571	+ load_data.total_traffic[max_device_index, experid_from_max_to_min]
		572	+ - load_data.total_traffic[min_device_index, experid_from_min_to_max]
		573	+ )
467		574
468	# 计算模拟移动后的两个设备间的负载差，和之前的负载差相比，计算收益，收益大于0则执行移动	575	# 计算模拟移动后的两个设备间的负载差，和之前的负载差相比，计算收益，收益大于0则执行移动
469	new_delta_load = abs((max_load - move_expert_load) - (min_load + move_expert_load)) // 2	576	new_delta_load = abs((max_load - move_expert_load) - (min_load + move_expert_load)) // 2
470	gain = delta_load - new_delta_load	577	gain = delta_load - new_delta_load
471	if gain > 0:	578	if gain > 0:
472	# 执行移动前进行防御性检查	579	# 执行移动前进行防御性检查
473	- if (max_device_index in device_to_expert and	580	+ if max_device_index in device_to_expert and min_device_index in device_to_expert:
474	- min_device_index in device_to_expert):
475	# 执行移动	581	# 执行移动
476	device_to_expert[max_device_index][experid_from_max_to_min_index] = experid_from_min_to_max	582	device_to_expert[max_device_index][experid_from_max_to_min_index] = experid_from_min_to_max
477	device_to_expert[min_device_index][experid_from_min_to_max_index] = experid_from_max_to_min	583	device_to_expert[min_device_index][experid_from_min_to_max_index] = experid_from_max_to_min
478	else:	584	else:
		585	+ logger.error(
		G G guowenna16月3日专家交换失败前新增 logger.error，打印期望设备、实际 device_to_expert key、可能原因和排查方案。原来直接抛 IndexError，现场信息不足。
		586	+ "[MindIE-SD/eplb] Expert exchange failed. "
		587	+ "issue=device index is missing from device_to_expert, expected_devices=(%s,%s), "
		588	+ "actual_devices=%s. possible_cause=EPLB placement state is inconsistent before exchange. "
		589	+ "Troubleshooting: inspect device_to_expert construction and load report rank coverage.",
		590	+ max_device_index,
		591	+ min_device_index,
		592	+ list(device_to_expert.keys()),
		593	+ )
479	raise IndexError("Device or expert index out of bounds")	594	raise IndexError("Device or expert index out of bounds")
480		595
481	device_loads[max_device_index] -= move_expert_load	596	device_loads[max_device_index] -= move_expert_load
	@@ -490,7 +605,8 @@ class ExpertExchangeService(EPLBService):
490	f"device {min_device_index} load "	605	f"device {min_device_index} load "
491	f"reduced by {load_data.total_traffic[min_device_index, experid_from_min_to_max]}, "	606	f"reduced by {load_data.total_traffic[min_device_index, experid_from_min_to_max]}, "
492	f"device {max_device_index} load "	607	f"device {max_device_index} load "
493	- f"increased {load_data.total_traffic[min_device_index, experid_from_min_to_max]}. ")	608	+ f"increased {load_data.total_traffic[min_device_index, experid_from_min_to_max]}. "
		609	+ )
494	logger.debug(f"--- Latest Max-min device load diff: {new_delta_load * 2}")	610	logger.debug(f"--- Latest Max-min device load diff: {new_delta_load * 2}")
495		611
496	# 记录1、移动专家数，2、更新内存，3、更新专家交换映射矩阵	612	# 记录1、移动专家数，2、更新内存，3、更新专家交换映射矩阵
	@@ -501,10 +617,12 @@ class ExpertExchangeService(EPLBService):
501	load_data.used_mems[max_device_index] += self.expert_mems[experid_from_min_to_max]	617	load_data.used_mems[max_device_index] += self.expert_mems[experid_from_min_to_max]
502	load_data.used_mems[max_device_index] -= self.expert_mems[experid_from_max_to_min]	618	load_data.used_mems[max_device_index] -= self.expert_mems[experid_from_max_to_min]
503		619
504	- i = experid_from_max_to_min_index + (max_device_index	620	+ i = experid_from_max_to_min_index + (
505	- * (load_data.expert_trans_tensor.shape[0] // self.num_devices))	621	+ max_device_index * (load_data.expert_trans_tensor.shape[0] // self.num_devices)
506	- j = experid_from_min_to_max_index + (min_device_index	622	+ )
507	- * (load_data.expert_trans_tensor.shape[0] // self.num_devices))	623	+ j = experid_from_min_to_max_index + (
		624	+ min_device_index * (load_data.expert_trans_tensor.shape[0] // self.num_devices)
		625	+ )
508	load_data.expert_trans_tensor[:, [i, j]] = load_data.expert_trans_tensor[:, [j, i]]	626	load_data.expert_trans_tensor[:, [i, j]] = load_data.expert_trans_tensor[:, [j, i]]
509	logger.debug(f"--- Current load status: {device_loads} ---")	627	logger.debug(f"--- Current load status: {device_loads} ---")
510	return device_to_expert	628	return device_to_expert
	@@ -525,7 +643,7 @@ def process_final_placement(results, num_experts):
525	if expert in device_experts:	643	if expert in device_experts:
526	device_indices[device_index].append(device_index)	644	device_indices[device_index].append(device_index)
527	else:	645	else:
528	- device_indices[device_index].append(random.choice(revert_list[expert]))	646	+ device_indices[device_index].append(random.choice(revert_list[expert])) # nosec B311
		G G guowenna16月3日这里是 EPLB 路由选择，不是安全随机，避免 bandit 误报。
529	local_expert_list = []	647	local_expert_list = []
530		648
531	for _, val in current_placement.items():	649	for _, val in current_placement.items():
	@@ -588,16 +706,19 @@ def eplb_greedy(**kwargs):
588		706
589	# 定义每个场景关联的算法服务	707	# 定义每个场景关联的算法服务
590	handlers = {	708	handlers = {
591	- 'A2A': A2ARedundantExpertService(world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote,	709	+ 'A2A': A2ARedundantExpertService(
592	- max_move, load_balance_threshold),	710	+ world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote, max_move, load_balance_threshold
593	- 'AG': AGRedundantExpertService(world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote,	711	+ ),
594	- max_move, load_balance_threshold),	712	+ 'AG': AGRedundantExpertService(
595	- 'EX': ExpertExchangeService(world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote,	713	+ world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote, max_move, load_balance_threshold
596	- max_move, load_balance_threshold)	714	+ ),
		715	+ 'EX': ExpertExchangeService(
		716	+ world_size, expert_num, expert_mems, device_mems, cost_local, cost_remote, max_move, load_balance_threshold
		717	+ ),
597	}	718	}
598	algorithm_service = handlers.get(algorithm_type)	719	algorithm_service = handlers.get(algorithm_type)
599	result = algorithm_service.placement_greedy(current_pattern, device_to_expert)	720	result = algorithm_service.placement_greedy(current_pattern, device_to_expert)
600	output = process_final_placement(result, expert_num)	721	output = process_final_placement(result, expert_num)
601	device_indices, local_expert_indices, local_expert_list, expert_trans_tensor = output	722	device_indices, local_expert_indices, local_expert_list, expert_trans_tensor = output
602	- logger.info(f"current_placement:{local_expert_list}")	723	+ logger.debug("[MindIE-SD/eplb] Current expert placement computed. local_expert_list=%s.", local_expert_list)
		G G guowenna16月3日原因是这类布局详情属于调试日志，且要符合 MindIE 模块前缀和参数化日志规范。
603	- return update, device_indices, local_expert_indices, local_expert_list, expert_trans_tensor	724	+ return update, device_indices, local_expert_indices, local_expert_list, expert_trans_tensor

Mmindiesd/eplb/task_manager.py+15-26

	@@ -35,46 +35,37 @@ def parse_module(module):
35	return dispatcher_list, expert_load_collector_list	35	return dispatcher_list, expert_load_collector_list
36		36
37		37
38	-def expert_info_transfer_pool(	38	+def expert_info_transfer_pool(module, instruction_queue, upload_queue, device):
39	- module,
40	- instruction_queue,
41	- upload_queue,
42	- device
43	- ):
44	dispatcher_list, expert_load_collector_list = parse_module(module)	39	dispatcher_list, expert_load_collector_list = parse_module(module)
45	transfer_stream = torch_npu.npu.Stream(device)	40	transfer_stream = torch_npu.npu.Stream(device)
46		41
47	for idx, collector in enumerate(expert_load_collector_list):	42	for idx, collector in enumerate(expert_load_collector_list):
48	- collector.task_transfer = ProfileTaskTransfer(	43	+ collector.task_transfer = ProfileTaskTransfer(instruction_queue, idx, collector.lb_interval)
49	- instruction_queue,
50	- idx,
51	- collector.lb_interval
52	- )
53		44
54	while True:	45	while True:
55	instruction = instruction_queue.get()	46	instruction = instruction_queue.get()
56	if instruction is None or instruction == 'exit':	47	if instruction is None or instruction == 'exit':
57	- logger.info(f"[ExpertInfoTransferPool] Get exit instruction")	48	+ logger.debug("[MindIE-SD/eplb] Expert info transfer pool received exit instruction.")
58	break	49	break
59	if isinstance(instruction, TaskPayload):	50	if isinstance(instruction, TaskPayload):
60	handler_function = TASK_DISPATCHER.get(instruction.task_type, handle_unknown_task)	51	handler_function = TASK_DISPATCHER.get(instruction.task_type, handle_unknown_task)
61	handler_function(instruction, upload_queue, expert_load_collector_list, dispatcher_list, transfer_stream)	52	handler_function(instruction, upload_queue, expert_load_collector_list, dispatcher_list, transfer_stream)
62	else:	53	else:
63	- logger.debug(f"Unknown instruction: {instruction}")	54	+ logger.debug("[MindIE-SD/eplb] Unknown instruction ignored. instruction=%s.", instruction)
64		55
65		56
66	-def connect_to_schedule_manager(	57	+def connect_to_schedule_manager(rank_in_group, ip, port, auth_key):
67	- rank_in_group,
68	- ip,
69	- port,
70	- auth_key
71	- ):
72	addr = (ip, port)	58	addr = (ip, port)
73	manager = get_manager_client(addr, auth_key)	59	manager = get_manager_client(addr, auth_key)
74	manager.connect()	60	manager.connect()
75	- logger.info(f"Connected to schedule manager, rank_in_group: {rank_in_group}")	61	+ logger.debug(
76	- instruction_queue = manager.get_instruction_queues(rank=rank_in_group)	62	+ "[MindIE-SD/eplb] Connected to schedule manager. rank_in_group=%s, manager_addr=%s:%s.",
77	- upload_queue = manager.get_upload_queues(rank=rank_in_group)	63	+ rank_in_group,
		64	+ ip,
		65	+ port,
		66	+ )
		67	+ instruction_queue = manager.get_instruction_queues(rank=rank_in_group) # pylint: disable=no-member
		68	+ upload_queue = manager.get_upload_queues(rank=rank_in_group) # pylint: disable=no-member
78	return instruction_queue, upload_queue	69	return instruction_queue, upload_queue
79		70
80		71
	@@ -89,9 +80,7 @@ def construct_expert_info_transfer_pool(**kwargs):
89	if instruction_queue is None or upload_queue is None:	80	if instruction_queue is None or upload_queue is None:
90	return None, None	81	return None, None
91	worker = threading.Thread(	82	worker = threading.Thread(
92	- target=expert_info_transfer_pool,	83	+ target=expert_info_transfer_pool, args=(module, instruction_queue, upload_queue, device), daemon=True
93	- args=(module, instruction_queue, upload_queue, device),
94	- daemon=True
95	)	84	)
96	worker.start()	85	worker.start()
97	- return worker, instruction_queue	86	+ return worker, instruction_queue

Mmindiesd/eplb/task_transfer.py+16-27

	@@ -17,22 +17,14 @@ from .task_payload import TaskType, TaskPayload
17		17
18		18
19	class ProfileTaskTransfer:	19	class ProfileTaskTransfer:
20	- def __init__(	20	+ def __init__(self, task_queue: queue.Queue, moe_layer_idx: int, lb_interval: int = 1):
21	- self,
22	- task_queue: queue.Queue,
23	- moe_layer_idx: int,
24	- lb_interval: int = 1
25	- ):
26	self.instruction_queue = task_queue	21	self.instruction_queue = task_queue
27	self.moe_layer_idx = moe_layer_idx	22	self.moe_layer_idx = moe_layer_idx
28	self.lb_interval = lb_interval	23	self.lb_interval = lb_interval
29	self.flag = 0	24	self.flag = 0
30		25
31	def profile_emit_task(self):	26	def profile_emit_task(self):
32	- task_payload = TaskPayload(	27	+ task_payload = TaskPayload(task_type=TaskType.PROFILE, moe_layer_idx=self.moe_layer_idx)
33	- task_type=TaskType.PROFILE,
34	- moe_layer_idx=self.moe_layer_idx
35	- )
36	if self.instruction_queue:	28	if self.instruction_queue:
37	self.flag += 1	29	self.flag += 1
38	if self.flag != self.lb_interval:	30	if self.flag != self.lb_interval:
	@@ -41,38 +33,35 @@ class ProfileTaskTransfer:
41	try:	33	try:
42	self.instruction_queue.put_nowait(task_payload)	34	self.instruction_queue.put_nowait(task_payload)
43	except queue.Full:	35	except queue.Full:
44	- logger.info(f"[Warning] instruction_queue full!!!")	36	+ logger.warning(
		G G guowenna16月3日 instruction_queue 满时从泛泛 info 改为 warning，补充 moe_layer_idx、lb_interval、根因和排查建议。
45	- pass	37	+ "[MindIE-SD/eplb] EPLB profile task enqueue failed. "
		38	+ "issue=instruction_queue is full, moe_layer_idx=%s, lb_interval=%s, expected=queue has free slot. "
		39	+ "possible_cause=scheduler consumes profile tasks slower than workers produce them. "
		40	+ "Troubleshooting: check scheduler process state, queue consumer thread, and EPLB interval configuration.",
		41	+ self.moe_layer_idx,
		42	+ self.lb_interval,
		43	+ )
46		44
47		45
48	class UpdateTaskTransfer:	46	class UpdateTaskTransfer:
49	- def __init__(	47	+ def __init__(self, task_queue: queue.Queue, moe_layer_idx):
50	- self,
51	- task_queue: queue.Queue,
52	- moe_layer_idx
53	- ):
54	self.instruction_queue = task_queue	48	self.instruction_queue = task_queue
55	self.moe_layer_idx = moe_layer_idx	49	self.moe_layer_idx = moe_layer_idx
56		50
57	def update_emit_task(	51	def update_emit_task(
58	- self,	52	+ self, device_indices_list, local_expert_indices_list, local_expert_list, expert_trans_tensor, world_size
59	- device_indices_list,	53	+ ):
60	- local_expert_indices_list,
61	- local_expert_list,
62	- expert_trans_tensor,
63	- world_size
64	- ):
65	for rank in range(world_size):	54	for rank in range(world_size):
66	layout_command = {	55	layout_command = {
67	'device_indices': device_indices_list[rank],	56	'device_indices': device_indices_list[rank],
68	'local_expert_indices': local_expert_indices_list[rank],	57	'local_expert_indices': local_expert_indices_list[rank],
69	'local_expert_list': local_expert_list[rank],	58	'local_expert_list': local_expert_list[rank],
70	- 'expert_trans_tensor': expert_trans_tensor	59	+ 'expert_trans_tensor': expert_trans_tensor,
71	}	60	}
72	task_payload = TaskPayload(	61	task_payload = TaskPayload(
73	task_type=TaskType.UPDATE_LAYOUT,	62	task_type=TaskType.UPDATE_LAYOUT,
74	worker_rank=rank,	63	worker_rank=rank,
75	moe_layer_idx=self.moe_layer_idx,	64	moe_layer_idx=self.moe_layer_idx,
76	- data=layout_command	65	+ data=layout_command,
77	)	66	)
78	- self.instruction_queue[rank].put(task_payload)	67	+ self.instruction_queue[rank].put(task_payload)

Mmindiesd/layers/flash_attn/attention_forward.py+6-2

	@@ -127,9 +127,13 @@ def check_input_params(input_params):
127	def get_manual_attention_op_type(attn_param, op_type):	127	def get_manual_attention_op_type(attn_param, op_type):
128	if is_a5_device() and op_type in _A5_DEPRECATED_OP_TYPES:	128	if is_a5_device() and op_type in _A5_DEPRECATED_OP_TYPES:
129	logger.warning(	129	logger.warning(
130	- "'%s' is not supported on A5 devices and has been routed to 'fused_attn_score'. "	130	+ "[MindIE-SD/flash_attn] Manual attention operator remapped for A5. "
131	- "Please switch to 'fused_attn_score' (or remove the manual op_type) to silence this warning.",	131	+ "issue=manual op_type is not supported on A5, expected_op_type=fused_attn_score, actual_op_type=%s, "
		132	+ "q_seqlen=%s, kv_seqlen=%s. possible_cause=the requested operator is deprecated on A5 devices. "
		133	+ "Troubleshooting: set op_type='fused_attn_score' or remove the manual op_type setting.",
132	op_type,	134	op_type,
		135	+ attn_param.q_seqlen,
		136	+ attn_param.kv_seqlen,
133	)	137	)
134	return "fused_attn_score"	138	return "fused_attn_score"
135		139

Mmindiesd/layers/flash_attn/attention_func.py+15-4

	@@ -79,9 +79,14 @@ def get_attention_function_static(attn_param):
79		79
80	if is_a5_device() and op_type in _A5_DEPRECATED_OP_TYPES:	80	if is_a5_device() and op_type in _A5_DEPRECATED_OP_TYPES:
81	logger.warning(	81	logger.warning(
82	- "Static-table op_type '%s' is not supported on A5 devices and has been routed to "	82	+ "[MindIE-SD/flash_attn] Static-table attention operator remapped for A5. "
83	- "'fused_attn_score'. Please refresh the static table entry to silence this warning.",	83	+ "issue=static-table op_type is not supported on A5, expected_op_type=fused_attn_score, "
		84	+ "actual_op_type=%s, q_seqlen=%s, kv_seqlen=%s. "
		85	+ "possible_cause=the static attention table contains a deprecated A5 operator. "
		86	+ "Troubleshooting: refresh the static table entry to use fused_attn_score.",
84	op_type,	87	op_type,
		88	+ attn_param.q_seqlen,
		89	+ attn_param.kv_seqlen,
85	)	90	)
86	op_type = "fused_attn_score"	91	op_type = "fused_attn_score"
87	return get_attention_function(attn_param, op_type, layout)	92	return get_attention_function(attn_param, op_type, layout)
	@@ -96,7 +101,10 @@ def get_attention_function(attn_param, op_type, layout):
96	elif npu_device == NPUDevice.A5:	101	elif npu_device == NPUDevice.A5:
97	op_registry = device_a5_op.get_all()	102	op_registry = device_a5_op.get_all()
98	else:	103	else:
99	- raise ParametersInvalid("Platform invalid. Please check env.")	104	+ raise ParametersInvalid(
		105	+ f"Platform invalid. expected one of {[item.name for item in NPUDevice if item != NPUDevice.UNDEFINED]}, "
		106	+ f"actual={npu_device}. Please check env."
		107	+ )
100		108
101	if op_type not in op_registry:	109	if op_type not in op_registry:
102	raise ParametersInvalid(	110	raise ParametersInvalid(
	@@ -132,7 +140,10 @@ def get_attention_function_runtime(attn_param, query, key, value, attn_mask=None
132	elif npu_device == NPUDevice.A5:	140	elif npu_device == NPUDevice.A5:
133	all_op = device_a5_op.get_all()	141	all_op = device_a5_op.get_all()
134	else:	142	else:
135	- raise ParametersInvalid("Platform invalid.")	143	+ raise ParametersInvalid(
		144	+ f"Platform invalid. expected one of {[item.name for item in NPUDevice if item != NPUDevice.UNDEFINED]}, "
		145	+ f"actual={npu_device}."
		146	+ )
136		147
137	func_lists = get_test_func_lists(attn_param, all_op)	148	func_lists = get_test_func_lists(attn_param, all_op)
138	if len(func_lists) == 0:	149	if len(func_lists) == 0:

Mmindiesd/layers/flash_attn/sparse_flash_attn.py+11-8

	@@ -55,17 +55,20 @@ def _resolve_sparse_type_for_a5(sparse_type, inner_precise):
55		55
56	if sparse_type == "rf_v2":	56	if sparse_type == "rf_v2":
57	if inner_precise != A5_V3_INNER_PRECISE:	57	if inner_precise != A5_V3_INNER_PRECISE:
58	- logger.info(	58	+ logger.debug(
59	- "sparse_type='rf_v2' is routed to 'rf_v3' on A5 devices, and inner_precise has been "	59	+ "[MindIE-SD/flash_attn] Sparse attention type remapped for A5. "
60	- "overridden from %s to %s as required by the v3 operator. "	60	+ "sparse_type=rf_v2, actual_sparse_type=rf_v3, expected_inner_precise=%s, actual_inner_precise=%s. "
61	- "Please switch to sparse_type='rf_v3' explicitly to silence this notice.",	61	+ "possible_cause=rf_v2 is deprecated on A5 and rf_v3 requires a fixed inner_precise value. "
		62	+ "Troubleshooting: set sparse_type='rf_v3' and inner_precise=%s explicitly.",
		63	+ A5_V3_INNER_PRECISE,
62	inner_precise,	64	inner_precise,
63	A5_V3_INNER_PRECISE,	65	A5_V3_INNER_PRECISE,
64	)	66	)
65	else:	67	else:
66	- logger.info(	68	+ logger.debug(
67	- "sparse_type='rf_v2' is routed to 'rf_v3' on A5 devices. "	69	+ "[MindIE-SD/flash_attn] Sparse attention type remapped for A5. "
68	- "Please switch to sparse_type='rf_v3' explicitly to silence this notice."	70	+ "sparse_type=rf_v2, actual_sparse_type=rf_v3. possible_cause=rf_v2 is deprecated on A5. "
		71	+ "Troubleshooting: set sparse_type='rf_v3' explicitly."
69	)	72	)
70	return "rf_v3", A5_V3_INNER_PRECISE	73	return "rf_v3", A5_V3_INNER_PRECISE
71		74
	@@ -226,7 +229,7 @@ def sparse_attention(
226	sparse_size=block_size,	229	sparse_size=block_size,
227	)	230	)
228	elif sparse_type is None:	231	elif sparse_type is None:
229	- out = torch_npu.npu_fusion_attention(	232	+ out = torch_npu.npu_fusion_attention( # pylint: disable=no-member
230	q,	233	q,
231	k,	234	k,
232	v,	235	v,

Mmindiesd/layers/moe/moe.py+2-2

	@@ -72,8 +72,8 @@ def _log_moe_config_once(dispatcher_cls, tokens_full, reduce_results):
72	return	72	return
73	dispatcher_name = "dynamic" if dispatcher_cls.__name__ == "DynamicDispatcher" else "static"	73	dispatcher_name = "dynamic" if dispatcher_cls.__name__ == "DynamicDispatcher" else "static"
74	comm_type = get_moe_comm_type().value	74	comm_type = get_moe_comm_type().value
75	- logger.info(	75	+ logger.debug(
76	- "MindIE-SD moe config: dispatcher=%s, comm_type=%s, tokens_full=%s, reduce_results=%s.",	76	+ "[MindIE-SD/moe] MoE config resolved. dispatcher=%s, comm_type=%s, tokens_full=%s, reduce_results=%s.",
77	dispatcher_name,	77	dispatcher_name,
78	comm_type,	78	comm_type,
79	tokens_full,	79	tokens_full,

Mmindiesd/layers/register_ops.py+22-15

	@@ -12,10 +12,11 @@
12	from pathlib import Path	12	from pathlib import Path
13	from functools import wraps	13	from functools import wraps
14	import os	14	import os
15	-from typing import Dict, Callable	15	+from typing import Callable
16	import torch	16	import torch
17	from torch.library import Library	17	from torch.library import Library
18	from ..utils import file_utils, ParametersInvalid, is_npu_available	18	from ..utils import file_utils, ParametersInvalid, is_npu_available
		19	+from ..utils.logs.logging import logger
19		20
20		21
21	MINDIE_NS = "mindiesd" # 固定命名空间，与 torch.ops.mindiesd 对应	22	MINDIE_NS = "mindiesd" # 固定命名空间，与 torch.ops.mindiesd 对应
	@@ -37,10 +38,7 @@ def _load_mindie_ops_library() -> None:
37	ops_path = file_utils.standardize_path(str(ops_path))	38	ops_path = file_utils.standardize_path(str(ops_path))
38	ops_file = os.path.join(ops_path, "libPTAExtensionOPS.so")	39	ops_file = os.path.join(ops_path, "libPTAExtensionOPS.so")
39		40
40	- file_utils.check_file_safety(	41	+ file_utils.check_file_safety(ops_file, permission_mode=file_utils.BINARY_FILE_PERMISSION)
41	- ops_file,
42	- permission_mode=file_utils.BINARY_FILE_PERMISSION
43	- )
44	torch.ops.load_library(ops_file)	42	torch.ops.load_library(ops_file)
45		43
46		44
	@@ -70,22 +68,18 @@ if torch.__version__.startswith("2.1"):
70		68
71	def _compatible_register_fake(op_name: str):	69	def _compatible_register_fake(op_name: str):
72	"""Compatibility wrapper for PyTorch 2.1 fake registration."""	70	"""Compatibility wrapper for PyTorch 2.1 fake registration."""
		71	+
73	def decorator(fake_func: Callable):	72	def decorator(fake_func: Callable):
74	@wraps(fake_func)	73	@wraps(fake_func)
75	def wrapper(args, *kwargs):	74	def wrapper(args, *kwargs):
76	# Ensure all tensor inputs are on Meta device (required for PyTorch 2.1)	75	# Ensure all tensor inputs are on Meta device (required for PyTorch 2.1)
77	- args = [	76	+ args = [a.to(device="meta") if isinstance(a, torch.Tensor) else a for a in args]
78	- a.to(device="meta") if isinstance(a, torch.Tensor) else a	77	+ kwargs = {k: v.to(device="meta") if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
79	- for a in args
80	- ]
81	- kwargs = {
82	- k: v.to(device="meta") if isinstance(v, torch.Tensor) else v
83	- for k, v in kwargs.items()
84	- }
85	return fake_func(args, *kwargs)	78	return fake_func(args, *kwargs)
86		79
87	_lib.impl(op_name, wrapper, "Meta")	80	_lib.impl(op_name, wrapper, "Meta")
88	return fake_func	81	return fake_func
		82	+
89	return decorator	83	return decorator
90	else:	84	else:
91	# PyTorch 2.2+ 使用 register_fake 或 impl_abstract	85	# PyTorch 2.2+ 使用 register_fake 或 impl_abstract
	@@ -99,7 +93,6 @@ else:
99	return _native_register_fake(op_name)	93	return _native_register_fake(op_name)
100		94
101		95
102	-
103	def register_mindie_fake_op(op_name: str):	96	def register_mindie_fake_op(op_name: str):
104	"""Decorator to register a fake implementation for a MindIE operator.	97	"""Decorator to register a fake implementation for a MindIE operator.
105		98
	@@ -115,14 +108,28 @@ def register_mindie_fake_op(op_name: str):
115	Decorator function that registers the fake implementation.	108	Decorator function that registers the fake implementation.
116	"""	109	"""
117	if not is_npu_available():	110	if not is_npu_available():
		111	+
118	def dummy_decorator(func):	112	def dummy_decorator(func):
119	return func	113	return func
		114	+
120	return dummy_decorator	115	return dummy_decorator
121		116
122	if not check_mindie_operator_exists(op_name):	117	if not check_mindie_operator_exists(op_name):
		118	+ logger.error(
		119	+ "[MindIE-SD/layers] MindIE custom operator registration failed. "
		120	+ "issue=operator is not found in torch.ops.%s, op_name=%s, expected=%s::%s exists. "
		121	+ "possible_cause=custom operator shared library was not loaded or TORCH_LIBRARY registration is missing. "
		122	+ "Troubleshooting: check libPTAExtensionOPS.so path, ASCEND_CUSTOM_OPP_PATH, operator build output, "
		123	+ "and torch.ops.%s registry.",
		124	+ MINDIE_NS,
		125	+ op_name,
		126	+ MINDIE_NS,
		127	+ op_name,
		128	+ MINDIE_NS,
		129	+ )
123	raise RuntimeError(	130	raise RuntimeError(
124	f"MindIE operator {MINDIE_NS}::{op_name} not found! "	131	f"MindIE operator {MINDIE_NS}::{op_name} not found! "
125	"Ensure the SO library is loaded and the operator is registered with TORCH_LIBRARY."	132	"Ensure the SO library is loaded and the operator is registered with TORCH_LIBRARY."
126	)	133	)
127		134
128	- return _compatible_register_fake(f"{MINDIE_NS}::{op_name}")	135	+ return _compatible_register_fake(f"{MINDIE_NS}::{op_name}")

Mmindiesd/offload.py+56-0

	@@ -14,6 +14,20 @@ import itertools
14	import torch	14	import torch
15	from torch.nn import ModuleList	15	from torch.nn import ModuleList
16		16
		17	+from .utils.logs.logging import logger
		18	+
		19	+
		20	+def _log_offload_param_error(issue, expected, actual, troubleshooting):
		21	+ logger.error(
		22	+ "[MindIE-SD/offload] Offload parameter validation failed. "
		23	+ "issue=%s, expected=%s, actual=%s. possible_cause=caller passed invalid offload configuration. "
		24	+ "Troubleshooting: %s",
		25	+ issue,
		26	+ expected,
		27	+ actual,
		28	+ troubleshooting,
		29	+ )
		30	+
17		31
18	def enable_offload(model, blocks, min_reserved_blocks_count=2):	32	def enable_offload(model, blocks, min_reserved_blocks_count=2):
19	"""	33	"""
	@@ -87,26 +101,68 @@ def enable_offload(model, blocks, min_reserved_blocks_count=2):
87	>>>	101	>>>
88	"""	102	"""
89	if not isinstance(model, torch.nn.Module):	103	if not isinstance(model, torch.nn.Module):
		104	+ _log_offload_param_error(
		105	+ "model type mismatch",
		106	+ "torch.nn.Module",
		107	+ type(model).__name__,
		108	+ "pass a torch.nn.Module instance as model",
		109	+ )
90	raise TypeError(f"model must be torch.nn.Module type, current type: {type(model).__name__}")	110	raise TypeError(f"model must be torch.nn.Module type, current type: {type(model).__name__}")
91		111
92	if not isinstance(blocks, ModuleList):	112	if not isinstance(blocks, ModuleList):
		113	+ _log_offload_param_error(
		114	+ "blocks type mismatch",
		115	+ "torch.nn.ModuleList",
		116	+ type(blocks).__name__,
		117	+ "pass model blocks as torch.nn.ModuleList",
		118	+ )
93	raise TypeError(f"blocks must be ModuleList, current type: {type(blocks).__name__}")	119	raise TypeError(f"blocks must be ModuleList, current type: {type(blocks).__name__}")
94		120
95	if not blocks:	121	if not blocks:
		122	+ _log_offload_param_error(
		123	+ "blocks is empty",
		124	+ "len(blocks)>0",
		125	+ len(blocks),
		126	+ "provide at least one block before enabling offload",
		127	+ )
96	raise ValueError("blocks cannot be empty list")	128	raise ValueError("blocks cannot be empty list")
97		129
98	for i, block in enumerate(blocks):	130	for i, block in enumerate(blocks):
99	if not isinstance(block, torch.nn.Module):	131	if not isinstance(block, torch.nn.Module):
		132	+ _log_offload_param_error(
		133	+ "block type mismatch",
		134	+ "torch.nn.Module",
		135	+ f"blocks[{i}]={type(block).__name__}",
		136	+ "ensure every item in blocks is a torch.nn.Module",
		137	+ )
100	raise TypeError(f"blocks[{i}] must be torch.nn.Module type, current type: {type(block).__name__}")	138	raise TypeError(f"blocks[{i}] must be torch.nn.Module type, current type: {type(block).__name__}")
101		139
102	if not isinstance(min_reserved_blocks_count, int):	140	if not isinstance(min_reserved_blocks_count, int):
		141	+ _log_offload_param_error(
		142	+ "min_reserved_blocks_count type mismatch",
		143	+ "int",
		144	+ type(min_reserved_blocks_count).__name__,
		145	+ "pass an integer min_reserved_blocks_count",
		146	+ )
103	raise TypeError(	147	raise TypeError(
104	f"min_reserved_blocks_count must be int type, current type: {type(min_reserved_blocks_count).__name__}"	148	f"min_reserved_blocks_count must be int type, current type: {type(min_reserved_blocks_count).__name__}"
105	)	149	)
106	if min_reserved_blocks_count < 0:	150	if min_reserved_blocks_count < 0:
		151	+ _log_offload_param_error(
		152	+ "min_reserved_blocks_count is negative",
		153	+ "min_reserved_blocks_count>=0",
		154	+ min_reserved_blocks_count,
		155	+ "set min_reserved_blocks_count to a non-negative integer",
		156	+ )
107	raise ValueError(f"min_reserved_blocks_count must be >= 0, current value: {min_reserved_blocks_count}")	157	raise ValueError(f"min_reserved_blocks_count must be >= 0, current value: {min_reserved_blocks_count}")
108		158
109	if min_reserved_blocks_count >= len(blocks):	159	if min_reserved_blocks_count >= len(blocks):
		160	+ _log_offload_param_error(
		161	+ "min_reserved_blocks_count exceeds block count",
		162	+ "min_reserved_blocks_count<len(blocks)",
		163	+ f"min_reserved_blocks_count={min_reserved_blocks_count}, len(blocks)={len(blocks)}",
		164	+ "reduce min_reserved_blocks_count or provide more blocks",
		165	+ )
110	raise ValueError(	166	raise ValueError(
111	f"min_reserved_blocks_count must be < len(blocks), "	167	f"min_reserved_blocks_count must be < len(blocks), "
112	f"current value: {min_reserved_blocks_count}, blocks length: {len(blocks)}"	168	f"current value: {min_reserved_blocks_count}, blocks length: {len(blocks)}"

Mmindiesd/quantization/quantize.py+49-30

	@@ -16,14 +16,22 @@ from typing import Dict
16	from collections import OrderedDict	16	from collections import OrderedDict
17	from functools import wraps	17	from functools import wraps
18	import torch	18	import torch
19	-import torch.nn as nn	19	+from torch import nn
20	import safetensors	20	import safetensors
21	from .mode import QuantAlgorithm	21	from .mode import QuantAlgorithm
22	from .config import QuantConfig, LayerQuantConfig, TimestepPolicyConfig	22	from .config import QuantConfig, LayerQuantConfig, TimestepPolicyConfig
23	-from .mode import W4A4_LIST,W8A8_LIST	23	+from .mode import W4A4_LIST, W8A8_LIST
24	from .utils import replace_rank_suffix, get_quant_weight, extract_constructor_args, MAX_WEIGHT_SIZE	24	from .utils import replace_rank_suffix, get_quant_weight, extract_constructor_args, MAX_WEIGHT_SIZE
25	-from .layer import (W4A4QuantLinear, W4A4MXFP4DualQuantLinear, W8A8QuantLinear, W8A8TimeStepQuantLinear,	25	+from .layer import (
26	- WeightQuantLinear, FP8RotateQuantFA, W8A8MXFP8QuantLinear, W4A4MXFP4QuantLinear)	26	+ W4A4QuantLinear,
		27	+ W4A4MXFP4DualQuantLinear,
		28	+ W8A8QuantLinear,
		29	+ W8A8TimeStepQuantLinear,
		30	+ WeightQuantLinear,
		31	+ FP8RotateQuantFA,
		32	+ W8A8MXFP8QuantLinear,
		33	+ W4A4MXFP4QuantLinear,
		34	+)
27	from ..utils import ParametersInvalid, ConfigError	35	from ..utils import ParametersInvalid, ConfigError
28	from ..utils import file_utils	36	from ..utils import file_utils
29	from ..utils.logs.logging import logger	37	from ..utils.logs.logging import logger
	@@ -35,7 +43,7 @@ def get_key_patterns(layer_name):
35	f'{layer_name}.weight',	43	f'{layer_name}.weight',
36	f'{layer_name}',	44	f'{layer_name}',
37	f'{layer_name}.fa_q.scale',	45	f'{layer_name}.fa_q.scale',
38	- f'{layer_name}.quant_type'	46	+ f'{layer_name}.quant_type',
39	]	47	]
40	return key_patterns	48	return key_patterns
41		49
	@@ -47,9 +55,7 @@ def weight_quantize(name, layer, cfg, quant_weights, **kwargs):
47		55
48		56
49	def w8a16_quantize(name, layer, cfg, quant_weights, **kwargs):	57	def w8a16_quantize(name, layer, cfg, quant_weights, **kwargs):
50	- quant_map = OrderedDict([	58	+ quant_map = OrderedDict([(nn.Linear, WeightQuantLinear)])
51	- (nn.Linear, WeightQuantLinear)
52	- ])
53		59
54	# 如果模型指定了类的匹配规则，优先匹配模型指定的	60	# 如果模型指定了类的匹配规则，优先匹配模型指定的
55	user_dict = kwargs.get('map', None)	61	user_dict = kwargs.get('map', None)
	@@ -129,8 +135,13 @@ def smooth_quantize_w8a8(name, layer, cfg, quant_weights, **kwargs):
129	else:	135	else:
130	init_params[bias] = False	136	init_params[bias] = False
131		137
132	- if cfg.quant_algo in [QuantAlgorithm.W8A8_DYNAMIC, QuantAlgorithm.W8A8_MXFP8, QuantAlgorithm.W4A4_DYNAMIC,	138	+ if cfg.quant_algo in [
133	- QuantAlgorithm.W4A4_MXFP4_DUALSCALE, QuantAlgorithm.W4A4_MXFP4_DYNAMIC]:	139	+ QuantAlgorithm.W8A8_DYNAMIC,
		140	+ QuantAlgorithm.W8A8_MXFP8,
		141	+ QuantAlgorithm.W4A4_DYNAMIC,
		142	+ QuantAlgorithm.W4A4_MXFP4_DUALSCALE,
		143	+ QuantAlgorithm.W4A4_MXFP4_DYNAMIC,
		144	+ ]:
134	init_params['is_dynamic'] = True	145	init_params['is_dynamic'] = True
135		146
136	init_params['weights'] = quant_weights	147	init_params['weights'] = quant_weights
	@@ -158,7 +169,6 @@ def smooth_quantize(name, layer, cfg, quant_weights, **kwargs):
158	def add_fa_quant(layer, cfg, prefix, quant_weights):	169	def add_fa_quant(layer, cfg, prefix, quant_weights):
159	if cfg.quant_algo in [QuantAlgorithm.FP8_DYNAMIC]:	170	if cfg.quant_algo in [QuantAlgorithm.FP8_DYNAMIC]:
160	layer.fa_quant = FP8RotateQuantFA(prefix, quant_weights)	171	layer.fa_quant = FP8RotateQuantFA(prefix, quant_weights)
161	- return
162		172
163		173
164	def get_layer_quant_mode(name, layer, cfg):	174	def get_layer_quant_mode(name, layer, cfg):
	@@ -198,14 +208,15 @@ def modify_graph(model, modified_layers):
198	def get_cfg_and_weights(quant_des_path):	208	def get_cfg_and_weights(quant_des_path):
199	quant_des_path, filename, rank = replace_rank_suffix(quant_des_path)	209	quant_des_path, filename, rank = replace_rank_suffix(quant_des_path)
200	quant_algo_str = "quant_algo"	210	quant_algo_str = "quant_algo"
201	- with file_utils.safe_open(quant_des_path, "r", encoding="utf-8",	211	+ with file_utils.safe_open(
202	- permission_mode=file_utils.CONFIG_FILE_PERMISSION) as reader:	212	+ quant_des_path, "r", encoding="utf-8", permission_mode=file_utils.CONFIG_FILE_PERMISSION
		213	+ ) as reader:
203	data = reader.read()	214	data = reader.read()
204	quant_des_dict = json.loads(data, strict=False)	215	quant_des_dict = json.loads(data, strict=False)
205	- logger.info(f"Quant Description Filename:{filename}")	216	+ logger.debug("[MindIE-SD/quantization] Quant description loaded. filename=%s.", filename)
206		217
207	if not quant_des_dict:	218	if not quant_des_dict:
208	- raise ParametersInvalid(f"quant_des_dict is none!")	219	+ raise ParametersInvalid("quant_des_dict is none!")
209	exclude_layers = [k for k, v in quant_des_dict.items() if v == "FLOAT"]	220	exclude_layers = [k for k, v in quant_des_dict.items() if v == "FLOAT"]
210	valid_values = {item.value for item in QuantAlgorithm} # 预计算有效值集合	221	valid_values = {item.value for item in QuantAlgorithm} # 预计算有效值集合
211	quantized_layers = {	222	quantized_layers = {
	@@ -215,7 +226,7 @@ def get_cfg_and_weights(quant_des_path):
215	}	226	}
216	quant_algo = quant_des_dict.get("model_quant_type", None)	227	quant_algo = quant_des_dict.get("model_quant_type", None)
217	if quant_algo is None:	228	if quant_algo is None:
218	- raise ParametersInvalid(f"quant_algo must be the type of QuantAlgorithm.")	229	+ raise ParametersInvalid("quant_algo must be the type of QuantAlgorithm.")
219		230
220	quant_config = {"quant_algo": quant_algo}	231	quant_config = {"quant_algo": quant_algo}
221	quant_config.update({'exclude_layers': tuple(exclude_layers)})	232	quant_config.update({'exclude_layers': tuple(exclude_layers)})
	@@ -233,10 +244,11 @@ def get_cfg_and_weights(quant_des_path):
233	weight_name = f'quant_model_weight_{quant_algo.lower()}.safetensors'	244	weight_name = f'quant_model_weight_{quant_algo.lower()}.safetensors'
234	quant_weight_path = os.path.join(quant_weight_dir, weight_name)	245	quant_weight_path = os.path.join(quant_weight_dir, weight_name)
235	quant_weight_path = file_utils.standardize_path(quant_weight_path)	246	quant_weight_path = file_utils.standardize_path(quant_weight_path)
236	- file_utils.check_file_safety(quant_weight_path,	247	+ file_utils.check_file_safety(
237	- permission_mode=file_utils.MODELDATA_FILE_PERMISSION, max_file_size=MAX_WEIGHT_SIZE)	248	+ quant_weight_path, permission_mode=file_utils.MODELDATA_FILE_PERMISSION, max_file_size=MAX_WEIGHT_SIZE
		249	+ )
238	quant_weights = safetensors.safe_open(quant_weight_path, framework="pytorch")	250	quant_weights = safetensors.safe_open(quant_weight_path, framework="pytorch")
239	- logger.info(f"Quant Weight Path:{quant_weight_path}")	251	+ logger.debug("[MindIE-SD/quantization] Quant weight file loaded. path=%s.", quant_weight_path)
240		252
241	return cfg, quant_weights	253	return cfg, quant_weights
242		254
	@@ -256,8 +268,9 @@ def validate_quantize_params(func):
256		268
257	timestep_config = kwargs.get('timestep_config')	269	timestep_config = kwargs.get('timestep_config')
258	if timestep_config is not None and not isinstance(timestep_config, TimestepPolicyConfig):	270	if timestep_config is not None and not isinstance(timestep_config, TimestepPolicyConfig):
259	- raise ParametersInvalid(f"Timestep_config must be the type of TimestepPolicyConfig,"	271	+ raise ParametersInvalid(
260	- "but currently got {type(timestep_config)}.")	272	+ f"Timestep_config must be the type of TimestepPolicyConfig, but currently got {type(timestep_config)}."
		273	+ )
261		274
262	dtype = kwargs.get('dtype', torch.bfloat16)	275	dtype = kwargs.get('dtype', torch.bfloat16)
263	if not isinstance(dtype, torch.dtype) or dtype not in (torch.float16, torch.bfloat16):	276	if not isinstance(dtype, torch.dtype) or dtype not in (torch.float16, torch.bfloat16):
	@@ -265,9 +278,11 @@ def validate_quantize_params(func):
265		278
266	module_map = kwargs.get('map', None)	279	module_map = kwargs.get('map', None)
267	if module_map is not None:	280	if module_map is not None:
268	- if not isinstance(module_map, Dict) or \	281	+ if (
269	- not all(isinstance(v, nn.Module) for v in module_map.values()) or \	282	+ not isinstance(module_map, Dict)
270	- not all(isinstance(k, nn.Module) for k in module_map.keys()):	283	+ or not all(isinstance(v, nn.Module) for v in module_map.values())
		284	+ or not all(isinstance(k, nn.Module) for k in module_map.keys())
		285	+ ):
271	raise ParametersInvalid("The data type of map must be dictionary, and its KVType must be nn.Module.")	286	raise ParametersInvalid("The data type of map must be dictionary, and its KVType must be nn.Module.")
272		287
273	return func(model, quant_des_path, **kwargs)	288	return func(model, quant_des_path, **kwargs)
	@@ -301,7 +316,7 @@ def quantize(model, quant_des_path, **kwargs):
301	return model	316	return model
302		317
303	modified_layers = []	318	modified_layers = []
304	- rank = int(os.getenv("RANK", 0))	319	+ rank = int(os.getenv("RANK", "0"))
305		320
306	for name, layer in model.named_modules():	321	for name, layer in model.named_modules():
307	# 跳过回退层	322	# 跳过回退层
	@@ -324,21 +339,25 @@ def quantize(model, quant_des_path, **kwargs):
324	if layer_quant_mode.contains_activation_and_weight_quant():	339	if layer_quant_mode.contains_activation_and_weight_quant():
325	quant_layer, is_modified = smooth_quantize(name, layer, layer_quant_cfg, quant_weights, **kwargs)	340	quant_layer, is_modified = smooth_quantize(name, layer, layer_quant_cfg, quant_weights, **kwargs)
326	if is_modified:	341	if is_modified:
327	- logger.debug(f"W8A8 Quant layer name:%s, Quant class name:%s.", name, quant_layer.__class__.__name__)	342	+ logger.debug("W8A8 Quant layer name:%s, Quant class name:%s.", name, quant_layer.__class__.__name__)
328	modified_layers.append((name, quant_layer))	343	modified_layers.append((name, quant_layer))
329	elif layer_quant_mode.check_weight_only_mode():	344	elif layer_quant_mode.check_weight_only_mode():
330	quant_layer, is_modified = weight_quantize(name, layer, layer_quant_cfg, quant_weights, **kwargs)	345	quant_layer, is_modified = weight_quantize(name, layer, layer_quant_cfg, quant_weights, **kwargs)
331	if is_modified:	346	if is_modified:
332	- logger.debug(f"Weight Quant layer name:%s, Quant class name:%s.", name, quant_layer.__class__.__name__)	347	+ logger.debug("Weight Quant layer name:%s, Quant class name:%s.", name, quant_layer.__class__.__name__)
333	modified_layers.append((name, quant_layer))	348	modified_layers.append((name, quant_layer))
334	elif layer_quant_mode.contains_fa_quantization():	349	elif layer_quant_mode.contains_fa_quantization():
335	add_fa_quant(layer, layer_quant_cfg, name, quant_weights)	350	add_fa_quant(layer, layer_quant_cfg, name, quant_weights)
336	if rank == 0:	351	if rank == 0:
337	- logger.info(f"FA Quant layer name:%s, Quant class name:%s, Quant algo:%s.",	352	+ logger.debug(
338	- name, layer.__class__.__name__, layer_quant_cfg.quant_algo)	353	+ "FA Quant layer name:%s, Quant class name:%s, Quant algo:%s.",
		354	+ name,
		355	+ layer.__class__.__name__,
		356	+ layer_quant_cfg.quant_algo,
		357	+ )
339		358
340	# 执行改图	359	# 执行改图
341	modify_graph(model, modified_layers)	360	modify_graph(model, modified_layers)
342	torch.npu.empty_cache()	361	torch.npu.empty_cache()
343		362
344	- return model	363	+ return model

Mmindiesd/share_memory.py+28-34

	@@ -10,32 +10,25 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
13	-import time
14	import atexit	13	import atexit
15	-from dataclasses import dataclass
16	-import logging
17	from typing import Optional, Union	14	from typing import Optional, Union
18		15
19	import zmq	16	import zmq
20	import torch	17	import torch
21	-import torch_npu	18	+from .utils.exception import ParametersInvalid
22	-from .utils.exception import TorchError, ParametersInvalid	19	+from .utils.logs.logging import logger
23	-
24	-logger = logging.getLogger(__name__)
25		20
26		21
27	class ShareMemoryManager:	22	class ShareMemoryManager:
28	- def __init__(self,	23	+ def __init__(
29	- instance_world_size: int,	24	+ self, instance_world_size: int, instance_id: int, master_addr: str = "127.0.0.1", base_port: int = 5555
30	- instance_id: int,	25	+ ):
31	- master_addr: str = "127.0.0.1",
32	- base_port: int = 5555):
33	self.instance_world_size = instance_world_size	26	self.instance_world_size = instance_world_size
34	self.instance_id = instance_id	27	self.instance_id = instance_id
35	self.device_id = torch.npu.current_device()	28	self.device_id = torch.npu.current_device()
36	self.master_addr = master_addr	29	self.master_addr = master_addr
37	self.base_port = base_port	30	self.base_port = base_port
38	- self.is_master = (instance_id == 0)	31	+ self.is_master = instance_id == 0
39		32
40	self.pub_port = self.base_port + self.device_id + 100	33	self.pub_port = self.base_port + self.device_id + 100
41	self.rep_port = self.pub_port + 1	34	self.rep_port = self.pub_port + 1
	@@ -72,7 +65,7 @@ class ShareMemoryManager:
72	self.rep_socket.send(b"ACK")	65	self.rep_socket.send(b"ACK")
73	ready_count += 1	66	ready_count += 1
74	except zmq.Again as e:	67	except zmq.Again as e:
75	- raise TimeoutError(f"Master timeout waiting for child processes ready") from e	68	+ raise TimeoutError("Master timeout waiting for child processes ready") from e
76		69
77	logger.debug("Master broadcasting handle on tcp://%s:%s", self.master_addr, self.pub_port)	70	logger.debug("Master broadcasting handle on tcp://%s:%s", self.master_addr, self.pub_port)
78	self.pub_socket.send_pyobj(handle)	71	self.pub_socket.send_pyobj(handle)
	@@ -80,8 +73,7 @@ class ShareMemoryManager:
80	return handle	73	return handle
81	else:	74	else:
82	logger.debug(	75	logger.debug(
83	- "Device %s subscribing to handle on tcp://%s:%s",	76	+ "Device %s subscribing to handle on tcp://%s:%s", self.device_id, self.master_addr, self.pub_port
84	- self.device_id, self.master_addr, self.pub_port
85	)	77	)
86	self.req_socket.send(b"READY")	78	self.req_socket.send(b"READY")
87	self.req_socket.recv()	79	self.req_socket.recv()
	@@ -98,10 +90,9 @@ ZMQ_CONTEXT = zmq.Context.instance()
98	manager: Optional[ShareMemoryManager] = None	90	manager: Optional[ShareMemoryManager] = None
99		91
100		92
101	-def init_share_memory(instance_world_size: int,	93	+def init_share_memory(
102	- instance_id: int,	94	+ instance_world_size: int, instance_id: int, master_addr: str = "127.0.0.1", base_port: int = 5555
103	- master_addr: str = "127.0.0.1",	95	+) -> ShareMemoryManager:
104	- base_port: int = 5555) -> ShareMemoryManager:
105	"""	96	"""
106	设置共享内存管理器实例	97	设置共享内存管理器实例
107	Args:	98	Args:
	@@ -118,23 +109,21 @@ def init_share_memory(instance_world_size: int,
118		109
119		110
120	def get_share_memory_manager() -> ShareMemoryManager:	111	def get_share_memory_manager() -> ShareMemoryManager:
121	- global manager
122	if not manager:	112	if not manager:
123	- raise ParametersInvalid("ShareMemoryManager has not been initialized."	113	+ raise ParametersInvalid("ShareMemoryManager has not been initialized.Please call init_share_memory first.")
124	- "Please call init_share_memory first.")
125	return manager	114	return manager
126		115
127		116
128	-def _check_device_and_dtype(module: torch.nn.Module,	117	+def _check_device_and_dtype(
129	- target_device: Optional[torch.device],	118	+ module: torch.nn.Module, target_device: Optional[torch.device], target_dtype: Optional[torch.dtype]
130	- target_dtype: Optional[torch.dtype]):	119	+):
131	cur_device = next(module.parameters()).device	120	cur_device = next(module.parameters()).device
132	cpu_device = torch.device("cpu")	121	cpu_device = torch.device("cpu")
133	meta_device = torch.device("meta")	122	meta_device = torch.device("meta")
134		123
135	if cur_device == cpu_device and (target_device is None or target_device == cpu_device):	124	if cur_device == cpu_device and (target_device is None or target_device == cpu_device):
136	return True, torch.nn.Module.to(module, target_device, target_dtype), None, None	125	return True, torch.nn.Module.to(module, target_device, target_dtype), None, None
137	- if cur_device == meta_device or target_device == meta_device:	126	+ if meta_device in (cur_device, target_device):
138	return True, torch.nn.Module.to(module, target_device, target_dtype), None, None	127	return True, torch.nn.Module.to(module, target_device, target_dtype), None, None
139		128
140	device_id = torch.npu.current_device()	129	device_id = torch.npu.current_device()
	@@ -150,9 +139,9 @@ def _check_device_and_dtype(module: torch.nn.Module,
150	return False, None, npu_device, device_id	139	return False, None, npu_device, device_id
151		140
152		141
153	-def share_memory(module: torch.nn.Module,	142	+def share_memory(
154	- device: Optional[Union[str, torch.device]] = None,	143	+ module: torch.nn.Module, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None
155	- dtype: Optional[torch.dtype] = None) -> torch.nn.Module:	144	+) -> torch.nn.Module:
156	"""	145	"""
157	Args:	146	Args:
158	module (torch.nn.Module): 待迁移的模型实例（必传，强类型校验）	147	module (torch.nn.Module): 待迁移的模型实例（必传，强类型校验）
	@@ -173,8 +162,13 @@ def share_memory(module: torch.nn.Module,
173	target_dtype = dtype	162	target_dtype = dtype
174		163
175	cur_device = next(module.parameters()).device	164	cur_device = next(module.parameters()).device
176	- logger.debug("%s from device '%s' to device '%s' and dtype '%s'",	165	+ logger.debug(
177	- type(module).__name__, cur_device, target_device, target_dtype)	166	+ "%s from device '%s' to device '%s' and dtype '%s'",
		167	+ type(module).__name__,
		168	+ cur_device,
		169	+ target_device,
		170	+ target_dtype,
		171	+ )
178		172
179	should_fallback, fallback_result, _, _ = _check_device_and_dtype(module, target_device, target_dtype)	173	should_fallback, fallback_result, _, _ = _check_device_and_dtype(module, target_device, target_dtype)
180	if should_fallback:	174	if should_fallback:
	@@ -192,7 +186,7 @@ def share_memory(module: torch.nn.Module,
192	else:	186	else:
193	new_tensor = torch.empty_like(param, device=target_device)	187	new_tensor = torch.empty_like(param, device=target_device)
194	recv_storage = sm_manager.broadcast_handle(None)	188	recv_storage = sm_manager.broadcast_handle(None)
195	- rebuild_storage = torch.UntypedStorage._new_shared_npu(*recv_storage)	189	+ rebuild_storage = torch.UntypedStorage._new_shared_npu(*recv_storage) # pylint: disable=no-member
196	new_tensor.set_(rebuild_storage)	190	new_tensor.set_(rebuild_storage)
197	param.data = new_tensor.view(param.shape)	191	param.data = new_tensor.view(param.shape)
198		192
	@@ -200,4 +194,4 @@ def share_memory(module: torch.nn.Module,
200	return module	194	return module
201		195
202		196
203	-atexit.register(ShareMemoryManager.cleanup)	197	+atexit.register(ShareMemoryManager.cleanup)

Mmindiesd/utils/exception.py+47-6

	@@ -13,36 +13,77 @@
13	from .logs.logging import logger	13	from .logs.logging import logger
14		14
15		15
		16	+def _log_exception(component, title, message, cause, troubleshooting):
		17	+ logger.error(
		18	+ "%s %s. issue=%s. possible_cause=%s. Troubleshooting: %s",
		19	+ component,
		20	+ title,
		21	+ message,
		22	+ cause,
		23	+ troubleshooting,
		24	+ )
		25	+
		26	+
16	class ParametersInvalid(Exception):	27	class ParametersInvalid(Exception):
17	def __init__(self, message):	28	def __init__(self, message):
18	self.message = f"[MIE06E000001] Parameters invalid. {message}"	29	self.message = f"[MIE06E000001] Parameters invalid. {message}"
19	- logger.error(self.message)	30	+ _log_exception(
		31	+ "[MindIE-SD/utils]",
		32	+ "Parameter validation failed",
		33	+ self.message,
		34	+ "an input parameter does not match the required type, range, shape, or supported value list",
		35	+ "compare the actual parameter in the message with the expected value and fix the caller input",
		36	+ )
20	super().__init__(self.message)	37	super().__init__(self.message)
21		38
22		39
23	class ConfigError(Exception):	40	class ConfigError(Exception):
24	def __init__(self, message):	41	def __init__(self, message):
25	self.message = f"[MIE06E000002] Config parameter err. {message}"	42	self.message = f"[MIE06E000002] Config parameter err. {message}"
26	- logger.error(self.message)	43	+ _log_exception(
		44	+ "[MindIE-SD/utils]",
		45	+ "Configuration validation failed",
		46	+ self.message,
		47	+ "a configuration item is missing, invalid, or inconsistent with the runtime requirement",
		48	+ "check the configuration file, environment variables, and the expected value shown in the message",
		49	+ )
27	super().__init__(self.message)	50	super().__init__(self.message)
28		51
29		52
30	class TorchError(Exception):	53	class TorchError(Exception):
31	def __init__(self, message):	54	def __init__(self, message):
32	self.message = f"[MIE06E000003] Torch exec err. {message}"	55	self.message = f"[MIE06E000003] Torch exec err. {message}"
33	- logger.error(self.message)	56	+ _log_exception(
		57	+ "[MindIE-SD/utils]",
		58	+ "Torch execution failed",
		59	+ self.message,
		60	+ "the torch or torch_npu operator failed during execution",
		61	+ "check the operator input shape, dtype, device placement, and the CANN/torch_npu error stack",
		62	+ )
34	super().__init__(self.message)	63	super().__init__(self.message)
35		64
36		65
37	class ModelInitError(Exception):	66	class ModelInitError(Exception):
38	def __init__(self, message):	67	def __init__(self, message):
39	self.message = f"[MIE06E000004] Model init err. {message}"	68	self.message = f"[MIE06E000004] Model init err. {message}"
40	- logger.error(self.message)	69	+ _log_exception(
		70	+ "[MindIE-SD/utils]",
		71	+ "Model initialization failed",
		72	+ self.message,
		73	+ "model weights, configuration, or runtime resources are not ready",
		74	+ "check model path, weight files, configuration values, NPU memory, and initialization stack",
		75	+ )
41	super().__init__(self.message)	76	super().__init__(self.message)
42		77
43		78
44	class ModelExecError(Exception):	79	class ModelExecError(Exception):
45	def __init__(self, message):	80	def __init__(self, message):
46	self.message = f"[MIE06E000005] Model exec err. {message}"	81	self.message = f"[MIE06E000005] Model exec err. {message}"
47	- logger.error(self.message)	82	+ _log_exception(
48	- super().__init__(self.message)	83	+ "[MindIE-SD/utils]",
		84	+ "Model execution failed",
		85	+ self.message,
		86	+ "the model forward, scheduler, or custom operator path failed during execution",
		87	+ "check the request parameters, tensor shape and dtype, scheduler state, and CANN operator error stack",
		88	+ )
		89	+ super().__init__(self.message)

Mmindiesd/utils/get_platform.py+6-1

	@@ -44,7 +44,12 @@ def get_npu_device() -> NPUDevice:
44	else:	44	else:
45	PLATFORM = NPUDevice.UNDEFINED	45	PLATFORM = NPUDevice.UNDEFINED
46	except RuntimeError as exc:	46	except RuntimeError as exc:
47	- logger.warning("Failed to get NPU SoC version: %s", exc)	47	+ logger.warning(
		48	+ "[MindIE-SD/utils] NPU SoC version query failed. issue=torch_npu failed to return SoC version, "
		49	+ "actual_error=%s. possible_cause=NPU driver, CANN, or device environment is unavailable. "
		50	+ "Troubleshooting: check npu-smi info, CANN environment variables, and torch_npu installation.",
		51	+ exc,
		52	+ )
48	PLATFORM = NPUDevice.UNDEFINED	53	PLATFORM = NPUDevice.UNDEFINED
49	return PLATFORM	54	return PLATFORM
50		55

Mmindiesd/utils/logs/logging.py+97-35

	@@ -29,7 +29,9 @@ BACKUP_OWNER_SHIP = 0o440
29	FILE_OWNER_SHIP = 0o640	29	FILE_OWNER_SHIP = 0o640
30	PATH_OWNER_SHIP = 0o750	30	PATH_OWNER_SHIP = 0o750
31	MB = 1024 * 1024	31	MB = 1024 * 1024
32	-MAX_LOG_STRING_LEN = 256	32	+MAX_LOG_STRING_LEN = 2048
		33	+LOG_COMPONENT = "[MindIE-SD/log]"
		34	+
33		35
34	def get_pid():	36	def get_pid():
35	return os.getpid()	37	return os.getpid()
	@@ -44,16 +46,32 @@ def check_owner_permission(file_path, max_mode) -> bool:
44	file_owner = os.stat(file_path).st_uid	46	file_owner = os.stat(file_path).st_uid
45	cur_owner = get_uid()	47	cur_owner = get_uid()
46	if file_owner != cur_owner:	48	if file_owner != cur_owner:
47	- logging.warning("File doesn't belong to current user.")	49	+ logging.warning(
		50	+ "%s File owner validation failed. issue=file owner mismatch, path=%s, expected_uid=%s, actual_uid=%s. "
		51	+ "Possible cause: the log file or directory was created by another user. "
		52	+ "Troubleshooting: check file ownership and change it to the current service user before enabling log output.",
		53	+ LOG_COMPONENT,
		54	+ file_path,
		55	+ cur_owner,
		56	+ file_owner,
		57	+ )
48	return False	58	return False
49		59
50	# check permission	60	# check permission
51	- file_mode = os.stat(file_path).st_mode & 0o777 # use 777 as mask to get 3-digit octal number	61	+ file_mode = os.stat(file_path).st_mode & 0o777 # use 777 as mask to get 3-digit octal number
52	- file_mode_bin = bin(file_mode)[2:].zfill(9) # transeform into 9-bit binary number	62	+ file_mode_bin = bin(file_mode)[2:].zfill(9) # transeform into 9-bit binary number
53	- max_mode_bin = bin(max_mode)[2:].zfill(9) # transeform into 9-bit binary number	63	+ max_mode_bin = bin(max_mode)[2:].zfill(9) # transeform into 9-bit binary number
54	- for i in range(9): # 9 means 9-bit binary number, checking every bit	64	+ for i in range(9): # 9 means 9-bit binary number, checking every bit
55	- if file_mode_bin[i] > max_mode_bin[i]: # 2 means the head of binary number '0b'	65	+ if file_mode_bin[i] > max_mode_bin[i]: # 2 means the head of binary number '0b'
56	- logging.warning("The permission of file is higher than %s.", oct(max_mode))	66	+ logging.warning(
		67	+ "%s File permission validation failed. issue=permission is higher than allowed, path=%s, "
		68	+ "expected_mode<=%s, actual_mode=%s. Possible cause: log file permission is too open. "
		69	+ "Troubleshooting: use chmod to reduce the file permission and retry.",
		70	+ LOG_COMPONENT,
		71	+ file_path,
		72	+ oct(max_mode),
		73	+ oct(file_mode),
		74	+ )
57	return False	75	return False
58		76
59	return True	77	return True
	@@ -67,7 +85,13 @@ def check_path(file_path, checking_conf=False):
67	# check if the path is symbolic link	85	# check if the path is symbolic link
68	trimmed_path = file_path.rstrip("/")	86	trimmed_path = file_path.rstrip("/")
69	if os.path.islink(trimmed_path):	87	if os.path.islink(trimmed_path):
70	- logging.warning("File path is a soft link.")	88	+ logging.warning(
		89	+ "%s Log path validation failed. issue=symbolic link is not allowed, path=%s. "
		90	+ "Possible cause: MINDIE_LOG_PATH points to a symlink. "
		91	+ "Troubleshooting: set MINDIE_LOG_PATH to a real directory path.",
		92	+ LOG_COMPONENT,
		93	+ trimmed_path,
		94	+ )
71	return False	95	return False
72		96
73	if checking_conf:	97	if checking_conf:
	@@ -81,6 +105,7 @@ class MindIELogFileHandler(BaseRotatingHandler):
81	Adapt from logging's TimedRotatingHandler and RotationFileHandler to combine both of their features.	105	Adapt from logging's TimedRotatingHandler and RotationFileHandler to combine both of their features.
82	Beside, add more detail about controlling log files' owner ships and rotation.	106	Beside, add more detail about controlling log files' owner ships and rotation.
83	"""	107	"""
		108	+
84	def __init__(self, real_log_path, max_file_num, max_file_size, rotate_cycle_num, rotate_cycle):	109	def __init__(self, real_log_path, max_file_num, max_file_size, rotate_cycle_num, rotate_cycle):
85	encoding = io.text_encoding(None)	110	encoding = io.text_encoding(None)
86	now_time_str = time.strftime("_%Y%m%d%H%M%S", time.localtime())	111	now_time_str = time.strftime("_%Y%m%d%H%M%S", time.localtime())
	@@ -94,7 +119,7 @@ class MindIELogFileHandler(BaseRotatingHandler):
94	self._max_file_num = max_file_num	119	self._max_file_num = max_file_num
95	self._rotate_cycle_num = rotate_cycle_num	120	self._rotate_cycle_num = rotate_cycle_num
96	self._rotate_cycle = rotate_cycle	121	self._rotate_cycle = rotate_cycle
97	- self._next_rollover = self._get_rollover_timepoint() # use time() since it is easier to compute	122	+ self._next_rollover = self._get_rollover_timepoint() # use time() since it is easier to compute
98		123
99	# Be aware that, at this point, _cur_log_file is not created yet since the use of delay mode.	124	# Be aware that, at this point, _cur_log_file is not created yet since the use of delay mode.
100	log_path_files = os.listdir(real_log_path)	125	log_path_files = os.listdir(real_log_path)
	@@ -104,7 +129,7 @@ class MindIELogFileHandler(BaseRotatingHandler):
104	time_str = self._get_time_str(it)	129	time_str = self._get_time_str(it)
105	if time_str and real_file_path.startswith(real_log_path) and os.path.exists(real_file_path):	130	if time_str and real_file_path.startswith(real_log_path) and os.path.exists(real_file_path):
106	log_files.append((real_file_path, time_str))	131	log_files.append((real_file_path, time_str))
107	- self._history_files = sorted(log_files, key=lambda x : x[1])	132	+ self._history_files = sorted(log_files, key=lambda x: x[1])
108	# Deal with history file number by deleting oldest one	133	# Deal with history file number by deleting oldest one
109	self._delete_file_by_number()	134	self._delete_file_by_number()
110	self._delete_file_by_time()	135	self._delete_file_by_time()
	@@ -162,7 +187,7 @@ class MindIELogFileHandler(BaseRotatingHandler):
162	cur_time_str = self._get_time_str(cur_log_name)	187	cur_time_str = self._get_time_str(cur_log_name)
163	self._history_files.append((self._cur_log_file, cur_time_str))	188	self._history_files.append((self._cur_log_file, cur_time_str))
164		189
165	- # delete oldest file by file number constaint	190	+ # delete oldest file by file number constraint
166	self._delete_file_by_number()	191	self._delete_file_by_number()
167		192
168	# delete oldest file by time	193	# delete oldest file by time
	@@ -203,7 +228,7 @@ class MindIELogFileHandler(BaseRotatingHandler):
203	# rollover timepoint is everyday's midnight, no log file will have log that cross two days.	228	# rollover timepoint is everyday's midnight, no log file will have log that cross two days.
204	now = time.time()	229	now = time.time()
205	tomorrow = datetime.fromtimestamp(now) + timedelta(days=1)	230	tomorrow = datetime.fromtimestamp(now) + timedelta(days=1)
206	- tomorrow_midnight = datetime(tomorrow.year, tomorrow.month, tomorrow.day, 0, 0, 0) # midnight is 00:00	231	+ tomorrow_midnight = datetime(tomorrow.year, tomorrow.month, tomorrow.day, 0, 0, 0) # midnight is 00:00
207	rollover_timepoint = int(time.mktime(tomorrow_midnight.timetuple()))	232	rollover_timepoint = int(time.mktime(tomorrow_midnight.timetuple()))
208	return rollover_timepoint	233	return rollover_timepoint
209		234
	@@ -267,6 +292,14 @@ class MindIELogFileHandler(BaseRotatingHandler):
267	case "yearly":	292	case "yearly":
268	return self._check_year(log_date, cur_date)	293	return self._check_year(log_date, cur_date)
269	case _:	294	case _:
		295	+ logging.warning(
		296	+ "%s Log rotation validation failed. issue=unknown rotate cycle, expected=%s, actual=%s. "
		297	+ "Possible cause: MINDIE_LOG_ROTATE contains an unsupported cycle. "
		298	+ "Troubleshooting: set rotate cycle to daily, weekly, monthly, or yearly.",
		299	+ LOG_COMPONENT,
		300	+ "daily\|weekly\|monthly\|yearly",
		301	+ self._rotate_cycle,
		302	+ )
270	raise ValueError(f"Unknown rotate cycle: {self._rotate_cycle}")	303	raise ValueError(f"Unknown rotate cycle: {self._rotate_cycle}")
271		304
272	def _remove_oldest_log(self):	305	def _remove_oldest_log(self):
	@@ -276,8 +309,12 @@ class MindIELogFileHandler(BaseRotatingHandler):
276		309
277	def _open(self):	310	def _open(self):
278	create_flags = os.O_RDWR \| os.O_CREAT	311	create_flags = os.O_RDWR \| os.O_CREAT
279	- open_func = os.fdopen(os.open(self._cur_log_file, create_flags, FILE_OWNER_SHIP),	312	+ open_func = os.fdopen(
280	- self.mode, encoding=self.encoding, errors=self.errors)	313	+ os.open(self._cur_log_file, create_flags, FILE_OWNER_SHIP),
		314	+ self.mode,
		315	+ encoding=self.encoding,
		316	+ errors=self.errors,
		317	+ )
281	return open_func	318	return open_func
282		319
283		320
	@@ -294,6 +331,14 @@ def str_to_loglevel(level_str):
294	case "CRITICAL":	331	case "CRITICAL":
295	return logging.CRITICAL	332	return logging.CRITICAL
296	case _:	333	case _:
		334	+ logging.warning(
		335	+ "%s Log level validation failed. issue=unknown log level, expected=%s, actual=%s. "
		336	+ "Possible cause: MINDIE_LOG_LEVEL contains an unsupported value. "
		337	+ "Troubleshooting: set log level to critical, error, warn, info, debug, or null.",
		338	+ LOG_COMPONENT,
		339	+ "critical\|error\|warn\|info\|debug\|null",
		340	+ level_str,
		341	+ )
297	raise ValueError(f"Unknown log level: {level_str}")	342	raise ValueError(f"Unknown log level: {level_str}")
298		343
299		344
	@@ -301,7 +346,7 @@ class LoggerFormatter(logging.Formatter):
301	def formatTime(self, record, datefmt=None):	346	def formatTime(self, record, datefmt=None):
302	dt = datetime.fromtimestamp(record.created).astimezone()	347	dt = datetime.fromtimestamp(record.created).astimezone()
303	formatted_time = dt.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3] + " " + dt.strftime('%z %Z')	348	formatted_time = dt.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3] + " " + dt.strftime('%z %Z')
304	- return f"{formatted_time[:23]}{formatted_time[24:27]}:{formatted_time[27:]}" # change 0800 to 08:00	349	+ return f"{formatted_time[:23]}{formatted_time[24:27]}:{formatted_time[27:]}" # change 0800 to 08:00
305		350
306	def format(self, record):	351	def format(self, record):
307	original = logging.Formatter.format(self, record)	352	original = logging.Formatter.format(self, record)
	@@ -312,15 +357,19 @@ class LoggerFormatter(logging.Formatter):
312	if len(message) > MAX_LOG_STRING_LEN:	357	if len(message) > MAX_LOG_STRING_LEN:
313	message = message[:MAX_LOG_STRING_LEN]	358	message = message[:MAX_LOG_STRING_LEN]
314	invalid_chars = {	359	invalid_chars = {
315	- '\f', '\r', '\b', '\t', '\v', '\n',	360	+ '\f',
316	- '\u000A', '\u000D', '\u000C', '\u000B',	361	+ '\r',
317	- '\u0008', '\u007F', '\u0009'	362	+ '\b',
		363	+ '\t',
		364	+ '\v',
		365	+ '\n',
		366	+ '\u007f',
318	}	367	}
319	for char in invalid_chars:	368	for char in invalid_chars:
320	message = message.replace(char, "")	369	message = message.replace(char, "")
321	message = re.sub(R"[ ]+", " ", message)	370	message = re.sub(R"[ ]+", " ", message)
322	else:	371	else:
323	- message = f'log is None!'	372	+ message = 'log is None!'
324	return message	373	return message
325		374
326		375
	@@ -339,17 +388,23 @@ def create_directory_with_permissions(real_log_path, permission) -> bool:
339	try:	388	try:
340	os.makedirs(current_path, mode=permission, exist_ok=True)	389	os.makedirs(current_path, mode=permission, exist_ok=True)
341	except Exception:	390	except Exception:
342	- logging.warning("Failed to create log directory.")	391	+ logging.warning(
		392	+ "%s Log directory creation failed. issue=failed to create directory, path=%s, expected_mode=%s. "
		393	+ "Possible cause: parent path permission is insufficient or the path is invalid. "
		394	+ "Troubleshooting: check parent directory permission, disk status, and MINDIE_LOG_PATH.",
		395	+ LOG_COMPONENT,
		396	+ current_path,
		397	+ oct(permission),
		398	+ )
343	return False	399	return False
344	return True	400	return True
345		401
346		402
347	def init_logger():	403	def init_logger():
348	- global logger
349	log_level = str_to_loglevel(ENV.component_log_level)	404	log_level = str_to_loglevel(ENV.component_log_level)
350	logger.setLevel(log_level)	405	logger.setLevel(log_level)
351	if ENV.disable_log:	406	if ENV.disable_log:
352	- logger.disabled=True	407	+ logger.disabled = True
353	return logger	408	return logger
354		409
355	if ENV.component_log_verbose in POSITIVE_BOOLEAN:	410	if ENV.component_log_verbose in POSITIVE_BOOLEAN:
	@@ -357,9 +412,7 @@ def init_logger():
357	'%(asctime)s [%(process)d] [%(thread)d] [MindIE-SD] [%(levelname)s] %(filename)s:%(lineno)d: %(message)s'	412	'%(asctime)s [%(process)d] [%(thread)d] [MindIE-SD] [%(levelname)s] %(filename)s:%(lineno)d: %(message)s'
358	)	413	)
359	else:	414	else:
360	- formatter = LoggerFormatter(	415	+ formatter = LoggerFormatter('%(asctime)s [MindIE-SD] [%(levelname)s] %(message)s')
361	- '%(asctime)s [%(levelname)s]: %(message)s'
362	- )
363		416
364	if ENV.component_log_stdout in POSITIVE_BOOLEAN:	417	if ENV.component_log_stdout in POSITIVE_BOOLEAN:
365	print_handler = logging.StreamHandler(stream=sys.stdout)	418	print_handler = logging.StreamHandler(stream=sys.stdout)
	@@ -370,11 +423,11 @@ def init_logger():
370	if ENV.component_log_to_file in POSITIVE_BOOLEAN:	423	if ENV.component_log_to_file in POSITIVE_BOOLEAN:
371	# check and standarlize the path	424	# check and standarlize the path
372	log_base_path = ENV.mindie_log_path	425	log_base_path = ENV.mindie_log_path
373	- log_base_path = os.path.expanduser(log_base_path) # expand '~'	426	+ log_base_path = os.path.expanduser(log_base_path) # expand '~'
374	# relative path	427	# relative path
375	if not log_base_path.startswith("/"):	428	if not log_base_path.startswith("/"):
376	log_base_path = os.path.join(MINDIE_DEFAULTS_LOG_PATH, log_base_path)	429	log_base_path = os.path.join(MINDIE_DEFAULTS_LOG_PATH, log_base_path)
377	- log_base_path = os.path.expanduser(log_base_path) # expand '~'	430	+ log_base_path = os.path.expanduser(log_base_path) # expand '~'
378		431
379	real_log_path = ""	432	real_log_path = ""
380	if check_path(log_base_path):	433	if check_path(log_base_path):
	@@ -384,18 +437,27 @@ def init_logger():
384	debug_log_path = os.path.join(real_log_path, "debug")	437	debug_log_path = os.path.join(real_log_path, "debug")
385	need_add_handler = create_directory_with_permissions(debug_log_path, PATH_OWNER_SHIP)	438	need_add_handler = create_directory_with_permissions(debug_log_path, PATH_OWNER_SHIP)
386	if need_add_handler:	439	if need_add_handler:
387	- file_handler = MindIELogFileHandler(debug_log_path,	440	+ file_handler = MindIELogFileHandler(
388	- max_file_num=ENV.rotate_max_file_num,	441	+ debug_log_path,
389	- max_file_size=ENV.rotate_max_file_size*MB,	442	+ max_file_num=ENV.rotate_max_file_num,
390	- rotate_cycle_num=ENV.rotate_cycle_num,	443	+ max_file_size=ENV.rotate_max_file_size * MB,
391	- rotate_cycle=ENV.rotate_cycle)	444	+ rotate_cycle_num=ENV.rotate_cycle_num,
		445	+ rotate_cycle=ENV.rotate_cycle,
		446	+ )
392	file_handler.setFormatter(formatter)	447	file_handler.setFormatter(formatter)
393	file_handler.setLevel(log_level)	448	file_handler.setLevel(log_level)
394	logger.addHandler(file_handler)	449	logger.addHandler(file_handler)
395	else:	450	else:
396	- logging.warning("The log file path is invalid or does not exist. The log cannot be saved!")	451	+ logging.warning(
		452	+ "%s Log file path validation failed. issue=invalid log path, path=%s. "
		453	+ "Possible cause: path is empty, too long, a symlink, or inaccessible. "
		454	+ "Troubleshooting: check MINDIE_LOG_PATH and directory permission; logs will only be printed to stdout.",
		455	+ LOG_COMPONENT,
		456	+ log_base_path,
		457	+ )
397	logger.propagate = False	458	logger.propagate = False
		459	+ return logger
398		460
399		461
400	logger = logging.getLogger('mindie-sd')	462	logger = logging.getLogger('mindie-sd')
401	-init_logger()	463	+init_logger()