test: verify community 2.10.0 features and fixes-pytorch-AtomGit

已合并

test: verify community 2.10.0 features and fixes #36918

kuhn7创建于 5月27日从kuhn7/pytorch： v2.12.0_distributed_testcase_verify 合入到 Ascend/pytorch： v2.12.0

test: verify community 2.10.0 features and fixes #36918

已合并

kuhn7创建于 5月27日

从kuhn7/pytorch v2.12.0_distributed_testcase_verify合入到Ascend/pytorch v2.12.0

共 4 个文件变更+1020-40

Atest/distributed/_composable/fsdp/test_fully_shard_training.py+146-0

@@ -0,0 +1,146 @@
	1	+# Owner(s): ["oncall: distributed"]
	2	+
	3	+import copy
	4	+from collections.abc import Callable
	5	+from typing import Optional
	6	+
	7	+import torch
	8	+import torch.distributed as dist
	9	+import torch.nn as nn
	10	+
	11	+from torch.distributed.fsdp import (
	12	+ fully_shard,
	13	+ share_comm_ctx,
	14	+)
	15	+from torch.distributed.fsdp._fully_shard._fsdp_collectives import (
	16	+ foreach_all_gather,
	17	+ foreach_reduce,
	18	+)
	19	+from torch.testing._internal.common_fsdp import (
	20	+ check_sharded_parity,
	21	+ MLP,
	22	+ patch_foreach_all_gather,
	23	+ patch_foreach_reduce,
	24	+)
	25	+from torch.testing._internal.common_fsdp import get_devtype
	26	+from torch.testing._internal.common_utils import run_tests
	27	+
	28	+from torch_npu.testing._internal.common_fsdp import FSDPNPUTest
	29	+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
	30	+
	31	+
	32	+device_type = torch.device(get_devtype())
	33	+
	34	+
	35	+class TestFullyShardShareCommContext(FSDPNPUTest):
	36	+ @property
	37	+ def world_size(self) -> int:
	38	+ return min(torch.get_device_module(device_type).device_count(), 2)
	39	+
	40	+ @skipIfUnsupportMultiNPU(2)
	41	+ def test_share_comm_context(self):
	42	+ torch.manual_seed(42)
	43	+ n_layers = 3
	44	+ lin_dim = 16
	45	+ model = nn.Sequential(
	46	+ *[MLP(lin_dim, torch.device("cpu")) for _ in range(n_layers)]
	47	+ )
	48	+ ref_model = copy.deepcopy(model).to(device_type)
	49	+ for layer in model:
	50	+ fully_shard(layer)
	51	+ layer._get_fsdp_state()._lazy_init()
	52	+ share_comm_ctx(list(model))
	53	+
	54	+ torch.manual_seed(42 + self.rank + 1)
	55	+ inp = torch.randn(4, 3, lin_dim, device=device_type.type)
	56	+ ref_loss = ref_model(inp).sum()
	57	+
	58	+ all_gather_streams = set()
	59	+ reduce_scatter_streams = set()
	60	+
	61	+ from torch.distributed.fsdp._fully_shard._fsdp_api import (
	62	+ AllGather,
	63	+ ReduceScatter,
	64	+ )
	65	+ from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
	66	+
	67	+ orig_foreach_all_gather = foreach_all_gather
	68	+
	69	+ def foreach_all_gather_with_assert(
	70	+ fsdp_params: list[FSDPParam],
	71	+ group: dist.ProcessGroup,
	72	+ async_op: bool,
	73	+ all_gather_copy_in_stream: torch.Stream,
	74	+ all_gather_stream: torch.Stream,
	75	+ device: torch.device,
	76	+ all_gather_comm: AllGather,
	77	+ ):
	78	+ nonlocal all_gather_streams
	79	+ all_gather_streams.add(all_gather_stream)
	80	+ return orig_foreach_all_gather(
	81	+ fsdp_params,
	82	+ group,
	83	+ async_op,
	84	+ all_gather_copy_in_stream,
	85	+ all_gather_stream,
	86	+ device,
	87	+ all_gather_comm,
	88	+ )
	89	+
	90	+ orig_foreach_reduce = foreach_reduce
	91	+
	92	+ @torch.no_grad()
	93	+ def foreach_reduce_with_assert(
	94	+ fsdp_params: list[FSDPParam],
	95	+ unsharded_grads: list[torch.Tensor],
	96	+ reduce_scatter_group: dist.ProcessGroup,
	97	+ reduce_scatter_stream: torch.Stream,
	98	+ reduce_scatter_comm: ReduceScatter,
	99	+ orig_dtype: Optional[torch.dtype],
	100	+ reduce_dtype: Optional[torch.dtype],
	101	+ device: torch.device,
	102	+ gradient_divide_factor: Optional[float],
	103	+ all_reduce_group: Optional[dist.ProcessGroup], # not `None` iff HSDP
	104	+ all_reduce_stream: torch.Stream,
	105	+ all_reduce_grads: bool,
	106	+ partial_reduce_output: Optional[torch.Tensor], # only used for HSDP
	107	+ all_reduce_hook: Optional[Callable[[torch.Tensor], None]],
	108	+ force_sum_reduction_for_comms: bool = False,
	109	+ ):
	110	+ nonlocal reduce_scatter_streams
	111	+ reduce_scatter_streams.add(reduce_scatter_stream)
	112	+ return orig_foreach_reduce(
	113	+ fsdp_params,
	114	+ unsharded_grads,
	115	+ reduce_scatter_group,
	116	+ reduce_scatter_stream,
	117	+ reduce_scatter_comm,
	118	+ orig_dtype,
	119	+ reduce_dtype,
	120	+ device,
	121	+ gradient_divide_factor,
	122	+ all_reduce_group,
	123	+ all_reduce_stream,
	124	+ all_reduce_grads,
	125	+ partial_reduce_output,
	126	+ all_reduce_hook,
	127	+ force_sum_reduction_for_comms,
	128	+ )
	129	+
	130	+ with (
	131	+ patch_foreach_all_gather(foreach_all_gather_with_assert),
	132	+ patch_foreach_reduce(foreach_reduce_with_assert),
	133	+ ):
	134	+ loss = model(inp).sum()
	135	+ self.assertEqual(ref_loss, loss)
	136	+ ref_loss.backward()
	137	+ loss.backward()
	138	+ for param in ref_model.parameters():
	139	+ dist.all_reduce(param.grad, op=dist.ReduceOp.AVG)
	140	+ self.assertEqual(len(all_gather_streams), 1)
	141	+ self.assertEqual(len(reduce_scatter_streams), 1)
	142	+ check_sharded_parity(self, ref_model, model)
	143	+
	144	+
	145	+if __name__ == "__main__":
	146	+ run_tests()

Mtest/distributed/_tensor/test_utils.py+131-36

	@@ -1,14 +1,24 @@
1	import itertools	1	import itertools
		2	+from contextlib import nullcontext
		3	+
2	import torch	4	import torch
		5	+import torch.distributed as dist
		6	+from torch.distributed._local_tensor import (
		7	+ local_tensor_mode,
		8	+ LocalTensor,
		9	+ LocalTensorMode,
		10	+)
3	from torch.distributed._tensor import distribute_tensor	11	from torch.distributed._tensor import distribute_tensor
4	-from torch.distributed._tensor._utils import (	12	+from torch.distributed.tensor._utils import (
5	- compute_local_shape,
6	compute_local_shape_and_global_offset,	13	compute_local_shape_and_global_offset,
7	)	14	)
8	-from torch.distributed._tensor.device_mesh import DeviceMesh	15	+from torch.distributed.device_mesh import init_device_mesh
9	-from torch.distributed._tensor.placement_types import Replicate, Shard	16	+from torch.distributed.tensor import DeviceMesh
		17	+from torch.distributed.tensor._utils import ExplicitRedistributionContext
		18	+from torch.distributed.tensor.debug import CommDebugMode
		19	+from torch.distributed.tensor.placement_types import Replicate, Shard
10		20
11	-from torch.testing._internal.common_utils import run_tests	21	+from torch.testing._internal.common_utils import run_tests, TestCase
12	from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase	22	from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase
13		23
14	import torch_npu	24	import torch_npu
	@@ -20,37 +30,6 @@ class UtilTest(DTensorTestBase):
20	def world_size(self):	30	def world_size(self):
21	return 8	31	return 8
22		32
23	- @skipIfUnsupportMultiNPU(4)
24	- @with_comms
25	- def test_compute_local_shape_2d_uneven(self):
26	- # mesh: 4 * 2
27	- mesh_tensor = torch.arange(self.world_size).reshape(4, 2)
28	- mesh = DeviceMesh(self.device_type, mesh_tensor)
29	- size = torch.Size([7, 7])
30	- rank_coordinates = mesh.get_coordinate()
31	-
32	- # replicate, shard
33	- placements2 = [Replicate(), Shard(0)]
34	- local_size2 = compute_local_shape(size, mesh, placements2)
35	- if rank_coordinates[1] < 1:
36	- self.assertEqual(local_size2, torch.Size([4, 7]))
37	- else:
38	- self.assertEqual(local_size2, torch.Size([3, 7]))
39	-
40	- # shard, shard
41	- placements3 = [Shard(0), Shard(1)]
42	- local_size3 = compute_local_shape(size, mesh, placements3)
43	- # first dim
44	- if rank_coordinates[0] < 3:
45	- self.assertEqual(local_size3[0], 2)
46	- else:
47	- self.assertEqual(local_size3[0], 1)
48	- # second dim
49	- if rank_coordinates[1] < 1:
50	- self.assertEqual(local_size3[1], 4)
51	- else:
52	- self.assertEqual(local_size3[1], 3)
53	-
54	@skipIfUnsupportMultiNPU(4)	33	@skipIfUnsupportMultiNPU(4)
55	@with_comms	34	@with_comms
56	def test_compute_local_shape_and_global_offset_1D(self):	35	def test_compute_local_shape_and_global_offset_1D(self):
	@@ -111,5 +90,121 @@ class UtilTest(DTensorTestBase):
111	)	90	)
112		91
113		92
		93	+class LocalTensorTestBase(TestCase):
		94	+ def assertEqual(self, lhs, rhs, **kwargs):
		95	+ mode = local_tensor_mode()
		96	+ with nullcontext() if mode is None else mode.disable():
		97	+ if isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor):
		98	+ assert isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)
		99	+ super().assertEqual(lhs._ranks, rhs._ranks)
		100	+ for r in lhs._ranks:
		101	+ super().assertEqual(
		102	+ lhs._local_tensors[r],
		103	+ rhs._local_tensors[r],
		104	+ lambda m: f"rank {r}: {m}",
		105	+ )
		106	+ elif isinstance(lhs, LocalTensor) or isinstance(rhs, LocalTensor):
		107	+ lhs, rhs = (lhs, rhs) if isinstance(lhs, LocalTensor) else (rhs, lhs)
		108	+ for r in lhs._ranks:
		109	+ super().assertEqual(
		110	+ lhs._local_tensors[r], rhs, lambda m: f"rank {r}: {m}"
		111	+ )
		112	+ else:
		113	+ return super().assertEqual(lhs, rhs, **kwargs)
		114	+
		115	+ @property
		116	+ def world_size(self):
		117	+ raise NotImplementedError("override world-size in your subclass")
		118	+
		119	+ def build_device_mesh(self) -> DeviceMesh:
		120	+ return init_device_mesh("cpu", (self.world_size,))
		121	+
		122	+ def setUp(self):
		123	+ super().setUp()
		124	+ torch.distributed.init_process_group(
		125	+ # TODO: test other ranks too
		126	+ "fake",
		127	+ rank=0,
		128	+ world_size=self.world_size,
		129	+ )
		130	+
		131	+ def tearDown(self):
		132	+ super().tearDown()
		133	+ try:
		134	+ dist.destroy_process_group()
		135	+ except AssertionError:
		136	+ pass
		137	+
		138	+
		139	+class TestExplicitRedistribute(LocalTensorTestBase):
		140	+ @property
		141	+ def world_size(self):
		142	+ return 4
		143	+
		144	+ def test_explicit_matmul(self):
		145	+ with LocalTensorMode(self.world_size):
		146	+ device_mesh = self.build_device_mesh()
		147	+ dim = 128
		148	+ x = torch.randn(8, dim, requires_grad=True)
		149	+ A = torch.randn(dim, dim, requires_grad=True)
		150	+
		151	+ # Prepare DTensors
		152	+ dx = distribute_tensor(x, device_mesh, [Shard(0)])
		153	+ dA = distribute_tensor(A, device_mesh, [Shard(0)])
		154	+
		155	+ # implicit redistribute works as usual by default
		156	+ with CommDebugMode() as comm_mode:
		157	+ torch.matmul(dx, dA)
		158	+ self.assertEqual(comm_mode.get_total_counts(), 1)
		159	+
		160	+ # explicit redistribute works too
		161	+ with ExplicitRedistributionContext():
		162	+ with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"):
		163	+ torch.matmul(dx, dA)
		164	+ with ExplicitRedistributionContext(mode="warn"):
		165	+ with self.assertLogs(
		166	+ torch.distributed.tensor._utils.logger, level="WARN"
		167	+ ) as captured:
		168	+ torch.matmul(dx, dA)
		169	+ self.assertEqual(len(captured.output), 1)
		170	+ self.assertRegex(
		171	+ captured.output[0],
		172	+ r"WARNING:.*Implicit redistribution occurred",
		173	+ )
		174	+ # TODO enable this once fixing the issue that op_info.schema is None in some calls to
		175	+ # redistribute_local_tensor
		176	+ # self.assertRegex(
		177	+ # captured.output[0],
		178	+ # r".aten\.mm\.default.",
		179	+ # )
		180	+
		181	+ # explicit redistribute allows manual redistribute
		182	+ with ExplicitRedistributionContext():
		183	+ dA_repl = dA.redistribute(device_mesh, [Replicate()])
		184	+ torch.matmul(dx, dA_repl)
		185	+
		186	+ dx = distribute_tensor(x, device_mesh, [Shard(0)])
		187	+ dA = distribute_tensor(A, device_mesh, [Replicate()])
		188	+ with ExplicitRedistributionContext(strict=True):
		189	+ dY = torch.matmul(dx, dA_repl)
		190	+ loss = dY.sum()
		191	+
		192	+ # we now see the error during backwards
		193	+ with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"):
		194	+ loss.backward(retain_graph=True)
		195	+
		196	+ with ExplicitRedistributionContext(strict=False):
		197	+ # but since it's a 'free' redistribute, we can still do it under non-strict mode
		198	+ loss.backward(retain_graph=True)
		199	+
		200	+ with ExplicitRedistributionContext(enable=False):
		201	+ # and we can disable
		202	+ loss.backward(retain_graph=True)
		203	+
		204	+ # and re-enable
		205	+ with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"):
		206	+ loss.backward(retain_graph=True)
		207	+
		208	+
114	if __name__ == "__main__":	209	if __name__ == "__main__":
115	run_tests()	210	run_tests()

Atest/distributed/test_local_tensor.py+735-0

@@ -0,0 +1,735 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+
	4	+from contextlib import nullcontext
	5	+
	6	+import torch
	7	+import torch.distributed as dist
	8	+from torch.distributed._local_tensor import (
	9	+ local_tensor_mode,
	10	+ LocalIntNode,
	11	+ LocalRunnerMode,
	12	+ LocalTensor,
	13	+ LocalTensorMode,
	14	+ maybe_disable_local_tensor_mode,
	15	+)
	16	+from torch.distributed.tensor import (
	17	+ DeviceMesh,
	18	+ distribute_tensor,
	19	+ init_device_mesh,
	20	+ Partial,
	21	+ Replicate,
	22	+ Shard,
	23	+ zeros,
	24	+)
	25	+from torch.testing._internal.common_utils import run_tests, TestCase
	26	+from torch.testing._internal.distributed._tensor.common_dtensor import reduce_local_int
	27	+
	28	+
	29	+class LocalTensorTestBase(TestCase):
	30	+ def assertEqual(self, lhs, rhs, **kwargs):
	31	+ mode = local_tensor_mode()
	32	+ with nullcontext() if mode is None else mode.disable():
	33	+ if isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor):
	34	+ assert isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)
	35	+ super().assertEqual(lhs._ranks, rhs._ranks)
	36	+ for r in lhs._ranks:
	37	+ super().assertEqual(
	38	+ lhs._local_tensors[r],
	39	+ rhs._local_tensors[r],
	40	+ lambda m: f"rank {r}: {m}",
	41	+ )
	42	+ elif isinstance(lhs, LocalTensor) or isinstance(rhs, LocalTensor):
	43	+ lhs, rhs = (lhs, rhs) if isinstance(lhs, LocalTensor) else (rhs, lhs)
	44	+ for r in lhs._ranks:
	45	+ super().assertEqual(
	46	+ lhs._local_tensors[r], rhs, lambda m: f"rank {r}: {m}"
	47	+ )
	48	+ else:
	49	+ return super().assertEqual(lhs, rhs, **kwargs)
	50	+
	51	+ @property
	52	+ def world_size(self):
	53	+ raise NotImplementedError("override world-size in your subclass")
	54	+
	55	+ def build_device_mesh(self) -> DeviceMesh:
	56	+ return init_device_mesh("cpu", (self.world_size,))
	57	+
	58	+
	59	+class LocalTensorRankTest(LocalTensorTestBase):
	60	+ def setUp(self):
	61	+ super().setUp()
	62	+
	63	+ def tearDown(self):
	64	+ super().tearDown()
	65	+ if dist.is_initialized():
	66	+ dist.destroy_process_group()
	67	+
	68	+ def run(self, result=None):
	69	+ # save the original test method
	70	+ test_name = self.id().split(".")[-1]
	71	+ original_test = getattr(self, test_name)
	72	+
	73	+ # replace the original test with new test that loops ranks
	74	+ def rank_loop_wrapper():
	75	+ for rank in range(self.world_size):
	76	+ if dist.is_initialized():
	77	+ dist.destroy_process_group()
	78	+ torch.distributed.init_process_group(
	79	+ "fake", rank=rank, world_size=self.world_size
	80	+ )
	81	+ original_test()
	82	+
	83	+ setattr(self, test_name, rank_loop_wrapper)
	84	+ return super().run(result)
	85	+
	86	+ @property
	87	+ def rank(self):
	88	+ assert dist.is_initialized(), "Process group is not initialized!"
	89	+ return dist.get_rank()
	90	+
	91	+
	92	+class LocalTensorWorldTest(LocalTensorTestBase):
	93	+ def setUp(self):
	94	+ super().setUp()
	95	+ torch.distributed.init_process_group(
	96	+ "fake",
	97	+ rank=0,
	98	+ world_size=self.world_size,
	99	+ )
	100	+
	101	+ def tearDown(self):
	102	+ super().tearDown()
	103	+ try:
	104	+ dist.destroy_process_group()
	105	+ except AssertionError:
	106	+ pass
	107	+
	108	+
	109	+class TestLocalTensorWorld2(LocalTensorWorldTest):
	110	+ world_size = 2
	111	+
	112	+ def test_local_tensor_dtype_consistency(self):
	113	+ """Test that LocalTensor enforces dtype consistency."""
	114	+ device = torch.device("cpu")
	115	+ shape = (2, 3)
	116	+
	117	+ inconsistent_tensors = {
	118	+ 0: torch.randn(shape, dtype=torch.float32, device=device),
	119	+ 1: torch.randn(
	120	+ shape, dtype=torch.float64, device=device
	121	+ ), # Different dtype
	122	+ }
	123	+
	124	+ with self.assertRaises(AssertionError):
	125	+ LocalTensor(inconsistent_tensors)
	126	+
	127	+ def test_local_tensor_creation_fails_with_grad_tensors(self):
	128	+ """Test that LocalTensor creation fails when local tensors have requires_grad=True."""
	129	+ device = torch.device("cpu")
	130	+ shape = (2, 3)
	131	+ dtype = torch.float32
	132	+
	133	+ # Create sample local tensors for different ranks
	134	+ local_tensors = {
	135	+ 0: torch.randn(shape, dtype=dtype, device=device, requires_grad=True),
	136	+ 1: torch.randn(shape, dtype=dtype, device=device, requires_grad=True),
	137	+ }
	138	+
	139	+ with self.assertRaises(AssertionError):
	140	+ LocalTensor(local_tensors)
	141	+
	142	+ def test_local_tensor_mode(self):
	143	+ """Test LocalTensorMode functionality."""
	144	+ device = torch.device("cpu")
	145	+ shape = (2, 3)
	146	+ dtype = torch.float32
	147	+
	148	+ # Create identical local tensors for consistency tests
	149	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	150	+ identical_local_tensors = {
	151	+ 0: base_tensor.clone(),
	152	+ 1: base_tensor.clone(),
	153	+ }
	154	+
	155	+ lt = LocalTensor(identical_local_tensors)
	156	+
	157	+ with LocalTensorMode(lt._ranks):
	158	+ result = lt + 1.0
	159	+ self.assertIsInstance(result, LocalTensor)
	160	+
	161	+ regular = torch.ones(2, 2)
	162	+ regular_result = regular + 1.0
	163	+ self.assertIsInstance(regular, LocalTensor)
	164	+ self.assertIsInstance(regular_result, LocalTensor)
	165	+
	166	+ def test_empty_local_tensors(self):
	167	+ """Test behavior with empty local tensors dict."""
	168	+ with self.assertRaises(ValueError):
	169	+ LocalTensor({})
	170	+
	171	+ def test_scalar_mul_reduction_bug(self):
	172	+ with LocalTensorMode(self.world_size):
	173	+ mesh = self.build_device_mesh()
	174	+
	175	+ tensor = torch.tensor([10, 10]).float()
	176	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	177	+ y = dt.sum() * 1 # noqa: F841
	178	+
	179	+ tensor = torch.arange(10).reshape(10, 1).float().requires_grad_()
	180	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	OO O openLiBingCI5月27日此条代码评论区间+175至+180 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI5月27日此条代码评论区间+175至+180 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	181	+
	182	+ print(dt.sum() * 1, dt.sum() * 2, dt.sum() * 3)
	183	+
	184	+ def test_uneven_sharding_mean_bug(self):
	185	+ with LocalTensorMode(self.world_size):
	186	+ mesh = self.build_device_mesh()
	187	+ tensor = torch.arange(12).reshape(-1, 4).float()
	188	+
	189	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	190	+
	191	+ mean = dt.mean()
	192	+ self.assertEqual(mean.placements, [Replicate()])
	193	+ full = mean.full_tensor()
	194	+ self.assertEqual(tensor.mean(), full)
	195	+
	196	+ def test_uneven_sharding_prod(self):
	197	+ with LocalTensorMode(self.world_size):
	198	+ mesh = self.build_device_mesh()
	199	+ tensor = (torch.arange(12) + 1).reshape(-1, 4).float()
	200	+
	201	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	202	+
	203	+ x = dt.prod()
	204	+ full = x.full_tensor()
	205	+ self.assertEqual(tensor.prod(), full)
	206	+
	207	+ def test_even_sharding_mean_is_partial(self):
	208	+ with LocalTensorMode(self.world_size):
	209	+ mesh = self.build_device_mesh()
	210	+ tensor = torch.arange(16).reshape(4, 4).float()
	211	+
	212	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	213	+
	214	+ mean = dt.mean()
	215	+ full = mean.full_tensor()
	216	+ self.assertEqual(tensor.mean(), full)
	217	+ self.assertEqual(mean.placements, [Partial("avg")])
	218	+
	219	+ def test_local_int_node_sym_min(self):
	220	+ node1 = LocalIntNode({0: 10, 1: 5})
	221	+ node2 = LocalIntNode({0: 7, 1: 8})
	222	+ result = node1.sym_min(node2)
	223	+ self.assertEqual(result._local_ints[0], 7)
	224	+ self.assertEqual(result._local_ints[1], 5)
	225	+
	226	+ def test_local_int_node_le(self):
	227	+ node1 = LocalIntNode({0: 3, 1: 4})
	228	+ node2 = LocalIntNode({0: 10, 1: 10})
	229	+ result = node1.le(node2)
	230	+ self.assertTrue(bool(result))
	231	+
	232	+ def test_data_ptr_raises(self):
	233	+ """data_ptr() on a LocalTensor should raise instead of returning 0."""
	234	+ local_tensors = {
	235	+ 0: torch.randn(4),
	236	+ 1: torch.randn(4),
	237	+ }
	238	+ lt = LocalTensor(local_tensors)
	239	+ with self.assertRaises(RuntimeError):
	240	+ lt.data_ptr()
	241	+
	242	+ # Native ops still work (dispatched per-rank via __torch_dispatch__)
	243	+ with LocalTensorMode(lt._ranks):
	244	+ result = lt + lt
	245	+ self.assertIsInstance(result, LocalTensor)
	246	+
	247	+ def test_sym_and_sym_or(self):
	248	+ node1 = LocalIntNode({0: 3, 1: 4})
	249	+ node2 = LocalIntNode({0: 10, 1: 10})
	250	+ sym_true = torch.SymBool(node1.le(node2))
	251	+ self.assertTrue(bool(sym_true & True))
	252	+ self.assertFalse(bool(sym_true & False))
	253	+ self.assertTrue(bool(sym_true \| False))
	254	+
	255	+
	256	+class TestLocalTensorRankWorld2(LocalTensorRankTest):
	257	+ world_size = 2
	258	+
	259	+ def test_flatten_unflatten(self):
	260	+ """Test that LocalTensor can be flattened and unflattened correctly."""
	261	+ device = torch.device("cpu")
	262	+ dtype = torch.float32
	263	+ # test samples
	264	+ test_cases = [
	265	+ {
	266	+ i: torch.randn(2, 3, dtype=dtype, device=device)
	267	+ for i in range(self.world_size)
	268	+ },
	269	+ {
	270	+ 0: torch.randn(2, 3, dtype=dtype, device=device),
	271	+ 1: torch.randn(3, 3, dtype=dtype, device=device),
	272	+ },
	273	+ ]
	274	+
	275	+ for local_tensors in test_cases:
	276	+ lt = LocalTensor(local_tensors)
	277	+ attrs, spec = lt.__tensor_flatten__()
	278	+ inner_tensors = {attr: getattr(lt, attr) for attr in attrs}
	279	+ lt_reconstruct = LocalTensor.__tensor_unflatten__(
	280	+ inner_tensors, spec, lt.size(), lt.stride()
	281	+ )
	282	+
	283	+ self.assertEqual(
	284	+ lt._local_tensors[self.rank], lt_reconstruct._local_tensors[self.rank]
	285	+ )
	286	+
	287	+ def test_basic_arithmetic_operations(self):
	288	+ """Test basic arithmetic operations on LocalTensors."""
	289	+ device = torch.device("cpu")
	290	+ shape = (2, 3)
	291	+ dtype = torch.float32
	292	+
	293	+ # Create identical local tensors for consistency tests
	294	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	295	+ identical_local_tensors = {
	296	+ 0: base_tensor.clone(),
	297	+ 1: base_tensor.clone(),
	298	+ }
	299	+
	300	+ lt1 = LocalTensor(identical_local_tensors)
	301	+ lt2 = LocalTensor(identical_local_tensors)
	302	+
	303	+ # Test addition
	304	+ result_add = lt1 + lt2
	305	+ self.assertIsInstance(result_add, LocalTensor)
	306	+ self.assertEqual(len(result_add._local_tensors), 2)
	307	+
	308	+ # Verify the operation was applied to each local tensor
	309	+ expected = (
	310	+ identical_local_tensors[self.rank] + identical_local_tensors[self.rank]
	311	+ )
	312	+ self.assertEqual(result_add._local_tensors[self.rank], expected)
	313	+
	314	+ # Test multiplication
	315	+ result_mul = lt1 * 2.0
	316	+ self.assertIsInstance(result_mul, LocalTensor)
	317	+ expected = identical_local_tensors[self.rank] * 2.0
	318	+ self.assertEqual(result_mul._local_tensors[self.rank], expected)
	319	+
	320	+ def test_view_ops(self):
	321	+ """Test that view operations work correctly on LocalTensor (standard subclass style)."""
	322	+ device = torch.device("cpu")
	323	+ base_tensor = torch.arange(8, device=device).reshape(2, 4).float()
	324	+ local_tensors = {
	325	+ 0: base_tensor.clone(),
	326	+ 1: base_tensor.clone(),
	327	+ }
	328	+ lt = LocalTensor(local_tensors)
	329	+
	330	+ test_cases = [
	331	+ (torch.flip, (lt,), {"dims": [1]}),
	332	+ (torch.fliplr, (lt,), {}),
	333	+ (torch.flipud, (lt,), {}),
	334	+ (torch.flatten, (lt,), {}),
	335	+ (torch.ravel, (lt,), {}),
	336	+ (torch.reshape, (lt, (8,)), {}),
	337	+ (torch.transpose, (lt, 0, 1), {}),
	338	+ (torch.t, (lt,), {}),
	339	+ (torch.squeeze, (lt,), {}),
	340	+ (torch.unsqueeze, (lt, 0), {}),
	341	+ ]
	342	+
	343	+ for op_func, args, kwargs in test_cases:
	344	+ with self.subTest(op=op_func.__name__):
	345	+ result = op_func(args, *kwargs)
	346	+ self.assertIsInstance(result, LocalTensor)
	347	+
	348	+ ref_args = tuple(
	349	+ local_tensors[self.rank] if a is lt else a for a in args
	350	+ )
	351	+ expected = op_func(ref_args, *kwargs)
	352	+ self.assertEqual(result._local_tensors[self.rank], expected)
	353	+
	354	+ def test_mixed_operations_with_regular_tensors(self):
	355	+ """Test operations between LocalTensors and regular tensors."""
	356	+ device = torch.device("cpu")
	357	+ shape = (2, 3)
	358	+ dtype = torch.float32
	359	+
	360	+ # Create identical local tensors for consistency tests
	361	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	362	+ identical_local_tensors = {
	363	+ 0: base_tensor.clone(),
	364	+ 1: base_tensor.clone(),
	365	+ }
	366	+
	367	+ lt = LocalTensor(identical_local_tensors)
	368	+ regular_tensor = torch.ones_like(identical_local_tensors[0])
	369	+
	370	+ # Test LocalTensor + regular tensor
	371	+ result = lt + regular_tensor
	372	+ self.assertIsInstance(result, LocalTensor)
	373	+
	374	+ expected = identical_local_tensors[self.rank] + regular_tensor
	375	+ self.assertEqual(result._local_tensors[self.rank], expected)
	376	+
	377	+ def test_collectives_within_local_tensor_mode(self):
	378	+ """Test that collective operations work within LocalTensorMode context."""
	379	+ test_tensors = {
	380	+ 0: torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
	381	+ 1: torch.tensor([[5.0, 6.0], [7.0, 8.0]]),
	382	+ }
	383	+ lt = LocalTensor(test_tensors)
	384	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	385	+
	386	+ with LocalTensorMode(lt._ranks):
	387	+ # Test all_reduce within mode
	388	+ lt_sum = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	389	+ dist.all_reduce(lt_sum, group=fake_pg)
	390	+
	391	+ expected_sum = torch.tensor([[6.0, 8.0], [10.0, 12.0]])
	392	+ self.assertEqual(lt_sum._local_tensors[self.rank], expected_sum)
	393	+
	394	+ # Test broadcast within mode
	395	+ lt_broadcast = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	396	+ dist.broadcast(lt_broadcast, src=self.rank, group=fake_pg)
	397	+
	398	+ # test current rank to other ranks.
	399	+ for _rank in test_tensors:
	400	+ if _rank == self.rank:
	401	+ continue
	402	+ self.assertEqual(
	403	+ lt_broadcast._local_tensors[_rank], test_tensors[self.rank]
	404	+ )
	405	+
	406	+ # Test that regular operations still work
	407	+ result = lt + 1.0
	408	+ self.assertIsInstance(result, LocalTensor)
	409	+
	410	+
	411	+class TestLocalTensorRankWorld3(LocalTensorRankTest):
	412	+ world_size = 3
	413	+
	414	+ def test_collective_reduction_operations(self):
	415	+ """Test different reduction operations for all_reduce."""
	416	+ # Create different tensors for each rank with simple values for testing
	417	+ test_tensors = {
	418	+ 0: torch.tensor([[1.0, 4.0], [2.0, 5.0]]),
	419	+ 1: torch.tensor([[2.0, 1.0], [3.0, 6.0]]),
	420	+ 2: torch.tensor([[3.0, 2.0], [1.0, 4.0]]),
	421	+ }
	422	+
	423	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	424	+
	425	+ # Test SUM reduction
	426	+ lt_sum = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	427	+ dist.all_reduce(lt_sum, op=dist.ReduceOp.SUM, group=fake_pg)
	428	+ expected_sum = torch.tensor([[6.0, 7.0], [6.0, 15.0]]) # Sum of all tensors
	429	+ self.assertEqual(lt_sum._local_tensors[self.rank], expected_sum)
	430	+
	431	+ # Test MAX reduction
	432	+ lt_max = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	433	+ dist.all_reduce(lt_max, op=dist.ReduceOp.MAX, group=fake_pg)
	434	+ expected_max = torch.tensor([[3.0, 4.0], [3.0, 6.0]]) # Max across all tensors
	435	+ self.assertEqual(lt_max._local_tensors[self.rank], expected_max)
	436	+
	437	+ # Test MIN reduction
	438	+ lt_min = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	439	+ dist.all_reduce(lt_min, op=dist.ReduceOp.MIN, group=fake_pg)
	440	+ expected_min = torch.tensor([[1.0, 1.0], [1.0, 4.0]]) # Min across all tensors
	441	+ self.assertEqual(lt_min._local_tensors[self.rank], expected_min)
	442	+
	443	+ def test_all_reduce_collective(self):
	444	+ """Test that all_reduce collective operation works correctly with LocalTensor."""
	445	+ # Create different tensors for each rank
	446	+ different_tensors = {
	447	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	448	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	449	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	450	+ }
	451	+
	452	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	453	+
	454	+ # Test all_reduce with SUM (default)
	455	+ lt_sum = LocalTensor({k: v.clone() for k, v in different_tensors.items()})
	456	+ lt_sum = lt_sum + 1
	457	+ dist.all_reduce(lt_sum, group=fake_pg)
	458	+
	459	+ # Verify all ranks have the sum of all tensors (after adding 1 to each)
	460	+ expected_sum = torch.tensor([[114.0, 225.0, 336.0], [447.0, 558.0, 669.0]])
	461	+ self.assertEqual(lt_sum._local_tensors[self.rank], expected_sum)
	462	+
	463	+ def test_broadcast_collective(self):
	464	+ """Test that broadcast collective operation works correctly with LocalTensor."""
	465	+ # Create different tensors for each rank
	466	+ different_tensors = {
	467	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	468	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	469	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	470	+ }
	471	+
	472	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	473	+
	474	+ # Test broadcast from current rank
	475	+ lt_broadcast = LocalTensor({k: v.clone() for k, v in different_tensors.items()})
	476	+ dist.broadcast(lt_broadcast, src=self.rank, group=fake_pg)
	477	+
	478	+ # Verify all other ranks have current rank's original tensor
	479	+ expected_broadcast = different_tensors[self.rank]
	480	+ for rank in different_tensors:
	481	+ if rank == self.rank:
	482	+ continue
	483	+ self.assertEqual(lt_broadcast._local_tensors[rank], expected_broadcast)
	484	+
	485	+ def test_all_gather_collective(self):
	486	+ """Test that all_gather collective operation works correctly with LocalTensor."""
	487	+ # Create different tensors for each rank
	488	+ different_tensors = {
	489	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	490	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	491	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	492	+ }
	493	+
	494	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	495	+
	496	+ # Test all_gather
	497	+ lt_gather = LocalTensor(different_tensors)
	498	+ tensor_list = [torch.zeros_like(lt_gather) for _ in range(3)]
	499	+
	500	+ dist.all_gather(tensor_list, lt_gather, group=fake_pg)
	501	+
	502	+ # Verify each position in tensor_list contains the corresponding rank's tensor
	503	+ self.assertEqual(tensor_list[self.rank], different_tensors[self.rank])
	504	+
	505	+ def test_all_to_all_single_collective(self):
	506	+ """Test that all_to_all_single collective operation works correctly with LocalTensor."""
	507	+ from torch.distributed._functional_collectives import all_to_all_single
	508	+
	509	+ # Create different tensors for each rank
	510	+ # Each rank will split its tensor and send parts to other ranks
	511	+ different_tensors = {
	512	+ 0: torch.tensor(
	513	+ [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
	514	+ ), # rank 0 sends [0,0], [0,0], [0,0] to ranks 0,1,2
	515	+ 1: torch.tensor(
	516	+ [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
	517	+ ), # rank 1 sends [1,1], [1,1], [1,1] to ranks 0,1,2
	518	+ 2: torch.tensor(
	519	+ [2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
	520	+ ), # rank 2 sends [2,2], [2,2], [2,2] to ranks 0,1,2
	521	+ }
	522	+
	523	+ # Each rank splits its input into 3 parts of size 2 each
	524	+ input_split_sizes = [2, 2, 2]
	525	+ # Each rank receives 3 parts of size 2 each from all ranks
	526	+ output_split_sizes = [2, 2, 2]
	527	+
	528	+ with LocalTensorMode(self.world_size):
	529	+ lt_input = LocalTensor(different_tensors)
	530	+
	531	+ # Test all_to_all_single using functional collectives API
	532	+ result = all_to_all_single(
	533	+ lt_input,
	534	+ output_split_sizes=output_split_sizes,
	535	+ input_split_sizes=input_split_sizes,
	536	+ group=torch.distributed.distributed_c10d._get_default_group(),
	537	+ )
	538	+
	539	+ result = result.wait()
	540	+ # Verify result is a LocalTensor
	541	+ self.assertIsInstance(result, LocalTensor)
	542	+
	543	+ # After all_to_all_single:
	544	+ # rank 0 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	545	+ # rank 1 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	546	+ # rank 2 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	547	+ expected_output = torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0, 2.0])
	548	+ self.assertEqual(result._local_tensors[self.rank], expected_output)
	549	+
	550	+
	551	+class TestLocalTensorWorld3(LocalTensorWorldTest):
	552	+ world_size = 3
	553	+
	554	+ def test_reduce_scatter_tensor_collective(self):
	555	+ """Test that reduce_scatter_tensor collective operation works correctly with LocalTensor."""
	556	+ # Create different tensors for each rank
	557	+ different_tensors = {
	558	+ 0: torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
	559	+ 1: torch.tensor([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]),
	560	+ 2: torch.tensor([[100.0, 200.0], [300.0, 400.0], [500.0, 600.0]]),
	561	+ }
	562	+
	563	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	564	+
	565	+ # Test reduce_scatter_tensor
	566	+ with LocalTensorMode(self.world_size):
	567	+ lt_reduce_scatter = LocalTensor(different_tensors)
	568	+ lt_reduce_scatter_size = lt_reduce_scatter.size()
	569	+ lt_output_tensor = torch.zeros(
	570	+ lt_reduce_scatter_size[0] // fake_pg.size(),
	571	+ *lt_reduce_scatter_size[1:],
	572	+ dtype=lt_reduce_scatter.dtype,
	573	+ device=lt_reduce_scatter.device,
	574	+ )
	575	+
	576	+ dist.reduce_scatter_tensor(
	577	+ lt_output_tensor, lt_reduce_scatter, group=fake_pg
	578	+ )
	579	+
	580	+ expected_output = LocalTensor(
	581	+ {
	582	+ 0: torch.tensor([[111.0, 222.0]]),
	583	+ 1: torch.tensor([[333.0, 444.0]]),
	584	+ 2: torch.tensor([[555.0, 666.0]]),
	585	+ }
	586	+ )
	587	+ print(lt_output_tensor)
	588	+ self.assertEqual(lt_output_tensor, expected_output)
	589	+
	590	+ def test_all_gather_into_tensor_collective(self):
	591	+ """Test that all_gather_into_tensor collective operation works correctly with LocalTensor."""
	592	+ # Create different tensors for each rank
	593	+ different_tensors = {
	594	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	595	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	596	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	597	+ }
	598	+
	599	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	600	+
	601	+ # Test all_gather_into_tensor
	602	+ with LocalTensorMode(self.world_size):
	603	+ lt_gather = LocalTensor(different_tensors)
	604	+ lt_gather_size = lt_gather.size()
	605	+ lt_output_tensor = torch.zeros(
	606	+ lt_gather_size[0] * fake_pg.size(),
	607	+ *lt_gather_size[1:],
	608	+ dtype=lt_gather.dtype,
	609	+ device=lt_gather.device,
	610	+ )
	611	+
	612	+ dist.all_gather_into_tensor(lt_output_tensor, lt_gather, group=fake_pg)
	613	+
	614	+ expected_output = torch.cat(list(different_tensors.values()))
	615	+
	616	+ self.assertEqual(lt_output_tensor, expected_output)
	617	+
	618	+
	619	+class TestLocalTensorWorld4(LocalTensorWorldTest):
	620	+ world_size = 4
	621	+
	622	+ def test_dtensor_cat(self):
	623	+ with LocalTensorMode(self.world_size):
	624	+ device_mesh = self.build_device_mesh()
	625	+
	626	+ t1 = torch.arange(16).view(4, 4).float()
	627	+ d1 = distribute_tensor(t1, device_mesh, [Replicate()])
	628	+ t2 = (torch.arange(16) + 16).view(4, 4).float()
	629	+ d2 = distribute_tensor(t2, device_mesh, [Shard(0)])
	630	+
	631	+ local_res = torch.cat([t1, t2], dim=-1)
	632	+ dist_res = torch.cat([d1, d2], dim=-1)
	633	+ full_tensor = dist_res.full_tensor()
	634	+ self.assertEqual(full_tensor, local_res)
	635	+
	636	+
	637	+class TestLocalTensorWorld8(LocalTensorWorldTest):
	638	+ world_size = 8
	639	+
	640	+ def test_dtensor_addmm(self):
	641	+ with LocalTensorMode(self.world_size):
	642	+ device_mesh = self.build_device_mesh()
	643	+
	644	+ shard_spec = [Shard(0)]
	645	+ replica_spec = [Replicate()]
	646	+
	647	+ tensor_to_shard = torch.randn(12, 8)
	648	+ mat1 = distribute_tensor(tensor_to_shard, device_mesh, shard_spec)
	649	+ tensor_to_replicate = torch.randn(8, 4)
	650	+ mat2 = distribute_tensor(tensor_to_replicate, device_mesh, replica_spec)
	651	+ input_tensor = torch.randn(4)
	652	+ input = distribute_tensor(input_tensor, device_mesh, replica_spec)
	653	+
	654	+ dist_res = torch.addmm(input, mat1, mat2)
	655	+ local_res = torch.addmm(input_tensor, tensor_to_shard, tensor_to_replicate)
	656	+ full_tensor = dist_res.full_tensor()
	657	+ self.assertEqual(full_tensor, local_res)
	658	+
	659	+
	660	+from torch.distributed._local_tensor._c10d import local_p2p_op, wait_all
	661	+
	662	+
	663	+class TestLocalRunner(LocalTensorWorldTest):
	664	+ world_size = 6
	665	+
	666	+ @staticmethod
	667	+ def _get_pp_peer(pp_index, mesh, dim, dir):
	668	+ with maybe_disable_local_tensor_mode():
	669	+ pp_meshes = mesh._get_all_submeshes(dim)
	670	+ pp_ret = {}
	671	+ for pp_mesh in pp_meshes:
	672	+ global_rank = pp_mesh.mesh[pp_index].item()
	673	+ global_peer = pp_mesh.mesh[(pp_index + dir) % pp_mesh.size()].item()
	674	+ pp_ret[global_rank] = global_peer
	675	+
	676	+ return torch.SymInt(LocalIntNode(pp_ret))
	677	+
	678	+ def _run_dp_pp(
	679	+ self,
	680	+ mesh: DeviceMesh,
	681	+ pp_index: int,
	682	+ actual: list[torch.Tensor \| None],
	683	+ expected: list[torch.Tensor \| None],
	684	+ ) -> None:
	685	+ ltm = LocalTensorMode(mesh.size())
	686	+ with ltm:
	687	+ dp_mesh = mesh["dp"]
	688	+ pp_mesh = mesh["pp"]
	689	+
	690	+ x = torch.rand(2, 4)
	691	+ xd = distribute_tensor(x, dp_mesh, [Shard(0)])
	692	+ xd = xd * 2
	693	+ x = x * 2
	694	+
	695	+ yd = zeros(*xd.shape, device_mesh=dp_mesh, placements=[Shard(0)])
	696	+
	697	+ if pp_index != pp_mesh.size(0) - 1:
	698	+ # Send to next pp rank
	699	+ pp_next_rank = TestLocalRunner._get_pp_peer(pp_index, mesh, "pp", +1)
	700	+ local_p2p_op(pp_next_rank, xd, dist.isend)
	701	+ expected[pp_index + 1] = ltm.tensor_map(
	702	+ x,
	703	+ lambda r, t: t
	704	+ if reduce_local_int(pp_next_rank, lambda vals: r in vals.values())
	705	+ else torch.zeros_like(t),
	706	+ )
	707	+
	708	+ if pp_index != 0:
	709	+ # Receive from prev pp rank
	710	+ pp_prev_rank = TestLocalRunner._get_pp_peer(pp_index, mesh, "pp", -1)
	711	+ rw = local_p2p_op(pp_prev_rank, yd, dist.irecv)
	712	+ wait_all(rw)
	713	+
	714	+ y = yd.full_tensor()
	715	+ actual[pp_index] = y
	716	+
	717	+ def test_dp_pp(self):
	718	+ pp_size = 3
	719	+ mesh = init_device_mesh(
	720	+ "cpu", (self.world_size // pp_size, pp_size), mesh_dim_names=("dp", "pp")
	721	+ )
	722	+ actual: list[torch.Tensor \| None] = [None] * pp_size
	723	+ expected: list[torch.Tensor \| None] = [None] * pp_size
	724	+ with LocalRunnerMode(
	725	+ self.world_size,
	726	+ pp_size,
	727	+ lambda pp_index: self._run_dp_pp(mesh, pp_index, actual, expected),
	728	+ ):
	729	+ pass
	730	+
	731	+ self.assertEqual(actual, expected)
	732	+
	733	+
	734	+if __name__ == "__main__":
	735	+ run_tests()

Mtorch_npu/npu/random.py+8-4

	@@ -2,7 +2,7 @@ from typing import Iterable, Union
2	import torch	2	import torch
3		3
4	import torch_npu	4	import torch_npu
5	-from . import _lazy_init, _lazy_call, device_count, current_device	5	+from . import _lazy_init, _lazy_call, device_count, current_device, is_initialized
6		6
7	__all__ = ['get_rng_state', 'set_rng_state',	7	__all__ = ['get_rng_state', 'set_rng_state',
8	'get_rng_state_all', 'set_rng_state_all',	8	'get_rng_state_all', 'set_rng_state_all',
	@@ -49,8 +49,12 @@ def set_rng_state(new_state: torch.Tensor, device: Union[int, str, torch.device]
49	device (torch.device or int, optional): The device to set the RNG state.	49	device (torch.device or int, optional): The device to set the RNG state.
50	Default: ``'npu'`` (i.e., ``torch.device('npu')``, the current NPU device).	50	Default: ``'npu'`` (i.e., ``torch.device('npu')``, the current NPU device).
51	"""	51	"""
52	- with torch._C._DisableFuncTorch():	52	+ if not is_initialized():
53	- new_state_copy = new_state.clone(memory_format=torch.contiguous_format)	53	+ with torch._C._DisableFuncTorch():
		54	+ # Clone the state because the callback will be triggered
		55	+ # later when NPU is lazy initialized.
		56	+ new_state = new_state.clone(memory_format=torch.contiguous_format)
		57	+
54	if isinstance(device, str):	58	if isinstance(device, str):
55	device = torch.device(device)	59	device = torch.device(device)
56	elif isinstance(device, int):	60	elif isinstance(device, int):
	@@ -61,7 +65,7 @@ def set_rng_state(new_state: torch.Tensor, device: Union[int, str, torch.device]
61	if idx is None:	65	if idx is None:
62	idx = current_device()	66	idx = current_device()
63	default_generator = torch_npu.npu.default_generators[idx]	67	default_generator = torch_npu.npu.default_generators[idx]
64	- default_generator.set_state(new_state_copy)	68	+ default_generator.set_state(new_state)
65		69
66	_lazy_call(cb)	70	_lazy_call(cb)
67		71