The context manager style 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 i
(model: nn.Module,
dest_folder: Optional[str] = None,
filter: Optional[Filter] = None)
| 263 | |
| 264 | @contextmanager |
| 265 | def debug_mode(model: nn.Module, |
| 266 | dest_folder: Optional[str] = None, |
| 267 | filter: Optional[Filter] = None): |
| 268 | """ |
| 269 | The context manager style interface to enable debugger on model. |
| 270 | If filter is provided, it will be used to filter out satisfied module to register hook. |
| 271 | If filter is not provided, all modules will be registered with hooks. |
| 272 | Example: |
| 273 | from tensorrt_llm._torch.debug.debug_hook import debug_mode |
| 274 | model_config = ModelConfig(pretrained_config=llama_config, |
| 275 | attn_backend=backend) |
| 276 | llama = LlamaForCausalLM(model_config).to(dtype).to(device) |
| 277 | llama.load_weights(hf_llama.state_dict()) |
| 278 | with torch.inference_mode() and debug_mode(llama, r"tensor_dump"): |
| 279 | attn_metadata.prepare() |
| 280 | logits = llama.forward(input_ids=input_ids, |
| 281 | position_ids=position_ids, |
| 282 | attn_metadata=attn_metadata) |
| 283 | Args: |
| 284 | model (nn.Module): the model to enable debug hook. |
| 285 | dest_folder: the working directory set to debug context to set where the hook dumped data/info. |
| 286 | filter: a filter to decide what modules will be registered with debug hook. |
| 287 | Returns: |
| 288 | None |
| 289 | """ |
| 290 | try: |
| 291 | enable_debug(model, dest_folder, filter) |
| 292 | register_tensor_dump_hook() |
| 293 | yield model |
| 294 | finally: |
| 295 | disable_debug() |
| 296 | |
| 297 | |
| 298 | def get_forward_arg_names(module: nn.Module): |