The function style to interface to enable debugger on model. If filter is provided, it will be used to filter out satisfied module to register hook. If filter is not provided, all modules will be registered with hooks. Example: from tensorrt_llm._torch.debug.debug_hook impor
(model: nn.Module,
dest_folder: Optional[str] = None,
filter: Optional[Filter] = None)
| 190 | |
| 191 | |
| 192 | def enable_debug(model: nn.Module, |
| 193 | dest_folder: Optional[str] = None, |
| 194 | filter: Optional[Filter] = None): |
| 195 | """ |
| 196 | The function style to interface to enable debugger on model. |
| 197 | If filter is provided, it will be used to filter out satisfied module to register hook. |
| 198 | If filter is not provided, all modules will be registered with hooks. |
| 199 | Example: |
| 200 | from tensorrt_llm._torch.debug.debug_hook import enable_debug |
| 201 | model_config = ModelConfig(pretrained_config=llama_config, |
| 202 | attn_backend=backend) |
| 203 | llama = LlamaForCausalLM(model_config).to(dtype).to(device) |
| 204 | llama.load_weights(hf_llama.state_dict()) |
| 205 | with torch.inference_mode(): |
| 206 | enable_debug(llama, r"tensor_dump"): |
| 207 | attn_metadata.prepare() |
| 208 | logits = llama.forward(input_ids=input_ids, |
| 209 | position_ids=position_ids, |
| 210 | attn_metadata=attn_metadata) |
| 211 | |
| 212 | Note: this method need user to disable debug by calling disable_debug |
| 213 | Args: |
| 214 | model (nn.Module): the model to enable debug hook. |
| 215 | dest_folder: the working directory set to debug context to set where the hook dumped data/info. |
| 216 | filter: a filter to decide what modules will be registered with debug hook. |
| 217 | Returns: |
| 218 | None |
| 219 | """ |
| 220 | debug_ctx = get_current_debug_ctx() |
| 221 | assert debug_ctx is None, "DebugContext shall be None when enable debugger context." |
| 222 | debug_ctx = DebuggerContext(dest_folder) |
| 223 | set_current_debug_ctx(debug_ctx) |
| 224 | |
| 225 | debug_ctx.get_current_modules_tree().clear() |
| 226 | debug_ctx.get_module_indices_tree().clear() |
| 227 | for name, submodule in model.named_modules(): |
| 228 | if name == "": |
| 229 | continue |
| 230 | |
| 231 | if submodule not in debug_ctx.forward_hook_handles: |
| 232 | do_hook = filter(submodule) if filter is not None else True |
| 233 | if do_hook: |
| 234 | debug_ctx.forward_hook_handles[ |
| 235 | submodule] = submodule.register_forward_hook( |
| 236 | after_forward, with_kwargs=True, always_call=True) |
| 237 | |
| 238 | if submodule not in debug_ctx.forward_pre_hook_handles: |
| 239 | do_hook = filter(submodule) if filter is not None else True |
| 240 | if do_hook: |
| 241 | debug_ctx.forward_pre_hook_handles[ |
| 242 | submodule] = submodule.register_forward_pre_hook( |
| 243 | pre_forward, with_kwargs=True) |
| 244 | |
| 245 | |
| 246 | def disable_debug(): |
no test coverage detected