Star847
1326
代码
代码
Issues1504
Pull Requests2268
流水线
项目讨论
Wiki
分析
Star847
1326
已合并
Add test for _EmptyStateDictLoadPlanner on NPU #42765
costan创建于 7月25日
Add test for _EmptyStateDictLoadPlanner on NPU #42765
已合并
costan创建于 7月25日
从ljfnjj1234567/pytorchEmptyStateDictLoadPlanner-v2.12.0合入到Ascend/pytorchv2.12.0
共 1 个文件变更+325-0
Atest/distributed/checkpoint/test_empty_state_dict_load_planner.py+325-0
@@ -0,0 +1,325 @@
1+# Copyright (c) 2026 Huawei Technologies Co., Ltd
2+# All rights reserved.
3+#
4+# Licensed under the BSD 3-Clause License (the "License");
5+# you may not use this file except in compliance with the License.
6+# You may obtain a copy of the License at
7+#
8+# https://opensource.org/licenses/BSD-3-Clause
9+#
10+# Unless required by applicable law or agreed to in writing, software
11+# distributed under the License is distributed on an "AS IS" BASIS,
12+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+# See the License for the specific language governing permissions and
14+# limitations under the License.
15+ 
16+"""
17+Add validation cases for torch.distributed.checkpoint APIs.
18+ 
19+This file validates _EmptyStateDictLoadPlanner behavior.
20+ 
21+Reasons:
22+1. PyTorch community lacks direct API validations for
23+ _EmptyStateDictLoadPlanner.
24+2. _EmptyStateDictLoadPlanner only handles checkpoint planning logic,
25+ which is device-independent.
26+3. These tests intentionally avoid NPU initialization and distributed
27+ environment requirements.
28+ 
29+This file can be extended with more planner validation cases.
30+"""
31+ 
32+import torch
33+ 
34+from torch.distributed.checkpoint.default_planner import (
35+ _EmptyStateDictLoadPlanner,
36+)
37+from torch.distributed.checkpoint.metadata import (
38+ ChunkStorageMetadata,
39+ Metadata,
40+ TensorProperties,
41+ TensorStorageMetadata,
42+)
43+from torch.distributed.checkpoint.planner import LoadPlan, LoadItemType
44+from torch.testing._internal.common_utils import TestCase, run_tests
45+ 
46+ 
47+def _make_tensor_md(size, dtype=torch.float32):
48+ """Helper to create TensorStorageMetadata for testing."""
49+ return TensorStorageMetadata(
50+ properties=TensorProperties.create_from_tensor(
51+ torch.empty(1, dtype=dtype)
52+ ),
53+ size=torch.Size(size),
54+ chunks=[
55+ ChunkStorageMetadata(
56+ offsets=torch.Size([0] * len(size)),
57+ sizes=torch.Size(size),
58+ )
59+ ],
60+ )
61+ 
62+ 
63+def _make_metadata(state_dict_metadata, *, planner_data=None):
64+ """Helper to create Metadata with a non-None planner_data default."""
65+ if planner_data is None:
66+ planner_data = {k: [k] for k in state_dict_metadata}
67+ return Metadata(
68+ state_dict_metadata=state_dict_metadata,
69+ planner_data=planner_data,
70+ )
71+ 
72+ 
73+class TestEmptyStateDictLoadPlanner(TestCase):
74+ 
75+ # ---------------------------------------------------------------
76+ # Construction
77+ # ---------------------------------------------------------------
78+ 
79+ def test_init_without_keys(self):
80+ """Verify planner initialization without specifying keys."""
81+ planner = _EmptyStateDictLoadPlanner()
82+ self.assertIsNone(planner.keys)
83+ 
84+ def test_init_with_keys(self):
85+ """Verify planner initialization with specified keys."""
86+ keys = ["model"]
87+ planner = _EmptyStateDictLoadPlanner(keys=keys)
88+ self.assertEqual(planner.keys, keys)
89+ 
90+ # ---------------------------------------------------------------
91+ # set_up_planner -- error paths
92+ # ---------------------------------------------------------------
93+ 
94+ def test_set_up_planner_with_non_empty_state_dict(self):
95+ """_EmptyStateDictLoadPlanner requires an empty state_dict.
96+ 
97+ A non-empty state_dict should trigger assertion failure.
98+ """
99+ planner = _EmptyStateDictLoadPlanner()
100+ with self.assertRaises(AssertionError):
101+ planner.set_up_planner(
102+ {"model": torch.tensor([1])},
103+ metadata=None,
104+ )
105+ 
106+ def test_set_up_planner_without_metadata(self):
107+ """Metadata is required when setting up planner.
108+ 
109+ Missing metadata should trigger assertion failure.
110+ """
111+ planner = _EmptyStateDictLoadPlanner()
112+ with self.assertRaises(AssertionError):
113+ planner.set_up_planner({}, metadata=None)
114+ 
115+ # ---------------------------------------------------------------
116+ # set_up_planner -- success paths
117+ # ---------------------------------------------------------------
118+ 
119+ def test_set_up_planner_success(self):
120+ """Empty state_dict with valid metadata: state_dict is populated.
121+ 
122+ _EmptyStateDictLoadPlanner rebuilds the state_dict from metadata,
123+ creating empty tensors for each TensorStorageMetadata entry.
124+ After set_up_planner, the planner's internal state (metadata,
125+ state_dict) should be properly initialized.
126+ """
127+ planner = _EmptyStateDictLoadPlanner()
128+ metadata = _make_metadata(
129+ state_dict_metadata={
130+ "a": _make_tensor_md((2, 3)),
131+ "b": _make_tensor_md((4,), dtype=torch.float64),
132+ }
133+ )
134+ sd = {}
135+ planner.set_up_planner(sd, metadata)
136+ 
137+ self.assertEqual(len(sd), 2)
138+ self.assertIn("a", sd)
139+ self.assertIn("b", sd)
140+ self.assertEqual(sd["a"].shape, torch.Size((2, 3)))
141+ self.assertEqual(sd["a"].dtype, torch.float32)
142+ self.assertEqual(sd["b"].shape, torch.Size((4,)))
143+ self.assertEqual(sd["b"].dtype, torch.float64)
144+ self.assertIs(planner.metadata, metadata)
145+ self.assertIsNotNone(planner.state_dict)
146+ 
147+ def test_set_up_planner_with_is_coordinator(self):
148+ """set_up_planner should accept and store is_coordinator flag."""
149+ planner = _EmptyStateDictLoadPlanner()
150+ metadata = _make_metadata(
151+ state_dict_metadata={"a": _make_tensor_md((1,))}
152+ )
153+ sd = {}
154+ planner.set_up_planner(sd, metadata, is_coordinator=True)
155+ self.assertTrue(planner.is_coordinator)
156+ 
157+ # ---------------------------------------------------------------
158+ # keys filtering in set_up_planner
159+ # ---------------------------------------------------------------
160+ 
161+ def test_keys_none_loads_all(self):
162+ """When keys=None, every key in metadata is loaded."""
163+ planner = _EmptyStateDictLoadPlanner(keys=None)
164+ metadata = _make_metadata(
165+ state_dict_metadata={
166+ "x": _make_tensor_md((1,)),
167+ "y": _make_tensor_md((2,)),
168+ "z": _make_tensor_md((3,)),
169+ }
170+ )
171+ sd = {}
172+ planner.set_up_planner(sd, metadata)
173+ self.assertEqual(len(sd), 3)
174+ self.assertIn("x", sd)
175+ self.assertIn("y", sd)
176+ self.assertIn("z", sd)
177+ 
178+ def test_keys_filter_loads_subset(self):
179+ """When keys is a specific set, only those keys are loaded."""
180+ planner = _EmptyStateDictLoadPlanner(keys={"x", "z"})
181+ metadata = _make_metadata(
182+ state_dict_metadata={
183+ "x": _make_tensor_md((1,)),
184+ "y": _make_tensor_md((2,)),
185+ "z": _make_tensor_md((3,)),
186+ },
187+ planner_data={"x": ["x"], "y": ["y"], "z": ["z"]},
188+ )
189+ sd = {}
190+ planner.set_up_planner(sd, metadata)
191+ self.assertEqual(len(sd), 2)
192+ self.assertIn("x", sd)
193+ self.assertIn("z", sd)
194+ self.assertNotIn("y", sd)
195+ 
196+ def test_keys_filter_loads_nothing_when_no_match(self):
197+ """When keys matches nothing, state_dict stays empty."""
198+ planner = _EmptyStateDictLoadPlanner(keys={"nonexistent"})
199+ metadata = _make_metadata(
200+ state_dict_metadata={
201+ "a": _make_tensor_md((1,)),
202+ "b": _make_tensor_md((2,)),
203+ },
204+ planner_data={"a": ["a"], "b": ["b"]},
205+ )
206+ sd = {}
207+ planner.set_up_planner(sd, metadata)
208+ self.assertEqual(len(sd), 0)
209+ 
210+ def test_keys_filter_with_planner_data(self):
211+ """keys filter works when metadata has planner_data (nested paths).
212+ 
213+ The planner should match keys against both the storage key and
214+ the unflattened path components from planner_data.
215+ """
216+ planner = _EmptyStateDictLoadPlanner(keys={"model.layer.weight"})
217+ metadata = _make_metadata(
218+ state_dict_metadata={
219+ "0": _make_tensor_md((3, 4)),
220+ },
221+ planner_data={
222+ "0": ["model", "layer", "weight"],
223+ },
224+ )
225+ sd = {}
226+ planner.set_up_planner(sd, metadata)
227+ self.assertIn("model", sd)
228+ self.assertIn("layer", sd["model"])
229+ self.assertIn("weight", sd["model"]["layer"])
230+ self.assertEqual(
231+ sd["model"]["layer"]["weight"].shape, torch.Size((3, 4))
232+ )
233+ 
234+ # ---------------------------------------------------------------
235+ # create_local_plan after set_up_planner
236+ # ---------------------------------------------------------------
237+ 
238+ def test_create_local_plan_after_setup(self):
239+ """After successful set_up_planner, create_local_plan returns a
240+ LoadPlan with the expected number of ReadItems."""
241+ planner = _EmptyStateDictLoadPlanner()
242+ metadata = _make_metadata(
243+ state_dict_metadata={
244+ "a": _make_tensor_md((2, 3)),
245+ "b": _make_tensor_md((4,)),
246+ }
247+ )
248+ sd = {}
249+ planner.set_up_planner(sd, metadata)
250+ 
251+ local_plan = planner.create_local_plan()
252+ self.assertIsInstance(local_plan, LoadPlan)
253+ self.assertGreaterEqual(len(local_plan.items), 2)
254+ item_types = {item.type for item in local_plan.items}
255+ self.assertIn(LoadItemType.TENSOR, item_types)
256+ 
257+ def test_create_local_plan_empty_when_keys_filter_all(self):
258+ """When keys filter removes all metadata entries, local plan is empty."""
259+ planner = _EmptyStateDictLoadPlanner(keys={"nonexistent"})
260+ metadata = _make_metadata(
261+ state_dict_metadata={
262+ "a": _make_tensor_md((1,)),
263+ },
264+ planner_data={"a": ["a"]},
265+ )
266+ sd = {}
267+ planner.set_up_planner(sd, metadata)
268+ 
269+ local_plan = planner.create_local_plan()
270+ self.assertIsInstance(local_plan, LoadPlan)
271+ self.assertEqual(len(local_plan.items), 0)
272+ 
273+ # ---------------------------------------------------------------
274+ # create_global_plan
275+ # ---------------------------------------------------------------
276+ 
277+ def test_create_global_plan(self):
278+ """create_global_plan should return a list of LoadPlan for each rank."""
279+ planner = _EmptyStateDictLoadPlanner()
280+ metadata = _make_metadata(
281+ state_dict_metadata={"a": _make_tensor_md((1,))}
282+ )
283+ sd = {}
284+ planner.set_up_planner(sd, metadata)
285+ local_plan = planner.create_local_plan()
286+ 
287+ global_plans = planner.create_global_plan([local_plan])
288+ self.assertIsInstance(global_plans, list)
289+ self.assertEqual(len(global_plans), 1)
290+ self.assertIsInstance(global_plans[0], LoadPlan)
291+ 
292+ def test_create_global_plan_multiple_ranks(self):
293+ """create_global_plan with plans from multiple ranks."""
294+ planner = _EmptyStateDictLoadPlanner()
295+ metadata = _make_metadata(
296+ state_dict_metadata={"a": _make_tensor_md((1,))}
297+ )
298+ sd = {}
299+ planner.set_up_planner(sd, metadata)
300+ p1 = planner.create_local_plan()
301+ p2 = planner.create_local_plan()
302+ 
303+ global_plans = planner.create_global_plan([p1, p2])
304+ self.assertEqual(len(global_plans), 2)
305+ 
306+ # ---------------------------------------------------------------
307+ # finish_plan
308+ # ---------------------------------------------------------------
309+ 
310+ def test_finish_plan_passthrough(self):
311+ """finish_plan returns the plan unchanged (identity)."""
312+ planner = _EmptyStateDictLoadPlanner()
313+ metadata = _make_metadata(
314+ state_dict_metadata={"a": _make_tensor_md((1,))}
315+ )
316+ sd = {}
317+ planner.set_up_planner(sd, metadata)
318+ local_plan = planner.create_local_plan()
319+ 
320+ finished_plan = planner.finish_plan(local_plan)
321+ self.assertIs(finished_plan, local_plan)
322+ 
323+ 
324+if __name__ == "__main__":
325+ run_tests()