| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager |
| 5 | |
| 6 | import torch |
| 7 | import torch.cuda.nvtx as nvtx |
| 8 | |
| 9 | |
| 10 | def print_tensor(tensor_obj, prefix, tensor_list=None): |
| 11 | """Descends iterators that contains Tensors and prints the Tensor. |
| 12 | Recursive function that descends iterator type arguments until |
| 13 | it finds a Tensor object. |
| 14 | """ |
| 15 | if tensor_list is None: |
| 16 | tensor_list = [] |
| 17 | |
| 18 | if isinstance(tensor_obj, (list, tuple)): |
| 19 | for ten in tensor_obj: |
| 20 | tensor_list = print_tensor(ten, prefix, tensor_list) |
| 21 | elif isinstance(tensor_obj, torch.Tensor): |
| 22 | tensor_dims = list(tensor_obj.size()) |
| 23 | tensor_list.append(tensor_dims) |
| 24 | return tensor_list |
| 25 | |
| 26 | |
| 27 | def process_layer_params(module_obj): |
| 28 | """Extract the static parameters from LLM and VLM relevant layer types.""" |
| 29 | param_info = {} |
| 30 | # Extract parameters for layers commonly used in LLMs and VLMs |
| 31 | if isinstance(module_obj, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)): |
| 32 | conv_params = {} |
| 33 | conv_params["in_chan"] = module_obj.in_channels |
| 34 | conv_params["out_chan"] = module_obj.out_channels |
| 35 | conv_params["filter_dim"] = module_obj.kernel_size |
| 36 | conv_params["stride"] = module_obj.stride |
| 37 | conv_params["padding"] = module_obj.padding |
| 38 | conv_params["dilation"] = module_obj.dilation |
| 39 | conv_params["transposed"] = module_obj.transposed |
| 40 | conv_params["output_padding"] = module_obj.output_padding |
| 41 | conv_params["groups"] = module_obj.groups |
| 42 | conv_params["padding_mode"] = module_obj.padding_mode |
| 43 | param_info = conv_params |
| 44 | elif isinstance( |
| 45 | module_obj, |
| 46 | ( |
| 47 | torch.nn.ConvTranspose1d, |
| 48 | torch.nn.ConvTranspose2d, |
| 49 | torch.nn.ConvTranspose3d, |
| 50 | ), |
| 51 | ): |
| 52 | convtranspose_params = {} |
| 53 | convtranspose_params["in_chan"] = module_obj.in_channels |
| 54 | convtranspose_params["out_chan"] = module_obj.out_channels |
| 55 | convtranspose_params["filter_dim"] = module_obj.kernel_size |
| 56 | convtranspose_params["stride"] = module_obj.stride |
| 57 | convtranspose_params["padding"] = module_obj.padding |
| 58 | convtranspose_params["dilation"] = module_obj.dilation |
| 59 | convtranspose_params["transposed"] = module_obj.transposed |
| 60 | convtranspose_params["output_padding"] = module_obj.output_padding |
| 61 | convtranspose_params["groups"] = module_obj.groups |
| 62 | convtranspose_params["padding_mode"] = module_obj.padding_mode |
| 63 | param_info = convtranspose_params |
| 64 | elif isinstance( |
| 65 | module_obj, (torch.nn.MaxPool1d, torch.nn.MaxPool2d, torch.nn.MaxPool3d) |
| 66 | ): |
| 67 | |
| 68 | def _handle_int_or_tuple(parameter): |
| 69 | if isinstance(parameter, tuple): |
| 70 | return list(parameter) |
| 71 | elif isinstance(parameter, int): |
| 72 | return [parameter, parameter] |
| 73 | |
| 74 | pooling_params = {} |
| 75 | pooling_params["filter_dim"] = _handle_int_or_tuple(module_obj.kernel_size) |
| 76 | pooling_params["stride"] = _handle_int_or_tuple(module_obj.stride) |
| 77 | pooling_params["padding"] = _handle_int_or_tuple(module_obj.padding) |
| 78 | pooling_params["dilation"] = _handle_int_or_tuple(module_obj.dilation) |
| 79 | param_info = pooling_params |
| 80 | elif isinstance( |
| 81 | module_obj, (torch.nn.AvgPool1d, torch.nn.AvgPool2d, torch.nn.AvgPool3d) |
| 82 | ): |
| 83 | pooling_params = {} |
| 84 | pooling_params["filter_dim"] = [ |
| 85 | module_obj.kernel_size, |
| 86 | module_obj.kernel_size, |
| 87 | ] |
| 88 | pooling_params["stride"] = [module_obj.stride, module_obj.stride] |
| 89 | pooling_params["padding"] = [module_obj.padding, module_obj.padding] |
| 90 | pooling_params["ceil_mode"] = module_obj.ceil_mode |
| 91 | pooling_params["count_include_pad"] = module_obj.count_include_pad |
| 92 | param_info = pooling_params |
| 93 | elif isinstance( |
| 94 | module_obj, |
| 95 | ( |
| 96 | torch.nn.AdaptiveAvgPool1d, |
| 97 | torch.nn.AdaptiveAvgPool2d, |
| 98 | torch.nn.AdaptiveAvgPool3d, |
| 99 | ), |
| 100 | ): |
| 101 | pooling_params = {} |
| 102 | pooling_params["output_size"] = [ |
| 103 | module_obj.output_size, |
| 104 | module_obj.output_size, |
| 105 | ] |
| 106 | param_info = pooling_params |
| 107 | elif isinstance(module_obj, torch.nn.Linear): |
| 108 | param_info["in_features"] = module_obj.in_features |
| 109 | param_info["out_features"] = module_obj.out_features |
| 110 | elif isinstance( |
| 111 | module_obj, |
| 112 | (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d), |
| 113 | ): |
| 114 | param_info["num_features"] = module_obj.num_features |
| 115 | param_info["epsilon"] = module_obj.eps |
| 116 | param_info["momentum"] = module_obj.momentum |
| 117 | elif isinstance(module_obj, torch.nn.ReLU): |
| 118 | param_info["in_place"] = module_obj.inplace |
| 119 | elif isinstance(module_obj, torch.nn.Dropout): |
| 120 | param_info["p"] = module_obj.p |
| 121 | param_info["in_place"] = module_obj.inplace |
| 122 | elif isinstance(module_obj, torch.nn.Embedding): |
| 123 | param_info["num_embeddings"] = module_obj.num_embeddings |
| 124 | param_info["embedding_dim"] = module_obj.embedding_dim |
| 125 | elif isinstance( |
| 126 | module_obj, |
| 127 | ( |
| 128 | torch.nn.Upsample, |
| 129 | torch.nn.UpsamplingNearest2d, |
| 130 | torch.nn.UpsamplingBilinear2d, |
| 131 | ), |
| 132 | ): |
| 133 | param_info["scale_factor"] = module_obj.scale_factor |
| 134 | |
| 135 | return param_info |
| 136 | |
| 137 | |
| 138 | def construct_marker_dict_and_push( |
| 139 | module_name, module_obj, in_tensor, kwargs=None, out_tensor=None |
| 140 | ): |
| 141 | marker_dict = {} |
| 142 | marker_dict["Module"] = module_name |
| 143 | |
| 144 | ## Get trainable parameters like weights and bias |
| 145 | module_params = module_obj.named_parameters(recurse=False) |
| 146 | for idx, (param_name, param_obj) in enumerate(module_params): |
| 147 | if idx == 0: |
| 148 | marker_dict["TrainableParams"] = {} |
| 149 | marker_dict["TrainableParams"][param_name] = list(param_obj.size()) |
| 150 | |
| 151 | in_tensor_list = print_tensor(in_tensor, "Input") |
| 152 | if in_tensor_list: |
| 153 | marker_dict["Inputs"] = in_tensor_list |
| 154 | |
| 155 | out_tensor_list = print_tensor(out_tensor, "Output") |
| 156 | if out_tensor_list: |
| 157 | marker_dict["Outputs"] = out_tensor_list |
| 158 | |
| 159 | ## Get Kwargs like input_ids and positions for the top module |
| 160 | if kwargs: |
| 161 | for key, value in kwargs.items(): |
| 162 | if isinstance(value, (torch.Tensor, list, tuple)): |
| 163 | tensor_list = print_tensor(value, key) |
| 164 | if tensor_list: |
| 165 | marker_dict[key] = tensor_list |
| 166 | |
| 167 | param_info = process_layer_params(module_obj) |
| 168 | if param_info: |
| 169 | marker_dict["StaticParams"] = param_info |
| 170 | nvtx.range_push("{}".format(marker_dict)) |
| 171 | |
| 172 | |
| 173 | class ResultHolder: |
| 174 | """Holder for storing results from within a context manager.""" |
| 175 | |
| 176 | result = None |
| 177 | |
| 178 | |
| 179 | @contextmanager |
| 180 | def layerwise_nvtx_marker_context(module_name, module_obj, in_tensor=None, kwargs=None): |
| 181 | """Context manager for NVTX markers that automatically pushes on enter |
| 182 | and pops on exit. |
| 183 | |
| 184 | Example: |
| 185 | with nvtx_marker_context("Module:MyModule", module, in_tensor=args, |
| 186 | kwargs=kwargs) as ctx: |
| 187 | ctx.result = module(*args, **kwargs) |
| 188 | return ctx.result |
| 189 | |
| 190 | """ |
| 191 | holder = ResultHolder() |
| 192 | |
| 193 | # Push input marker |
| 194 | construct_marker_dict_and_push( |
| 195 | module_name, |
| 196 | module_obj, |
| 197 | in_tensor=in_tensor, |
| 198 | kwargs=kwargs, |
| 199 | ) |
| 200 | try: |
| 201 | yield holder |
| 202 | finally: |
| 203 | # Pop input marker |
| 204 | nvtx.range_pop() |
| 205 | # Push and pop output marker |
| 206 | output_name = module_name.replace("(input)", "(output)") |
| 207 | construct_marker_dict_and_push( |
| 208 | output_name, |
| 209 | module_obj, |
| 210 | in_tensor=None, |
| 211 | kwargs=None, |
| 212 | out_tensor=holder.result, |
| 213 | ) |
| 214 | nvtx.range_pop() |
| 215 | |
| 216 | |
| 217 | class PytHooks: |
| 218 | """This module contains all the code needed to enable forward hooks |
| 219 | in a pytorch network. |
| 220 | |
| 221 | To register the hooks for a given network, the user needs to instantiate |
| 222 | a PytHook object. Then call the register_hooks method. |
| 223 | |
| 224 | Example: |
| 225 | my_hook = PytHook() |
| 226 | my_hook.register_hooks(my_network_model) |
| 227 | |
| 228 | """ |
| 229 | |
| 230 | def __init__(self): |
| 231 | """Initialize module variables.""" |
| 232 | super().__init__() |
| 233 | self.module_to_name_map = {} |
| 234 | |
| 235 | def module_fwd_hook(self, module_obj, in_tensor, out_tensor): |
| 236 | """Callback function that ends the NVTX marker. |
| 237 | Records the module name and tensor information. |
| 238 | Called after the module executes the forward method. |
| 239 | """ |
| 240 | nvtx.range_pop() |
| 241 | module_name = self.module_to_name_map.get(module_obj, "unknown") |
| 242 | construct_marker_dict_and_push( |
| 243 | module_name, module_obj, in_tensor=None, kwargs=None, out_tensor=out_tensor |
| 244 | ) |
| 245 | nvtx.range_pop() |
| 246 | return |
| 247 | |
| 248 | def module_fwd_pre_hook(self, module_obj, in_tensor, kwargs): |
| 249 | """Creates an NVTX marker with the module name in it. |
| 250 | This function is called before the module executes. |
| 251 | """ |
| 252 | module_name = self.module_to_name_map.get(module_obj, "unknown") |
| 253 | construct_marker_dict_and_push( |
| 254 | module_name, module_obj, in_tensor=in_tensor, kwargs=kwargs, out_tensor=None |
| 255 | ) |
| 256 | return |
| 257 | |
| 258 | def register_hooks(self, network_model, module_prefix="top"): |
| 259 | """User level function that activates all the hooks. |
| 260 | The user needs to call this method from the network source code. |
| 261 | The code descends all the modules in the network and registers their |
| 262 | respective hooks. |
| 263 | """ |
| 264 | # Module types to skip (simple operations that don't need detailed profiling) |
| 265 | skip_types = ( |
| 266 | torch.nn.Identity, |
| 267 | torch.nn.Dropout, |
| 268 | torch.nn.Dropout1d, |
| 269 | torch.nn.Dropout2d, |
| 270 | torch.nn.Dropout3d, |
| 271 | ) |
| 272 | |
| 273 | for name, module in network_model.named_modules(prefix=module_prefix): |
| 274 | # Skip certain module types to reduce profiling overhead |
| 275 | if isinstance(module, skip_types): |
| 276 | continue |
| 277 | |
| 278 | module.register_forward_pre_hook(self.module_fwd_pre_hook, with_kwargs=True) |
| 279 | module.register_forward_hook(self.module_fwd_hook) |
| 280 | if module not in self.module_to_name_map: |
| 281 | self.module_to_name_map[module] = name |
| 282 | else: |
| 283 | raise ValueError("Module instance {} is not unique ".format(module)) |
| 284 | return |