Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / random.py: 72%
209 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""RNG state management for distributed tensor operations.
17Provides utilities for tracking and synchronizing random number generator states
18across multiple devices in distributed training scenarios.
19"""
21__all__ = [
22 "is_rng_supported_mesh",
23 "manual_seed",
24 "OffsetBasedRNGTracker",
25]
27import contextlib
28import warnings
29from logging import getLogger
30import typing
31from typing import Optional
32import functools
33import operator
35from hyper_parallel.core.dtensor.placement_types import Shard
36from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
37from hyper_parallel.platform import get_platform
39platform = get_platform()
40DTensorBase = platform.DTensorBase
41Tensor = platform.tensor
43logger = getLogger(__name__)
46def is_rng_supported_mesh(device_mesh: Optional[DeviceMesh] = None) -> bool:
47 """Check if the device mesh supports DTensor random operations.
49 Currently, DTensor random operations are only supported on CUDA and CUDA-like
50 devices. Users should call this function before using DTensor random APIs to
51 verify compatibility.
53 Args:
54 device_mesh: Optional :class:`DeviceMesh` to check (same semantics as PyTorch
55 ``torch.distributed.tensor``). If omitted, checks the active platform device
56 handle only.
58 Returns:
59 bool: ``True`` if the device mesh supports DTensor random operations,
60 ``False`` otherwise.
61 """
62 if device_mesh is not None and device_mesh.device_type == "cpu":
63 warnings.warn(
64 f"DTensor random operators may not have complete support on {device_mesh.device_type} device mesh",
65 stacklevel=2,
66 )
67 return False
68 device_handle = platform.get_device_handle()
69 if device_handle and hasattr(device_handle, "set_rng_state"):
70 return True
71 if device_mesh is not None:
72 warnings.warn(
73 f"DTensor random operators may not have complete support on {device_mesh.device_type} device mesh",
74 stacklevel=2,
75 )
76 return False
79class _PhiloxState:
80 """
81 Convenience accessor for interpreting the packed bits of (seed: uint64, offset: uint64) in the philox state,
82 which for some reason is actually exposed as a size-16 uint8 tensor.
84 The state is always moved to .cpu since it is necessary for it to be on CPU before applying it back to a generator.
85 """
87 def __init__(self, state: Tensor):
88 self._state = state.to("cpu")
90 @property
91 def state(self):
92 """Return the underlying RNG state tensor (CPU uint8)."""
93 return self._state
95 @property
96 def offset(self) -> int:
97 """Return the offset value (last 8 bytes) of the Philox RNG state."""
98 return int(self._state[8:].view(dtype=platform.tensor_dtype.int64).item())
100 @offset.setter
101 def offset(self, offset: int) -> None:
102 """Set the offset value of the Philox RNG state."""
103 offset_tensor = Tensor([offset], dtype=platform.tensor_dtype.uint64).view(
104 platform.tensor_dtype.uint8
105 ) # device?
106 self._state[8:] = offset_tensor
108 @property
109 def seed(self) -> int:
110 """Return the seed value (first 8 bytes) of the Philox RNG state."""
111 return int(self._state[:8].view(dtype=platform.tensor_dtype.uint64).item())
113 @seed.setter
114 def seed(self, seed: int) -> None:
115 """Set the seed value of the Philox RNG state."""
116 seed_tensor = Tensor([seed], dtype=platform.tensor_dtype.uint64).view(
117 platform.tensor_dtype.uint8
118 )# device
119 self._state[:8] = seed_tensor
122class _RNGStateTracker:
123 """
124 Tracks and manages RNG states for DTensor random operations.
126 Maintains a mapping from operation tags to RNG state tensors (ByteTensor),
127 providing standardized interfaces for state access and modification.
129 The core method `_distribute_region` establishes the proper RNG context
130 when DTensor executes random operators across distributed devices.
131 """
133 def __init__(self, device):
134 self._device = device
135 self._device_handle = platform.get_device_handle()
136 if not self._device_handle:
137 raise RuntimeError(
138 f"{self.__class__.__name__} instantiation requires the presence of "
139 )
140 self._use_distribute_region = True
142 @property
143 def distribute_region_enabled(self) -> bool:
144 """Return whether the RNG distribute region is enabled for distributed random operations."""
145 return self._use_distribute_region
147 @distribute_region_enabled.setter
148 def distribute_region_enabled(self, value) -> None:
149 """Set whether the RNG distribute region is enabled."""
150 self._use_distribute_region = value
152 def _distribute_region(
153 self, device_mesh, placements, global_shape, generator = None
154 ):
155 pass
157 def _manual_seed(self, parallel_seed: int) -> None:
158 pass
161class OffsetBasedRNGTracker(_RNGStateTracker):
162 """
163 This subclass of ``_RNGStateTracker`` defines the default policy of how RNG states
164 should be shared and synchronized among all ranks to respect the semantics of DTensor
165 random operators.
166 """
168 def __init__(
169 self,
170 run_state_sync: bool = True,
171 ):
172 super().__init__(_resolve_device())
173 rng_state = self._get_device_state()
174 if run_state_sync:
175 # synchronize RNG state using rank 0's current one
176 platform.broadcast(rng_state, 0)
177 my_rng_state = self._get_device_state()
178 if not all(my_rng_state == rng_state):
179 logger.warning(
180 "DTensor is synchronizing RNG states of every rank with the state from rank 0. "
181 "This behavior is deprecated. "
182 "Please call ``manual_seed(seed, device_mesh)`` from "
183 "``hyper_parallel.core.dtensor.random`` on every rank that participates in SPMD DTensor "
184 "operations with the same seed. If using Pipeline Parallelism, each pipelining state would use "
185 "a different seed, but all ranks belonging to one pipeline stage would use the same seed."
186 )
187 self._set_device_state(rng_state)
189 def _manual_seed(self, parallel_seed: int) -> None:
190 """Set default RNG seed (``platform.manual_seed``); same idea as PyTorch DTensor."""
191 platform.manual_seed(parallel_seed)
193 def _get_device_state(self):
194 rng_state = self._device_handle.get_rng_state().to(self._device)
195 return rng_state
197 def _set_device_state(self, state: Tensor):
198 # It seems that the underlying generator wants a cpu tensor but the dtensor code expects `_get_device_state`
199 # to convert to a 'device' tensor, probably because we may use it with our backend comms for sync/debug
200 # for now, we just convert back to cpu here to make sure it always works.
201 self._device_handle.set_rng_state(state.to("cpu"))
203 @contextlib.contextmanager
204 def _distribute_region(
205 self, device_mesh, placements, global_shape, generator = None
206 ):
208 # regular (non-LocalTensor) mode
209 if generator is not None:
210 # This is a little hacky, but for any user-passed generator, we store its state under a unique key,
211 # not because we need to keep a copy of it but because its the easiest way to make it work with the
212 # existing set/get APIs. We also ensure we remove it from rng_states after each _distribute_region.
213 state = _PhiloxState(generator.get_state())
214 else:
215 state = _PhiloxState(self._get_device_state())
217 if self.distribute_region_enabled:
218 old_offset = state.offset
219 self._set_pre_op_offset(state, device_mesh, placements, global_shape)
220 with fork_rng(
221 devices=[self._device], device_type=platform.device_type()
222 ):
223 self._device_handle.set_rng_state(state.state)
224 try:
225 yield # execute the region code
226 finally:
227 # update offset to synchronize among ranks
228 self._set_post_op_offset(state, global_shape, old_offset)
230 else:
231 yield
233 if generator is not None:
234 # ensure we (a) propagate the state advancement back to the user's RNG so its visible and impacts any future
235 # usage of that RNG (dtensor or non-dtensor), (b) drop it from our own cache so that if the user updates
236 # the seed value in their rng and uses it with DTensor again, we always use the latest value
237 generator.set_state(state.state)
238 else:
239 self._set_device_state(state.state)
241 def compute_offset_incr(self, device_mesh, placements, global_shape) -> int:
242 """Compute the per-shard RNG offset increment for the current rank.
244 Based on the shard linear index and local shard size, computes how much to
245 advance the offset so that each shard gets a unique portion of the random stream.
247 Args:
248 device_mesh (DeviceMesh): The device mesh describing the device topology.
249 placements (Sequence[Placement]): The placement strategy for each mesh dimension.
250 global_shape: input global shape
252 Returns:
253 int: The offset increment, 4-byte aligned.
254 """
255 mesh_coordinate = device_mesh.get_coordinate()
256 shard_idx_by_dim, total_num_shards_by_dim = _calc_shard_info(
257 mesh_coordinate,
258 device_mesh,
259 placements,
260 global_shape,
261 reverse_repeated_shards=True,
262 )
263 shard_linear_idx = self._calc_shard_linear_idx(
264 shard_idx_by_dim, total_num_shards_by_dim
265 )
266 local_size_on_rank_0 = _calc_first_shard_size(device_mesh, placements, global_shape)
267 local_size = functools.reduce(operator.mul, local_size_on_rank_0, 1)
268 return (shard_linear_idx * local_size + 3) // 4 * 4
270 def _set_pre_op_offset(self, state: _PhiloxState, device_mesh, placements, global_shape) -> None:
271 """Set the starting random number generator (RNG) offset for the local shard
272 on the current process before operation execution.The offset value begins from
273 the current accumulated position and increments by the local shard size until
274 covering the total elements of the global distributed tensor. Multiple processes
275 holding replicas of the same shard will share identical starting offset values.
277 Args:
278 state (`Tensor`): The generator state to modify
279 device_mesh (DeviceMesh): The device mesh describing the device topology.
280 placements (Sequence[Placement]): The placement strategy for each mesh dimension.
281 Each element should be a Placement object (Shard, Replicate, Partial, etc.).
282 global_shape: input global shape
284 Returns:
285 None
287 .. warning::
288 The current implementation does not consider memory layout contiguity.
290 Example:
291 take a DTensor of shape [8, 16] as an example. Assume that the DTensor
292 is placed on a device mesh with placements ([Shard(1), Replicate(), Shard(0)]),
293 and the mesh is:
294 [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
295 ``mesh.get_coordinate()`` provides the coordinate of the current rank
296 in the mesh. For example, the coordinate of rank 5 is (1, 0, 1).
298 Another concept to introduce besides rank coordinate is shard coordinate.
299 Each rank holds a local shard of the DTensor. In the example, the DTensor
300 is partitioned into 4 [4, 8] shards. The first shard has 2 replicas and
301 rank 0 (coord (0, 0, 0)) and rank 2 (coord (0, 1, 0)) have 1 replica each.
302 That being said, the local shard on rank 0 and rank 2 correspond to the same
303 shard of the DTensor. To denote each DTensor shard, we use a shard coordinate
304 (in the example, it will be a tuple (i, j) where shard (i, j) has the slice
305 DTensor[4 * i : 4 * (i + 1), 8 * j : 8 * (j + 1)], 0 <= i < 2, 0 <= j < 2).
307 Once we have rank coordinate and shard coordinate, we can calculate on each rank
308 what shard of the DTensor the rank holds, with the help of dim_map. The dim_map
309 of the above DTensor is [2, 0] so the shard coordinate of a rank with rank coord
310 (x, y, z) is simply (z, x) by taking(rank_coord[dim_map[0]],rank_coord[dim_map[1]]).
311 Following this calculation,
312 rank 0 and rank 2 holds the shard of coord (0, 0);
313 rank 1 and rank 3 holds the shard of coord (0, 1);
314 rank 4 and rank 6 holds the shard of coord (1, 0);
315 rank 5 and rank 7 holds the shard of coord (1, 1);
317 The last value to calculate before obtaining the starting offset is the shard linear index.
318 The starting offset for each rank will be its shard_linear_index * local_tensor_numel.
319 """
320 current_offset = state.offset
321 offset_incr = self.compute_offset_incr(device_mesh, placements, global_shape)
322 state.offset = current_offset + offset_incr
324 def _set_post_op_offset(
325 self, state: _PhiloxState, global_shape, old_offset: int
326 ) -> None:
327 """Sets the RNG to a synchronized state after running the local random op.
328 Restores the random number generator to a globally consistent state following
329 local shard execution. Each process must advance its offset by the total element
330 count of the distributed tensor, measured from the offset value recorded before
331 the operation began.
333 Args:
334 state (`Tensor`): The generator state to modify.
335 global_shape: The global shape of the distributed tensor.
336 old_offset (int): The RNG offset before the operation.
338 Returns:
339 None
340 """
341 numel = functools.reduce(operator.mul, global_shape, 1)
342 numel = (numel + 3) // 4 * 4
343 state.offset = old_offset + numel
345 def _calc_shard_linear_idx(
346 self, shard_coord: list[int], shard_size: list[int]
347 ) -> int:
348 return _calc_shard_linear_idx(shard_coord, shard_size)
351def _calc_first_shard_size(device_mesh, placements, global_shape) -> list[int]:
352 """Calculate the size of the first shard on rank 0.
354 Args:
355 device_mesh: The device mesh describing the device topology.
356 placements: Sequence of Placement objects (Shard, Replicate, etc.).
357 global_shape: input global shape
359 Returns:
360 list[int]: Shape of rank 0's local shard.
361 """
362 local_size_on_rank_0 = list(global_shape)
363 for idx, placement in enumerate(placements):
364 if isinstance(placement, Shard):
365 mesh_dim_size = device_mesh.size(idx)
366 shard_dim = placement.dim
367 current_dim_size = local_size_on_rank_0[shard_dim]
368 local_size_on_rank_0[shard_dim], _ = local_shard_size_and_offset(
369 current_dim_size,
370 mesh_dim_size,
371 0,
372 )
373 return local_size_on_rank_0
376def _calc_shard_info(
377 mesh_coordinate,
378 device_mesh,
379 placements,
380 global_shape,
381 reverse_repeated_shards=False,
382):
383 """Calculate shard information for a specific rank."""
384 mesh_size = device_mesh.mesh_shape
385 # note: dim_map does not allow double sharding which is the FSDP(fully_shard)+TP
386 # case. Replace the custom logic with dim_map once we support it.
387 dim_map = [-1] * len(global_shape)
388 for i, placement in enumerate(placements):
389 if isinstance(placement, Shard):
390 shard_dim = placement.dim
391 if dim_map[shard_dim] == -1:
392 dim_map[shard_dim] = [i]
393 else:
394 mesh_dim_list = dim_map[shard_dim]
395 if not isinstance(mesh_dim_list, list):
396 raise TypeError(f"Expected mesh_dim_list to be a list, got {type(mesh_dim_list)}")
397 mesh_dim_list.append(i)
399 # Compute shard coordinate:
400 # The coordinate on each tensor dim is a tuple (idx, range)
401 # If a DTensor is partitioned on its dim i into n shards, and the current rank
402 # holds the j-th, then its shard coordinate will be (idx=j, range=n) on dim i
403 if mesh_coordinate is None:
404 raise ValueError("mesh_coordinate must not be None")
405 shard_idx_by_dim = []
406 total_num_shards_by_dim = [] # total number of shards on each tensor dim
407 for mesh_dim in dim_map:
408 shard_idx = 0
409 total_num_shards = 1
410 # the tensor dim is sharded on more than 1 mesh dim
411 if isinstance(mesh_dim, list):
412 repeated_mesh_dims = mesh_dim
413 if reverse_repeated_shards:
414 # RNG follows the sequential placement application: the last
415 # placement subdivides each earlier shard and becomes major.
416 repeated_mesh_dims = list(reversed(repeated_mesh_dims))
417 rank_coord = [mesh_coordinate[d] for d in repeated_mesh_dims]
418 num_shards = [mesh_size[d] for d in repeated_mesh_dims]
419 # compute the shard idx and total number of shards
420 for idx, size in zip(rank_coord, num_shards):
421 shard_idx = shard_idx * size + idx
422 total_num_shards *= size
424 shard_idx_by_dim.append(shard_idx)
425 total_num_shards_by_dim.append(total_num_shards)
426 return shard_idx_by_dim, total_num_shards_by_dim
429def _calc_shard_linear_idx(shard_coord: list[int], shard_size: list[int]) -> int:
430 # compute shard linear index
431 shard_linear_idx = 0
432 shard_coord_stride = 1
433 for idx, size in zip(reversed(shard_coord), reversed(shard_size)):
434 shard_linear_idx += idx * shard_coord_stride
435 shard_coord_stride *= size
437 return shard_linear_idx
440def _resolve_device():
441 device_handle = platform.get_device_handle()
442 device_idx = platform.get_rank() % platform.device_count(device_handle)
444 def get_device(device_idx):
445 return platform.device(device_idx)
447 return get_device(device_idx)
450def manual_seed(seed: int, device_mesh: DeviceMesh) -> None:
451 """Set the seed for generating random numbers on the calling rank (PyTorch DTensor parity).
453 Ensures the global RNG used by DTensor random ops is initialized consistently. Lazily
454 creates the :class:`OffsetBasedRNGTracker` used by shard dispatch with
455 ``run_state_sync=False`` so ranks are not synchronized from rank 0's prior RNG state.
457 Args:
458 seed: Desired RNG seed (must be agreed across ranks in the mesh for SPMD).
459 device_mesh: Mesh that must include the current process rank.
461 Raises:
462 RuntimeError: If the current rank is not part of ``device_mesh`` (undefined DTensor
463 RNG behavior in that case).
465 Warning:
466 Does not validate that ``seed`` matches across ranks; callers must ensure SPMD
467 consistency. Pipeline parallel: use one seed per pipeline stage group as in PyTorch.
468 """
469 if not is_rng_supported_mesh(device_mesh):
470 warnings.warn(
471 "DTensor manual_seed() may not have complete support "
472 f"on {device_mesh.device_type} device mesh",
473 stacklevel=2,
474 )
475 return
477 # Local import avoids import cycle: _op_dispatch imports this module at load time.
478 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER # pylint: disable=C0415
480 if _OP_DISPATCHER._rng_tracker is None:
481 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
483 if device_mesh.get_coordinate() is None:
484 raise RuntimeError(
485 "manual_seed requires the current rank to be a part of the device mesh "
486 "otherwise DTensor RNG state on the rank will not be initialized and "
487 "the behavior of DTensor random ops is undefined."
488 )
490 platform.manual_seed(seed)
493def local_shard_size_and_offset(
494 curr_local_size: int,
495 num_chunks: int,
496 rank,
497):
498 """
499 Given the size of the current local tensor (which may already be sharded on some dimensions),
500 computes the new local shard size and offset given the desired number of chunks
501 (num_chunks is generally equal to the size of the current sharding dim).
503 Note: new local shard offset is relative to the current sharded tensor, not the global tensor.
504 See `_utils.compute_local_shape_and_global_offset` for computing global offset.
506 Returns (new local shard size, offset)
508 """
509 # Compute the chunk size inline
510 if curr_local_size % num_chunks == 0:
511 full_chunk_size = curr_local_size // num_chunks
512 shard_starting_idx = full_chunk_size * rank
513 return full_chunk_size, shard_starting_idx
515 # uneven sharding case
516 full_chunk_size = (curr_local_size + num_chunks - 1) // num_chunks
517 shard_starting_idx = full_chunk_size * rank
519 if curr_local_size < shard_starting_idx:
520 return 0, typing.cast(int, curr_local_size)
521 local_shard_size = (
522 min(curr_local_size, shard_starting_idx + full_chunk_size)
523 - shard_starting_idx
524 )
525 return local_shard_size, shard_starting_idx
528_fork_rng_warned_already = False
531@contextlib.contextmanager
532def fork_rng(
533 devices=None,
534 enabled=True,
535 device_type="npu",
536):
537 """
538 Forks the RNG, so that when you return, the RNG is reset
539 to the state that it was previously in.
541 Args:
542 devices (iterable of Device IDs): devices for which to fork
543 the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates
544 on all devices, but will emit a warning if your machine has a lot
545 of devices, since this function will run very slowly in that case.
546 If you explicitly specify devices, this warning will be suppressed
547 enabled (bool): if ``False``, the RNG is not forked. This is a convenience
548 argument for easily disabling the context manager without having
549 to delete it and unindent your Python code under it.
550 device_type (str): device type str, default is `npu`. As for supported device,
551 see details in :ref:`accelerator<accelerators>`
552 """
554 device_mod = platform.get_device_handle()
555 if device_mod is None:
556 raise RuntimeError(
557 f"{platform} has no module of `{device_type}`, you should register "
558 )
559 global _fork_rng_warned_already
561 if not enabled:
562 yield
563 return
565 if devices is None:
566 num_devices = platform.device_count(device_mod)
567 if num_devices > 1 and not _fork_rng_warned_already:
568 _fork_rng_warned_already = True
569 devices = list(range(num_devices))
570 else:
571 # Protect against user passing us a generator; we need to traverse this
572 # multiple times but a generator will be exhausted upon first traversal
573 devices = list(devices)
575 cpu_rng_state = platform.get_rng_state()
576 device_rng_states = [platform.get_rng_state(device, device_mod) for device in devices]
578 try:
579 yield
580 finally:
581 platform.set_rng_state(cpu_rng_state)
582 for device, device_rng_state in zip(devices, device_rng_states):
583 platform.set_rng_state(device_rng_state, device, device_mod)