Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / random.py: 72%
205 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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, device_mesh, placements, global_shape
258 )
259 shard_linear_idx = self._calc_shard_linear_idx(
260 shard_idx_by_dim, total_num_shards_by_dim
261 )
262 local_size_on_rank_0 = _calc_first_shard_size(device_mesh, placements, global_shape)
263 local_size = functools.reduce(operator.mul, local_size_on_rank_0, 1)
264 return (shard_linear_idx * local_size + 3) // 4 * 4
266 def _set_pre_op_offset(self, state: _PhiloxState, device_mesh, placements, global_shape) -> None:
267 """Set the starting random number generator (RNG) offset for the local shard
268 on the current process before operation execution.The offset value begins from
269 the current accumulated position and increments by the local shard size until
270 covering the total elements of the global distributed tensor. Multiple processes
271 holding replicas of the same shard will share identical starting offset values.
273 Args:
274 state (`Tensor`): The generator state to modify
275 device_mesh (DeviceMesh): The device mesh describing the device topology.
276 placements (Sequence[Placement]): The placement strategy for each mesh dimension.
277 Each element should be a Placement object (Shard, Replicate, Partial, etc.).
278 global_shape: input global shape
280 Returns:
281 None
283 .. warning::
284 The current implementation does not consider memory layout contiguity.
286 Example:
287 take a DTensor of shape [8, 16] as an example. Assume that the DTensor
288 is placed on a device mesh with placements ([Shard(1), Replicate(), Shard(0)]),
289 and the mesh is:
290 [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
291 ``mesh.get_coordinate()`` provides the coordinate of the current rank
292 in the mesh. For example, the coordinate of rank 5 is (1, 0, 1).
294 Another concept to introduce besides rank coordinate is shard coordinate.
295 Each rank holds a local shard of the DTensor. In the example, the DTensor
296 is partitioned into 4 [4, 8] shards. The first shard has 2 replicas and
297 rank 0 (coord (0, 0, 0)) and rank 2 (coord (0, 1, 0)) have 1 replica each.
298 That being said, the local shard on rank 0 and rank 2 correspond to the same
299 shard of the DTensor. To denote each DTensor shard, we use a shard coordinate
300 (in the example, it will be a tuple (i, j) where shard (i, j) has the slice
301 DTensor[4 * i : 4 * (i + 1), 8 * j : 8 * (j + 1)], 0 <= i < 2, 0 <= j < 2).
303 Once we have rank coordinate and shard coordinate, we can calculate on each rank
304 what shard of the DTensor the rank holds, with the help of dim_map. The dim_map
305 of the above DTensor is [2, 0] so the shard coordinate of a rank with rank coord
306 (x, y, z) is simply (z, x) by taking(rank_coord[dim_map[0]],rank_coord[dim_map[1]]).
307 Following this calculation,
308 rank 0 and rank 2 holds the shard of coord (0, 0);
309 rank 1 and rank 3 holds the shard of coord (0, 1);
310 rank 4 and rank 6 holds the shard of coord (1, 0);
311 rank 5 and rank 7 holds the shard of coord (1, 1);
313 The last value to calculate before obtaining the starting offset is the shard linear index.
314 The starting offset for each rank will be its shard_linear_index * local_tensor_numel.
315 """
316 current_offset = state.offset
317 offset_incr = self.compute_offset_incr(device_mesh, placements, global_shape)
318 state.offset = current_offset + offset_incr
320 def _set_post_op_offset(
321 self, state: _PhiloxState, global_shape, old_offset: int
322 ) -> None:
323 """Sets the RNG to a synchronized state after running the local random op.
324 Restores the random number generator to a globally consistent state following
325 local shard execution. Each process must advance its offset by the total element
326 count of the distributed tensor, measured from the offset value recorded before
327 the operation began.
329 Args:
330 state (`Tensor`): The generator state to modify.
331 global_shape: The global shape of the distributed tensor.
332 old_offset (int): The RNG offset before the operation.
334 Returns:
335 None
336 """
337 numel = functools.reduce(operator.mul, global_shape, 1)
338 numel = (numel + 3) // 4 * 4
339 state.offset = old_offset + numel
341 def _calc_shard_linear_idx(
342 self, shard_coord: list[int], shard_size: list[int]
343 ) -> int:
344 return _calc_shard_linear_idx(shard_coord, shard_size)
347def _calc_first_shard_size(device_mesh, placements, global_shape) -> list[int]:
348 """Calculate the size of the first shard on rank 0.
350 Args:
351 device_mesh: The device mesh describing the device topology.
352 placements: Sequence of Placement objects (Shard, Replicate, etc.).
353 global_shape: input global shape
355 Returns:
356 list[int]: Shape of rank 0's local shard.
357 """
358 local_size_on_rank_0 = list(global_shape)
359 for idx, placement in enumerate(placements):
360 if isinstance(placement, Shard):
361 mesh_dim_size = device_mesh.size(idx)
362 shard_dim = placement.dim
363 local_size_on_rank_0[shard_dim], _ = local_shard_size_and_offset(
364 global_shape[shard_dim],
365 mesh_dim_size,
366 0,
367 )
368 return local_size_on_rank_0
371def _calc_shard_info(
372 mesh_coordinate, device_mesh, placements, global_shape
373):
374 """Calculate shard information for a specific rank."""
375 mesh_size = device_mesh.mesh_shape
376 # note: dim_map does not allow double sharding which is the FSDP(fully_shard)+TP
377 # case. Replace the custom logic with dim_map once we support it.
378 dim_map = [-1] * len(global_shape)
379 for i, placement in enumerate(placements):
380 if isinstance(placement, Shard):
381 shard_dim = placement.dim
382 if dim_map[shard_dim] == -1:
383 dim_map[shard_dim] = [i]
384 else:
385 mesh_dim_list = dim_map[shard_dim]
386 if not isinstance(mesh_dim_list, list):
387 raise TypeError(f"Expected mesh_dim_list to be a list, got {type(mesh_dim_list)}")
388 mesh_dim_list.append(i)
390 # Compute shard coordinate:
391 # The coordinate on each tensor dim is a tuple (idx, range)
392 # If a DTensor is partitioned on its dim i into n shards, and the current rank
393 # holds the j-th, then its shard coordinate will be (idx=j, range=n) on dim i
394 if mesh_coordinate is None:
395 raise ValueError("mesh_coordinate must not be None")
396 shard_idx_by_dim = []
397 total_num_shards_by_dim = [] # total number of shards on each tensor dim
398 for mesh_dim in dim_map:
399 shard_idx = 0
400 total_num_shards = 1
401 # the tensor dim is sharded on more than 1 mesh dim
402 if isinstance(mesh_dim, list):
403 rank_coord = [mesh_coordinate[d] for d in mesh_dim]
404 num_shards = [mesh_size[d] for d in mesh_dim]
405 # compute the shard idx and total number of shards
406 for idx, size in zip(rank_coord, num_shards):
407 shard_idx = shard_idx * size + idx
408 total_num_shards *= size
410 shard_idx_by_dim.append(shard_idx)
411 total_num_shards_by_dim.append(total_num_shards)
412 return shard_idx_by_dim, total_num_shards_by_dim
415def _calc_shard_linear_idx(shard_coord: list[int], shard_size: list[int]) -> int:
416 # compute shard linear index
417 shard_linear_idx = 0
418 shard_coord_stride = 1
419 for idx, size in zip(reversed(shard_coord), reversed(shard_size)):
420 shard_linear_idx += idx * shard_coord_stride
421 shard_coord_stride *= size
423 return shard_linear_idx
426def _resolve_device():
427 device_handle = platform.get_device_handle()
428 device_idx = platform.get_rank() % platform.device_count(device_handle)
430 def get_device(device_idx):
431 return platform.device(device_idx)
433 return get_device(device_idx)
436def manual_seed(seed: int, device_mesh: DeviceMesh) -> None:
437 """Set the seed for generating random numbers on the calling rank (PyTorch DTensor parity).
439 Ensures the global RNG used by DTensor random ops is initialized consistently. Lazily
440 creates the :class:`OffsetBasedRNGTracker` used by shard dispatch with
441 ``run_state_sync=False`` so ranks are not synchronized from rank 0's prior RNG state.
443 Args:
444 seed: Desired RNG seed (must be agreed across ranks in the mesh for SPMD).
445 device_mesh: Mesh that must include the current process rank.
447 Raises:
448 RuntimeError: If the current rank is not part of ``device_mesh`` (undefined DTensor
449 RNG behavior in that case).
451 Warning:
452 Does not validate that ``seed`` matches across ranks; callers must ensure SPMD
453 consistency. Pipeline parallel: use one seed per pipeline stage group as in PyTorch.
454 """
455 if not is_rng_supported_mesh(device_mesh):
456 warnings.warn(
457 "DTensor manual_seed() may not have complete support "
458 f"on {device_mesh.device_type} device mesh",
459 stacklevel=2,
460 )
461 return
463 # Local import avoids import cycle: _op_dispatch imports this module at load time.
464 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER # pylint: disable=C0415
466 if _OP_DISPATCHER._rng_tracker is None:
467 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
469 if device_mesh.get_coordinate() is None:
470 raise RuntimeError(
471 "manual_seed requires the current rank to be a part of the device mesh "
472 "otherwise DTensor RNG state on the rank will not be initialized and "
473 "the behavior of DTensor random ops is undefined."
474 )
476 platform.manual_seed(seed)
479def local_shard_size_and_offset(
480 curr_local_size: int,
481 num_chunks: int,
482 rank,
483):
484 """
485 Given the size of the current local tensor (which may already be sharded on some dimensions),
486 computes the new local shard size and offset given the desired number of chunks
487 (num_chunks is generally equal to the size of the current sharding dim).
489 Note: new local shard offset is relative to the current sharded tensor, not the global tensor.
490 See `_utils.compute_local_shape_and_global_offset` for computing global offset.
492 Returns (new local shard size, offset)
494 """
495 # Compute the chunk size inline
496 if curr_local_size % num_chunks == 0:
497 full_chunk_size = curr_local_size // num_chunks
498 shard_starting_idx = full_chunk_size * rank
499 return full_chunk_size, shard_starting_idx
501 # uneven sharding case
502 full_chunk_size = (curr_local_size + num_chunks - 1) // num_chunks
503 shard_starting_idx = full_chunk_size * rank
505 if curr_local_size < shard_starting_idx:
506 return 0, typing.cast(int, curr_local_size)
507 local_shard_size = (
508 min(curr_local_size, shard_starting_idx + full_chunk_size)
509 - shard_starting_idx
510 )
511 return local_shard_size, shard_starting_idx
514_fork_rng_warned_already = False
517@contextlib.contextmanager
518def fork_rng(
519 devices=None,
520 enabled=True,
521 device_type="npu",
522):
523 """
524 Forks the RNG, so that when you return, the RNG is reset
525 to the state that it was previously in.
527 Args:
528 devices (iterable of Device IDs): devices for which to fork
529 the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates
530 on all devices, but will emit a warning if your machine has a lot
531 of devices, since this function will run very slowly in that case.
532 If you explicitly specify devices, this warning will be suppressed
533 enabled (bool): if ``False``, the RNG is not forked. This is a convenience
534 argument for easily disabling the context manager without having
535 to delete it and unindent your Python code under it.
536 device_type (str): device type str, default is `npu`. As for supported device,
537 see details in :ref:`accelerator<accelerators>`
538 """
540 device_mod = platform.get_device_handle()
541 if device_mod is None:
542 raise RuntimeError(
543 f"{platform} has no module of `{device_type}`, you should register "
544 )
545 global _fork_rng_warned_already
547 if not enabled:
548 yield
549 return
551 if devices is None:
552 num_devices = platform.device_count(device_mod)
553 if num_devices > 1 and not _fork_rng_warned_already:
554 _fork_rng_warned_already = True
555 devices = list(range(num_devices))
556 else:
557 # Protect against user passing us a generator; we need to traverse this
558 # multiple times but a generator will be exhausted upon first traversal
559 devices = list(devices)
561 cpu_rng_state = platform.get_rng_state()
562 device_rng_states = [platform.get_rng_state(device, device_mod) for device in devices]
564 try:
565 yield
566 finally:
567 platform.set_rng_state(cpu_rng_state)
568 for device, device_rng_state in zip(devices, device_rng_states):
569 platform.set_rng_state(device_rng_state, device, device_mod)