chore: import upstream snapshot with attribution
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,179 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
from paddle.base import core
|
||||
|
||||
from .custom_streams import ( # noqa: F401
|
||||
Event,
|
||||
Stream,
|
||||
create_event,
|
||||
create_stream,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import CPUPlace
|
||||
|
||||
_CPUPlaceLike: TypeAlias = (
|
||||
CPUPlace
|
||||
| str # some string like "iluvatar_gpu" "metax_gpu:0", etc.
|
||||
| int # some int like 0, 1, etc.
|
||||
)
|
||||
|
||||
|
||||
def device_count() -> int:
|
||||
'''
|
||||
Return the number of GPUs available.
|
||||
|
||||
Returns:
|
||||
int: the number of GPUs available.
|
||||
|
||||
Note:
|
||||
This function returns 0 when compiled without CUDA support.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.device.device_count()
|
||||
|
||||
'''
|
||||
return 0
|
||||
|
||||
|
||||
def get_rng_state(
|
||||
device: _CPUPlaceLike | None = None,
|
||||
) -> core.GeneratorState:
|
||||
r'''
|
||||
Get the random state for the default generator.
|
||||
|
||||
Returns:
|
||||
Tensor: The random state tensor.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||||
>>> import paddle
|
||||
>>> paddle.device.get_rng_state()
|
||||
|
||||
'''
|
||||
return core.default_cpu_generator().get_state()
|
||||
|
||||
|
||||
def set_rng_state(
|
||||
new_state: core.GeneratorState, device: _CPUPlaceLike | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Set the random number generator state of the specified device.
|
||||
|
||||
Args:
|
||||
new_state (core.GeneratorState): The desired RNG state to set.
|
||||
This should be a state object previously obtained from ``get_rng_state()``.
|
||||
device (DeviceLike, optional): The device to set the RNG state for.
|
||||
If not specified, uses the current default device (as returned by ``paddle.framework._current_expected_place_()``).
|
||||
Can be a device object, integer device ID, or device string.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> # Save RNG state
|
||||
>>> state = paddle.device.get_rng_state()
|
||||
>>> # Do some random operations
|
||||
>>> x = paddle.randn([2, 3])
|
||||
>>> # Restore RNG state
|
||||
>>> paddle.device.set_rng_state(state)
|
||||
"""
|
||||
core.default_cpu_generator().set_state(new_state)
|
||||
|
||||
|
||||
def manual_seed(seed: int) -> None:
|
||||
r"""Set the seed for generating random numbers for the current Device.
|
||||
|
||||
.. warning::
|
||||
If you are working with a multi-Device model, this function is insufficient
|
||||
to get determinism. To seed all Devices, use :func:`manual_seed_all`.
|
||||
|
||||
Sets the seed for global default generator, which manages the random number generation.
|
||||
|
||||
Args:
|
||||
seed(int): The random seed to set.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.manual_seed(102)
|
||||
>>> # paddle.cuda.manual_seed(102) is equivalent to paddle.device.manual_seed(102)
|
||||
>>> paddle.cuda.manual_seed(102)
|
||||
|
||||
"""
|
||||
seed = int(seed)
|
||||
core.default_cpu_generator().manual_seed(seed)
|
||||
|
||||
|
||||
def max_memory_allocated(device: _CPUPlaceLike | None = None) -> int:
|
||||
r"""
|
||||
The API max_memory_allocated is not supported in CPU PaddlePaddle.
|
||||
Please reinstall PaddlePaddle with GPU or XPU support to call this API.
|
||||
"""
|
||||
raise ValueError(
|
||||
"The API paddle.device.max_memory_allocated is not supported in CPU PaddlePaddle. "
|
||||
"Please reinstall PaddlePaddle with GPU or XPU support to call this API."
|
||||
)
|
||||
|
||||
|
||||
def max_memory_reserved(device: _CPUPlaceLike | None = None) -> int:
|
||||
r"""
|
||||
The API max_memory_reserved is not supported in CPU PaddlePaddle.
|
||||
Please reinstall PaddlePaddle with GPU or XPU support to call this API.
|
||||
"""
|
||||
raise ValueError(
|
||||
"The API paddle.device.max_memory_reserved is not supported in CPU PaddlePaddle. "
|
||||
"Please reinstall PaddlePaddle with GPU or XPU support to call this API."
|
||||
)
|
||||
|
||||
|
||||
def reset_max_memory_allocated(device: _CPUPlaceLike | None = None) -> None:
|
||||
r"""
|
||||
The API reset_max_memory_allocated is not supported in CPU PaddlePaddle.
|
||||
Please reinstall PaddlePaddle with GPU or XPU support to call this API.
|
||||
"""
|
||||
raise ValueError(
|
||||
"The API paddle.device.reset_max_memory_allocated is not supported in CPU PaddlePaddle. "
|
||||
"Please reinstall PaddlePaddle with GPU or XPU support to call this API."
|
||||
)
|
||||
|
||||
|
||||
def reset_max_memory_reserved(device: _CPUPlaceLike | None = None) -> None:
|
||||
r"""
|
||||
The API reset_max_memory_reserved is not supported in CPU PaddlePaddle.
|
||||
Please reinstall PaddlePaddle with GPU or XPU support to call this API.
|
||||
"""
|
||||
raise ValueError(
|
||||
"The API paddle.device.reset_max_memory_reserved is not supported in CPU PaddlePaddle. "
|
||||
"Please reinstall PaddlePaddle with GPU or XPU support to call this API."
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,500 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections import deque
|
||||
from enum import Enum
|
||||
|
||||
import paddle
|
||||
from paddle.base import log_helper
|
||||
|
||||
from .graphs import CUDAGraph
|
||||
|
||||
# CUDAGraphedLayer Debug tools
|
||||
enable_debug_print = bool(
|
||||
int(os.getenv('PADDLE_DEBUG_ENABLE_CUDAGRAPH_LAYER_LOGGING', '0'))
|
||||
)
|
||||
debug_cudagraphedlayer_fallback_to_default = bool(
|
||||
int(os.getenv('PADDLE_DEBUG_CUDAGRAPHEDLAYER_FALLBACK_TO_DEFAULT', '0'))
|
||||
)
|
||||
|
||||
logger = log_helper.get_logger(
|
||||
__name__, logging.INFO, fmt='[%(levelname)s] %(message)s'
|
||||
)
|
||||
|
||||
|
||||
def debug_print(x):
|
||||
if not enable_debug_print:
|
||||
return
|
||||
logger.info(x)
|
||||
|
||||
|
||||
def print_tensor(
|
||||
t,
|
||||
name="Unnamed",
|
||||
print_meta=True,
|
||||
print_ptr=False,
|
||||
print_hash=True,
|
||||
hash=None,
|
||||
):
|
||||
output = []
|
||||
if name:
|
||||
output.append(name)
|
||||
if hash is None:
|
||||
hash = lambda t: float((t.astype('float32') * 1000).sum())
|
||||
|
||||
if t is None:
|
||||
debug_print(f"{name} is None")
|
||||
elif isinstance(t, paddle.Tensor):
|
||||
if print_meta:
|
||||
output.append(f"shape = {t.shape}")
|
||||
output.append(f"place = {t.place}")
|
||||
if print_ptr:
|
||||
output.append(f"ptr = {hex(t.data_ptr())}")
|
||||
if print_hash:
|
||||
output.append(f"hash = {hash(t)}")
|
||||
debug_print(" | ".join(output))
|
||||
|
||||
|
||||
def printer(x, banner="printer"):
|
||||
if not enable_debug_print:
|
||||
return
|
||||
debug_print(banner.center(100, "-"))
|
||||
recursive_apply(print_tensor, x)
|
||||
|
||||
|
||||
# We need this function, for any kind of inputs with iterables
|
||||
# we recursively apply the function to the leave nodes
|
||||
def recursive_apply(function, input_var):
|
||||
if isinstance(input_var, list):
|
||||
return [recursive_apply(function, item) for item in input_var]
|
||||
elif isinstance(input_var, tuple):
|
||||
return tuple(recursive_apply(function, item) for item in input_var)
|
||||
elif isinstance(input_var, dict):
|
||||
return {
|
||||
key: recursive_apply(function, value)
|
||||
for key, value in input_var.items()
|
||||
}
|
||||
else:
|
||||
return function(input_var)
|
||||
|
||||
|
||||
def detach_tensor(tensor):
|
||||
# Detach an individual tensor and preserve its 'stop_gradient' property
|
||||
if isinstance(tensor, paddle.Tensor):
|
||||
detached_tensor = tensor.detach()
|
||||
detached_tensor.stop_gradient = tensor.stop_gradient
|
||||
return detached_tensor
|
||||
return tensor
|
||||
|
||||
|
||||
# We try our best to flatten the input to list of tensors
|
||||
# example: args = ((t1,t2),(t3,(t4,t5))) -> [t1, t2, t3, t4, t5]
|
||||
def recursive_flatten(target):
|
||||
ret = []
|
||||
|
||||
def append(arg):
|
||||
if isinstance(arg, paddle.Tensor):
|
||||
# [NOTE] sometimes unnecessary tensors, such as the constant `mask` tensor in the PP layer, is passed into subsequent layers.
|
||||
# When a tensor is marked with `stop_gradient=True`, it indicates that it does not contribute to gradient calculations,
|
||||
# suggesting it's unrelated to the main computational process.
|
||||
# Therefore, I try to eliminate the copying of such tensors in the to optimize performance.
|
||||
# if not arg.stop_gradient:
|
||||
# [NOTE] However, `stop_gradient=True` propagation rules within the framework appear to be flawed, so directly eliminate stop_gradient may cause bug
|
||||
ret.append(arg)
|
||||
|
||||
recursive_apply(append, target)
|
||||
return ret
|
||||
|
||||
|
||||
# input any kind of args / kwargs structure, output list of tensor
|
||||
def recursive_flatten_args_kwargs(args, kwargs):
|
||||
return [
|
||||
*recursive_flatten(args),
|
||||
*recursive_flatten(tuple(kwargs.values())),
|
||||
]
|
||||
|
||||
|
||||
detach = lambda x: recursive_apply(detach_tensor, x)
|
||||
|
||||
|
||||
def get_grad_tensor(x):
|
||||
"""Returns the gradient of a Paddle Tensor if it's a tensor; otherwise, returns the input."""
|
||||
if isinstance(x, paddle.Tensor):
|
||||
if x.stop_gradient:
|
||||
return None
|
||||
else:
|
||||
return x.grad
|
||||
return None
|
||||
|
||||
|
||||
# CUDA Graph with Static Input and Output
|
||||
class CUDAGraphWithStaticInputOutput:
|
||||
def __init__(self, num_warmup_steps):
|
||||
self.num_warmup_steps = num_warmup_steps
|
||||
self.graph = CUDAGraph()
|
||||
|
||||
self.has_recorded = False
|
||||
|
||||
self.has_preserved_inputs = False
|
||||
self.args_static = None
|
||||
self.kwargs_static = None
|
||||
|
||||
# inputs is the recursively flattened args and kwargs
|
||||
self.inputs_static = None
|
||||
self.outputs_static = None
|
||||
|
||||
def preserve_or_copy(self, args, kwargs):
|
||||
"""
|
||||
For the CUDA Graph, it is crucial that the buffer remains address-stable,
|
||||
meaning that the buffer addresses for any inputs to the CUDA Graph should not change.
|
||||
One solution to achieve this is to preserve all input tensors.
|
||||
|
||||
This function attempts to recursively flatten the input arguments and keyword arguments
|
||||
to identify all tensors passed to the layer (though it may still miss some due to other implicit
|
||||
ways inputs can be passed to a layer). It then preserves references to these input tensors
|
||||
as `self.inputs_static` so that the buffer pointers can be reused later.
|
||||
|
||||
When this method is called subsequently, it copies the values back to the preserved input tensors
|
||||
to ensure the buffers are reused.
|
||||
"""
|
||||
if not self.has_preserved_inputs:
|
||||
self.has_preserved_inputs = True
|
||||
self.args_static = args
|
||||
self.kwargs_static = kwargs
|
||||
self.inputs_static = recursive_flatten_args_kwargs(
|
||||
self.args_static, self.kwargs_static
|
||||
)
|
||||
else:
|
||||
inputs = recursive_flatten_args_kwargs(args, kwargs)
|
||||
for x_static, x in zip(self.inputs_static, inputs):
|
||||
x_static.copy_(x, True)
|
||||
|
||||
def record(self, f, *args, **kwargs):
|
||||
self.preserve_or_copy(args, kwargs)
|
||||
|
||||
self.graph.capture_begin()
|
||||
self.outputs_static = f(*self.args_static, **self.kwargs_static)
|
||||
self.graph.capture_end()
|
||||
debug_print(
|
||||
"[CUDAGraph] Record-Replay Start (Graph is replayed for the first time)"
|
||||
)
|
||||
self.graph.replay()
|
||||
|
||||
self.has_recorded = True
|
||||
|
||||
return self.outputs_static
|
||||
|
||||
def set_output_static(self, outputs_static):
|
||||
self.outputs_static = outputs_static
|
||||
|
||||
def replay(self, *args, **kwargs):
|
||||
if not self.has_recorded:
|
||||
raise RuntimeError("Graph should be recorded first")
|
||||
|
||||
self.preserve_or_copy(args, kwargs)
|
||||
|
||||
debug_print("[CUDAGraph] Replay Start")
|
||||
self.graph.replay()
|
||||
return self.outputs_static
|
||||
|
||||
def save(self, name):
|
||||
logging.info(f"save graph to {name}")
|
||||
self.graph.print_to_dot_files(name)
|
||||
|
||||
|
||||
# CUDA Graph Layer Status Enumeration
|
||||
class CUDAGraphLayerStatus(Enum):
|
||||
"""Enum to represent the status of a CUDA Graph Layer."""
|
||||
|
||||
WARMUP = 1
|
||||
RECORD = 2
|
||||
CUDAGRAPH = 3
|
||||
|
||||
|
||||
class CUDAGraphForwardBackward:
|
||||
def __init__(self, num_warmup_steps):
|
||||
self.forward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
|
||||
self.backward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
|
||||
self.status = CUDAGraphLayerStatus.RECORD
|
||||
|
||||
def capture_end(self):
|
||||
self.status = CUDAGraphLayerStatus.CUDAGRAPH
|
||||
|
||||
def is_record_step(self):
|
||||
return self.status == CUDAGraphLayerStatus.RECORD
|
||||
|
||||
def is_cuda_graph_step(self):
|
||||
return self.status == CUDAGraphLayerStatus.CUDAGRAPH
|
||||
|
||||
|
||||
class CUDAGraphContext:
|
||||
"""
|
||||
Manages the context for CUDA graph execution in layers. This includes handling
|
||||
the state of CUDA graph layers, managing forward and backward graphs, and
|
||||
tracking the execution steps.
|
||||
"""
|
||||
|
||||
def __init__(self, layer, num_warmup_steps):
|
||||
"""
|
||||
Initializes the CUDA graph context.
|
||||
:param layer: The layer to be used in the CUDA graph.
|
||||
:param num_warmup_steps: Number of warmup steps before recording starts.
|
||||
"""
|
||||
self.layer = layer
|
||||
self.num_warmup_steps = num_warmup_steps
|
||||
|
||||
# The state of context is in either WARMUP or CUDAGRAPH
|
||||
self._step = 0
|
||||
self.status = CUDAGraphLayerStatus.WARMUP
|
||||
|
||||
# Queue to support 1f1b/interleaved scheduler, assuming FIFO order
|
||||
# data queue
|
||||
self.data_queue = deque()
|
||||
|
||||
# graph queue
|
||||
self.graph_queue = deque()
|
||||
|
||||
# Graph Operations
|
||||
def get_graph(self):
|
||||
if len(self.graph_queue) == 0:
|
||||
return CUDAGraphForwardBackward(self.num_warmup_steps)
|
||||
else:
|
||||
return self.graph_queue.popleft()
|
||||
|
||||
def reuse_graph(self, g):
|
||||
self.graph_queue.append(g)
|
||||
|
||||
# Tensor Queue Operations
|
||||
def push_data(self, args):
|
||||
self.data_queue.append(args)
|
||||
|
||||
def pop_data(self):
|
||||
return self.data_queue.popleft()
|
||||
|
||||
# Finite State Machine of Layer State
|
||||
def warmup_step(self):
|
||||
self._step += 1
|
||||
if self._step == self.num_warmup_steps:
|
||||
self.status = CUDAGraphLayerStatus.CUDAGRAPH
|
||||
|
||||
def is_warmup_step(self):
|
||||
return self.status == CUDAGraphLayerStatus.WARMUP
|
||||
|
||||
def is_cuda_graph_step(self):
|
||||
return self.status == CUDAGraphLayerStatus.CUDAGRAPH
|
||||
|
||||
|
||||
def select_y_with_grad(ys, dys):
|
||||
# [TODO] when there is multiple output tensor, we support only one y that allows backward
|
||||
y, dy = None, None
|
||||
if isinstance(ys, paddle.Tensor):
|
||||
y, dy = ys, dys[0]
|
||||
elif isinstance(ys, (list, tuple)):
|
||||
for v, dv in zip(ys, dys):
|
||||
if isinstance(v, paddle.Tensor) and (not v.stop_gradient):
|
||||
y, dy = v, dv
|
||||
break
|
||||
assert isinstance(y, paddle.Tensor) and isinstance(dy, paddle.Tensor)
|
||||
return y, dy
|
||||
|
||||
|
||||
# we get the output of the backward from the detached inputs after the backward is calculated
|
||||
# we save it to the graph itself
|
||||
def get_args_grad(inputs):
|
||||
grad_inputs, detached_grad_inputs = inputs
|
||||
args_grad = []
|
||||
for x, detached_x in zip(grad_inputs, detached_grad_inputs):
|
||||
# if required grad
|
||||
if not x.stop_gradient:
|
||||
if detached_x.grad is None:
|
||||
# if input requires grad but we don't have grad, we just allocate some zeros
|
||||
# x.stop_gradient = True
|
||||
args_grad.append(paddle.zeros(detached_x.shape))
|
||||
# args_grad.append(None)
|
||||
else:
|
||||
args_grad.append(detached_x.grad)
|
||||
else:
|
||||
args_grad.append(None)
|
||||
return tuple(args_grad)
|
||||
|
||||
|
||||
class _CUDAGraphedLayer(paddle.autograd.PyLayer):
|
||||
"""
|
||||
A custom layer that integrates CUDA Graph recording and execution into PaddlePaddle's autograd system.
|
||||
It handles forward and backward operations differently based on the CUDA graph layer status.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, context, arg_tuple, *grad_inputs):
|
||||
"""
|
||||
Handles the forward pass of the layer. It operates differently based on the
|
||||
context's status: warmup, recording, or CUDA graph step.
|
||||
"""
|
||||
args, kwargs = arg_tuple
|
||||
# Detach all inputs from the computational graph
|
||||
args = detach(args)
|
||||
kwargs = detach(kwargs)
|
||||
|
||||
detached_grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
|
||||
inputs = (grad_inputs, detached_grad_inputs)
|
||||
|
||||
printer(detached_grad_inputs, "Forward input")
|
||||
if (
|
||||
context.is_warmup_step()
|
||||
or debug_cudagraphedlayer_fallback_to_default
|
||||
):
|
||||
debug_print("[CUDAGraph] Forward Step (Default)")
|
||||
|
||||
with paddle.enable_grad():
|
||||
y = context.layer(*args, **kwargs)
|
||||
|
||||
context.push_data((CUDAGraphLayerStatus.WARMUP, None, inputs, y))
|
||||
else:
|
||||
graph = context.get_graph()
|
||||
if graph.is_record_step():
|
||||
# In record step, record the forward pass in CUDA graph
|
||||
debug_print(f"[CUDAGraph] Forward Step (Record) id {id(graph)}")
|
||||
|
||||
def forward(*args, **kwargs):
|
||||
with paddle.enable_grad():
|
||||
return context.layer(*args, **kwargs)
|
||||
|
||||
y = graph.forward_graph.record(forward, *args, **kwargs)
|
||||
|
||||
context.push_data(
|
||||
(CUDAGraphLayerStatus.RECORD, graph, inputs, y)
|
||||
)
|
||||
else:
|
||||
debug_print(f"[CUDAGraph] Forward Step (Graph) id {id(graph)}")
|
||||
y = graph.forward_graph.replay(*args, **kwargs)
|
||||
|
||||
context.push_data(
|
||||
(CUDAGraphLayerStatus.CUDAGRAPH, graph, None, y)
|
||||
)
|
||||
|
||||
debug_print("[CUDAGraph] Forward Step End")
|
||||
|
||||
ctx.save_for_backward(context)
|
||||
printer(y, "Forward output")
|
||||
return detach(y)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *dys):
|
||||
"""
|
||||
Handles the backward pass of the layer. Similar to forward, it handles
|
||||
backward based on the context's status: warmup, record, or CUDAGraph.
|
||||
"""
|
||||
(context,) = ctx.saved_tensor()
|
||||
|
||||
(status, graph, inputs, ys) = context.pop_data()
|
||||
y, dy = select_y_with_grad(ys, dys)
|
||||
|
||||
printer((y, dy), "Backward input")
|
||||
|
||||
if status == CUDAGraphLayerStatus.WARMUP:
|
||||
debug_print("[CUDAGraph] Backward Step (Default)")
|
||||
|
||||
# In warmup step, perform standard backward operation
|
||||
y.backward(dy)
|
||||
args_grad = get_args_grad(inputs)
|
||||
|
||||
context.warmup_step()
|
||||
elif status == CUDAGraphLayerStatus.RECORD:
|
||||
debug_print(f"[CUDAGraph] Backward Step (Record) id {id(graph)}")
|
||||
|
||||
# In record step, record the backward pass in CUDA graph
|
||||
def backward(y, dy):
|
||||
y.backward(dy)
|
||||
|
||||
graph.backward_graph.record(backward, y, dy)
|
||||
|
||||
# [NOTE] the get_args_grad should not put inside backward
|
||||
# the args_grad should be calculated after graph is replayed
|
||||
args_grad = get_args_grad(inputs)
|
||||
graph.backward_graph.set_output_static(args_grad)
|
||||
graph.capture_end()
|
||||
|
||||
context.reuse_graph(graph)
|
||||
elif status == CUDAGraphLayerStatus.CUDAGRAPH:
|
||||
debug_print(f"[CUDAGraph] Backward Step (Graph) id {id(graph)}")
|
||||
|
||||
# In CUDA graph step, replay the recorded graph for backward pass
|
||||
args_grad = graph.backward_graph.replay(y, dy)
|
||||
context.reuse_graph(graph)
|
||||
else:
|
||||
raise RuntimeError("Unknown cuda graph status")
|
||||
|
||||
debug_print("[CUDAGraph] Backward Step End")
|
||||
|
||||
printer(args_grad, "Backward output")
|
||||
return args_grad
|
||||
|
||||
|
||||
class CUDAGraphedLayer(paddle.nn.Layer):
|
||||
"""
|
||||
CUDAGraphedLayer: A PaddlePaddle Layer to convert an eager mode model to utilize CUDA Graphs.
|
||||
|
||||
CUDA Graphs provide a way to capture kernel-level operations of a model and play
|
||||
them back efficiently, allowing for potential speedups in repetitive computations,
|
||||
such as those during training iterations. This layer is a wrapper that enables
|
||||
the usage of CUDA Graphs with PaddlePaddle models.
|
||||
|
||||
Overview:
|
||||
- The layer encapsulates another layer (the model to be converted).
|
||||
- During the first few (num_warmup_steps) iterations, the layer operates in
|
||||
eager mode without any CUDA Graphs.
|
||||
- After the warmup steps, the layer captures the forward and backward computations
|
||||
and replays them using CUDA Graphs in subsequent iterations.
|
||||
|
||||
Usage:
|
||||
model = Model()
|
||||
graphed_model = CUDAGraphedLayer(model)
|
||||
|
||||
Parameters:
|
||||
- layer (paddle.nn.Layer): The PaddlePaddle model/layer to be converted.
|
||||
- num_warmup_steps (int): The number of iterations before the CUDA Graph
|
||||
capture begins. Default is 3.
|
||||
|
||||
Notes:
|
||||
- Restrictions:
|
||||
* CPU-GPU Synchronization: Operations that synchronize the CPU with the GPU, like device to host transfers, are not allowed.
|
||||
* CPU Work: Any operations on the CPU within the captured graph are not recorded.
|
||||
* Memory Address (Pointer) Consistency: Replays consistently read from and write to identical virtual memory addresses.
|
||||
* Dynamic Operations:
|
||||
- Control Flow: Dynamic control flows, especially those based on CPU data like if/else statements, are prohibited.
|
||||
- Tensor Shapes: Dynamic tensor shapes are not supported.
|
||||
|
||||
- Allowed Operations:
|
||||
* CUDA RNG Operations: CUDA-based Random Number Generation operations are allowed.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: paddle.nn.Layer, num_warmup_steps=3):
|
||||
super().__init__()
|
||||
self.context = CUDAGraphContext(layer, num_warmup_steps)
|
||||
self.add_sublayer(f"Graphed {type(layer).__name__}", layer)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
# We collect them into a list of tensor that required grad
|
||||
grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
|
||||
return _CUDAGraphedLayer.apply(
|
||||
self.context, (args, kwargs), *grad_inputs
|
||||
)
|
||||
|
||||
def is_warmup_step(self):
|
||||
return self.context.is_warmup_step()
|
||||
|
||||
def is_cuda_graph_step(self):
|
||||
return self.context.is_cuda_graph_step()
|
||||
@@ -0,0 +1,276 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from typing import NoReturn, overload
|
||||
|
||||
from paddle.base.core import (
|
||||
CUDAPlace,
|
||||
CustomPlace,
|
||||
XPUPlace,
|
||||
get_all_custom_device_type,
|
||||
is_compiled_with_cuda,
|
||||
is_compiled_with_custom_device,
|
||||
is_compiled_with_rocm,
|
||||
is_compiled_with_xpu,
|
||||
)
|
||||
|
||||
|
||||
def check_compiled_with_custom_device():
|
||||
custom_device_flag = False
|
||||
custom_devices_types = get_all_custom_device_type()
|
||||
for device_type in custom_devices_types:
|
||||
if is_compiled_with_custom_device(device_type):
|
||||
custom_device_flag = True
|
||||
break
|
||||
return custom_device_flag
|
||||
|
||||
|
||||
if (
|
||||
is_compiled_with_cuda()
|
||||
or is_compiled_with_rocm()
|
||||
or check_compiled_with_custom_device()
|
||||
or is_compiled_with_xpu()
|
||||
):
|
||||
from paddle.base.core import CUDAGraph as CoreCUDAGraph
|
||||
|
||||
def is_cuda_graph_supported():
|
||||
return True
|
||||
else:
|
||||
CoreCUDAGraph = None
|
||||
|
||||
def is_cuda_graph_supported():
|
||||
return False
|
||||
|
||||
|
||||
def current_expected_place():
|
||||
for device in get_all_custom_device_type():
|
||||
selected_devices = os.getenv(f"FLAGS_selected_{device}s", "0").split(
|
||||
","
|
||||
)
|
||||
device_id = int(selected_devices[0])
|
||||
return CustomPlace(device, device_id)
|
||||
return None
|
||||
|
||||
|
||||
ALL_MODES = ["global", "thread_local", "relaxed"]
|
||||
cuda_graph_id = 0
|
||||
|
||||
|
||||
class CUDAGraph:
|
||||
"""
|
||||
The native Paddle constructor takes ``place``, ``mode``, ``pool_id`` and
|
||||
``enable_replace``; the PyTorch-compatible ``keep_graph`` keyword is
|
||||
accepted as well. ``capture_begin`` additionally accepts the PyTorch
|
||||
keywords ``pool`` and ``capture_error_mode`` so the same instance can be
|
||||
driven from either API style.
|
||||
"""
|
||||
|
||||
@overload
|
||||
def __init__(self, keep_graph: bool, /) -> None: ...
|
||||
|
||||
@overload
|
||||
def __init__(
|
||||
self,
|
||||
place: CUDAPlace | XPUPlace | CustomPlace | None = None,
|
||||
mode: str = "thread_local",
|
||||
pool_id: int | None = None,
|
||||
enable_replace: bool = False,
|
||||
*,
|
||||
keep_graph: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
place=None,
|
||||
mode="thread_local",
|
||||
pool_id=None,
|
||||
enable_replace=False,
|
||||
*,
|
||||
keep_graph: bool = False,
|
||||
):
|
||||
assert CoreCUDAGraph is not None, (
|
||||
"CUDA Graph is only supported on PaddlePaddle compiled with NVIDIA GPU."
|
||||
)
|
||||
|
||||
if isinstance(place, bool):
|
||||
if keep_graph is not False:
|
||||
raise TypeError(
|
||||
"keep_graph is specified both positionally and by keyword"
|
||||
)
|
||||
keep_graph = place
|
||||
place = None
|
||||
|
||||
self._graph = None
|
||||
if place is None and check_compiled_with_custom_device():
|
||||
place = current_expected_place()
|
||||
elif place is None:
|
||||
if is_compiled_with_cuda():
|
||||
device_id = int(os.environ.get('FLAGS_selected_gpus', 0))
|
||||
place = CUDAPlace(device_id)
|
||||
elif is_compiled_with_xpu():
|
||||
device_id = int(os.environ.get('FLAGS_selected_xpus', 0))
|
||||
place = XPUPlace(device_id)
|
||||
else:
|
||||
raise RuntimeError("Not Supported devices")
|
||||
|
||||
self._place = place
|
||||
assert mode in ALL_MODES
|
||||
self._mode = ALL_MODES.index(mode)
|
||||
self._pool_id = pool_id
|
||||
self._enable_replace = enable_replace
|
||||
self._keep_graph = keep_graph
|
||||
self._debug_mode = False
|
||||
|
||||
def capture_begin(
|
||||
self, pool: int | None = None, capture_error_mode: str | None = None
|
||||
) -> None:
|
||||
"""Begin capturing CUDA work on the current stream.
|
||||
|
||||
Args:
|
||||
pool (int, optional): A memory pool token from
|
||||
:func:`paddle.cuda.graph_pool_handle` or another graph's
|
||||
:meth:`pool`. When provided, this graph shares the indicated
|
||||
memory pool. Overrides ``pool_id`` from the constructor.
|
||||
capture_error_mode (str, optional): One of ``'global'``,
|
||||
``'thread_local'``, ``'relaxed'`` (see :data:`ALL_MODES`).
|
||||
When ``None`` (default) the constructor's ``mode`` is used;
|
||||
otherwise it overrides the constructor for this capture and a
|
||||
:class:`UserWarning` is emitted to flag the precedence.
|
||||
Invalid values raise :class:`ValueError`.
|
||||
"""
|
||||
if pool is not None:
|
||||
self._pool_id = pool
|
||||
elif self._pool_id is None:
|
||||
self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()
|
||||
|
||||
if capture_error_mode is None:
|
||||
mode = self._mode
|
||||
else:
|
||||
if capture_error_mode not in ALL_MODES:
|
||||
raise ValueError(
|
||||
f"capture_error_mode must be one of {ALL_MODES}, "
|
||||
f"but got {capture_error_mode!r}."
|
||||
)
|
||||
mode = ALL_MODES.index(capture_error_mode)
|
||||
if mode != self._mode:
|
||||
warnings.warn(
|
||||
f"capture_error_mode={capture_error_mode!r} differs from "
|
||||
f"the constructor mode={ALL_MODES[self._mode]!r}; the "
|
||||
f"explicit capture_error_mode takes precedence for this "
|
||||
f"capture.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
CoreCUDAGraph.begin_capture_with_pool_id(
|
||||
self._place, mode, self._pool_id, self._enable_replace
|
||||
)
|
||||
|
||||
def capture_end(self):
|
||||
self._graph = CoreCUDAGraph.end_capture()
|
||||
|
||||
def _require_captured(self) -> None:
|
||||
"""Raise a clear error if no graph has been captured yet.
|
||||
|
||||
``self._graph`` is only populated by :meth:`capture_end`; methods that
|
||||
consume it (``replay`` / ``reset`` / ``debug_dump`` / ...) would
|
||||
otherwise raise ``AttributeError`` on ``NoneType`` when called too
|
||||
early. Centralizing the check produces a single, actionable message.
|
||||
"""
|
||||
if self._graph is None:
|
||||
raise RuntimeError(
|
||||
"CUDAGraph has not been captured yet. "
|
||||
"Call capture_begin/capture_end first."
|
||||
)
|
||||
|
||||
def instantiate(self) -> CoreCUDAGraph:
|
||||
"""Return the instantiated core CUDA graph held by this wrapper.
|
||||
|
||||
Paddle builds the executable graph eagerly inside :meth:`capture_end`,
|
||||
so by the time this is called the graph is already instantiated. It is
|
||||
kept for source compatibility with ``torch.cuda.CUDAGraph.instantiate``
|
||||
and returns the held core :class:`~paddle.base.core.CUDAGraph` produced
|
||||
by :meth:`capture_end`.
|
||||
"""
|
||||
self._require_captured()
|
||||
return self._graph
|
||||
|
||||
def replay(self):
|
||||
self._require_captured()
|
||||
self._graph.replay()
|
||||
|
||||
def reset(self):
|
||||
self._require_captured()
|
||||
self._graph.reset()
|
||||
|
||||
def pool(self) -> int:
|
||||
"""Return an opaque integer token representing this graph's memory pool.
|
||||
|
||||
The token can be passed as the ``pool`` argument to another graph's
|
||||
:meth:`capture_begin` (or to :class:`paddle.cuda.graph`) so the two
|
||||
graphs share the same memory pool.
|
||||
"""
|
||||
if self._pool_id is None:
|
||||
self._pool_id = CoreCUDAGraph.gen_new_memory_pool_id()
|
||||
return self._pool_id
|
||||
|
||||
def enable_debug_mode(self) -> None:
|
||||
"""Enable debug mode so that :meth:`debug_dump` is permitted."""
|
||||
self._debug_mode = True
|
||||
|
||||
def debug_dump(self, debug_path) -> None:
|
||||
"""Dump the captured graph to ``debug_path`` for inspection.
|
||||
|
||||
:meth:`enable_debug_mode` must be called first.
|
||||
"""
|
||||
if not self._debug_mode:
|
||||
raise RuntimeError(
|
||||
"debug_dump requires debug mode to be enabled first. "
|
||||
"Call enable_debug_mode() before debug_dump()."
|
||||
)
|
||||
self._require_captured()
|
||||
self.print_to_dot_files(debug_path)
|
||||
|
||||
def raw_cuda_graph(self) -> NoReturn:
|
||||
"""Paddle does not expose the raw ``cudaGraph_t`` handle."""
|
||||
raise NotImplementedError(
|
||||
"raw_cuda_graph is not yet supported in Paddle CUDAGraph. "
|
||||
"The underlying cudaGraph_t handle is not exposed by the Python "
|
||||
"binding."
|
||||
)
|
||||
|
||||
def raw_cuda_graph_exec(self) -> NoReturn:
|
||||
"""Paddle does not expose the raw ``cudaGraphExec_t`` handle."""
|
||||
raise NotImplementedError(
|
||||
"raw_cuda_graph_exec is not yet supported in Paddle CUDAGraph. "
|
||||
"The underlying cudaGraphExec_t handle is not exposed by the "
|
||||
"Python binding."
|
||||
)
|
||||
|
||||
def print_to_dot_files(self, dirname, flags=None):
|
||||
if not isinstance(dirname, (str, bytes)):
|
||||
dirname = dirname.name
|
||||
os.makedirs(name=dirname, exist_ok=True)
|
||||
assert os.path.isdir(dirname), (
|
||||
f"The dirname {dirname} should be a directory"
|
||||
)
|
||||
if flags is None:
|
||||
flags = 2047 # only all information. It can be any integer inside [1, 2048)
|
||||
self._graph.print_to_dot_files(dirname, flags)
|
||||
|
||||
def replace_input_ptrs(self, old_ptrs, new_ptrs):
|
||||
self._graph.replace_input_ptrs(old_ptrs, new_ptrs)
|
||||
@@ -0,0 +1,540 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
# --- Constants ---
|
||||
KB = 1024
|
||||
MB = 1024 * 1024
|
||||
GB = 1024 * 1024 * 1024
|
||||
|
||||
|
||||
# --- Formatting Helpers ---
|
||||
def format_size(size_bytes):
|
||||
if size_bytes == 0:
|
||||
return "0 B"
|
||||
if size_bytes < MB:
|
||||
return f"{size_bytes / KB:.2f} KB"
|
||||
if size_bytes < GB:
|
||||
return f"{size_bytes / MB:.2f} MB"
|
||||
return f"{size_bytes / GB:.2f} GB"
|
||||
|
||||
|
||||
def print_table(title, headers, rows):
|
||||
if not rows:
|
||||
return
|
||||
# Calculate widths
|
||||
col_widths = [len(str(h)) for h in headers]
|
||||
for row in rows:
|
||||
for i, cell in enumerate(row):
|
||||
if i < len(col_widths):
|
||||
col_widths[i] = max(col_widths[i], len(str(cell)))
|
||||
col_widths = [w + 2 for w in col_widths]
|
||||
|
||||
# Build lines
|
||||
row_fmt = "|" + "|".join([f"{{:^{w}}}" for w in col_widths]) + "|"
|
||||
header_sep = "+" + "+".join(["=" * w for w in col_widths]) + "+"
|
||||
inner_sep = "+" + "+".join(["-" * w for w in col_widths]) + "+"
|
||||
|
||||
print(f"\n### {title}")
|
||||
print(header_sep)
|
||||
print(
|
||||
"|" + "|".join([f"{h:^{w}}" for h, w in zip(headers, col_widths)]) + "|"
|
||||
)
|
||||
print(header_sep)
|
||||
|
||||
for i, row in enumerate(rows):
|
||||
print(row_fmt.format(*[str(c) for c in row]))
|
||||
if (
|
||||
title == "Block Size Distribution"
|
||||
and (i + 1) % 2 == 0
|
||||
and i != len(rows) - 1
|
||||
):
|
||||
print(inner_sep)
|
||||
elif title != "Block Size Distribution":
|
||||
print(inner_sep)
|
||||
if title == "Block Size Distribution":
|
||||
print(header_sep)
|
||||
|
||||
|
||||
class MemoryAnalysisTool:
|
||||
def __init__(self):
|
||||
raise TypeError("Utility class should not be instantiated.")
|
||||
|
||||
@classmethod
|
||||
def vmm_max_free_size(
|
||||
self, device_id: int | None = None
|
||||
) -> tuple[int, int]:
|
||||
name = 'paddle.device.cuda.vmm_max_free_size'
|
||||
if not (core.is_compiled_with_cuda()):
|
||||
raise ValueError(
|
||||
f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
|
||||
)
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
return core.vmm_max_free_size(device_id)
|
||||
|
||||
@classmethod
|
||||
def vmm_free_block_info(
|
||||
self,
|
||||
device_id: int | None = None,
|
||||
) -> list[list[tuple[int, int]]]:
|
||||
name = 'paddle.device.cuda.vmm_free_block_info'
|
||||
if not (core.is_compiled_with_cuda()):
|
||||
raise ValueError(
|
||||
f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
|
||||
)
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
return core.vmm_free_block_info(device_id)
|
||||
|
||||
@classmethod
|
||||
def all_block_info(
|
||||
self,
|
||||
device_id: int | None = None,
|
||||
) -> list[list[tuple[int, int, bool]]]:
|
||||
name = 'paddle.device.cuda.all_block_info'
|
||||
if not (core.is_compiled_with_cuda()):
|
||||
raise ValueError(
|
||||
f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
|
||||
)
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
info = core.all_block_info(device_id)
|
||||
return [list(chunk) for chunk in info]
|
||||
|
||||
@classmethod
|
||||
def vmm_all_block_info(
|
||||
self,
|
||||
device_id: int | None = None,
|
||||
) -> list[list[tuple[int, int, bool]]]:
|
||||
return self.all_block_info(device_id)
|
||||
|
||||
@classmethod
|
||||
def vmm_large_all_block_info(
|
||||
self,
|
||||
device_id: int | None = None,
|
||||
) -> list[list[tuple[int, int, bool]]]:
|
||||
name = 'paddle.device.cuda.vmm_large_all_block_info'
|
||||
if not (core.is_compiled_with_cuda()):
|
||||
raise ValueError(
|
||||
f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
|
||||
)
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
return core.large_pool_block_info(device_id)
|
||||
|
||||
@classmethod
|
||||
def vmm_small_all_block_info(
|
||||
self,
|
||||
device_id: int | None = None,
|
||||
) -> list[list[tuple[int, int, bool]]]:
|
||||
name = 'paddle.device.cuda.vmm_small_all_block_info'
|
||||
if not (core.is_compiled_with_cuda()):
|
||||
raise ValueError(
|
||||
f"The API {name} is not supported in CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU support to call this API."
|
||||
)
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
return core.small_pool_block_info(device_id)
|
||||
|
||||
@classmethod
|
||||
def memory_summary(self, device_id: int | None = None) -> None:
|
||||
device_id = (
|
||||
device_id
|
||||
if device_id is not None
|
||||
else core.get_cuda_current_device_id()
|
||||
)
|
||||
nvidia_smi_AVAILABLE = False
|
||||
try:
|
||||
# import nvidia_smi, pip install nvidia-ml-py3
|
||||
import nvidia_smi
|
||||
|
||||
nvidia_smi_AVAILABLE = True
|
||||
except ImportError:
|
||||
nvidia_smi_AVAILABLE = False
|
||||
|
||||
THRESHOLDS = [
|
||||
1 * MB,
|
||||
10 * MB,
|
||||
50 * MB,
|
||||
100 * MB,
|
||||
200 * MB,
|
||||
400 * MB,
|
||||
600 * MB,
|
||||
800 * MB,
|
||||
1 * GB,
|
||||
2 * GB,
|
||||
3 * GB,
|
||||
]
|
||||
RANGE_HEADERS = [
|
||||
"[0B,1M)",
|
||||
"[1M,10M)",
|
||||
"[10M,50M)",
|
||||
"[50M,100M)",
|
||||
"[100M,200M)",
|
||||
"[200M,400M)",
|
||||
"[400M,600M)",
|
||||
"[600M,800M)",
|
||||
"[800M,1G)",
|
||||
"[1G,2G)",
|
||||
"[2G,3G)",
|
||||
"[3G,+INF)",
|
||||
]
|
||||
|
||||
allocator_lists = self.all_block_info(device_id=device_id)
|
||||
# --- Feature 1: Global Summary with NVML & Rates ---
|
||||
|
||||
# 1.1 Get Paddle Stats
|
||||
mem_allocated = paddle.device.cuda.memory_allocated()
|
||||
max_mem_allocated = paddle.device.cuda.max_memory_allocated()
|
||||
mem_reserved = paddle.device.cuda.memory_reserved()
|
||||
max_mem_reserved = paddle.device.cuda.max_memory_reserved()
|
||||
|
||||
# 1.2 Calculate Rates (Utilization of the Reserved Pool)
|
||||
# Rate = How much of the reserved pool is actually holding tensor data?
|
||||
max_alloc_rate = (
|
||||
((mem_reserved - max_mem_allocated) / mem_reserved)
|
||||
if mem_reserved > 0
|
||||
else 0.0
|
||||
)
|
||||
|
||||
# 1.3 Get Physical Usage via nvidia_smi
|
||||
phy_used_str = "N/A"
|
||||
if nvidia_smi_AVAILABLE:
|
||||
try:
|
||||
nvidia_smi.nvmlInit()
|
||||
|
||||
handle = nvidia_smi.nvmlDeviceGetHandleByIndex(device_id)
|
||||
info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
|
||||
phy_used_str = format_size(info.used)
|
||||
phy_total_str = format_size(info.total)
|
||||
# nvidia_smi.nvmlShutdown() # Optional, depends on lifecycle
|
||||
except Exception as e:
|
||||
phy_used_str = "Err"
|
||||
phy_total_str = "Err"
|
||||
else:
|
||||
print(
|
||||
"Place install nvidia-smi to check real memory usage, pip install command: `pip install nvidia-ml-py3`"
|
||||
)
|
||||
phy_used_str = "No nvidia_smi"
|
||||
phy_total_str = "No nvidia_smi"
|
||||
|
||||
global_headers = [
|
||||
"Allocators",
|
||||
"Allocated",
|
||||
"Max Alloc",
|
||||
"Reserved",
|
||||
"Max Reserved",
|
||||
"Max Frag Rate",
|
||||
"Phy GPU Used / Total",
|
||||
]
|
||||
|
||||
global_rows = [
|
||||
[
|
||||
len(allocator_lists),
|
||||
format_size(mem_allocated),
|
||||
format_size(max_mem_allocated),
|
||||
format_size(mem_reserved),
|
||||
format_size(max_mem_reserved),
|
||||
f"{max_alloc_rate:.2%}",
|
||||
phy_used_str + ' / ' + phy_total_str,
|
||||
]
|
||||
]
|
||||
|
||||
print_table("Global Memory Snapshot", global_headers, global_rows)
|
||||
|
||||
# --- 2. Allocator Analysis ---
|
||||
summary_rows = []
|
||||
dist_rows = []
|
||||
|
||||
for idx, blocks in enumerate(allocator_lists):
|
||||
allocator_name = f"Allocator_{idx}"
|
||||
|
||||
# A. Basic Counting
|
||||
total_blocks = len(blocks)
|
||||
free_blocks = 0
|
||||
total_size = 0
|
||||
free_size = 0
|
||||
max_free_size = 0
|
||||
max_used_size = 0
|
||||
buckets = [[0, 0] for _ in range(len(RANGE_HEADERS))]
|
||||
|
||||
for size, addr, is_free in blocks:
|
||||
total_size += size
|
||||
if is_free:
|
||||
free_blocks += 1
|
||||
free_size += size
|
||||
max_free_size = max(max_free_size, size)
|
||||
else:
|
||||
max_used_size = max(max_used_size, size)
|
||||
|
||||
# Bucket Mapping
|
||||
b_idx = len(THRESHOLDS)
|
||||
for i, t in enumerate(THRESHOLDS):
|
||||
if size < t:
|
||||
b_idx = i
|
||||
break
|
||||
buckets[b_idx][0 if is_free else 1] += 1
|
||||
|
||||
used_blocks = total_blocks - free_blocks
|
||||
used_size = total_size - free_size
|
||||
|
||||
# B. Summary Row (Total -> Used -> Free)
|
||||
summary_rows.append(
|
||||
[
|
||||
allocator_name,
|
||||
total_blocks,
|
||||
used_blocks,
|
||||
free_blocks,
|
||||
format_size(total_size),
|
||||
format_size(used_size),
|
||||
format_size(free_size),
|
||||
format_size(max_used_size),
|
||||
format_size(max_free_size),
|
||||
]
|
||||
)
|
||||
|
||||
# D. Distribution Rows
|
||||
dist_rows.append(
|
||||
[allocator_name, "Free Blocks"] + [b[0] for b in buckets]
|
||||
)
|
||||
dist_rows.append(
|
||||
[allocator_name, "Used Blocks"] + [b[1] for b in buckets]
|
||||
)
|
||||
|
||||
# --- 3. Render Outputs ---
|
||||
sum_headers = [
|
||||
"ID",
|
||||
"Tot Blks",
|
||||
"Used Blks",
|
||||
"Free Blks",
|
||||
"Tot Size",
|
||||
"Used Size",
|
||||
"Free Size",
|
||||
"Max Used",
|
||||
"Max Free",
|
||||
]
|
||||
print_table("Allocator Summary Statistics", sum_headers, summary_rows)
|
||||
|
||||
dist_headers = ["Allocator ID", "Block Type", *RANGE_HEADERS]
|
||||
print_table("Block Size Distribution", dist_headers, dist_rows)
|
||||
|
||||
@classmethod
|
||||
def allocate_record_table(self, data, output_filepath: str = ""):
|
||||
if not data:
|
||||
print("No data to display.")
|
||||
return
|
||||
|
||||
print(f"Record data size: {len(data)}, start printing...")
|
||||
headers = [
|
||||
'Allocator_Instance',
|
||||
'Is_Allocate',
|
||||
'Seq_ID',
|
||||
'Req_Size',
|
||||
'Cur_Alloc',
|
||||
'Cur_Rsrv',
|
||||
]
|
||||
formatted_row = []
|
||||
all_lines = []
|
||||
all_lines.append("\t".join(headers))
|
||||
for row in data:
|
||||
formatted_row = [
|
||||
str(row[0]),
|
||||
"Allocate" if row[1] else "Free",
|
||||
str(row[2]),
|
||||
str(row[3]),
|
||||
str(row[4]),
|
||||
str(row[5]),
|
||||
]
|
||||
line = "\t".join(formatted_row)
|
||||
all_lines.append(line)
|
||||
|
||||
try:
|
||||
with open(output_filepath, 'w', encoding='utf-8') as f:
|
||||
f.write("\n".join(all_lines))
|
||||
print(f"Data successfully written to: {output_filepath}")
|
||||
except OSError as e:
|
||||
print(f"Error writing to file {output_filepath}: {e}")
|
||||
|
||||
@classmethod
|
||||
def allocate_record_plot(self, data, save_path: str = ""):
|
||||
try:
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib import ticker
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"matplotlib is required but not installed. Please install it using: pip install matplotlib"
|
||||
)
|
||||
|
||||
if not data:
|
||||
print("No data to plot.")
|
||||
return
|
||||
print(f"Record data size: {len(data)}, start plotting...")
|
||||
data_np = np.array(data)
|
||||
is_allocate = data_np[:, 1]
|
||||
filter_mask = is_allocate == 1
|
||||
data_np = data_np[filter_mask]
|
||||
allocator_instance = data_np[:, 0] # allocator_instance not used
|
||||
ids = data_np[:, 2]
|
||||
sizes = data_np[:, 3]
|
||||
allocated = data_np[:, 4]
|
||||
reserved = data_np[:, 5]
|
||||
|
||||
LOG_START_VALUE = 1
|
||||
plt.style.use('seaborn-v0_8-whitegrid')
|
||||
fig, (ax1, ax2) = plt.subplots(
|
||||
2,
|
||||
1,
|
||||
sharex=True,
|
||||
figsize=(16, 10),
|
||||
dpi=120,
|
||||
gridspec_kw={'height_ratios': [3, 1], 'hspace': 0},
|
||||
)
|
||||
|
||||
# allocated event plot
|
||||
ax1.plot(
|
||||
ids, sizes, color='#2ca02c', linestyle='-', linewidth=1, alpha=0.3
|
||||
)
|
||||
ax1.scatter(
|
||||
ids,
|
||||
sizes,
|
||||
color='#2ca02c',
|
||||
s=60,
|
||||
alpha=1.0,
|
||||
edgecolors='white',
|
||||
linewidth=0.5,
|
||||
label='Request Size',
|
||||
zorder=5,
|
||||
)
|
||||
|
||||
ax1.set_ylabel(
|
||||
'Request Size (Linear Scale)',
|
||||
fontsize=12,
|
||||
fontweight='bold',
|
||||
labelpad=10,
|
||||
)
|
||||
ax1.set_title(
|
||||
'Paddle GPU Memory Allocation Analysis',
|
||||
fontsize=16,
|
||||
fontweight='bold',
|
||||
pad=20,
|
||||
)
|
||||
|
||||
ax1.set_ylim(bottom=LOG_START_VALUE)
|
||||
ax1.tick_params(axis='x', length=0)
|
||||
plt.setp(ax1.get_xticklabels(), visible=False)
|
||||
|
||||
# memory allocated, reserved plot
|
||||
ax2.plot(
|
||||
ids,
|
||||
reserved,
|
||||
color='#d62728',
|
||||
linestyle='--',
|
||||
linewidth=1.5,
|
||||
alpha=0.8,
|
||||
label='Reserved (Pool)',
|
||||
)
|
||||
ax2.fill_between(ids, 0, reserved, color='#d62728', alpha=0.1)
|
||||
ax2.plot(
|
||||
ids,
|
||||
allocated,
|
||||
color='#1f77b4',
|
||||
linestyle='-',
|
||||
linewidth=2,
|
||||
alpha=0.9,
|
||||
label='Allocated (Used)',
|
||||
)
|
||||
ax2.fill_between(ids, 0, allocated, color='#1f77b4', alpha=0.15)
|
||||
|
||||
ax2.invert_yaxis()
|
||||
ax2.set_ylim(reserved.max() * 3.0, LOG_START_VALUE)
|
||||
# ax2.set_yscale('symlog', linthresh=1024 * 1024)
|
||||
ax2.set_ylabel(
|
||||
'Pool Status (Inverted)',
|
||||
fontsize=11,
|
||||
fontweight='bold',
|
||||
labelpad=10,
|
||||
)
|
||||
|
||||
ax2.set_xlabel('')
|
||||
ax2.tick_params(axis='x', which='both', length=0)
|
||||
plt.setp(ax2.get_xticklabels(), visible=False)
|
||||
|
||||
# y axis setting 0
|
||||
def y_axis_formatter(x, pos):
|
||||
val = abs(x)
|
||||
if val <= LOG_START_VALUE * 1.5:
|
||||
return '0'
|
||||
return format_size(val).replace(" ", "")
|
||||
|
||||
formatter = ticker.FuncFormatter(y_axis_formatter)
|
||||
ax1.yaxis.set_major_formatter(formatter)
|
||||
ax2.yaxis.set_major_formatter(formatter)
|
||||
|
||||
for ax in [ax1, ax2]:
|
||||
current_ticks = ax.get_yticks().tolist()
|
||||
if LOG_START_VALUE not in current_ticks:
|
||||
current_ticks.append(LOG_START_VALUE)
|
||||
ax.set_yticks(sorted(current_ticks))
|
||||
|
||||
# axis setting
|
||||
for ax in [ax1, ax2]:
|
||||
for spine in ax.spines.values():
|
||||
spine.set_edgecolor('black')
|
||||
spine.set_linewidth(1.5)
|
||||
ax.tick_params(
|
||||
axis='both', which='major', colors='black', width=1.0, length=5
|
||||
)
|
||||
|
||||
lines1, labels1 = ax1.get_legend_handles_labels()
|
||||
lines2, labels2 = ax2.get_legend_handles_labels()
|
||||
|
||||
ax1.legend(
|
||||
lines1 + lines2,
|
||||
labels1 + labels2,
|
||||
loc='upper right',
|
||||
fontsize=10,
|
||||
frameon=True,
|
||||
facecolor='white',
|
||||
framealpha=0.9,
|
||||
edgecolor='black',
|
||||
shadow=False,
|
||||
)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.subplots_adjust(hspace=0.05)
|
||||
|
||||
plt.savefig(save_path)
|
||||
plt.close()
|
||||
print(f"Analysis plot saved to: {save_path}")
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from paddle.base.core import CUDAEvent as Event, CUDAPlace, CUDAStream as Stream
|
||||
|
||||
|
||||
def create_stream(
|
||||
device_id: CUDAPlace | int | None = None,
|
||||
priority: int = 2,
|
||||
device_type: str | None = None, # Ignored for compatibility
|
||||
blocking: bool = False, # Ignored for compatibility
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create CUDA Stream
|
||||
"""
|
||||
return Stream(device_id, priority)
|
||||
|
||||
|
||||
def create_event(
|
||||
enable_timing: bool = False,
|
||||
blocking: bool = False,
|
||||
interprocess: bool = False,
|
||||
device_type: str | None = None,
|
||||
device_id: int = 0,
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create CUDA Event
|
||||
"""
|
||||
return Event(enable_timing, blocking, interprocess)
|
||||
@@ -0,0 +1,603 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
from .custom_streams import ( # noqa: F401
|
||||
Event,
|
||||
Stream,
|
||||
create_event,
|
||||
create_stream,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import CustomPlace
|
||||
|
||||
_CustomPlaceLike: TypeAlias = (
|
||||
CustomPlace
|
||||
| str # some string like "iluvatar_gpu" "metax_gpu:0", etc.
|
||||
| int # some int like 0, 1, etc.
|
||||
)
|
||||
|
||||
dev_types = core.get_all_custom_device_type()
|
||||
|
||||
dev_type = dev_types[0] if dev_types else None
|
||||
|
||||
if dev_type and not core.is_compiled_with_custom_device(dev_type):
|
||||
raise Exception(
|
||||
"No custom device available, please install paddle with custom device support"
|
||||
)
|
||||
if dev_type and dev_type in ['metax_gpu', 'iluvatar_gpu']:
|
||||
from .gpgpu_backend import get_device_properties
|
||||
else:
|
||||
from .default_backend import get_device_properties
|
||||
|
||||
__all__ = [
|
||||
'Stream',
|
||||
'Event',
|
||||
'device_count',
|
||||
'get_device_properties',
|
||||
'empty_cache',
|
||||
'max_memory_allocated',
|
||||
'max_memory_reserved',
|
||||
'reset_max_memory_allocated',
|
||||
'reset_max_memory_reserved',
|
||||
'memory_allocated',
|
||||
'memory_reserved',
|
||||
'current_stream',
|
||||
'synchronize',
|
||||
]
|
||||
|
||||
|
||||
def device_count(device_type: str | None = None) -> int:
|
||||
'''
|
||||
Return the number of custom devices available.
|
||||
|
||||
Args:
|
||||
device_type (str, optional): The type of custom device (e.g., 'npu', 'mlu', etc.).
|
||||
If None, returns the count of the first available custom device type.
|
||||
|
||||
Returns:
|
||||
int: the number of custom devices available.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.device_count()
|
||||
>>> paddle.device.device_count('npu')
|
||||
'''
|
||||
|
||||
if device_type:
|
||||
num = core.get_custom_device_count(device_type)
|
||||
else:
|
||||
num = core.get_custom_device_count(dev_type)
|
||||
|
||||
return num
|
||||
|
||||
|
||||
def empty_cache() -> None:
|
||||
'''
|
||||
Releases idle cached memory held by the allocator so that those can be used in other GPU
|
||||
application and visible in device-specific tools.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.empty_cache()
|
||||
'''
|
||||
core.device_empty_cache()
|
||||
|
||||
|
||||
def max_memory_allocated(device: _CustomPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the peak size of memory that is allocated to tensor of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Returns:
|
||||
int: The peak size of memory that is allocated to tensor of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.max_memory_allocated('npu:0')
|
||||
>>> paddle.device.max_memory_allocated('npu')
|
||||
>>> paddle.device.max_memory_allocated(0)
|
||||
>>> paddle.device.max_memory_allocated(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"max_memory_allocated only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
return core.device_memory_stat_peak_value("Allocated", device_id)
|
||||
|
||||
|
||||
def max_memory_reserved(device: _CustomPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the peak size of memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Returns:
|
||||
int: The peak size of memory that is held by the allocator of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.max_memory_reserved('npu:0')
|
||||
>>> paddle.device.max_memory_reserved('npu')
|
||||
>>> paddle.device.max_memory_reserved(0)
|
||||
>>> paddle.device.max_memory_reserved(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"max_memory_reserved only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
return core.device_memory_stat_peak_value("Reserved", device_id)
|
||||
|
||||
|
||||
def reset_max_memory_allocated(device: _CustomPlaceLike | None = None) -> None:
|
||||
'''
|
||||
Reset the peak size of memory that is allocated to tensor of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.reset_max_memory_allocated('npu:0')
|
||||
>>> paddle.device.reset_max_memory_allocated('npu')
|
||||
>>> paddle.device.reset_max_memory_allocated(0)
|
||||
>>> paddle.device.reset_max_memory_allocated(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"reset_max_memory_allocated only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
core.device_memory_stat_reset_peak_value("Allocated", device_id)
|
||||
|
||||
|
||||
def reset_max_memory_reserved(device: _CustomPlaceLike | None = None) -> None:
|
||||
'''
|
||||
Reset the peak size of memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.reset_max_memory_reserved('npu:0')
|
||||
>>> paddle.device.reset_max_memory_reserved('npu')
|
||||
>>> paddle.device.reset_max_memory_reserved(0)
|
||||
>>> paddle.device.reset_max_memory_reserved(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"reset_max_memory_reserved only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
core.device_memory_stat_reset_peak_value("Reserved", device_id)
|
||||
|
||||
|
||||
def memory_allocated(device: _CustomPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the current size of memory that is allocated to tensor of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Returns:
|
||||
int: The current size of memory that is allocated to tensor of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.memory_allocated('npu:0')
|
||||
>>> paddle.device.memory_allocated('npu')
|
||||
>>> paddle.device.memory_allocated(0)
|
||||
>>> paddle.device.memory_allocated(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"memory_allocated only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
return core.device_memory_stat_current_value("Allocated", device_id)
|
||||
|
||||
|
||||
def memory_reserved(device: _CustomPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the current size of memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Returns:
|
||||
int: The current size of memory that is held by the allocator of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.memory_reserved('npu:0')
|
||||
>>> paddle.device.memory_reserved('npu')
|
||||
>>> paddle.device.memory_reserved(0)
|
||||
>>> paddle.device.memory_reserved(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"memory_reserved only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
return core.device_memory_stat_current_value("Reserved", device_id)
|
||||
|
||||
|
||||
def current_stream(device: _CustomPlaceLike | None = None) -> core.CustomStream:
|
||||
'''
|
||||
Return the current stream by the device.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Returns:
|
||||
Stream: The stream to the device.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.current_stream('npu:0')
|
||||
>>> paddle.device.current_stream('npu')
|
||||
>>> paddle.device.current_stream(0)
|
||||
>>> paddle.device.current_stream(Paddle.CustomPlace('npu', 0))
|
||||
'''
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"current_stream only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
return core._get_current_custom_device_stream(dev_type, device_id)
|
||||
|
||||
|
||||
def synchronize(device: _CustomPlaceLike | None = None) -> None:
|
||||
"""
|
||||
Wait for the compute on the given device to finish.
|
||||
|
||||
Args:
|
||||
device(_CustomPlaceLike, optional): Support input like 'npu:0', 'mlu', int, or CustomPlace.
|
||||
If None, the device is the first available custom device with index 0.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.synchronize('npu:0')
|
||||
>>> paddle.device.synchronize('npu')
|
||||
>>> paddle.device.synchronize(0)
|
||||
>>> paddle.device.synchronize(Paddle.CustomPlace('npu', 0))
|
||||
"""
|
||||
device_id = 0
|
||||
|
||||
if device is None:
|
||||
device_id = 0
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_id = 0
|
||||
else:
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
device_id = int(device_id_str)
|
||||
elif isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_id = device.get_device_id()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The input: {device} is not expected. Because paddle.device."
|
||||
"synchronize only support str, int or CustomPlace. "
|
||||
"Please input appropriate device again! "
|
||||
"Example: 'npu:0'"
|
||||
)
|
||||
|
||||
core._synchronize_custom_device(dev_type, device_id)
|
||||
|
||||
|
||||
def get_rng_state(
|
||||
device: _CustomPlaceLike | None = None,
|
||||
) -> core.GeneratorState:
|
||||
r'''
|
||||
Get the random state for the default generator.
|
||||
|
||||
Returns:
|
||||
Tensor: The random state tensor.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||||
>>> import paddle
|
||||
>>> paddle.device.get_rng_state()
|
||||
|
||||
'''
|
||||
place = paddle.device.device_to_place(device)
|
||||
if isinstance(place, core.CPUPlace):
|
||||
return core.default_cpu_generator().get_state()
|
||||
return core.default_custom_device_generator(place).get_state()
|
||||
|
||||
|
||||
def set_rng_state(
|
||||
new_state: core.GeneratorState, device: _CustomPlaceLike | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Set the random number generator state of the specified device.
|
||||
|
||||
Args:
|
||||
new_state (core.GeneratorState): The desired RNG state to set.
|
||||
This should be a state object previously obtained from ``get_rng_state()``.
|
||||
device (DeviceLike, optional): The device to set the RNG state for.
|
||||
If not specified, uses the current default device (as returned by ``paddle.framework._current_expected_place_()``).
|
||||
Can be a device object, integer device ID, or device string.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> # Save RNG state
|
||||
>>> state = paddle.device.get_rng_state()
|
||||
>>> # Do some random operations
|
||||
>>> x = paddle.randn([2, 3])
|
||||
>>> # Restore RNG state
|
||||
>>> paddle.device.set_rng_state(state)
|
||||
"""
|
||||
place = paddle.device.device_to_place(device)
|
||||
if isinstance(place, core.CPUPlace):
|
||||
core.default_cpu_generator().set_state(new_state)
|
||||
else:
|
||||
core.default_custom_device_generator(place).set_state(new_state)
|
||||
|
||||
|
||||
def manual_seed(seed: int) -> None:
|
||||
r"""Set the seed for generating random numbers for the current Device.
|
||||
|
||||
.. warning::
|
||||
If you are working with a multi-Device model, this function is insufficient
|
||||
to get determinism. To seed all Devices, use :func:`manual_seed_all`.
|
||||
If current Device is CPU, this function will set the seed of the default CPU generator.
|
||||
|
||||
Sets the seed for global default generator, which manages the random number generation.
|
||||
|
||||
Args:
|
||||
seed(int): The random seed to set.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||||
>>> import paddle
|
||||
>>> paddle.device.manual_seed(102)
|
||||
>>> # paddle.cuda.manual_seed(102) is equivalent to paddle.device.manual_seed(102)
|
||||
|
||||
"""
|
||||
seed = int(seed)
|
||||
place = paddle.framework._current_expected_place()
|
||||
if isinstance(place, core.CPUPlace):
|
||||
core.default_cpu_generator().manual_seed(seed)
|
||||
else:
|
||||
core.default_custom_device_generator(place).manual_seed(seed)
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from paddle.base.core import (
|
||||
CustomDeviceEvent as Event,
|
||||
CustomDeviceStream as Stream,
|
||||
CustomPlace,
|
||||
)
|
||||
|
||||
|
||||
def create_stream(
|
||||
device_id: CustomPlace | int | None = None,
|
||||
priority: int = 2,
|
||||
device_type: str | None = None, # Ignored for compatibility
|
||||
blocking: bool = False, # Ignored for compatibility
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create Custom Stream
|
||||
"""
|
||||
return Stream(
|
||||
device_type,
|
||||
device_id,
|
||||
priority,
|
||||
blocking=blocking,
|
||||
)
|
||||
|
||||
|
||||
def create_event(
|
||||
enable_timing: bool = False,
|
||||
blocking: bool = False,
|
||||
interprocess: bool = False,
|
||||
device_type: str | None = None,
|
||||
device_id: int = 0,
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create Custom Event
|
||||
"""
|
||||
return Event(
|
||||
device_type,
|
||||
device_id,
|
||||
enable_timing,
|
||||
blocking,
|
||||
interprocess,
|
||||
)
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import CustomPlace
|
||||
from paddle.base.libpaddle import _customDeviceProperties
|
||||
|
||||
_CustomPlaceLike: TypeAlias = CustomPlace | str | int
|
||||
|
||||
__all__ = [
|
||||
'get_device_properties',
|
||||
]
|
||||
|
||||
|
||||
def get_device_properties(
|
||||
device: _CustomPlaceLike | None = None,
|
||||
) -> _customDeviceProperties:
|
||||
"""
|
||||
Return the properties of given custom device.
|
||||
|
||||
Args:
|
||||
device (CustomPlace|str|int|None, optional): The device, the id of the device or
|
||||
the string name of device like 'metax_gpu:x' which to get the properties of the
|
||||
device from. Notice that this api only supports gpgpu devices. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Returns:
|
||||
_customDeviceProperties: The properties of the device which include device name,
|
||||
major compute capability, minor compute capability, global memory available
|
||||
and the number of multiprocessors on the device.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:METAX_GPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.get_device_properties('metax_gpu:0')
|
||||
>>> paddle.device.get_device_properties(0)
|
||||
>>> paddle.device.get_device_properties(paddle.CustomPlace('metax_gpu', 0))
|
||||
"""
|
||||
raise RuntimeError(
|
||||
"get_device_properties is not supported for this device type. "
|
||||
"This function is only available for gpgpu devices."
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
from paddle.base import core
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import CustomPlace
|
||||
from paddle.base.libpaddle import _customDeviceProperties
|
||||
|
||||
_CustomPlaceLike: TypeAlias = CustomPlace | str | int
|
||||
|
||||
__all__ = [
|
||||
'get_device_properties',
|
||||
]
|
||||
|
||||
|
||||
def get_device_properties(
|
||||
device: _CustomPlaceLike | None = None,
|
||||
) -> _customDeviceProperties:
|
||||
"""
|
||||
Return the properties of given custom device.
|
||||
|
||||
Args:
|
||||
device (CustomPlace|str|int|None, optional): The device, the id of the device or
|
||||
the string name of device like 'metax_gpu:x' which to get the properties of the
|
||||
device from. Notice that this api only supports gpgpu backend. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Returns:
|
||||
_customDeviceProperties: The properties of the device which include device name,
|
||||
major compute capability, minor compute capability, global memory available
|
||||
and the number of multiprocessors on the device.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.get_device_properties('metax_gpu:0')
|
||||
>>> paddle.device.get_device_properties(0)
|
||||
>>> paddle.device.get_device_properties(paddle.CustomPlace('metax_gpu', 0))
|
||||
"""
|
||||
if device is not None:
|
||||
if isinstance(device, int):
|
||||
device_id = device
|
||||
# Use default custom device type
|
||||
dev_types = core.get_all_custom_device_type()
|
||||
if not dev_types:
|
||||
raise ValueError("No custom device types available")
|
||||
device_name = dev_types[0]
|
||||
elif isinstance(device, core.CustomPlace):
|
||||
device_name = device.get_device_type()
|
||||
device_id = device.get_device_id()
|
||||
elif isinstance(device, str):
|
||||
colon_idx = device.rfind(':')
|
||||
if colon_idx == -1:
|
||||
device_name = device
|
||||
device_id = 0
|
||||
else:
|
||||
device_name = device[:colon_idx]
|
||||
device_id_str = device[colon_idx + 1 :]
|
||||
|
||||
if not device_id_str.isdigit():
|
||||
raise ValueError(
|
||||
f"Invalid device ID '{device_id_str}'. "
|
||||
f"After colon must be digits only. "
|
||||
"Example: 'metax_gpu:0'"
|
||||
)
|
||||
|
||||
device_id = int(device_id_str)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The device type {device} is not expected. Because paddle.device."
|
||||
"get_device_properties only support int, str or CustomPlace. "
|
||||
"Please input appropriate device again!"
|
||||
)
|
||||
else:
|
||||
# Use default custom device type and device id
|
||||
dev_types = core.get_all_custom_device_type()
|
||||
if not dev_types:
|
||||
raise ValueError("No custom device types available")
|
||||
device_name = dev_types[0]
|
||||
device_id = 0
|
||||
|
||||
return core.get_device_properties(device_name, device_id)
|
||||
@@ -0,0 +1,689 @@
|
||||
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
from paddle.utils import deprecated
|
||||
|
||||
from .streams import Event, Stream, create_event, create_stream # noqa: F401
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import XPUPlace
|
||||
from paddle.base.libpaddle import _gpuDeviceProperties
|
||||
|
||||
_XPUPlaceLike: TypeAlias = (
|
||||
XPUPlace
|
||||
| str # some str like 'xpu:0', 'xpu:1', etc.
|
||||
| int # some int like 0, 1, etc.
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
'Stream',
|
||||
'Event',
|
||||
'synchronize',
|
||||
'device_count',
|
||||
'set_debug_level',
|
||||
'empty_cache',
|
||||
'max_memory_allocated',
|
||||
'max_memory_reserved',
|
||||
'reset_max_memory_allocated',
|
||||
'reset_max_memory_reserved',
|
||||
'memory_allocated',
|
||||
'memory_reserved',
|
||||
'memory_total', # memory managed by runtime, not paddle
|
||||
'memory_used', # memory managed by runtime, not paddle
|
||||
'get_device_properties',
|
||||
]
|
||||
|
||||
|
||||
def current_stream(device: _XPUPlaceLike | None = None) -> core.XPUStream:
|
||||
'''
|
||||
Return the current XPU stream by the device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace()|int|None, optional): The device or the ID of the device which want to get stream from.
|
||||
If device is None, the device is the current device. Default: None.
|
||||
|
||||
Returns:
|
||||
XPUStream: the stream to the device.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> s1 = paddle.device.xpu.current_stream()
|
||||
|
||||
>>> s2 = paddle.device.xpu.current_stream(0)
|
||||
|
||||
>>> s3 = paddle.device.xpu.current_stream(paddle.XPUPlace(0))
|
||||
|
||||
'''
|
||||
|
||||
device_id = -1
|
||||
|
||||
if device is not None:
|
||||
if isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.XPUPlace):
|
||||
device_id = device.get_device_id()
|
||||
elif isinstance(device, str):
|
||||
place = paddle.device._convert_to_place(device)
|
||||
device_id = place.get_device_id()
|
||||
else:
|
||||
raise ValueError("device type must be int or paddle.XPUPlace")
|
||||
|
||||
return core._xpu_get_current_stream(device_id)
|
||||
|
||||
|
||||
def extract_xpu_device_id(device: _XPUPlaceLike, op_name: str) -> int:
|
||||
'''
|
||||
Return the id of the given xpu device. It is just a utility that will not be exposed to users.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace or int or str): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The id of the given device. If device is None, return the id of current device.
|
||||
'''
|
||||
if device is None:
|
||||
return core.get_xpu_current_device_id()
|
||||
|
||||
if isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.XPUPlace):
|
||||
device_id = device.get_device_id()
|
||||
elif isinstance(device, str):
|
||||
if device.startswith('xpu:'):
|
||||
device_id = int(device[4:])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The current string {device} is not expected. Because {op_name} only support string which is like 'xpu:x'. "
|
||||
"Please input appropriate string again!"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The device type {device} is not expected. Because {op_name} only support int, str or paddle.XPUPlace. "
|
||||
"Please input appropriate device again!"
|
||||
)
|
||||
|
||||
assert device_id >= 0, (
|
||||
f"The device id must be not less than 0, but got id = {device_id}."
|
||||
)
|
||||
assert device_id < device_count(), (
|
||||
f"The device id {device_id} exceeds xpu card number {device_count()}"
|
||||
)
|
||||
return device_id
|
||||
|
||||
|
||||
@deprecated(
|
||||
since="2.5.0",
|
||||
update_to="paddle.device.synchronize",
|
||||
level=1,
|
||||
reason="synchronize in paddle.device.xpu will be removed in future",
|
||||
)
|
||||
def synchronize(device: _XPUPlaceLike | None = None) -> int:
|
||||
"""
|
||||
Wait for the compute on the given XPU device to finish.
|
||||
|
||||
Parameters:
|
||||
device(paddle.XPUPlace()|int, optional): The device or the ID of the device.
|
||||
If device is None, the device is the current device. Default: None.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
>>> paddle.device.xpu.synchronize()
|
||||
>>> paddle.device.xpu.synchronize(0)
|
||||
>>> paddle.device.xpu.synchronize(paddle.XPUPlace(0))
|
||||
|
||||
"""
|
||||
|
||||
device_id = -1
|
||||
|
||||
if device is not None:
|
||||
if isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.XPUPlace):
|
||||
device_id = device.get_device_id()
|
||||
elif isinstance(device, str):
|
||||
if device.startswith('xpu:'):
|
||||
device_id = int(device[4:])
|
||||
elif device == 'xpu':
|
||||
device_id = 0
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The current string {device} is not expected. Because paddle.device.cuda."
|
||||
"synchronize only support string which is like 'xpu:x' or 'xpu'. "
|
||||
"Please input appropriate string again!"
|
||||
)
|
||||
else:
|
||||
raise ValueError("device type must be int or paddle.XPUPlace")
|
||||
|
||||
return core._xpu_device_synchronize(device_id)
|
||||
|
||||
|
||||
def device_count() -> int:
|
||||
'''
|
||||
Return the number of XPUs available.
|
||||
|
||||
Returns:
|
||||
int: the number of XPUs available.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> paddle.device.xpu.device_count()
|
||||
|
||||
'''
|
||||
|
||||
num_xpus = (
|
||||
core.get_xpu_device_count()
|
||||
if hasattr(core, 'get_xpu_device_count')
|
||||
else 0
|
||||
)
|
||||
|
||||
return num_xpus
|
||||
|
||||
|
||||
def set_debug_level(level: int = 0) -> None:
|
||||
'''
|
||||
Set the debug level of XPUs' api. The default level is 0, which means no debug info.
|
||||
|
||||
Args:
|
||||
int: Debug level of XPUs available.
|
||||
Level 0x1 for trace (Print the invocation of the interface),
|
||||
0x10 for checksum (Print the checksum of the tensor),
|
||||
0x100 for dump (Save the tensor as a file in npy format),
|
||||
0x1000 for profiling (Record the execution time of each operator).
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> paddle.device.xpu.set_debug_level(1)
|
||||
'''
|
||||
name = "paddle.device.xpu.set_debug_level"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
else:
|
||||
core.set_xpu_debug_level(level)
|
||||
|
||||
|
||||
def empty_cache() -> None:
|
||||
'''
|
||||
Releases idle cached memory held by the allocator so that those can be used in other XPU
|
||||
application and visible in `xpu-smi`. In most cases you don't need to use this function,
|
||||
Paddle does not release the memory back to the OS when you remove Tensors on the XPU,
|
||||
Because it keeps xpu memory in a pool so that next allocations can be done much faster.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> tensor = paddle.randn([512, 512, 512], "float64")
|
||||
>>> del tensor
|
||||
>>> paddle.device.xpu.empty_cache()
|
||||
'''
|
||||
name = "paddle.device.xpu.empty_cache"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
else:
|
||||
core.xpu_empty_cache()
|
||||
|
||||
|
||||
def get_device_properties(
|
||||
device: _XPUPlaceLike | None = None,
|
||||
) -> _gpuDeviceProperties:
|
||||
'''
|
||||
Return the properties of given device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x' which to get the properties of the
|
||||
device from. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Returns:
|
||||
_gpuDeviceProperties: The properties of the device which include ASCII string
|
||||
identifying device, major compute capability, minor compute capability, global
|
||||
memory available and the number of multiprocessors on the device.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
>>> paddle.device.xpu.get_device_properties()
|
||||
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
||||
|
||||
>>> paddle.device.xpu.get_device_properties(0)
|
||||
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
||||
|
||||
>>> paddle.device.xpu.get_device_properties('xpu:0')
|
||||
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
||||
|
||||
>>> paddle.device.xpu.get_device_properties(paddle.XPUPlace(0))
|
||||
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
||||
|
||||
'''
|
||||
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
"The API paddle.device.xpu.get_device_properties is not supported in "
|
||||
"CPU-only PaddlePaddle. Please reinstall PaddlePaddle with XPU support "
|
||||
"to call this API."
|
||||
)
|
||||
|
||||
if device is not None:
|
||||
if isinstance(device, int):
|
||||
device_id = device
|
||||
elif isinstance(device, core.XPUPlace):
|
||||
device_id = device.get_device_id()
|
||||
elif isinstance(device, str):
|
||||
if device.startswith('xpu:'):
|
||||
device_id = int(device[4:])
|
||||
elif device == 'xpu':
|
||||
device_id = 0
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The current string {device} is not expected. Because paddle.device."
|
||||
"xpu.get_device_properties only support string which is like 'xpu:x' or 'xpu'. "
|
||||
"Please input appropriate string again!"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The device type {device} is not expected. Because paddle.device.xpu."
|
||||
"get_device_properties only support int, str or paddle.XPUPlace. "
|
||||
"Please input appropriate device again!"
|
||||
)
|
||||
else:
|
||||
device_id = -1
|
||||
|
||||
return core.get_device_properties(device_id)
|
||||
|
||||
|
||||
def max_memory_allocated(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the peak size of xpu memory that is allocated to tensor of the given device.
|
||||
|
||||
Note:
|
||||
The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may larger than the memory size that tensor actually need.
|
||||
For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The peak size of xpu memory that is allocated to tensor of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(paddle.XPUPlace(0))
|
||||
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(0)
|
||||
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.max_memory_allocated"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.device_memory_stat_peak_value("Allocated", device_id)
|
||||
|
||||
|
||||
def max_memory_reserved(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the peak size of XPU memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The peak size of XPU memory that is held by the allocator of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(paddle.XPUPlace(0))
|
||||
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(0)
|
||||
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.max_memory_reserved"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.device_memory_stat_peak_value("Reserved", device_id)
|
||||
|
||||
|
||||
def reset_max_memory_allocated(device: _XPUPlaceLike | None = None) -> None:
|
||||
'''
|
||||
Reset the peak size of XPU memory that is allocated to tensor of the given device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> paddle.device.xpu.reset_max_memory_allocated(paddle.XPUPlace(0))
|
||||
>>> paddle.device.xpu.reset_max_memory_allocated(0)
|
||||
>>> paddle.device.xpu.reset_max_memory_allocated("xpu:0")
|
||||
'''
|
||||
|
||||
name = "paddle.device.xpu.reset_max_memory_allocated"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
core.device_memory_stat_reset_peak_value("Allocated", device_id)
|
||||
|
||||
|
||||
def reset_max_memory_reserved(device: _XPUPlaceLike | None = None) -> None:
|
||||
'''
|
||||
Reset the peak size of XPU memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> paddle.device.xpu.reset_max_memory_reserved(paddle.XPUPlace(0))
|
||||
>>> paddle.device.xpu.reset_max_memory_reserved(0)
|
||||
>>> paddle.device.xpu.reset_max_memory_reserved("xpu:0")
|
||||
'''
|
||||
|
||||
name = "paddle.device.xpu.reset_max_memory_reserved"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
core.device_memory_stat_reset_peak_value("Reserved", device_id)
|
||||
|
||||
|
||||
def memory_allocated(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the current size of xpu memory that is allocated to tensor of the given device.
|
||||
|
||||
Note:
|
||||
The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may be larger than the memory size that tensor actually need.
|
||||
For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The current size of xpu memory that is allocated to tensor of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> memory_allocated_size = paddle.device.xpu.memory_allocated(paddle.XPUPlace(0))
|
||||
>>> memory_allocated_size = paddle.device.xpu.memory_allocated(0)
|
||||
>>> memory_allocated_size = paddle.device.xpu.memory_allocated("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.memory_allocated"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.device_memory_stat_current_value("Allocated", device_id)
|
||||
|
||||
|
||||
def memory_reserved(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the current size of XPU memory that is held by the allocator of the given device.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The current size of XPU memory that is held by the allocator of the given device, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> memory_reserved_size = paddle.device.xpu.memory_reserved(paddle.XPUPlace(0))
|
||||
>>> memory_reserved_size = paddle.device.xpu.memory_reserved(0)
|
||||
>>> memory_reserved_size = paddle.device.xpu.memory_reserved("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.memory_reserved"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.device_memory_stat_current_value("Reserved", device_id)
|
||||
|
||||
|
||||
def memory_total(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the total size of XPU memory of the given device that is held by the XPU Runtime.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The total size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> memory_total_size = paddle.device.xpu.memory_total(paddle.XPUPlace(0))
|
||||
>>> memory_total_size = paddle.device.xpu.memory_total(0)
|
||||
>>> memory_total_size = paddle.device.xpu.memory_total("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.memory_total"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.get_xpu_device_total_memory(device_id)
|
||||
|
||||
|
||||
def memory_used(device: _XPUPlaceLike | None = None) -> int:
|
||||
'''
|
||||
Return the used size of XPU memory of the given device that is held by the XPU Runtime.
|
||||
|
||||
Args:
|
||||
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
||||
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
||||
Default: None.
|
||||
|
||||
Return:
|
||||
int: The used size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('xpu')
|
||||
|
||||
>>> memory_used_size = paddle.device.xpu.memory_used(paddle.XPUPlace(0))
|
||||
>>> memory_used_size = paddle.device.xpu.memory_used(0)
|
||||
>>> memory_used_size = paddle.device.xpu.memory_used("xpu:0")
|
||||
'''
|
||||
name = "paddle.device.xpu.memory_used"
|
||||
if not core.is_compiled_with_xpu():
|
||||
raise ValueError(
|
||||
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
||||
)
|
||||
device_id = extract_xpu_device_id(device, op_name=name)
|
||||
return core.get_xpu_device_used_memory(device_id)
|
||||
|
||||
|
||||
def get_rng_state(device: _XPUPlaceLike | None = None) -> core.GeneratorState:
|
||||
'''
|
||||
Get the random state for the default generator.
|
||||
|
||||
Returns:
|
||||
Tensor: The random state tensor.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.get_rng_state()
|
||||
|
||||
'''
|
||||
place = paddle.device.device_to_place(device)
|
||||
if isinstance(place, core.CPUPlace):
|
||||
return core.default_cpu_generator().get_state()
|
||||
return core.default_xpu_generator(place.get_device_id()).get_state()
|
||||
|
||||
|
||||
def set_rng_state(
|
||||
new_state: core.GeneratorState, device: _XPUPlaceLike | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Set the random number generator state of the specified device.
|
||||
|
||||
Args:
|
||||
new_state (core.GeneratorState): The desired RNG state to set.
|
||||
This should be a state object previously obtained from ``get_rng_state()``.
|
||||
device (DeviceLike, optional): The device to set the RNG state for.
|
||||
If not specified, uses the current default device (as returned by ``paddle.framework._current_expected_place_()``).
|
||||
Can be a device object, integer device ID, or device string.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> # Save RNG state
|
||||
>>> state = paddle.device.get_rng_state()
|
||||
>>> # Do some random operations
|
||||
>>> x = paddle.randn([2, 3])
|
||||
>>> # Restore RNG state
|
||||
>>> paddle.device.set_rng_state(state)
|
||||
"""
|
||||
place = paddle.device.device_to_place(device)
|
||||
if isinstance(place, core.CPUPlace):
|
||||
core.default_cpu_generator().set_state(new_state)
|
||||
else:
|
||||
core.default_xpu_generator(place.get_device_id()).set_state(new_state)
|
||||
|
||||
|
||||
def manual_seed(seed: int) -> None:
|
||||
r"""Set the seed for generating random numbers for the current Device.
|
||||
|
||||
.. warning::
|
||||
If you are working with a multi-Device model, this function is insufficient
|
||||
to get determinism. To seed all Devices, use :func:`manual_seed_all`.
|
||||
If current Device is CPU, this function will set the seed of the default CPU generator.
|
||||
|
||||
Sets the seed for global default generator, which manages the random number generation.
|
||||
|
||||
Args:
|
||||
seed(int): The random seed to set.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:XPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.manual_seed(102)
|
||||
>>> # paddle.cuda.manual_seed(102) is equivalent to paddle.device.manual_seed(102)
|
||||
>>> paddle.cuda.manual_seed(102)
|
||||
|
||||
"""
|
||||
seed = int(seed)
|
||||
place = paddle.framework._current_expected_place_()
|
||||
if isinstance(place, core.CPUPlace):
|
||||
core.default_cpu_generator().manual_seed(seed)
|
||||
else:
|
||||
core.default_xpu_generator(place.get_device_id()).manual_seed(seed)
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from paddle.base.core import (
|
||||
XPUEvent as Event,
|
||||
XPUPlace,
|
||||
XPUStream as Stream,
|
||||
)
|
||||
|
||||
|
||||
def create_stream(
|
||||
device_id: XPUPlace | int | None = None,
|
||||
priority: int = 2,
|
||||
device_type: str | None = None, # Ignored for compatibility
|
||||
blocking: bool = False, # Ignored for compatibility
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create XPU Stream
|
||||
"""
|
||||
return Stream(device_id)
|
||||
|
||||
|
||||
def create_event(
|
||||
enable_timing: bool = False,
|
||||
blocking: bool = False,
|
||||
interprocess: bool = False,
|
||||
device_type: str | None = None,
|
||||
device_id: int = 0,
|
||||
):
|
||||
"""
|
||||
Factory Function, used to create XPU Event
|
||||
"""
|
||||
return Event()
|
||||
Reference in New Issue
Block a user