# 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()