fix: 修复 MegaMoe Level1 证据采集并收敛接口测试-hyper-parallel-AtomGit

已开启

fix: 修复 MegaMoe Level1 证据采集并收敛接口测试 #1368

Feiran Zhang创建于 9月9日从feiranzhang/hyper-parallel： megamoe-level1-test 合入到 mindspore/hyper-parallel： master

fix: 修复 MegaMoe Level1 证据采集并收敛接口测试 #1368

已开启

Feiran Zhang创建于 9月9日

从feiranzhang/hyper-parallel megamoe-level1-test合入到mindspore/hyper-parallel master

共 4 个文件变更+118-230

Mtests/torch/multicore/_mega_moe_utils.py+12-4文件内容审核中，请稍后刷新重试

Dtests/torch/multicore/_test_mega_moe_defaults.py+0-218

	@@ -1,218 +0,0 @@
1	-# Copyright 2026 Huawei Technologies Co., Ltd
2	-#
3	-# Licensed under the Apache License, Version 2.0 (the "License");
4	-# you may not use this file except in compliance with the License.
5	-# You may obtain a copy of the License at
6	-#
7	-# http://www.apache.org/licenses/LICENSE-2.0
8	-#
9	-# Unless required by applicable law or agreed to in writing, software
10	-# distributed under the License is distributed on an "AS IS" BASIS,
11	-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	-# See the License for the specific language governing permissions and
13	-# limitations under the License.
14	-# ============================================================================
15	-"""Default-capacity and Router-inclusive representative MegaMoe acceptance."""
16	-
17	-from __future__ import annotations
18	-
19	-import gc
20	-import os
21	-import statistics
22	-import time
23	-from dataclasses import asdict
24	-
25	-import torch
26	-import torch.distributed as dist
27	-
28	-from hyper_parallel.core.multicore import MegaMoeExperts
29	-from tests.torch.multicore import _test_mega_moe as baseline
30	-from tests.torch.multicore._mega_moe_utils import (
31	- assert_stable_memory,
32	- collect_memory,
33	- memory_sample,
34	- start_shmem_lifetime,
35	- write_evidence,
36	-)
37	-
38	-
39	-def _router_weight(shape: baseline.MoeShape) -> torch.nn.Parameter:
40	- """Use identical trainable Router weights in every A/B repetition."""
41	- generator = torch.Generator().manual_seed(706)
42	- weight = torch.randn(shape.num_experts, shape.hidden_size, generator=generator)
43	- return torch.nn.Parameter((weight / shape.hidden_size**0.5).to(baseline.DEVICE))
44	-
45	-
46	-def _route(hidden: torch.Tensor, weight: torch.Tensor, top_k: int) -> tuple:
47	- """Compute Router projection, softmax, Top-K and selected-weight normalization."""
48	- probabilities = (hidden.float() @ weight.T).softmax(dim=-1)
49	- values, indices = probabilities.topk(top_k, dim=-1)
50	- return indices.to(torch.int32), values / values.sum(dim=-1, keepdim=True)
51	-
52	-
53	-def _run_router_layer(layer: torch.nn.Module, shape: baseline.MoeShape) -> list[torch.Tensor]:
54	- """Validate a connected Router/expert graph, including Router dW and SGD."""
55	- hidden, upstream = baseline.make_data(shape)
56	- hidden.requires_grad_(True)
57	- router = _router_weight(shape)
58	- ids, weights = _route(hidden, router, shape.top_k)
59	- weights.retain_grad()
60	- output = baseline.forward_layer(layer, hidden, ids, weights)
61	- output.backward(upstream)
62	- optimizer = torch.optim.SGD([*layer.parameters(), router], lr=1e-3)
63	- optimizer.step()
64	- torch.npu.synchronize()
65	- compared = [output, hidden.grad, weights.grad, router.grad, router]
66	- compared.extend(baseline.expert_weight_gradients(layer))
67	- compared.extend(baseline.expert_weights(layer))
68	- for index, tensor in enumerate(compared):
69	- assert tensor is not None, f"rank={baseline.RANK}: missing Router graph tensor {index}."
70	- baseline.assert_finite(f"Router graph tensor {index}", tensor)
71	- return [tensor.detach().clone() for tensor in compared]
72	-
73	-
74	-def _validate_pair(shape: baseline.MoeShape, factor: float \| None, scope: str) -> None:
75	- """Compare real output, every gradient and one update at the measured shape."""
76	- start_shmem_lifetime()
77	- os.environ.pop("SYMMETRIC_MEMORY_HEAP_SIZE", None)
78	- mega, common = baseline.new_layers(shape, expert_capacity_factor=factor)
79	- try:
80	- if scope == "router_experts":
81	- reference = _run_router_layer(common, shape)
82	- actual = _run_router_layer(mega, shape)
83	- for index, (expected, observed) in enumerate(zip(reference, actual)):
84	- baseline.assert_close(f"Router graph tensor {index}", observed, expected)
85	- else:
86	- hidden, upstream = baseline.make_data(shape)
87	- ids, weights, _ = baseline.make_balanced_route(shape)
88	- reference = baseline.run_layer(common, hidden, ids, weights, None, upstream)
89	- actual = baseline.run_layer(mega, hidden, ids, weights, None, upstream)
90	- baseline.assert_results_close(reference, actual)
91	- finally:
92	- mega.close()
93	-
94	-
95	-def _measurement_layer(shape: baseline.MoeShape, backend: str, factor: float \| None) -> torch.nn.Module:
96	- """Retain only one parameter-aligned backend while measuring memory."""
97	- os.environ.pop("SYMMETRIC_MEMORY_HEAP_SIZE", None)
98	- if backend == "common":
99	- return baseline.new_common_moe(shape)
100	- start_shmem_lifetime()
101	- mega, common = baseline.new_layers(shape, expert_capacity_factor=factor)
102	- del common
103	- gc.collect()
104	- torch.npu.empty_cache()
105	- return mega
106	-
107	-
108	-def _iteration(
109	- layer: torch.nn.Module,
110	- router: torch.nn.Parameter,
111	- shape: baseline.MoeShape,
112	- inputs: tuple,
113	- scope: str,
114	-) -> tuple[float, float, int, int]:
115	- """Time host-to-completion F and F+B, including histogram construction."""
116	- source, upstream, ids, weights = inputs
117	- layer.zero_grad(set_to_none=True)
118	- router.grad = None
119	- hidden = source.detach().requires_grad_(True)
120	- selected_weights = weights.detach().requires_grad_(True)
121	- dist.barrier()
122	- torch.npu.synchronize()
123	- torch.npu.reset_peak_memory_stats()
124	- start = time.perf_counter()
125	- if scope == "router_experts":
126	- ids, selected_weights = _route(hidden, router, shape.top_k)
127	- output = baseline.forward_layer(layer, hidden, ids, selected_weights)
128	- torch.npu.synchronize()
129	- forward_ms = (time.perf_counter() - start) * 1000.0
130	- output.backward(upstream)
131	- torch.npu.synchronize()
132	- forward_backward_ms = (time.perf_counter() - start) * 1000.0
133	- peak_allocated = torch.npu.max_memory_allocated()
134	- peak_reserved = torch.npu.max_memory_reserved()
135	- gradients = [hidden.grad, *baseline.expert_weight_gradients(layer)]
136	- gradients.append(router.grad if scope == "router_experts" else selected_weights.grad)
137	- for index, gradient in enumerate(gradients):
138	- assert gradient is not None, f"rank={baseline.RANK}: missing measured gradient {index}."
139	- baseline.assert_finite(f"measured gradient {index}", gradient)
140	- return (
141	- baseline.rank_max_latency(forward_ms), baseline.rank_max_latency(forward_backward_ms),
142	- peak_allocated, peak_reserved,
143	- )
144	-
145	-
146	-def _measure_backend(
147	- shape: baseline.MoeShape, backend: str, factor: float \| None, scope: str,
148	-) -> dict:
149	- """Run fixed warmup and stable iterations with an isolated live backend."""
150	- layer = _measurement_layer(shape, backend, factor)
151	- router = _router_weight(shape)
152	- hidden, upstream = baseline.make_data(shape)
153	- ids, weights, _ = baseline.make_balanced_route(shape)
154	- inputs = (hidden, upstream, ids, weights)
155	- latencies, samples = [], []
156	- try:
157	- for step in range(8):
158	- latency = _iteration(layer, router, shape, inputs, scope)
159	- layer.zero_grad(set_to_none=True)
160	- router.grad = None
161	- if step >= 3:
162	- latencies.append(latency)
163	- samples.append(collect_memory())
164	- assert_stable_memory(samples)
165	- return {
166	- "backend": backend,
167	- "effective_capacity_factor": factor if backend == "mega_moe" else None,
168	- "peak_allocated_bytes": max(latency[2] for latency in latencies),
169	- "peak_reserved_bytes": max(latency[3] for latency in latencies),
170	- "stable_memory": samples,
171	- "rank_max_forward_ms": [latency[0] for latency in latencies],
172	- "rank_max_fwd_bwd_ms": [latency[1] for latency in latencies],
173	- "gradients_present_and_finite": True,
174	- }
175	- finally:
176	- if isinstance(layer, MegaMoeExperts):
177	- layer.close()
178	-
179	-
180	-def test_mega_moe_default_capacity_acceptance() -> None:
181	- """Record same-shape default/bounded A/B and Router-inclusive F/F+B costs."""
182	- shape = baseline.performance_shape()
183	- comparisons = []
184	- for factor in (None, 1.5):
185	- for scope in ("experts", "router_experts"):
186	- _validate_pair(shape, factor, scope)
187	- gc.collect()
188	- torch.npu.empty_cache()
189	- repetitions = []
190	- for backend in ("common", "mega_moe", "mega_moe", "common"):
191	- if baseline.RANK == 0:
192	- print(f"Level1 factor={factor} scope={scope} backend={backend}", flush=True)
193	- repetitions.append(_measure_backend(shape, backend, factor, scope))
194	- gc.collect()
195	- torch.npu.empty_cache()
196	- dist.barrier()
197	- medians = {
198	- backend: statistics.median(
199	- value for result in repetitions if result["backend"] == backend
200	- for value in result["rank_max_fwd_bwd_ms"]
201	- )
202	- for backend in ("common", "mega_moe")
203	- }
204	- comparisons.append({
205	- "comparison_capacity_factor": factor, "timing_scope": scope,
206	- "counts_supplied": False, "histogram_in_timing": True,
207	- "router_in_timing": scope == "router_experts", "warmup_steps": 3,
208	- "measured_steps": 5, "order": [result["backend"] for result in repetitions],
209	- "output_gradient_update_validation": True, "repetitions": repetitions,
210	- "mega_over_common_fwd_bwd": medians["mega_moe"] / medians["common"],
211	- "after_close_memory": memory_sample(),
212	- })
213	- write_evidence({
214	- "shape": asdict(shape), "dtype": "bfloat16", "ep_is_whole_world": True,
215	- "router_projection_dtype": "float32",
216	- "comparisons": comparisons,
217	- "timing_note": "Host wall time with a completion sync after F and B; setup and validation excluded.",
218	- })

Atests/torch/multicore/_test_mega_moe_router_interface.py+98-0

@@ -0,0 +1,98 @@
	1	+# Copyright 2026 Huawei Technologies Co., Ltd
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# http://www.apache.org/licenses/LICENSE-2.0
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+# ============================================================================
	15	+"""MegaMoe Router interface acceptance without caller-provided expert counts."""
	16	+
	17	+from __future__ import annotations
	18	+
	19	+import os
	20	+from dataclasses import asdict
	21	+
	22	+import torch
	23	+
	24	+from hyper_parallel.core.multicore import MegaMoeExperts
	25	+from tests.torch.multicore import _test_mega_moe as baseline
	26	+from tests.torch.multicore._mega_moe_utils import (
	27	+ start_shmem_lifetime,
	28	+ write_evidence,
	29	+)
	30	+
	31	+
	32	+def _router_weight(shape: baseline.MoeShape) -> torch.nn.Parameter:
	33	+ """Use identical trainable Router weights for the two expert backends."""
	34	+ generator = torch.Generator().manual_seed(706)
	35	+ weight = torch.randn(shape.num_experts, shape.hidden_size, generator=generator)
	36	+ return torch.nn.Parameter((weight / shape.hidden_size**0.5).to(baseline.DEVICE))
	37	+
	38	+
	39	+def _route(hidden: torch.Tensor, weight: torch.Tensor, top_k: int) -> tuple:
	40	+ """Compute Router projection, softmax, Top-K and selected-weight normalization."""
	41	+ probabilities = (hidden.float() @ weight.T).softmax(dim=-1)
	42	+ values, indices = probabilities.topk(top_k, dim=-1)
	43	+ return indices.to(torch.int32), values / values.sum(dim=-1, keepdim=True)
	44	+
	45	+
	46	+def _run_router_layer(layer: torch.nn.Module, shape: baseline.MoeShape) -> dict[str, torch.Tensor]:
	47	+ """Validate a connected Router/expert graph, including Router dW and SGD."""
	48	+ hidden, upstream = baseline.make_data(shape)
	49	+ hidden.requires_grad_(True)
	50	+ router = _router_weight(shape)
	51	+ ids, weights = _route(hidden, router, shape.top_k)
	52	+ weights.retain_grad()
	53	+ if isinstance(layer, MegaMoeExperts):
	54	+ output = layer(hidden, ids, weights)
	55	+ else:
	56	+ output = baseline.forward_layer(layer, hidden, ids, weights)
	57	+ output.backward(upstream)
	58	+ optimizer = torch.optim.SGD([*layer.parameters(), router], lr=1e-3)
	59	+ optimizer.step()
	60	+ torch.npu.synchronize()
	61	+ gate_up_grad, down_grad = baseline.expert_weight_gradients(layer)
	62	+ gate_up_weight, down_weight = baseline.expert_weights(layer)
	63	+ compared = {
	64	+ "output": output,
	65	+ "input_grad": hidden.grad,
	66	+ "route_weight_grad": weights.grad,
	67	+ "router_grad": router.grad,
	68	+ "router_weight": router,
	69	+ "gate_up_weight_grad": gate_up_grad,
	70	+ "down_weight_grad": down_grad,
	71	+ "gate_up_weight": gate_up_weight,
	72	+ "down_weight": down_weight,
	73	+ }
	74	+ for name, tensor in compared.items():
	75	+ assert tensor is not None, f"rank={baseline.RANK}: missing Router graph tensor {name}."
	76	+ baseline.assert_finite(f"Router graph {name}", tensor)
	77	+ return {name: tensor.detach().clone() for name, tensor in compared.items()}
	78	+
	79	+
	80	+def test_mega_moe_router_interface() -> None:
	81	+ """Validate omitted counts and Router gradients with one step per backend."""
	82	+ shape = baseline.performance_shape()
	83	+ start_shmem_lifetime()
	84	+ os.environ.pop("SYMMETRIC_MEMORY_HEAP_SIZE", None)
	85	+ mega, common = baseline.new_layers(shape)
	86	+ try:
	87	+ reference = _run_router_layer(common, shape)
	88	+ actual = _run_router_layer(mega, shape)
	89	+ for name, expected in reference.items():
	90	+ baseline.assert_close(f"Router graph {name}", actual[name], expected)
	91	+ finally:
	92	+ mega.close()
	93	+ write_evidence({
	94	+ "shape": asdict(shape), "dtype": "bfloat16", "ep_is_whole_world": True,
	95	+ "capacity_factor": None, "counts_supplied": False, "router_projection_dtype": "float32",
	96	+ "backend_steps": {"common": 1, "mega_moe": 1},
	97	+ "validated_tensors": list(reference), "output_gradient_update_validation": True,
	98	+ })

Mtests/torch/multicore/test_mega_moe.py+8-8

	@@ -24,13 +24,13 @@ from pathlib import Path
24	import pytest	24	import pytest
25		25
26	from tests.common.mark_utils import arg_mark	26	from tests.common.mark_utils import arg_mark
		27	+from tests.common.parallel_case import TorchCase, parallel_run
		28	+from tests.common.port_utils import allocate_port
27	from tests.torch.multicore._test_env import (	29	from tests.torch.multicore._test_env import (
28	multicore_adapter_is_available,	30	multicore_adapter_is_available,
29	prepare_multicore_test_environment,	31	prepare_multicore_test_environment,
30	without_inherited_rank_environment,	32	without_inherited_rank_environment,
31	)	33	)
32	-from tests.common.parallel_case import TorchCase, parallel_run
33	-from tests.common.port_utils import allocate_port
34		34
35	_WORKER = str(Path(__file__).resolve().parent / "_test_mega_moe.py")	35	_WORKER = str(Path(__file__).resolve().parent / "_test_mega_moe.py")
36	_PRECISION_WORLD_SIZE = 2	36	_PRECISION_WORLD_SIZE = 2
	@@ -153,7 +153,7 @@ def _run_acceptance_worker(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, case
153	_prepare_torch_multicore_test_environment()	153	_prepare_torch_multicore_test_environment()
154	result_dir = Path(os.getenv("HP_MEGA_MOE_EVIDENCE_DIR", str(tmp_path)))	154	result_dir = Path(os.getenv("HP_MEGA_MOE_EVIDENCE_DIR", str(tmp_path)))
155	result_path = result_dir / f"{case}.json"	155	result_path = result_dir / f"{case}.json"
156	- worker = "_test_mega_moe_resources.py" if cards == 2 else "_test_mega_moe_defaults.py"	156	+ worker = "_test_mega_moe_resources.py" if cards == 2 else "_test_mega_moe_router_interface.py"
157	monkeypatch.setenv("HP_MEGA_MOE_WORLD_SIZE", str(cards))	157	monkeypatch.setenv("HP_MEGA_MOE_WORLD_SIZE", str(cards))
158	monkeypatch.setenv("HP_MEGA_MOE_LEVEL1_RESULT", str(result_path))	158	monkeypatch.setenv("HP_MEGA_MOE_LEVEL1_RESULT", str(result_path))
159	monkeypatch.setenv("SHMEM_IP_PORT", f"tcp://127.0.0.1:{allocate_port()}")	159	monkeypatch.setenv("SHMEM_IP_PORT", f"tcp://127.0.0.1:{allocate_port()}")
	@@ -192,10 +192,10 @@ def test_mega_moe_shared_resources(monkeypatch: pytest.MonkeyPatch, tmp_path: Pa
192	card_mark="allcards",	192	card_mark="allcards",
193	essential_mark="unessential",	193	essential_mark="unessential",
194	)	194	)
195	-def test_mega_moe_default_interface(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:	195	+def test_mega_moe_router_interface(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
196	- """Feature: Default and explicit capacity without Router-provided counts.	196	+ """Feature: Router interface without caller-provided expert counts.
197		197
198	- Description: Compare identical representative shapes in ABBA order, including Router timing.	198	+ Description: Connect a trainable Router and run one common/MegaMoe step at the representative shape.
199	- Expectation: Outputs and gradients agree; stable latency, peak memory and software identity are recorded.	199	+ Expectation: Outputs, input/route/expert/Router gradients and one SGD update agree without supplied counts.
200	"""	200	"""
201	- _run_acceptance_worker(monkeypatch, tmp_path, "test_mega_moe_default_capacity_acceptance", 4)	201	+ _run_acceptance_worker(monkeypatch, tmp_path, "test_mega_moe_router_interface", 4)