Repository: pytorch
Issue: [RFC] DebugMode
Labels: triaged

### 🚀 The feature, motivation and pitch

## RFC: Debug Mode

Inspired by issue https://github.com/pytorch/pytorch/issues/93880, this RFC proposes to design a DEBUG mode for PyTorch eager mode.
A lot of CUDA kernels fail with device assert but understanding this assert maybe tricky. We rely on device side asserts so that the performance is not degraded due to syncs and extra checks. But this leads to a problem when user sees an error which is not easy to understand.

With debug mode, users would be able to get comprehensive error messages at the cost of decreased performance. This would be an opt-in mechanism which could be used via setting an ENVIRONMENT variable or used directly in code by running under a context manager. 


**We will provide a bunch of checks (especially for indexing ops) that users can utilise as well as the ability for the user to add their own checks.**


**NOTE** : Operators compiled by PT2 stack doesn't dispatch via the dispatcher (compiled functions directly called by Python) and hence we won't probably support checks on them in first iteration (Thoughts?).

### Usage with Environment Variable:

Enviroment variable `TORCH_EAGER_DEBUG` (emphasizing that eager ops would be checked) would be one way of running the under debug mode.

Examples
1. Running with debug mode
```bash
TORCH_EAGER_DEBUG=1 python my_script.py
```

2. Specifying the checks to run
```bash
TORCH_EAGER_DEBUG="inf, index" python my_script.py
```

### Usage with Context Manager:

We will also expose context manager which will run only the relevant code under debug mode.

```python
with DebugMode():
    y = model(x)
```

### Design

In terms of design, we want to enable user to register their own checks besides the ones provided by PyTorch.

#### 1. Context Manager which accepts user defined checks.

```python
class DebugMode(TorchDispatchMode):
    def __init__(self,
                 checks: Union[str, List[str]] = 'all',  # Selection of checks provided by PyTorch
                 # Additional pre-checks (callables) user can pass
                 # Arguments to the callable are op, op_args, op_kwargs
                 pre_checks: Optional[Union[Iterable[Callable], Callable]] = None,
                 # Additional post-checks (callables) user can pass
                 # Arguments to the callable are op, op_args, op_kwargs, result
                 post_checks: Optional[Union[Iterable[Callable], Callable]] = None,
                 # custom_error_handler for an uncaught error
                 custom_error_handler: Optional[Callable] = None) -> None:
```

Observation: With this design, it will probably not be simple to share checks between users (there can't be a third party package which you can just install and get new checks).

#### 2. Context Manager with registration mechanism.

```python
class DebugMode(TorchDispatchMode):
    def __init__(self,
                 # Selection of checks provided by PyTorch (as well as user defined checks)
                 checks: Union[str, List[str]] = 'all') -> None:
```

Registration helpers (class methods)
```python
# Arguments to the callable are op, op_args, op_kwargs
DebugMode.register_debug_mode_pre_check(name, callable)

# Arguments to the callable are op, op_args, op_kwargs, result
DebugMode.register_debug_mode_post_check(name, callable)

# Arguments to the callable are exception, op, op_args, op_kwargs
DebugMode.register_debug_mode_error_handler(callable)
```

With this design, users can easily share their checks (even as third party libraries).

#### 3. C++ Macro for Debug Checks in an Operator

```c++
// one debug check
TORCH_DEBUG_CHECK(TORCH_CHECK(condition, "some error message"));

// multiple debug check
TORCH_DEBUG_CHECK({
    TORCH_CHECK(condition1, "some error message");
    TORCH_CHECK(condition2, "some other error message");
    });
```

These checks will only be triggered if the asked by the user. User will use ENVIRONMENT variable or DebugMode context manager from above to turn on these checks. 

### Implementation

We will be leveraging the `TorchDispatchMode` to implement this utility. This will allow us to intercept and add checks for both forward and backward operators. For C++ macro, we will use an internal runtime global flag to turn the macro on and off.


**Draft PR [WIP]**: https://github.com/pytorch/pytorch/pull/95952

**Thanks @chillee, @zou3519 and @albanD for help and review of the design.**

### Alternatives

_No response_

### Additional context

_No response_

cc @ZainRizvi @kit1980 @huydhn @clee2000