Add distributed checkpoint support-pytorch-AtomGit

已合并

Add distributed checkpoint support #9796

chuboning创建于 2024年2月28日从refs/pull/9796/head 合入到 master

Add distributed checkpoint support #9796

已合并

chuboning创建于 2024年2月28日

从refs/pull/9796/head合入到master

共 6 个文件变更+626-31

Atest/distributed/checkpoint/test_checkpoint.py+387-0

@@ -0,0 +1,387 @@
	1	+import os
	2	+from typing import cast, List, Optional, Union
	3	+import torch
	4	+import torch.distributed as dist
	5	+import torch.nn
	6	+import torch.futures
	7	+from torch.futures import Future
	8	+from torch.distributed.checkpoint.storage import WriteResult
	9	+
	10	+from torch.distributed.checkpoint import (
	11	+ StorageReader,
	12	+ StorageWriter,
	13	+ CheckpointException,
	14	+ load_state_dict,
	15	+ save_state_dict,
	16	+)
	17	+from torch.distributed._shard import sharded_tensor
	18	+
	19	+from torch.distributed.checkpoint.default_planner import (
	20	+ _create_default_local_metadata,
	21	+)
	22	+
	23	+from torch.distributed.checkpoint.metadata import (
	24	+ BytesStorageMetadata,
	25	+ Metadata,
	26	+ TensorStorageMetadata,
	27	+)
	28	+
	29	+from torch.distributed.checkpoint.planner import (
	30	+ SavePlan,
	31	+ SavePlanner,
	32	+ LoadPlan,
	33	+ LoadPlanner,
	34	+)
	35	+
	36	+from torch.distributed._shard.sharded_tensor import (
	37	+ state_dict_hook,
	38	+ ShardedTensor,
	39	+)
	40	+from torch.distributed._shard.sharding_spec import ChunkShardingSpec
	41	+from torch.testing._internal.distributed._shard.sharded_tensor import (
	42	+ ShardedTensorTestBase,
	43	+)
	44	+
	45	+import torch_npu
	46	+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
	47	+from torch_npu.testing.testcase import run_tests
	48	+
	49	+
	50	+class TestModule(torch.nn.Module):
	51	+ def __init__(self) -> None:
	52	+ super().__init__()
	53	+ self.sharded: ShardedTensor = sharded_tensor.zeros(self.spec(), 4, 4)
	54	+ self.regular = torch.nn.Parameter(torch.ones(4, 4))
	55	+ self.extra_sharded: Optional[ShardedTensor] = None
	56	+ self.extra_param: Optional[torch.nn.Parameter] = None
	57	+ self._register_state_dict_hook(state_dict_hook)
	58	+
	59	+ def spec(self) -> ChunkShardingSpec:
	60	+ # pyre-fixme [28]: Unexpected keyword argument `dim` to call `dist._sharding_spec.api.ChunkShardingSpec.__init__`.
	61	+ return ChunkShardingSpec(
	62	+ dim=0,
	63	+ placements=[
	64	+ "rank:0/npu:0",
	65	+ "rank:1/npu:1",
	66	+ ],
	67	+ )
	68	+
	69	+
	70	+class TestDistributedCheckpointing(ShardedTensorTestBase):
	71	+ def destroy_pg(self) -> None:
	72	+ dist.barrier()
	73	+ dist.destroy_process_group()
	74	+
	75	+ @property
	76	+ def world_size(self) -> int:
	77	+ return 2
	78	+
	79	+ @with_comms
	80	+ @skipIfUnsupportMultiNPU(2)
	81	+ def test_tensor_metadata_with_missing_rank_spec(self) -> None:
	82	+ spec = ChunkShardingSpec(
	83	+ dim=0,
	84	+ placements=[
	85	+ "rank:1/npu:1",
	86	+ ],
	87	+ )
	88	+
	89	+ st = sharded_tensor.zeros(spec, 4, 4, dtype=torch.float64)
	90	+ mapping = {}
	91	+
	92	+ md = _create_default_local_metadata({"st": st})
	93	+
	94	+ st_md = md.state_dict_metadata["st"]
	95	+ self.assertEqual(1, len(st_md.chunks))
	96	+
	97	+ @with_comms
	98	+ @skipIfUnsupportMultiNPU(2)
	99	+ def test_default_metadata(self) -> None:
	100	+ device = f"npu:{dist.get_rank()}"
	101	+ spec = ChunkShardingSpec(
	102	+ dim=0,
	103	+ placements=[
	104	+ "rank:0/npu:0",
	105	+ "rank:1/npu:1",
	106	+ ],
	107	+ )
	108	+
	109	+ state_dict = {
	110	+ "sharded": sharded_tensor.rand(
	111	+ spec,
	112	+ (
	113	+ 10,
	114	+ 10,
	115	+ ),
	116	+ ),
	117	+ "replicated": torch.rand(4, device=device),
	118	+ "bytes": [1, 2, 3, 4],
	119	+ }
	120	+
	121	+ metadata = _create_default_local_metadata(state_dict)
	122	+ self.assertTrue("bytes" in metadata.state_dict_metadata)
	123	+ self.assertIsInstance(
	124	+ metadata.state_dict_metadata["bytes"], BytesStorageMetadata
	125	+ )
	126	+
	127	+ self.assertTrue("replicated" in metadata.state_dict_metadata)
	128	+ self.assertIsInstance(
	129	+ metadata.state_dict_metadata["replicated"], TensorStorageMetadata
	130	+ )
	131	+ md = metadata.state_dict_metadata["replicated"]
	132	+ self.assertEqual(md.size, state_dict["replicated"].size())
	133	+ self.assertEqual(md.properties.dtype, torch.float32)
	134	+ self.assertEqual(1, len(md.chunks))
	135	+
	136	+ self.assertTrue("sharded" in metadata.state_dict_metadata)
	137	+ self.assertIsInstance(
	138	+ metadata.state_dict_metadata["sharded"], TensorStorageMetadata
	139	+ )
	140	+ md = metadata.state_dict_metadata["sharded"]
	141	+ self.assertEqual(md.properties.dtype, torch.float32)
	142	+ self.assertEqual(md.size, state_dict["sharded"].size())
	143	+ self.assertEqual(2, len(md.chunks))
	144	+
	145	+
	146	+class TestStorageBase:
	147	+ def __init__(self, fail_conf):
	148	+ self.fail_conf = fail_conf
	149	+ self.rank = 0 if not dist.is_initialized() else dist.get_rank()
	150	+
	151	+ def _get_ranks(self, name):
	152	+ return self.fail_conf[name] if name in self.fail_conf else None
	153	+
	154	+ def _fail_rank(self, name):
	155	+ ranks = self._get_ranks(name)
	156	+ if ranks is not None and self.rank in ranks:
	157	+ raise ValueError(f"rank fail {self.rank} for {name}")
	158	+
	159	+ def _fail_rank_async(self, name, result=None):
	160	+ ranks = self._get_ranks(name)
	161	+ fut = Future()
	162	+ if ranks is not None and self.rank in ranks:
	163	+ fut.set_exception(ValueError(f"async rank fail {self.rank} for {name}"))
	164	+ else:
	165	+ fut.set_result(result)
	166	+ return fut
	167	+
	168	+
	169	+class FaultyStorageWriter(TestStorageBase, StorageWriter):
	170	+ def __init__(self, fail_conf):
	171	+ super().__init__(fail_conf)
	172	+
	173	+ def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None:
	174	+ return
	175	+
	176	+ def set_up_storage_writer(self, is_coordinator: bool) -> None:
	177	+ self._fail_rank("fail_set_up_storage_writer")
	178	+
	179	+ def prepare_local_plan(self, plan: SavePlan) -> SavePlan:
	180	+ self._fail_rank("fail_prepare_local_plan")
	181	+ return plan
	182	+
	183	+ def prepare_global_plan(self, plans: List[SavePlan]) -> List[SavePlan]:
	184	+ self._fail_rank("fail_prepare_global_plan")
	185	+ return plans
	186	+
	187	+ def write_data(
	188	+ self, plan: SavePlan, planner: SavePlanner
	189	+ ) -> Future[List[WriteResult]]:
	190	+ self._fail_rank("fail_write_data")
	191	+ return self._fail_rank_async("fail_write_data_async", [])
	192	+
	193	+ def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None:
	194	+ self._fail_rank("fail_finish")
	195	+
	196	+ @classmethod
	197	+ def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
	198	+ return True
	199	+
	200	+
	201	+class FaultyStorageReader(TestStorageBase, StorageReader):
	202	+ def __init__(self, metadata, fail_conf):
	203	+ super().__init__(fail_conf)
	204	+ self.metadata = metadata
	205	+
	206	+ def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None:
	207	+ return
	208	+
	209	+ def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None:
	210	+ self._fail_rank("fail_set_up_storage_reader")
	211	+
	212	+ def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan:
	213	+ self._fail_rank("fail_prepare_local_plan")
	214	+ return plan
	215	+
	216	+ def prepare_global_plan(self, plans: List[LoadPlan]) -> List[LoadPlan]:
	217	+ self._fail_rank("fail_prepare_global_plan")
	218	+ return plans
	219	+
	220	+ def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
	221	+ self._fail_rank("fail_read_data")
	222	+ return self._fail_rank_async("fail_read_data_async")
	223	+
	224	+ def read_metadata(self) -> Metadata:
	225	+ self._fail_rank("fail_read_metadata")
	226	+ return self.metadata
	227	+
	228	+ @classmethod
	229	+ def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
	230	+ return True
	231	+
	232	+
	233	+class TestDistributedFailure(ShardedTensorTestBase):
	234	+ def destroy_pg(self) -> None:
	235	+ dist.barrier()
	236	+ dist.destroy_process_group()
	237	+
	238	+ def get_spec(self):
	239	+ return ChunkShardingSpec(
	240	+ dim=0,
	241	+ placements=[
	242	+ f"rank:{r}/npu:{r}" for r in range(dist.get_world_size())
	243	+ ],
	244	+ )
	245	+
	246	+ @with_comms
	247	+ @skipIfUnsupportMultiNPU(4)
	248	+ def test_dummy_writer_works(self) -> None:
	249	+ state_dict = {
	250	+ "sharded": sharded_tensor.rand(self.get_spec(), 20, 20),
	251	+ "replicated": torch.rand(10, 10),
	252	+ "bytes": [1, 2, 3, 4],
	253	+ }
	254	+
	255	+ save_state_dict(state_dict, FaultyStorageWriter({}))
	256	+
	257	+ @with_comms
	258	+ @skipIfUnsupportMultiNPU(4)
	259	+ def test_dummy_reader_works(self) -> None:
	260	+ state_dict = {
	261	+ "sharded": sharded_tensor.rand(self.get_spec(), 20, 20),
	262	+ "replicated": torch.rand(10, 10),
	263	+ "bytes": [1, 2, 3, 4],
	264	+ }
	265	+ metadata = _create_default_local_metadata(state_dict)
	266	+
	267	+ load_state_dict(state_dict, FaultyStorageReader(metadata, {}))
	268	+
	269	+ def _test_dist_failure(self, callback, kwargs):
	270	+ bad_ranks = next(iter(kwargs.values())) if len(kwargs) > 0 else []
	271	+
	272	+ # Empty bad_ranks means it must work
	273	+ if len(bad_ranks) == 0:
	274	+ callback()
	275	+ else:
	276	+ with self.assertRaises(CheckpointException) as cm:
	277	+ callback()
	278	+ e = cast(CheckpointException, cm.exception)
	279	+ for rank, wrapped_ex in e.failures.items():
	280	+ ex = wrapped_ex[0]
	281	+ self.assertTrue(rank in bad_ranks, msg=f"{rank} did not fail")
	282	+ if not kwargs.get("ignore_exception_type", False):
	283	+ self.assertEqual(ValueError, type(ex), str(ex))
	284	+
	285	+ failed_ranks = e.failures.keys()
	286	+ for rank in bad_ranks:
	287	+ self.assertTrue(
	288	+ rank in failed_ranks,
	289	+ msg=f"{rank} was supposed to fail was fine",
	290	+ )
	291	+
	292	+ def _test_save(self, state_dict, coordinator=0, **kwargs):
	293	+ no_dist = not dist.is_initialized()
	294	+
	295	+ def _save():
	296	+ save_state_dict(
	297	+ state_dict,
	298	+ storage_writer=FaultyStorageWriter(kwargs),
	299	+ coordinator_rank=coordinator,
	300	+ no_dist=no_dist,
	301	+ )
	302	+
	303	+ self._test_dist_failure(_save, kwargs)
	304	+
	305	+ def _test_load(self, state_dict, coordinator=0, **kwargs):
	306	+ no_dist = not dist.is_initialized()
	307	+
	308	+ def _load():
	309	+ metadata = _create_default_local_metadata(state_dict)
	310	+ load_state_dict(
	311	+ state_dict,
	312	+ storage_reader=FaultyStorageReader(metadata, kwargs),
	313	+ coordinator_rank=coordinator,
	314	+ no_dist=no_dist,
	315	+ )
	316	+
	317	+ self._test_dist_failure(_load, kwargs)
	318	+
	319	+ @with_comms
	320	+ @skipIfUnsupportMultiNPU(4)
	321	+ def test_save_error_handling(self) -> None:
	322	+ state_dict = {
	323	+ "sharded": sharded_tensor.rand(self.get_spec(), 20, 20),
	324	+ "replicated": torch.rand(10, 10),
	325	+ "bytes": [1, 2, 3, 4],
	326	+ }
	327	+
	328	+ self._test_save(state_dict, fail_set_up_storage_writer=[0])
	329	+ self._test_save(state_dict, fail_finish=[0])
	330	+ self._test_save(state_dict, fail_prepare_global_plan=[0])
	331	+
	332	+ self._test_save(state_dict, fail_prepare_local_plan=[0])
	333	+ self._test_save(state_dict, fail_write_data=[2])
	334	+ self._test_save(state_dict, fail_write_data_async=[3])
	335	+
	336	+ self._test_save(state_dict, coordinator=1, fail_set_up_storage_writer=[1])
	337	+ self._test_save(state_dict, coordinator=1, fail_finish=[1])
	338	+
	339	+ def test_save_error_handling_no_dist(self) -> None:
	340	+ state_dict = {"replicated": torch.rand(10, 10), "bytes": [1, 2, 3, 4]}
	341	+
	342	+ self.assertFalse(dist.is_initialized())
	343	+
	344	+ self._test_save(state_dict, fail_set_up_storage_writer=[0])
	345	+ self._test_save(state_dict, fail_finish=[0])
	346	+ self._test_save(state_dict, fail_prepare_global_plan=[0])
	347	+
	348	+ self._test_save(state_dict, fail_prepare_local_plan=[0])
	349	+ self._test_save(state_dict, fail_write_data=[0])
	350	+ self._test_save(state_dict, fail_write_data_async=[0])
	351	+
	352	+ @with_comms
	353	+ @skipIfUnsupportMultiNPU(4)
	354	+ def test_load_error_handling(self) -> None:
	355	+ state_dict = {
	356	+ "sharded": sharded_tensor.rand(self.get_spec(), 20, 20),
	357	+ "replicated": torch.rand(10, 10),
	358	+ "bytes": [1, 2, 3, 4],
	359	+ }
	360	+
	361	+ self._test_load(state_dict)
	362	+ self._test_load(state_dict, fail_set_up_storage_reader=[0])
	363	+ self._test_load(state_dict, fail_prepare_global_plan=[0])
	364	+ self._test_load(state_dict, fail_read_metadata=[0])
	365	+ self._test_load(state_dict, fail_prepare_local_plan=[1])
	366	+ self._test_load(state_dict, fail_read_data=[3])
	367	+ self._test_load(state_dict, fail_read_data_async=[1])
	368	+
	369	+ self._test_load(state_dict, coordinator=3, fail_set_up_storage_reader=[0])
	370	+ self._test_load(state_dict, coordinator=1, fail_read_metadata=[3])
	371	+ self._test_load(state_dict, coordinator=2, fail_read_data=[0])
	372	+ self._test_load(state_dict, coordinator=3, fail_read_data_async=[2])
	373	+ self._test_load(state_dict, coordinator=1, fail_prepare_global_plan=[1])
	374	+
	375	+ def test_load_error_handling_no_dist(self) -> None:
	376	+ state_dict = {"replicated": torch.rand(10, 10), "bytes": [1, 2, 3, 4]}
	377	+ self._test_load(state_dict)
	378	+ self._test_load(state_dict, fail_set_up_storage_reader=[0])
	379	+ self._test_load(state_dict, fail_read_metadata=[0])
	380	+ self._test_load(state_dict, fail_prepare_local_plan=[0])
	381	+ self._test_load(state_dict, fail_prepare_global_plan=[0])
	382	+ self._test_load(state_dict, fail_read_data=[0])
	383	+ self._test_load(state_dict, fail_read_data_async=[0])
	384	+
	385	+
	386	+if __name__ == "__main__":
	387	+ run_tests()

Atest/distributed/checkpoint/test_fsdp_model_state.py+97-0

@@ -0,0 +1,97 @@
	1	+import torch
	2	+
	3	+import torch.distributed.checkpoint as dist_cp
	4	+import torch.distributed as dist
	5	+
	6	+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
	7	+from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
	8	+from torch.distributed.checkpoint.default_planner import (
	9	+ DefaultSavePlanner,
	10	+ DefaultLoadPlanner,
	11	+)
	12	+
	13	+from torch.testing._internal.distributed._tensor.common_dtensor import (
	14	+ DTensorTestBase,
	15	+)
	16	+from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
	17	+
	18	+import torch_npu
	19	+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
	20	+from torch_npu.testing.testcase import run_tests
	21	+
	22	+
	23	+class FsdpModelStateCheckpoint(DTensorTestBase):
	24	+ def _test_fsdp_model_state(self, process_group) -> None:
	25	+ CHECKPOINT_DIR = self.temp_dir
	26	+
	27	+ model = FSDP(torch.nn.Linear(8, 8, device="npu"))
	28	+ model(torch.rand(8, 8, device="npu")).sum().backward()
	29	+
	30	+ with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
	31	+ state_dict = {
	32	+ "model": model.state_dict(),
	33	+ }
	34	+
	35	+ dist_cp.save_state_dict(
	36	+ state_dict=state_dict,
	37	+ storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
	38	+ planner=DefaultSavePlanner(),
	39	+ )
	40	+
	41	+ model_2 = FSDP(
	42	+ torch.nn.Linear(8, 8, device="npu"), process_group=process_group
	43	+ )
	44	+
	45	+ with FSDP.summon_full_params(model):
	46	+ with FSDP.summon_full_params(model_2):
	47	+ self.assertNotEqual(model.weight, model_2.weight)
	48	+ self.assertNotEqual(model.bias, model_2.bias)
	49	+
	50	+ # now load the model and ensure the values are the same
	51	+ with FSDP.state_dict_type(model_2, StateDictType.SHARDED_STATE_DICT):
	52	+ state_dict = {
	53	+ "model": model_2.state_dict(),
	54	+ }
	55	+
	56	+ dist_cp.load_state_dict(
	57	+ state_dict=state_dict,
	58	+ storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
	59	+ planner=DefaultLoadPlanner(),
	60	+ )
	61	+ model_2.load_state_dict(state_dict["model"])
	62	+
	63	+ with FSDP.summon_full_params(model):
	64	+ with FSDP.summon_full_params(model_2):
	65	+ self.assertEqual(model.weight, model_2.weight)
	66	+ self.assertEqual(model.bias, model_2.bias)
	67	+
	68	+ @with_comms
	69	+ @skipIfUnsupportMultiNPU(4)
	70	+ @with_temp_dir
	71	+ def test_fsdp_model_state_no_resharding(self):
	72	+ self._test_fsdp_model_state(process_group=None)
	73	+
	74	+ def _create_new_dist_group(self):
	75	+ world_size = dist.get_world_size()
	76	+ group1 = [i for i in range(world_size) if i % 2 == 0]
	77	+ group2 = [i for i in range(world_size) if i % 2 != 0]
	78	+
	79	+ # create new fsdp group for resharding
	80	+ fsdp_0 = dist.new_group(ranks=group1)
	81	+ fsdp_1 = dist.new_group(ranks=group2)
	82	+ if dist.get_rank() % 2 == 0:
	83	+ my_fsdp = fsdp_0
	84	+ else:
	85	+ my_fsdp = fsdp_1
	86	+
	87	+ return my_fsdp
	88	+
	89	+ @with_comms
	90	+ @skipIfUnsupportMultiNPU(4)
	91	+ @with_temp_dir
	92	+ def test_fsdp_model_state_with_resharding(self):
	93	+ self._test_fsdp_model_state(process_group=self._create_new_dist_group())
	94	+
	95	+
	96	+if __name__ == "__main__":
	97	+ run_tests()

Atest/distributed/checkpoint/test_fsdp_optim_state.py+102-0

@@ -0,0 +1,102 @@
	1	+import torch
	2	+
	3	+import torch.distributed.checkpoint as dist_cp
	4	+import torch.distributed as dist
	5	+
	6	+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
	7	+from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
	8	+from torch.distributed.checkpoint.default_planner import (
	9	+ DefaultSavePlanner,
	10	+ DefaultLoadPlanner,
	11	+)
	12	+
	13	+from torch.testing._internal.distributed._tensor.common_dtensor import (
	14	+ DTensorTestBase,
	15	+)
	16	+from torch.testing._internal.distributed.checkpoint_utils import with_temp_dir
	17	+
	18	+import torch_npu
	19	+from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU
	20	+from torch_npu.testing.testcase import run_tests
	21	+
	22	+
	23	+class FsdpOptimStateCheckpoint(DTensorTestBase):
	24	+ @with_comms
	25	+ @skipIfUnsupportMultiNPU(4)
	26	+ @with_temp_dir
	27	+ def test_distributed_tensor_planner(self) -> None:
	28	+ CHECKPOINT_DIR = self.temp_dir
	29	+
	30	+ model = FSDP(torch.nn.Linear(8, 8, device="npu"))
	31	+ optim = torch.optim.Adam(model.parameters(), lr=0.1)
	32	+
	33	+ model(torch.rand(8, 8, device="npu")).sum().backward()
	34	+ optim.step()
	35	+
	36	+ with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
	37	+ state_dict = {
	38	+ "model": model.state_dict(),
	39	+ "optim": FSDP.optim_state_dict(model, optim),
	40	+ }
	41	+
	42	+ dist_cp.save_state_dict(
	43	+ state_dict=state_dict,
	44	+ storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
	45	+ planner=DefaultSavePlanner(),
	46	+ )
	47	+
	48	+ # now load the model and ensure the values are the same
	49	+ model_2 = FSDP(torch.nn.Linear(8, 8, device="npu"))
	50	+ optim_2 = torch.optim.Adam(model_2.parameters(), lr=0.1)
	51	+
	52	+ with FSDP.summon_full_params(model):
	53	+ with FSDP.summon_full_params(model_2):
	54	+ self.assertNotEqual(model.weight, model_2.weight)
	55	+ self.assertNotEqual(model.bias, model_2.bias)
	56	+
	57	+ # Adam lazily creates its state
	58	+ self.assertEqual(0, len(optim_2.state))
	59	+
	60	+ with FSDP.state_dict_type(model_2, StateDictType.SHARDED_STATE_DICT):
	61	+ state_dict = {
	62	+ "model": model_2.state_dict(),
	63	+ # cannot load the optimizer together with the model
	64	+ }
	65	+
	66	+ dist_cp.load_state_dict(
	67	+ state_dict=state_dict,
	68	+ storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
	69	+ planner=DefaultLoadPlanner(),
	70	+ )
	71	+ model_2.load_state_dict(state_dict["model"])
	72	+
	73	+ optim_state = torch.distributed.checkpoint.optimizer.load_sharded_optimizer_state_dict(
	74	+ model_state_dict=state_dict["model"],
	75	+ optimizer_key="optim",
	76	+ storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
	77	+ )
	78	+
	79	+ flattened_osd = FSDP.optim_state_dict_to_load(
	80	+ model_2, optim_2, optim_state["optim"]
	81	+ )
	82	+ optim_2.load_state_dict(flattened_osd)
	83	+
	84	+ with FSDP.summon_full_params(model):
	85	+ with FSDP.summon_full_params(model_2):
	86	+ self.assertEqual(model.weight, model_2.weight)
	87	+ self.assertEqual(model.bias, model_2.bias)
	88	+
	89	+ def opt_at(opt, idx):
	90	+ return list(iter(opt.state.values()))[idx]
	91	+
	92	+ # Adam lazily creates its state
	93	+ self.assertEqual(
	94	+ opt_at(optim, 0)["exp_avg"], opt_at(optim_2, 0)["exp_avg"]
	95	+ )
	96	+ self.assertEqual(
	97	+ opt_at(optim, 0)["exp_avg_sq"], opt_at(optim_2, 0)["exp_avg_sq"]
	98	+ )
	99	+
	100	+
	101	+if __name__ == "__main__":
	102	+ run_tests()

Mtorch_npu/csrc/aten/common/ToKernelNpu.cpp+38-31

	@@ -59,46 +59,53 @@ static inline at::Tensor to_impl_npu(
59	}	59	}
60		60
61	at::Tensor NPUNativeFunctions::to(	61	at::Tensor NPUNativeFunctions::to(
62	- const at::Tensor& self,	62	+ const at::Tensor &self,
63	c10::optional<at::ScalarType> dtype,	63	c10::optional<at::ScalarType> dtype,
64	c10::optional<c10::Layout> layout,	64	c10::optional<c10::Layout> layout,
65	c10::optional<c10::Device> device,	65	c10::optional<c10::Device> device,
66	c10::optional<bool> pin_memory,	66	c10::optional<bool> pin_memory,
67	bool non_blocking,	67	bool non_blocking,
68	bool copy,	68	bool copy,
69	- c10::optional<c10::MemoryFormat> optional_memory_format) {	69	+ c10::optional<c10::MemoryFormat> optional_memory_format)
70	- TORCH_CHECK(	70	+{
71	- !optional_memory_format.has_value(),	71	+ if (device.has_value() && device.value().is_cpu() && optional_memory_format.has_value()) {
72	- "NPU not support specify memory_format.");	72	+ TORCH_CHECK(
73	- c10::TensorOptions options_ = c10::TensorOptions().dtype(dtype)	73	+ optional_memory_format.value() == c10::MemoryFormat::Preserve \|\|
74	- .layout(layout)	74	+ optional_memory_format.value() == c10::MemoryFormat::Contiguous,
75	- .device(device);	75	+ "Only contiguous_format or preserve_format is supported.");
76	- TORCH_CHECK(	76	+ } else {
77	- !(options_.has_memory_format() && optional_memory_format.has_value()),	77	+ TORCH_CHECK(
78	- "Cannot set memory_format both in c10::TensorOptions and explicit argument; please delete "	78	+ !optional_memory_format.has_value(),
79	- "the redundant setter.");	79	+ "NPU not support specify memory_format.");
80	- auto options =	80	+ }
81	- options_.merge_in(c10::TensorOptions().memory_format(optional_memory_format));
82		81
83	- TORCH_CHECK(	82	+ c10::TensorOptions options_ = c10::TensorOptions().dtype(dtype).layout(layout).device(device);
84	- options.requires_grad_opt() == c10::nullopt,	83	+ TORCH_CHECK(
85	- "to(options) expects unset requires_grad flag, but got "	84	+ !(options_.has_memory_format() && optional_memory_format.has_value()),
86	- "options.requires_grad set as ",	85	+ "Cannot set memory_format both in c10::TensorOptions and explicit argument; please delete "
87	- options.requires_grad());	86	+ "the redundant setter.");
		87	+ auto options =
		88	+ options_.merge_in(c10::TensorOptions().memory_format(optional_memory_format));
88		89
89	- TORCH_CHECK(	90	+ TORCH_CHECK(
90	- !options.has_layout() \|\| self.layout() == options.layout(),	91	+ options.requires_grad_opt() == c10::nullopt,
91	- "to(options) doesn't support converting to a different layout, "	92	+ "to(options) expects unset requires_grad flag, but got "
92	- "but got self.layout being ",	93	+ "options.requires_grad set as ",
93	- self.layout(),	94	+ options.requires_grad());
94	- " and options.layout set as ",
95	- options.layout());
96		95
97	- if (options.has_device()) {	96	+ TORCH_CHECK(
98	- options = options.device(ensure_has_index(options.device()));	97	+ !options.has_layout() \|\| self.layout() == options.layout(),
99	- }	98	+ "to(options) doesn't support converting to a different layout, "
100	- auto specified_options = self.options().merge_in(options);	99	+ "but got self.layout being ",
101	- return to_impl_npu(self, specified_options, non_blocking, copy);	100	+ self.layout(),
		101	+ " and options.layout set as ",
		102	+ options.layout());
		103	+
		104	+ if (options.has_device()) {
		105	+ options = options.device(ensure_has_index(options.device()));
		106	+ }
		107	+ auto specified_options = self.options().merge_in(options);
		108	+ return to_impl_npu(self, specified_options, non_blocking, copy);
102	}	109	}
103		110
104	at::Tensor NPUNativeFunctions::to(	111	at::Tensor NPUNativeFunctions::to(

Mtorch_npu/testing/common_distributed.py+1-0

	@@ -76,6 +76,7 @@ def init_pg(backend: str = "hccl", world_size=1, rank=0, file_name="file://") ->
76	raise RuntimeError(f"Backend {backend} not supported!")	76	raise RuntimeError(f"Backend {backend} not supported!")
77		77
78	dist.init_process_group(	78	dist.init_process_group(
		79	+ backend=backend,
79	world_size=world_size,	80	world_size=world_size,
80	rank=rank, # pyre-ignore[16]	81	rank=rank, # pyre-ignore[16]
81	init_method=f"file://{file_name}", # pyre-ignore[16]	82	init_method=f"file://{file_name}", # pyre-ignore[16]

Mtorch_npu/utils/storage.py+1-0

	@@ -58,3 +58,4 @@ def add_storage_methods():
58	torch.storage.UntypedStorage.cpu = _cpu	58	torch.storage.UntypedStorage.cpu = _cpu
59	torch.storage.TypedStorage._deepcopy = _deepcopy	59	torch.storage.TypedStorage._deepcopy = _deepcopy
60	torch.storage.TypedStorage.resize_ = _resize	60	torch.storage.TypedStorage.resize_ = _resize
		61	+ torch.storage.TypedStorage._resize_ = _resize