Star846
1324
代码
代码
Issues1497
Pull Requests2281
流水线
项目讨论
Wiki
分析
Star846
1324
已合并
test: add NPU tests for BroadcastingTorchSaveReader APIs #35035
zjucn创建于 5月8日
test: add NPU tests for BroadcastingTorchSaveReader APIs #35035
已合并
zjucn创建于 5月8日
从zjucn/pytorchtest/2.7.1-broadcasting-reader-api-npu合入到Ascend/pytorchv2.7.1
共 4 个文件变更+171-1
Mdocs/zh/native_apis/pytorch_2-10-0/torch-distributed-checkpoint.md+5-0
@@ -60,6 +60,11 @@
60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|
61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|
62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|
63+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader|是|-|
64+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata|是|-|
65+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan|是|-|
66+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan|是|-|
67+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data|是|-|
63|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|68|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|
64|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|69|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|
65|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|70|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|
Mdocs/zh/native_apis/pytorch_2-7-1/torch-distributed-checkpoint.md+6-1
@@ -60,6 +60,11 @@
60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|
61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|
62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|
63+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader|是|-|
64+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata|是|-|
65+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan|是|-|
66+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan|是|-|
67+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data|是|-|
63|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|68|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|
64|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|69|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|
65-|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|70+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|
Mdocs/zh/native_apis/pytorch_2-9-0/torch-distributed-checkpoint.md+5-0
@@ -60,6 +60,11 @@
60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|60|torch.distributed.checkpoint.state_dict.set_model_state_dict|是|-|
61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|61|torch.distributed.checkpoint.state_dict.set_optimizer_state_dict|是|-|
62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|62|torch.distributed.checkpoint.state_dict.StateDictOptions|是|-|
63+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader|是|-|
64+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata|是|-|
65+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan|是|-|
66+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan|是|-|
67+|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data|是|-|
63|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|68|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset|是|-|
64|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|69|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader|是|-|
65|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|70|torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id|是|-|
Atest/distributed/checkpoint/test_broadcasting_reader_api.py+155-0
@@ -0,0 +1,155 @@
1+"""
2+1. PyTorch community lacks direct validation cases for some
3+ torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader APIs,
4+ so this file is added.
5+ 
6+2. This file validates the following APIs:
7+ torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata
8+ torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan
9+ torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan
10+ torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data
11+ (extendable)
12+"""
13+ 
14+import os
15+import tempfile
16+ 
17+import torch
18+import torch.distributed as dist
19+import torch.distributed.checkpoint as dcp
20+from torch.distributed.checkpoint.format_utils import (
21+ BroadcastingTorchSaveReader,
22+ DynamicMetaLoadPlanner,
23+)
24+from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex
25+from torch.distributed.checkpoint.planner import LoadItemType, LoadPlan, ReadItem
26+from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
27+from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
28+ 
29+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms
30+from torch_npu.testing.testcase import run_tests, TestCase
31+ 
32+ 
33+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
34+ 
35+ 
36+def _make_byteio_read_item(fqn: str = "payload") -> ReadItem:
37+ return ReadItem(
38+ type=LoadItemType.BYTE_IO,
39+ dest_index=MetadataIndex(fqn),
40+ dest_offsets=torch.Size((0,)),
41+ storage_index=MetadataIndex(fqn),
42+ storage_offsets=torch.Size((0,)),
43+ lengths=torch.Size((0,)),
44+ )
45+ 
46+ 
47+class DummyLoadPlanner:
48+ flatten_state_dict = False
49+ state_dict = {}
50+ 
51+ 
52+class TestBroadcastingTorchSaveReaderApi(TestCase):
53+ def test_read_metadata_returns_empty_metadata(self):
54+ reader = BroadcastingTorchSaveReader(checkpoint_id="unused.pt")
55+ 
56+ metadata = reader.read_metadata()
57+ 
58+ self.assertIsInstance(metadata, Metadata)
59+ self.assertEqual({}, metadata.state_dict_metadata)
60+ 
61+ def test_prepare_local_plan_returns_input_plan(self):
62+ reader = BroadcastingTorchSaveReader()
63+ plan = LoadPlan([], storage_data={"storage": 1}, planner_data={"planner": 2})
64+ 
65+ result = reader.prepare_local_plan(plan)
66+ 
67+ self.assertIs(result, plan)
68+ 
69+ def test_prepare_global_plan_returns_input_plans(self):
70+ reader = BroadcastingTorchSaveReader()
71+ plans = [
72+ LoadPlan([], storage_data={"rank": 0}),
73+ LoadPlan([], storage_data={"rank": 1}),
74+ ]
75+ 
76+ result = reader.prepare_global_plan(plans)
77+ 
78+ self.assertIs(result, plans)
79+ 
80+ def test_read_data_rejects_byte_io_items(self):
81+ with tempfile.TemporaryDirectory() as temp_dir:
82+ torch_path = os.path.join(temp_dir, "payload.pt")
83+ torch.save({"payload": ["step", 1]}, torch_path)
84+ 
85+ reader = BroadcastingTorchSaveReader(checkpoint_id=torch_path)
86+ reader.is_coordinator = True
87+ plan = LoadPlan([_make_byteio_read_item("payload")])
88+ 
89+ with self.assertRaisesRegex(RuntimeError, "only supports loading Tensors"):
90+ reader.read_data(plan, DummyLoadPlanner())
91+ 
92+ 
93+class TestBroadcastingTorchSaveReaderNpu(DTensorTestBase):
94+ @property
95+ def world_size(self) -> int:
96+ return 2
97+ 
98+ @with_comms
99+ @with_temp_dir
100+ @skipIfUnsupportMultiNPU(2)
101+ def test_read_data_loads_torch_save_tensor_to_npu_state_dict(self):
102+ source_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
103+ torch_path = os.path.join(self.temp_dir, "tensor.pt")
104+ if dist.get_rank() == 0:
105+ torch.save({"tensor": source_tensor}, torch_path)
106+ dist.barrier()
107+ 
108+ state_dict = {"tensor": torch.zeros(3, 4).to(device_type)}
109+ dcp.load(
110+ state_dict,
111+ planner=DynamicMetaLoadPlanner(),
112+ storage_reader=BroadcastingTorchSaveReader(),
113+ checkpoint_id=torch_path,
114+ )
115+ 
116+ self.assertEqual(source_tensor, state_dict["tensor"].cpu())
117+ 
118+ @with_comms
119+ @with_temp_dir
120+ @skipIfUnsupportMultiNPU(2)
121+ def test_read_data_handles_nested_state_dict(self):
122+ source_state_dict = {
123+ "model": {
124+ "weight": torch.arange(6, dtype=torch.float32).reshape(2, 3),
125+ "bias": torch.tensor([1.0, 2.0], dtype=torch.float32),
126+ }
127+ }
128+ torch_path = os.path.join(self.temp_dir, "nested.pt")
129+ if dist.get_rank() == 0:
130+ torch.save(source_state_dict, torch_path)
131+ dist.barrier()
132+ 
133+ state_dict = {
134+ "model": {
135+ "weight": torch.zeros(2, 3).to(device_type),
136+ "bias": torch.zeros(2).to(device_type),
137+ }
138+ }
139+ dcp.load(
140+ state_dict,
141+ planner=DynamicMetaLoadPlanner(),
142+ storage_reader=BroadcastingTorchSaveReader(),
143+ checkpoint_id=torch_path,
144+ )
145+ 
146+ self.assertEqual(
147+ source_state_dict["model"]["weight"], state_dict["model"]["weight"].cpu()
148+ )
149+ self.assertEqual(
150+ source_state_dict["model"]["bias"], state_dict["model"]["bias"].cpu()
151+ )
152+ 
153+ 
154+if __name__ == "__main__":
155+ run_tests()