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

已合并

test: verify community 2.10.0 features and fixes #36909

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

test: verify community 2.10.0 features and fixes #36909

已合并

kuhn7创建于 5月27日

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

共 4 个文件变更+888-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+603-0

@@ -0,0 +1,603 @@
	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	+)
	15	+from torch.distributed.tensor import (
	16	+ DeviceMesh,
	17	+ distribute_tensor,
	18	+ init_device_mesh,
	19	+ Partial,
	20	+ Replicate,
	21	+ Shard,
	22	+ zeros,
	23	+)
	24	+from torch.testing._internal.common_utils import run_tests, TestCase
	25	+from torch.testing._internal.distributed._tensor.common_dtensor import reduce_local_int
	26	+
	27	+
	28	+class LocalTensorTestBase(TestCase):
	29	+ def assertEqual(self, lhs, rhs, **kwargs):
	30	+ mode = local_tensor_mode()
	31	+ with nullcontext() if mode is None else mode.disable():
	32	+ if isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor):
	33	+ assert isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)
	34	+ super().assertEqual(lhs._ranks, rhs._ranks)
	35	+ for r in lhs._ranks:
	36	+ super().assertEqual(
	37	+ lhs._local_tensors[r],
	38	+ rhs._local_tensors[r],
	39	+ lambda m: f"rank {r}: {m}",
	40	+ )
	41	+ elif isinstance(lhs, LocalTensor) or isinstance(rhs, LocalTensor):
	42	+ lhs, rhs = (lhs, rhs) if isinstance(lhs, LocalTensor) else (rhs, lhs)
	43	+ for r in lhs._ranks:
	44	+ super().assertEqual(
	45	+ lhs._local_tensors[r], rhs, lambda m: f"rank {r}: {m}"
	46	+ )
	47	+ else:
	48	+ return super().assertEqual(lhs, rhs, **kwargs)
	49	+
	50	+ @property
	51	+ def world_size(self):
	52	+ raise NotImplementedError("override world-size in your subclass")
	53	+
	54	+ def build_device_mesh(self) -> DeviceMesh:
	55	+ return init_device_mesh("cpu", (self.world_size,))
	56	+
	57	+ def setUp(self):
	58	+ super().setUp()
	59	+ torch.distributed.init_process_group(
	60	+ # TODO: test other ranks too
	61	+ "fake",
	62	+ rank=0,
	63	+ world_size=self.world_size,
	64	+ )
	65	+
	66	+ def tearDown(self):
	67	+ super().tearDown()
	68	+ try:
	69	+ dist.destroy_process_group()
	70	+ except AssertionError:
	71	+ pass
	72	+
	73	+
	74	+class TestLocalTensorWorld2(LocalTensorTestBase):
	75	+ world_size = 2
	76	+
	77	+ def test_local_tensor_dtype_consistency(self):
	78	+ """Test that LocalTensor enforces dtype consistency."""
	79	+ device = torch.device("cpu")
	80	+ shape = (2, 3)
	81	+
	82	+ inconsistent_tensors = {
	83	+ 0: torch.randn(shape, dtype=torch.float32, device=device),
	84	+ 1: torch.randn(
	85	+ shape, dtype=torch.float64, device=device
	86	+ ), # Different dtype
	87	+ }
	88	+
	89	+ with self.assertRaises(AssertionError):
	90	+ LocalTensor(inconsistent_tensors)
	91	+
	92	+ def test_local_tensor_creation_fails_with_grad_tensors(self):
	93	+ """Test that LocalTensor creation fails when local tensors have requires_grad=True."""
	94	+ device = torch.device("cpu")
	95	+ shape = (2, 3)
	96	+ dtype = torch.float32
	97	+
	98	+ # Create sample local tensors for different ranks
	99	+ local_tensors = {
	100	+ 0: torch.randn(shape, dtype=dtype, device=device, requires_grad=True),
	101	+ 1: torch.randn(shape, dtype=dtype, device=device, requires_grad=True),
	102	+ }
	103	+
	104	+ with self.assertRaises(AssertionError):
	105	+ LocalTensor(local_tensors)
	106	+
	107	+ # TODO: test flatten/unflatten
	108	+
	109	+ def test_basic_arithmetic_operations(self):
	110	+ """Test basic arithmetic operations on LocalTensors."""
	111	+ device = torch.device("cpu")
	112	+ shape = (2, 3)
	113	+ dtype = torch.float32
	114	+
	115	+ # Create identical local tensors for consistency tests
	116	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	117	+ identical_local_tensors = {
	118	+ 0: base_tensor.clone(),
	119	+ 1: base_tensor.clone(),
	120	+ }
	121	+
	122	+ lt1 = LocalTensor(identical_local_tensors)
	123	+ lt2 = LocalTensor(identical_local_tensors)
	124	+
	125	+ # Test addition
	126	+ result_add = lt1 + lt2
	127	+ self.assertIsInstance(result_add, LocalTensor)
	128	+ self.assertEqual(len(result_add._local_tensors), 2)
	129	+
	130	+ # Verify the operation was applied to each local tensor
	131	+ for rank in identical_local_tensors:
	132	+ expected = identical_local_tensors[rank] + identical_local_tensors[rank]
	133	+ self.assertEqual(result_add._local_tensors[rank], expected)
	134	+
	135	+ # Test multiplication
	136	+ result_mul = lt1 * 2.0
	137	+ self.assertIsInstance(result_mul, LocalTensor)
	138	+ for rank in identical_local_tensors:
	139	+ expected = identical_local_tensors[rank] * 2.0
	140	+ self.assertEqual(result_mul._local_tensors[rank], expected)
	141	+
	142	+ # TODO: consider an op-info test; we don't actually need to cover all ops
	143	+ # but it will help make sure views and more exotic things are done
	144	+ # correctly (in standard subclass style)
	145	+
	146	+ def test_mixed_operations_with_regular_tensors(self):
	147	+ """Test operations between LocalTensors and regular tensors."""
	148	+ device = torch.device("cpu")
	149	+ shape = (2, 3)
	150	+ dtype = torch.float32
	151	+
	152	+ # Create identical local tensors for consistency tests
	153	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	154	+ identical_local_tensors = {
	155	+ 0: base_tensor.clone(),
	156	+ 1: base_tensor.clone(),
	157	+ }
	158	+
	159	+ lt = LocalTensor(identical_local_tensors)
	160	+ regular_tensor = torch.ones_like(identical_local_tensors[0])
	161	+
	162	+ # Test LocalTensor + regular tensor
	163	+ result = lt + regular_tensor
	164	+ self.assertIsInstance(result, LocalTensor)
	165	+
	166	+ for rank in identical_local_tensors:
	167	+ expected = identical_local_tensors[rank] + regular_tensor
	168	+ self.assertEqual(result._local_tensors[rank], expected)
	169	+
	170	+ def test_local_tensor_mode(self):
	171	+ """Test LocalTensorMode functionality."""
	172	+ device = torch.device("cpu")
	173	+ shape = (2, 3)
	174	+ dtype = torch.float32
	175	+
	176	+ # Create identical local tensors for consistency tests
	177	+ base_tensor = torch.randn(shape, dtype=dtype, device=device)
	178	+ identical_local_tensors = {
	179	+ 0: base_tensor.clone(),
	180	+ 1: base_tensor.clone(),
	181	+ }
	182	+
	183	+ lt = LocalTensor(identical_local_tensors)
	184	+
	185	+ with LocalTensorMode(lt._ranks):
	186	+ result = lt + 1.0
	187	+ self.assertIsInstance(result, LocalTensor)
	188	+
	189	+ regular = torch.ones(2, 2)
	190	+ regular_result = regular + 1.0
	191	+ self.assertIsInstance(regular, LocalTensor)
	192	+ self.assertIsInstance(regular_result, LocalTensor)
	193	+
	194	+ def test_empty_local_tensors(self):
	195	+ """Test behavior with empty local tensors dict."""
	196	+ # TODO: raise a better error here
	197	+ with self.assertRaises(StopIteration): # next() on empty iterator
	198	+ LocalTensor({})
	199	+
	200	+ def test_collectives_within_local_tensor_mode(self):
	201	+ """Test that collective operations work within LocalTensorMode context."""
	202	+ test_tensors = {
	203	+ 0: torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
	204	+ 1: torch.tensor([[5.0, 6.0], [7.0, 8.0]]),
	205	+ }
	206	+ lt = LocalTensor(test_tensors)
	207	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	208	+
	209	+ with LocalTensorMode(lt._ranks):
	210	+ # Test all_reduce within mode
	211	+ lt_sum = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	212	+ dist.all_reduce(lt_sum, group=fake_pg)
	213	+
	214	+ expected_sum = torch.tensor([[6.0, 8.0], [10.0, 12.0]])
	215	+ for rank in test_tensors:
	216	+ self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
	217	+
	218	+ # Test broadcast within mode
	219	+ lt_broadcast = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	220	+ dist.broadcast(lt_broadcast, src=0, group=fake_pg)
	221	+
	222	+ for rank in test_tensors:
	223	+ self.assertEqual(lt_broadcast._local_tensors[rank], test_tensors[0])
	224	+
	225	+ # Test that regular operations still work
	226	+ result = lt + 1.0
	227	+ self.assertIsInstance(result, LocalTensor)
	228	+
	229	+ def test_scalar_mul_reduction_bug(self):
	230	+ with LocalTensorMode(self.world_size):
	231	+ mesh = self.build_device_mesh()
	232	+
	233	+ tensor = torch.tensor([10, 10]).float()
	234	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	235	+ y = dt.sum() * 1 # noqa: F841
	236	+
	237	+ tensor = torch.arange(10).reshape(10, 1).float().requires_grad_()
	238	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	OO O openLiBingCI5月27日此条代码评论区间+233至+238 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI5月27日此条代码评论区间+233至+238 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	239	+
	240	+ print(dt.sum() * 1, dt.sum() * 2, dt.sum() * 3)
	241	+
	242	+ def test_uneven_sharding_mean_bug(self):
	243	+ with LocalTensorMode(self.world_size):
	244	+ mesh = self.build_device_mesh()
	245	+ tensor = torch.arange(12).reshape(-1, 4).float()
	246	+
	247	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	248	+
	249	+ mean = dt.mean()
	250	+ self.assertEqual(mean.placements, [Replicate()])
	251	+ full = mean.full_tensor()
	252	+ self.assertEqual(tensor.mean(), full)
	253	+
	254	+ def test_uneven_sharding_prod(self):
	255	+ with LocalTensorMode(self.world_size):
	256	+ mesh = self.build_device_mesh()
	257	+ tensor = (torch.arange(12) + 1).reshape(-1, 4).float()
	258	+
	259	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	260	+
	261	+ x = dt.prod()
	262	+ full = x.full_tensor()
	263	+ self.assertEqual(tensor.prod(), full)
	264	+
	265	+ def test_even_sharding_mean_is_partial(self):
	266	+ with LocalTensorMode(self.world_size):
	267	+ mesh = self.build_device_mesh()
	268	+ tensor = torch.arange(16).reshape(4, 4).float()
	269	+
	270	+ dt = distribute_tensor(tensor, device_mesh=mesh, placements=[Shard(0)])
	271	+
	272	+ mean = dt.mean()
	273	+ full = mean.full_tensor()
	274	+ self.assertEqual(tensor.mean(), full)
	275	+ self.assertEqual(mean.placements, [Partial("avg")])
	276	+
	277	+
	278	+class TestLocalTensorWorld3(LocalTensorTestBase):
	279	+ world_size = 3
	280	+
	281	+ def test_collective_reduction_operations(self):
	282	+ """Test different reduction operations for all_reduce."""
	283	+ # Create different tensors for each rank with simple values for testing
	284	+ test_tensors = {
	285	+ 0: torch.tensor([[1.0, 4.0], [2.0, 5.0]]),
	286	+ 1: torch.tensor([[2.0, 1.0], [3.0, 6.0]]),
	287	+ 2: torch.tensor([[3.0, 2.0], [1.0, 4.0]]),
	288	+ }
	289	+
	290	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	291	+
	292	+ # Test SUM reduction
	293	+ lt_sum = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	294	+ dist.all_reduce(lt_sum, op=dist.ReduceOp.SUM, group=fake_pg)
	295	+ expected_sum = torch.tensor([[6.0, 7.0], [6.0, 15.0]]) # Sum of all tensors
	296	+ for rank in test_tensors:
	297	+ self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
	298	+
	299	+ # Test MAX reduction
	300	+ lt_max = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	301	+ dist.all_reduce(lt_max, op=dist.ReduceOp.MAX, group=fake_pg)
	302	+ expected_max = torch.tensor([[3.0, 4.0], [3.0, 6.0]]) # Max across all tensors
	303	+ for rank in test_tensors:
	304	+ self.assertEqual(lt_max._local_tensors[rank], expected_max)
	305	+
	306	+ # Test MIN reduction
	307	+ lt_min = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
	308	+ dist.all_reduce(lt_min, op=dist.ReduceOp.MIN, group=fake_pg)
	309	+ expected_min = torch.tensor([[1.0, 1.0], [1.0, 4.0]]) # Min across all tensors
	310	+ for rank in test_tensors:
	311	+ self.assertEqual(lt_min._local_tensors[rank], expected_min)
	312	+
	313	+ def test_all_reduce_collective(self):
	314	+ """Test that all_reduce collective operation works correctly with LocalTensor."""
	315	+ # Create different tensors for each rank
	316	+ different_tensors = {
	317	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	318	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	319	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	320	+ }
	321	+
	322	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	323	+
	324	+ # Test all_reduce with SUM (default)
	325	+ lt_sum = LocalTensor({k: v.clone() for k, v in different_tensors.items()})
	326	+ lt_sum = lt_sum + 1
	327	+ dist.all_reduce(lt_sum, group=fake_pg)
	328	+
	329	+ # Verify all ranks have the sum of all tensors (after adding 1 to each)
	330	+ expected_sum = torch.tensor([[114.0, 225.0, 336.0], [447.0, 558.0, 669.0]])
	331	+ for rank in different_tensors:
	332	+ self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
	333	+
	334	+ def test_broadcast_collective(self):
	335	+ """Test that broadcast collective operation works correctly with LocalTensor."""
	336	+ # Create different tensors for each rank
	337	+ different_tensors = {
	338	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	339	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	340	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	341	+ }
	342	+
	343	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	344	+
	345	+ # Test broadcast from rank 1
	346	+ lt_broadcast = LocalTensor({k: v.clone() for k, v in different_tensors.items()})
	347	+ dist.broadcast(lt_broadcast, src=1, group=fake_pg)
	348	+
	349	+ # Verify all ranks have rank 1's original tensor
	350	+ expected_broadcast = different_tensors[1]
	351	+ for rank in different_tensors:
	352	+ self.assertEqual(lt_broadcast._local_tensors[rank], expected_broadcast)
	353	+
	354	+ def test_all_gather_collective(self):
	355	+ """Test that all_gather collective operation works correctly with LocalTensor."""
	356	+ # Create different tensors for each rank
	357	+ different_tensors = {
	358	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	359	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	360	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	361	+ }
	362	+
	363	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	364	+
	365	+ # Test all_gather
	366	+ lt_gather = LocalTensor(different_tensors)
	367	+ tensor_list = [torch.zeros_like(lt_gather) for _ in range(3)]
	368	+
	369	+ dist.all_gather(tensor_list, lt_gather, group=fake_pg)
	370	+
	371	+ # Verify each position in tensor_list contains the corresponding rank's tensor
	372	+ self.assertEqual(tensor_list[0], different_tensors[0])
	373	+ self.assertEqual(tensor_list[1], different_tensors[1])
	374	+ self.assertEqual(tensor_list[2], different_tensors[2])
	375	+
	376	+ def test_reduce_scatter_tensor_collective(self):
	377	+ """Test that reduce_scatter_tensor collective operation works correctly with LocalTensor."""
	378	+ # Create different tensors for each rank
	379	+ different_tensors = {
	380	+ 0: torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),
	381	+ 1: torch.tensor([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]),
	382	+ 2: torch.tensor([[100.0, 200.0], [300.0, 400.0], [500.0, 600.0]]),
	383	+ }
	384	+
	385	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	386	+
	387	+ # Test reduce_scatter_tensor
	388	+ with LocalTensorMode(self.world_size):
	389	+ lt_reduce_scatter = LocalTensor(different_tensors)
	390	+ lt_reduce_scatter_size = lt_reduce_scatter.size()
	391	+ lt_output_tensor = torch.zeros(
	392	+ lt_reduce_scatter_size[0] // fake_pg.size(),
	393	+ *lt_reduce_scatter_size[1:],
	394	+ dtype=lt_reduce_scatter.dtype,
	395	+ device=lt_reduce_scatter.device,
	396	+ )
	397	+
	398	+ dist.reduce_scatter_tensor(
	399	+ lt_output_tensor, lt_reduce_scatter, group=fake_pg
	400	+ )
	401	+
	402	+ expected_output = LocalTensor(
	403	+ {
	404	+ 0: torch.tensor([[111.0, 222.0]]),
	405	+ 1: torch.tensor([[333.0, 444.0]]),
	406	+ 2: torch.tensor([[555.0, 666.0]]),
	407	+ }
	408	+ )
	409	+ print(lt_output_tensor)
	410	+ self.assertEqual(lt_output_tensor, expected_output)
	411	+
	412	+ def test_all_gather_into_tensor_collective(self):
	413	+ """Test that all_gather_into_tensor collective operation works correctly with LocalTensor."""
	414	+ # Create different tensors for each rank
	415	+ different_tensors = {
	416	+ 0: torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	417	+ 1: torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]),
	418	+ 2: torch.tensor([[100.0, 200.0, 300.0], [400.0, 500.0, 600.0]]),
	419	+ }
	420	+
	421	+ fake_pg = torch.distributed.distributed_c10d._get_default_group()
	422	+
	423	+ # Test all_gather_into_tensor
	424	+ with LocalTensorMode(self.world_size):
	425	+ lt_gather = LocalTensor(different_tensors)
	426	+ lt_gather_size = lt_gather.size()
	427	+ lt_output_tensor = torch.zeros(
	428	+ lt_gather_size[0] * fake_pg.size(),
	429	+ *lt_gather_size[1:],
	430	+ dtype=lt_gather.dtype,
	431	+ device=lt_gather.device,
	432	+ )
	433	+
	434	+ dist.all_gather_into_tensor(lt_output_tensor, lt_gather, group=fake_pg)
	435	+
	436	+ expected_output = torch.cat(list(different_tensors.values()))
	437	+
	438	+ self.assertEqual(lt_output_tensor, expected_output)
	439	+
	440	+ def test_all_to_all_single_collective(self):
	441	+ """Test that all_to_all_single collective operation works correctly with LocalTensor."""
	442	+ from torch.distributed._functional_collectives import all_to_all_single
	443	+
	444	+ # Create different tensors for each rank
	445	+ # Each rank will split its tensor and send parts to other ranks
	446	+ different_tensors = {
	447	+ 0: torch.tensor(
	448	+ [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
	449	+ ), # rank 0 sends [0,0], [0,0], [0,0] to ranks 0,1,2
	450	+ 1: torch.tensor(
	451	+ [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
	452	+ ), # rank 1 sends [1,1], [1,1], [1,1] to ranks 0,1,2
	453	+ 2: torch.tensor(
	454	+ [2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
	455	+ ), # rank 2 sends [2,2], [2,2], [2,2] to ranks 0,1,2
	456	+ }
	457	+
	458	+ # Each rank splits its input into 3 parts of size 2 each
	459	+ input_split_sizes = [2, 2, 2]
	460	+ # Each rank receives 3 parts of size 2 each from all ranks
	461	+ output_split_sizes = [2, 2, 2]
	462	+
	463	+ with LocalTensorMode(self.world_size):
	464	+ lt_input = LocalTensor(different_tensors)
	465	+
	466	+ # Test all_to_all_single using functional collectives API
	467	+ result = all_to_all_single(
	468	+ lt_input,
	469	+ output_split_sizes=output_split_sizes,
	470	+ input_split_sizes=input_split_sizes,
	471	+ group=torch.distributed.distributed_c10d._get_default_group(),
	472	+ )
	473	+
	474	+ result = result.wait()
	475	+ # Verify result is a LocalTensor
	476	+ self.assertIsInstance(result, LocalTensor)
	477	+
	478	+ # After all_to_all_single:
	479	+ # rank 0 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	480	+ # rank 1 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	481	+ # rank 2 receives: [0,0] from rank 0, [1,1] from rank 1, [2,2] from rank 2 = [0,0,1,1,2,2]
	482	+ expected_output = torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0, 2.0])
	483	+
	484	+ for rank in different_tensors:
	485	+ self.assertEqual(result._local_tensors[rank], expected_output)
	486	+
	487	+
	488	+class TestLocalTensorWorld4(LocalTensorTestBase):
	489	+ world_size = 4
	490	+
	491	+ def test_dtensor_cat(self):
	492	+ with LocalTensorMode(self.world_size):
	493	+ device_mesh = self.build_device_mesh()
	494	+
	495	+ t1 = torch.arange(16).view(4, 4).float()
	496	+ d1 = distribute_tensor(t1, device_mesh, [Replicate()])
	497	+ t2 = (torch.arange(16) + 16).view(4, 4).float()
	498	+ d2 = distribute_tensor(t2, device_mesh, [Shard(0)])
	499	+
	500	+ local_res = torch.cat([t1, t2], dim=-1)
	501	+ dist_res = torch.cat([d1, d2], dim=-1)
	502	+ full_tensor = dist_res.full_tensor()
	503	+ self.assertEqual(full_tensor, local_res)
	504	+
	505	+
	506	+class TestLocalTensorWorld8(LocalTensorTestBase):
	507	+ world_size = 8
	508	+
	509	+ def test_dtensor_addmm(self):
	510	+ with LocalTensorMode(self.world_size):
	511	+ device_mesh = self.build_device_mesh()
	512	+
	513	+ shard_spec = [Shard(0)]
	514	+ replica_spec = [Replicate()]
	515	+
	516	+ tensor_to_shard = torch.randn(12, 8)
	517	+ mat1 = distribute_tensor(tensor_to_shard, device_mesh, shard_spec)
	518	+ tensor_to_replicate = torch.randn(8, 4)
	519	+ mat2 = distribute_tensor(tensor_to_replicate, device_mesh, replica_spec)
	520	+ input_tensor = torch.randn(4)
	521	+ input = distribute_tensor(input_tensor, device_mesh, replica_spec)
	522	+
	523	+ dist_res = torch.addmm(input, mat1, mat2)
	524	+ local_res = torch.addmm(input_tensor, tensor_to_shard, tensor_to_replicate)
	525	+ full_tensor = dist_res.full_tensor()
	526	+ self.assertEqual(full_tensor, local_res)
	527	+
	528	+
	529	+from torch.distributed._local_tensor._c10d import local_p2p_op, wait_all
	530	+
	531	+
	532	+class TestLocalRunner(LocalTensorTestBase):
	533	+ world_size = 6
	534	+
	535	+ @staticmethod
	536	+ def _get_pp_peer(pp_index, mesh, dim, dir):
	537	+ pp_meshes = mesh._get_all_submeshes(dim)
	538	+ pp_ret = {}
	539	+ for pp_mesh in pp_meshes:
	540	+ global_rank = pp_mesh.mesh[pp_index].item()
	541	+ global_peer = pp_mesh.mesh[(pp_index + dir) % pp_mesh.size()].item()
	542	+ pp_ret[global_rank] = global_peer
	543	+
	544	+ return torch.SymInt(LocalIntNode(pp_ret))
	545	+
	546	+ def _run_dp_pp(
	547	+ self,
	548	+ mesh: DeviceMesh,
	549	+ pp_index: int,
	550	+ actual: list[torch.Tensor \| None],
	551	+ expected: list[torch.Tensor \| None],
	552	+ ) -> None:
	553	+ ltm = LocalTensorMode(mesh.size())
	554	+ with ltm:
	555	+ dp_mesh = mesh["dp"]
	556	+ pp_mesh = mesh["pp"]
	557	+
	558	+ x = torch.rand(2, 4)
	559	+ xd = distribute_tensor(x, dp_mesh, [Shard(0)])
	560	+ xd = xd * 2
	561	+ x = x * 2
	562	+
	563	+ yd = zeros(*xd.shape, device_mesh=dp_mesh, placements=[Shard(0)])
	564	+
	565	+ if pp_index != pp_mesh.size(0) - 1:
	566	+ # Send to next pp rank
	567	+ pp_next_rank = TestLocalRunner._get_pp_peer(pp_index, mesh, "pp", +1)
	568	+ local_p2p_op(pp_next_rank, xd, dist.isend)
	569	+ expected[pp_index + 1] = ltm.tensor_map(
	570	+ x,
	571	+ lambda r, t: t
	572	+ if reduce_local_int(pp_next_rank, lambda vals: r in vals.values())
	573	+ else torch.zeros_like(t),
	574	+ )
	575	+
	576	+ if pp_index != 0:
	577	+ # Receive from prev pp rank
	578	+ pp_prev_rank = TestLocalRunner._get_pp_peer(pp_index, mesh, "pp", -1)
	579	+ rw = local_p2p_op(pp_prev_rank, yd, dist.irecv)
	580	+ wait_all(rw)
	581	+
	582	+ y = yd.full_tensor()
	583	+ actual[pp_index] = y
	584	+
	585	+ def test_dp_pp(self):
	586	+ pp_size = 3
	587	+ mesh = init_device_mesh(
	588	+ "cpu", (self.world_size // pp_size, pp_size), mesh_dim_names=("dp", "pp")
	589	+ )
	590	+ actual: list[torch.Tensor \| None] = [None] * pp_size
	591	+ expected: list[torch.Tensor \| None] = [None] * pp_size
	592	+ with LocalRunnerMode(
	593	+ self.world_size,
	594	+ pp_size,
	595	+ lambda pp_index: self._run_dp_pp(mesh, pp_index, actual, expected),
	596	+ ):
	597	+ pass
	598	+
	599	+ self.assertEqual(actual, expected)
	600	+
	601	+
	602	+if __name__ == "__main__":
	603	+ 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