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-10-25 05:45 +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"""
20# pylint: disable=C9006,C9007,W0603
22__all__ = [
23 "is_rng_supported_mesh",
24 "manual_seed",
25 "OffsetBasedRNGTracker",
26]
28import contextlib
29import warnings
30from logging import getLogger
31import typing
32from typing import Optional
33import functools
34import operator
36import torch
37import torch.distributed as dist
39from hyper_parallel.core.dtensor import _utils
40from hyper_parallel.core.dtensor.placement_types import Shard
41from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
43Tensor = torch.tensor
45logger = getLogger(__name__)
48def is_rng_supported_mesh(device_mesh: Optional[DeviceMesh] = None) -> bool:
49 """Check if the device mesh supports DTensor random operations.
51 Currently, DTensor random operations are only supported on CUDA and CUDA-like
52 devices. Users should call this function before using DTensor random APIs to
53 verify compatibility.
55 Args:
56 device_mesh: Optional :class:`DeviceMesh` to check (same semantics as PyTorch
57 ``torch.distributed.tensor``). If omitted, only the active device handle
58 is checked.
60 Returns:
61 bool: ``True`` if the device mesh supports DTensor random operations,
62 ``False`` otherwise.
63 """
64 if device_mesh is not None and device_mesh.device_type == "cpu":
65 warnings.warn(
66 f"DTensor random operators may not have complete support on {device_mesh.device_type} device mesh",
67 stacklevel=2,
68 )
69 return False
70 device_handle = _utils.get_device_handle()
71 if device_handle and hasattr(device_handle, "set_rng_state"):
72 return True
73 if device_mesh is not None:
74 warnings.warn(
75 f"DTensor random operators may not have complete support on {device_mesh.device_type} device mesh",
76 stacklevel=2,
77 )
78 return False
81class _PhiloxState:
82 """
83 Convenience accessor for interpreting the packed bits of (seed: uint64, offset: uint64) in the philox state,
84 which for some reason is actually exposed as a size-16 uint8 tensor.
86 The state is always moved to .cpu since it is necessary for it to be on CPU before applying it back to a generator.
87 """
89 def __init__(self, state: Tensor):
90 self._state = state.to("cpu")
92 @property
93 def state(self):
94 """Return the underlying RNG state tensor (CPU uint8)."""
95 return self._state
97 @property
98 def offset(self) -> int:
99 """Return the offset value (last 8 bytes) of the Philox RNG state."""
100 return int(self._state[8:].view(dtype=torch.int64).item())
102 @offset.setter
103 def offset(self, offset: int) -> None:
104 """Set the offset value of the Philox RNG state."""
105 offset_tensor = Tensor([offset], dtype=torch.uint64).view(
106 torch.uint8
107 ) # device?
108 self._state[8:] = offset_tensor
110 @property
111 def seed(self) -> int:
112 """Return the seed value (first 8 bytes) of the Philox RNG state."""
113 return int(self._state[:8].view(dtype=torch.uint64).item())
115 @seed.setter
116 def seed(self, seed: int) -> None:
117 """Set the seed value of the Philox RNG state."""
118 seed_tensor = Tensor([seed], dtype=torch.uint64).view(
119 torch.uint8
120 )# device
121 self._state[:8] = seed_tensor
124class _RNGStateTracker:
125 """
126 Tracks and manages RNG states for DTensor random operations.
128 Maintains a mapping from operation tags to RNG state tensors (ByteTensor),
129 providing standardized interfaces for state access and modification.
131 The core method `_distribute_region` establishes the proper RNG context
132 when DTensor executes random operators across distributed devices.
133 """
135 def __init__(self, device):
136 self._device = device
137 self._device_handle = _utils.get_device_handle()
138 if not self._device_handle:
139 raise RuntimeError(
140 f"{self.__class__.__name__} instantiation requires the presence of "
141 )
142 self._use_distribute_region = True
144 @property
145 def distribute_region_enabled(self) -> bool:
146 """Return whether the RNG distribute region is enabled for distributed random operations."""
147 return self._use_distribute_region
149 @distribute_region_enabled.setter
150 def distribute_region_enabled(self, value) -> None:
151 """Set whether the RNG distribute region is enabled."""
152 self._use_distribute_region = value
154 def _distribute_region(
155 self, device_mesh, placements, global_shape, generator = None
156 ):
157 pass
159 def _manual_seed(self, parallel_seed: int) -> None:
160 pass
163class OffsetBasedRNGTracker(_RNGStateTracker):
164 """
165 This subclass of ``_RNGStateTracker`` defines the default policy of how RNG states
166 should be shared and synchronized among all ranks to respect the semantics of DTensor
167 random operators.
168 """
170 def __init__(
171 self,
172 run_state_sync: bool = True,
173 ):
174 super().__init__(_resolve_device())
175 rng_state = self._get_device_state()
176 if run_state_sync:
177 # synchronize RNG state using rank 0's current one
178 dist.broadcast(rng_state, 0)
179 my_rng_state = self._get_device_state()
180 if not all(my_rng_state == rng_state):
181 logger.warning(
182 "DTensor is synchronizing RNG states of every rank with the state from rank 0. "
183 "This behavior is deprecated. "
184 "Please call ``manual_seed(seed, device_mesh)`` from "
185 "``hyper_parallel.core.dtensor.random`` on every rank that participates in SPMD DTensor "
186 "operations with the same seed. If using Pipeline Parallelism, each pipelining state would use "
187 "a different seed, but all ranks belonging to one pipeline stage would use the same seed."
188 )
189 self._set_device_state(rng_state)
191 def _manual_seed(self, parallel_seed: int) -> None:
192 """Set the default RNG seed (``torch.manual_seed``), as in PyTorch DTensor."""
193 torch.manual_seed(parallel_seed)
195 def _get_device_state(self):
196 rng_state = self._device_handle.get_rng_state().to(self._device)
197 return rng_state
199 def _set_device_state(self, state: Tensor):
200 # It seems that the underlying generator wants a cpu tensor but the dtensor code expects `_get_device_state`
201 # to convert to a 'device' tensor, probably because we may use it with our backend comms for sync/debug
202 # for now, we just convert back to cpu here to make sure it always works.
203 self._device_handle.set_rng_state(state.to("cpu"))
205 @contextlib.contextmanager
206 def _distribute_region(
207 self, device_mesh, placements, global_shape, generator = None
208 ):
210 # regular (non-LocalTensor) mode
211 if generator is not None:
212 # This is a little hacky, but for any user-passed generator, we store its state under a unique key,
213 # not because we need to keep a copy of it but because its the easiest way to make it work with the
214 # existing set/get APIs. We also ensure we remove it from rng_states after each _distribute_region.
215 state = _PhiloxState(generator.get_state())
216 else:
217 state = _PhiloxState(self._get_device_state())
219 if self.distribute_region_enabled:
220 old_offset = state.offset
221 self._set_pre_op_offset(state, device_mesh, placements, global_shape)
222 with fork_rng(
223 devices=[self._device], device_type=_utils.device_type()
224 ):
225 self._device_handle.set_rng_state(state.state)
226 try:
227 yield # execute the region code
228 finally:
229 # update offset to synchronize among ranks
230 self._set_post_op_offset(state, global_shape, old_offset)
232 else:
233 yield
235 if generator is not None:
236 # ensure we (a) propagate the state advancement back to the user's RNG so its visible and impacts any future
237 # usage of that RNG (dtensor or non-dtensor), (b) drop it from our own cache so that if the user updates
238 # the seed value in their rng and uses it with DTensor again, we always use the latest value
239 generator.set_state(state.state)
240 else:
241 self._set_device_state(state.state)
243 def compute_offset_incr(self, device_mesh, placements, global_shape) -> int:
244 """Compute the per-shard RNG offset increment for the current rank.
246 Based on the shard linear index and local shard size, computes how much to
247 advance the offset so that each shard gets a unique portion of the random stream.
249 Args:
250 device_mesh (DeviceMesh): The device mesh describing the device topology.
251 placements (Sequence[Placement]): The placement strategy for each mesh dimension.
252 global_shape: input global shape
254 Returns:
255 int: The offset increment, 4-byte aligned.
256 """
257 mesh_coordinate = device_mesh.get_coordinate()
258 shard_idx_by_dim, total_num_shards_by_dim = _calc_shard_info(
259 mesh_coordinate,
260 device_mesh,
261 placements,
262 global_shape,
263 reverse_repeated_shards=True,
264 )
265 shard_linear_idx = self._calc_shard_linear_idx(
266 shard_idx_by_dim, total_num_shards_by_dim
267 )
268 local_size_on_rank_0 = _calc_first_shard_size(device_mesh, placements, global_shape)
269 local_size = functools.reduce(operator.mul, local_size_on_rank_0, 1)
270 return (shard_linear_idx * local_size + 3) // 4 * 4
272 def _set_pre_op_offset(self, state: _PhiloxState, device_mesh, placements, global_shape) -> None:
273 """Set the starting random number generator (RNG) offset for the local shard
274 on the current process before operation execution.The offset value begins from
275 the current accumulated position and increments by the local shard size until
276 covering the total elements of the global distributed tensor. Multiple processes
277 holding replicas of the same shard will share identical starting offset values.
279 Args:
280 state (`Tensor`): The generator state to modify
281 device_mesh (DeviceMesh): The device mesh describing the device topology.
282 placements (Sequence[Placement]): The placement strategy for each mesh dimension.
283 Each element should be a Placement object (Shard, Replicate, Partial, etc.).
284 global_shape: input global shape
286 Returns:
287 None
289 .. warning::
290 The current implementation does not consider memory layout contiguity.
292 Example:
293 take a DTensor of shape [8, 16] as an example. Assume that the DTensor
294 is placed on a device mesh with placements ([Shard(1), Replicate(), Shard(0)]),
295 and the mesh is:
296 [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
297 ``mesh.get_coordinate()`` provides the coordinate of the current rank
298 in the mesh. For example, the coordinate of rank 5 is (1, 0, 1).
300 Another concept to introduce besides rank coordinate is shard coordinate.
301 Each rank holds a local shard of the DTensor. In the example, the DTensor
302 is partitioned into 4 [4, 8] shards. The first shard has 2 replicas and
303 rank 0 (coord (0, 0, 0)) and rank 2 (coord (0, 1, 0)) have 1 replica each.
304 That being said, the local shard on rank 0 and rank 2 correspond to the same
305 shard of the DTensor. To denote each DTensor shard, we use a shard coordinate
306 (in the example, it will be a tuple (i, j) where shard (i, j) has the slice
307 DTensor[4 * i : 4 * (i + 1), 8 * j : 8 * (j + 1)], 0 <= i < 2, 0 <= j < 2).
309 Once we have rank coordinate and shard coordinate, we can calculate on each rank
310 what shard of the DTensor the rank holds, with the help of dim_map. The dim_map
311 of the above DTensor is [2, 0] so the shard coordinate of a rank with rank coord
312 (x, y, z) is simply (z, x) by taking(rank_coord[dim_map[0]],rank_coord[dim_map[1]]).
313 Following this calculation,
314 rank 0 and rank 2 holds the shard of coord (0, 0);
315 rank 1 and rank 3 holds the shard of coord (0, 1);
316 rank 4 and rank 6 holds the shard of coord (1, 0);
317 rank 5 and rank 7 holds the shard of coord (1, 1);
319 The last value to calculate before obtaining the starting offset is the shard linear index.
320 The starting offset for each rank will be its shard_linear_index * local_tensor_numel.
321 """
322 current_offset = state.offset
323 offset_incr = self.compute_offset_incr(device_mesh, placements, global_shape)
324 state.offset = current_offset + offset_incr
326 def _set_post_op_offset(
327 self, state: _PhiloxState, global_shape, old_offset: int
328 ) -> None:
329 """Sets the RNG to a synchronized state after running the local random op.
330 Restores the random number generator to a globally consistent state following
331 local shard execution. Each process must advance its offset by the total element
332 count of the distributed tensor, measured from the offset value recorded before
333 the operation began.
335 Args:
336 state (`Tensor`): The generator state to modify.
337 global_shape: The global shape of the distributed tensor.
338 old_offset (int): The RNG offset before the operation.
340 Returns:
341 None
342 """
343 numel = functools.reduce(operator.mul, global_shape, 1)
344 numel = (numel + 3) // 4 * 4
345 state.offset = old_offset + numel
347 def _calc_shard_linear_idx(
348 self, shard_coord: list[int], shard_size: list[int]
349 ) -> int:
350 return _calc_shard_linear_idx(shard_coord, shard_size)
353def _calc_first_shard_size(device_mesh, placements, global_shape) -> list[int]:
354 """Calculate the size of the first shard on rank 0.
356 Args:
357 device_mesh: The device mesh describing the device topology.
358 placements: Sequence of Placement objects (Shard, Replicate, etc.).
359 global_shape: input global shape
361 Returns:
362 list[int]: Shape of rank 0's local shard.
363 """
364 local_size_on_rank_0 = list(global_shape)
365 for idx, placement in enumerate(placements):
366 if isinstance(placement, Shard):
367 mesh_dim_size = device_mesh.size(idx)
368 shard_dim = placement.dim
369 current_dim_size = local_size_on_rank_0[shard_dim]
370 local_size_on_rank_0[shard_dim], _ = local_shard_size_and_offset(
371 current_dim_size,
372 mesh_dim_size,
373 0,
374 )
375 return local_size_on_rank_0
378def _calc_shard_info(
379 mesh_coordinate,
380 device_mesh,
381 placements,
382 global_shape,
383 reverse_repeated_shards=False,
384):
385 """Calculate shard information for a specific rank."""
386 mesh_size = device_mesh.mesh_shape
387 # note: dim_map does not allow double sharding which is the FSDP(fully_shard)+TP
388 # case. Replace the custom logic with dim_map once we support it.
389 dim_map = [-1] * len(global_shape)
390 for i, placement in enumerate(placements):
391 if isinstance(placement, Shard):
392 shard_dim = placement.dim
393 if dim_map[shard_dim] == -1:
394 dim_map[shard_dim] = [i]
395 else:
396 mesh_dim_list = dim_map[shard_dim]
397 if not isinstance(mesh_dim_list, list):
398 raise TypeError(f"Expected mesh_dim_list to be a list, got {type(mesh_dim_list)}")
399 mesh_dim_list.append(i)
401 # Compute shard coordinate:
402 # The coordinate on each tensor dim is a tuple (idx, range)
403 # If a DTensor is partitioned on its dim i into n shards, and the current rank
404 # holds the j-th, then its shard coordinate will be (idx=j, range=n) on dim i
405 if mesh_coordinate is None:
406 raise ValueError("mesh_coordinate must not be None")
407 shard_idx_by_dim = []
408 total_num_shards_by_dim = [] # total number of shards on each tensor dim
409 for mesh_dim in dim_map:
410 shard_idx = 0
411 total_num_shards = 1
412 # the tensor dim is sharded on more than 1 mesh dim
413 if isinstance(mesh_dim, list):
414 repeated_mesh_dims = mesh_dim
415 if reverse_repeated_shards:
416 # RNG follows the sequential placement application: the last
417 # placement subdivides each earlier shard and becomes major.
418 repeated_mesh_dims = list(reversed(repeated_mesh_dims))
419 rank_coord = [mesh_coordinate[d] for d in repeated_mesh_dims]
420 num_shards = [mesh_size[d] for d in repeated_mesh_dims]
421 # compute the shard idx and total number of shards
422 for idx, size in zip(rank_coord, num_shards):
423 shard_idx = shard_idx * size + idx
424 total_num_shards *= size
426 shard_idx_by_dim.append(shard_idx)
427 total_num_shards_by_dim.append(total_num_shards)
428 return shard_idx_by_dim, total_num_shards_by_dim
431def _calc_shard_linear_idx(shard_coord: list[int], shard_size: list[int]) -> int:
432 # compute shard linear index
433 shard_linear_idx = 0
434 shard_coord_stride = 1
435 for idx, size in zip(reversed(shard_coord), reversed(shard_size)):
436 shard_linear_idx += idx * shard_coord_stride
437 shard_coord_stride *= size
439 return shard_linear_idx
442def _resolve_device():
443 device_handle = _utils.get_device_handle()
444 device_idx = dist.get_rank() % _utils.device_count(device_handle)
446 def get_device(device_idx):
447 return _utils.device(device_idx)
449 return get_device(device_idx)
452def manual_seed(seed: int, device_mesh: DeviceMesh) -> None:
453 """Set the seed for generating random numbers on the calling rank (PyTorch DTensor parity).
455 Ensures the global RNG used by DTensor random ops is initialized consistently. Lazily
456 creates the :class:`OffsetBasedRNGTracker` used by shard dispatch with
457 ``run_state_sync=False`` so ranks are not synchronized from rank 0's prior RNG state.
459 Args:
460 seed: Desired RNG seed (must be agreed across ranks in the mesh for SPMD).
461 device_mesh: Mesh that must include the current process rank.
463 Raises:
464 RuntimeError: If the current rank is not part of ``device_mesh`` (undefined DTensor
465 RNG behavior in that case).
467 Warning:
468 Does not validate that ``seed`` matches across ranks; callers must ensure SPMD
469 consistency. Pipeline parallel: use one seed per pipeline stage group as in PyTorch.
470 """
471 if not is_rng_supported_mesh(device_mesh):
472 warnings.warn(
473 "DTensor manual_seed() may not have complete support "
474 f"on {device_mesh.device_type} device mesh",
475 stacklevel=2,
476 )
477 return
479 # Local import avoids import cycle: _op_dispatch imports this module at load time.
480 from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER # pylint: disable=C0415
482 if _OP_DISPATCHER._rng_tracker is None:
483 _OP_DISPATCHER._rng_tracker = OffsetBasedRNGTracker(run_state_sync=False)
485 if device_mesh.get_coordinate() is None:
486 raise RuntimeError(
487 "manual_seed requires the current rank to be a part of the device mesh "
488 "otherwise DTensor RNG state on the rank will not be initialized and "
489 "the behavior of DTensor random ops is undefined."
490 )
492 torch.manual_seed(seed)
495def local_shard_size_and_offset(
496 curr_local_size: int,
497 num_chunks: int,
498 rank,
499):
500 """
501 Given the size of the current local tensor (which may already be sharded on some dimensions),
502 computes the new local shard size and offset given the desired number of chunks
503 (num_chunks is generally equal to the size of the current sharding dim).
505 Note: new local shard offset is relative to the current sharded tensor, not the global tensor.
506 See `_utils.compute_local_shape_and_global_offset` for computing global offset.
508 Returns (new local shard size, offset)
510 """
511 # Compute the chunk size inline
512 if curr_local_size % num_chunks == 0:
513 full_chunk_size = curr_local_size // num_chunks
514 shard_starting_idx = full_chunk_size * rank
515 return full_chunk_size, shard_starting_idx
517 # uneven sharding case
518 full_chunk_size = (curr_local_size + num_chunks - 1) // num_chunks
519 shard_starting_idx = full_chunk_size * rank
521 if curr_local_size < shard_starting_idx:
522 return 0, typing.cast(int, curr_local_size)
523 local_shard_size = (
524 min(curr_local_size, shard_starting_idx + full_chunk_size)
525 - shard_starting_idx
526 )
527 return local_shard_size, shard_starting_idx
530_fork_rng_warned_already = False
533@contextlib.contextmanager
534def fork_rng(
535 devices=None,
536 enabled=True,
537 device_type="npu",
538):
539 """
540 Forks the RNG, so that when you return, the RNG is reset
541 to the state that it was previously in.
543 Args:
544 devices (iterable of Device IDs): devices for which to fork
545 the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates
546 on all devices, but will emit a warning if your machine has a lot
547 of devices, since this function will run very slowly in that case.
548 If you explicitly specify devices, this warning will be suppressed
549 enabled (bool): if ``False``, the RNG is not forked. This is a convenience
550 argument for easily disabling the context manager without having
551 to delete it and unindent your Python code under it.
552 device_type (str): device type str, default is `npu`. As for supported device,
553 see details in :ref:`accelerator<accelerators>`
554 """
556 device_mod = _utils.get_device_handle()
557 if device_mod is None:
558 raise RuntimeError(
559 f"torch has no module of `{device_type}`, you should register it first."
560 )
561 global _fork_rng_warned_already
563 if not enabled:
564 yield
565 return
567 if devices is None:
568 num_devices = _utils.device_count(device_mod)
569 if num_devices > 1 and not _fork_rng_warned_already:
570 _fork_rng_warned_already = True
571 devices = list(range(num_devices))
572 else:
573 # Protect against user passing us a generator; we need to traverse this
574 # multiple times but a generator will be exhausted upon first traversal
575 devices = list(devices)
577 cpu_rng_state = _utils.get_rng_state()
578 device_rng_states = [_utils.get_rng_state(device, device_mod) for device in devices]
580 try:
581 yield
582 finally:
583 _utils.set_rng_state(cpu_rng_state)
584 for device, device_rng_state in zip(devices, device_rng_states):
585 _utils.set_rng_state(device_rng_state, device, device_mod)