# PyTorch 场景的精度数据采集

msprobe 工具主要通过在训练脚本内添加 dump 接口、启动训练的方式采集精度数据。

本工具提供固定的 API 支持列表，若需要删除或增加 dump 的 API，可以在 msprobe/pytorch/hook_module/support_wrap_ops.yaml 文件内手动修改，如下示例：

```yaml
functional:  # functional为算子类别，找到对应的类别，在该类别下按照下列格式删除或添加API
  - conv1d
  - conv2d
  - conv3d
```

删除API的场景：部分模型代码逻辑会存在API原生类型校验，工具执行dump操作时，对模型的API封装可能与模型的原生API类型不一致，此时可能引发校验失败，详见《[FAQ](FAQ.md)》中“异常情况”的第10和11条。

## 1 接口介绍

### 1.1 PrecisionDebugger

**功能说明**：通过加载 dump 配置文件的方式来确定 dump 操作的详细配置。

**原型**：

```Python
PrecisionDebugger(config_path=None, task=None, dump_path=None, level=None, model=None, step=None)
```

1. config_path：指定 dump 配置文件路径；model：指定具体的 torch.nn.Module，默认未配置，level 配置为"L0"或"mix"时，必须在该接口或 **start** 接口中配置该参数。其他参数均在 [config.json](../config.json) 文件中可配，详细配置可见 [config.json 介绍](./02.config_introduction.md)。
2. 此接口的参数均不是必要，且优先级高于 [config.json](../config.json) 文件中的配置，但可配置的参数相比 config.json 较少。

### 1.2 start

**功能说明**：启动精度数据采集。在模型初始化之后的位置添加。需要与 stop 函数一起添加在 for 循环内。

**原型**：

```Python
debugger.start(model=None)
```

1.  model：指定具体的 torch.nn.Module，默认未配置，level 配置为"L0"或"mix"时，必须在该接口或 **PrecisionDebugger** 接口中配置该参数。
本接口中的 model 比 PrecisionDebugger 中 model 参数优先级更高，会覆盖 PrecisionDebugger 中的 model 参数。

### 1.3 stop

**功能说明**：停止精度数据采集。在 **start** 函数之后的任意位置添加。若需要 dump 反向数据，则需要添加在反向计算代码（如，loss.backward）之后。使用示例可参见 [2.1 快速上手](#21-快速上手)和 [2.2 采集完整的前反向数据](#22-采集完整的前反向数据)。

**原型**：

```Python
debugger.stop()
```

### 1.4 forward_backward_dump_end

**功能说明**：停止精度数据采集。用于 dump 指定代码的前反向数据。在 **start** 函数之后，反向计算代码（如，loss.backward）之前的任意位置添加，可以采集 **start** 函数和该函数之间的前反向数据，可以通过调整 **start** 函数与该函数的位置，来指定需要 dump 的代码块。要求 **stop** 函数添加在反向计算代码（如，loss.backward）之后，此时该函数与 **stop** 函数之间的代码不会被 dump。使用示例可参见 [2.3 采集指定代码块的前反向数据](#23-采集指定代码块的前反向数据)

**原型**：

```Python
forward_backward_dump_end()
```

### 1.5 step

**功能说明**：更新 dump 参数。在最后一个 **stop** 函数后或一个 step 结束的位置添加。需要与 **start** 函数一起添加在 for 循环内。

**原型**：

```Python
debugger.step()
```

### 1.6 module_dump

**功能说明**：开启模块级精度数据dump。该接口为函数模块化接口，即只会dump输入的模块数据，不会dump子模块和模块内API的数据。
需要配合start、stop和step等接口使用。使用示例可参考[2.4 采集函数模块化数据](#24-采集函数模块化数据)

**原型**：

```Python
module_dump(module, module_name)
```
**参数说明**:

1. module: 网络中实例化好的nn.Module类对象。数据类型：torch.nn.Module。必选。
2. module_name：用户自定义该module名称，主要用于dump数据的命名。数据类型：str。必选。

### 1.7 module_dump_end

**功能说明**：结束模块级精度数据dump。在module_dump和module_dump_end之间的API或者Module，除了传入的module，其他数据均不会被dump。
在执行完module_dump_end后数据恢复正常的dump模式。

**原型**：

```Python
module_dump_end()
```

### 1.8 seed_all

**功能说明**：对于网络中随机性的固定和确定性计算，msprobe工具提供了seed_all接口用于固定网络中的随机性和开启确定性计算。

**原型**：
```python
seed_all(seed=1234, mode=False)
```

**参数说明**:

1. seed: 随机性种子。参数示例: seed=1000。默认值：1234。非必选
2. mode：确定性计算模式。可配置True或False。参数示例：mode=True。默认为False。非必选（注意：确定性计算会导致API执行性能降低，建议在发现模型多次执行结果不同的情况下开启）

seed_all 函数可固定随机数的范围如下表。

| API                                      | 固定随机数                  |
| ---------------------------------------- | --------------------------- |
| os.environ['PYTHONHASHSEED'] = str(seed) | 禁止 Python 中的 hash 随机化 |
| os.environ['HCCL_DETERMINISTIC'] = True  | 固定通信算子计算的确定性       |
| random.seed(seed)                        | 设置 random 随机生成器的种子  |
| np.random.seed(seed)                     | 设置 numpy 中随机生成器的种子 |
| torch.manual_seed(seed)                  | 设置当前 CPU 的随机种子       |
| torch.cuda.manual_seed(seed)             | 设置当前 GPU 的随机种子       |
| torch.cuda.manual_seed_all(seed)         | 设置所有 GPU 的随机种子       |
| torch_npu.npu.manual_seed(seed)          | 设置当前 NPU 的随机种子       |
| torch_npu.npu.manual_seed_all(seed)      | 设置所有 NPU 的随机种子       |
| torch.use_deterministic_algorithms(True) | CUDA/CANN 使能确定性计算（注意mode为True时才会调用该方法开启确定性     |
| torch.backends.cudnn.enable=False        | 关闭 cuDNN                   |
| torch.backends.cudnn.benchmark=False     | cuDNN 确定性地选择算法       |
| torch.backends.cudnn.deterministic=True  | cuDNN 仅使用确定性的卷积算法 |

需要保证 CPU 或 GPU 以及 NPU 的模型输入完全一致，dump 数据的比对才有意义，seed_all 并不能保证模型输入完全一致，如下表所示场景需要保证输入的一致性。

| 场景            | 固定方法      |
| --------------- | ------------- |
| 数据集的 shuffle | 关闭 shuffle。 |
| dropout         | 关闭 dropout。 |

关闭 shuffle 示例：

```python
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size = batch_size,
shuffle = False,
num_workers = num_workers
)
```

关闭 dropout：

在使用 `from msprobe.pytorch import PrecisionDebugger` 后，工具会自动将 `torch.nn.functional.dropout`、`torch.nn.functional.dropout2d`、`torch.nn.functional.dropout3d`、`torch.nn.Dropout`、`torch.nn.Dropout2d`、`torch.nn.Dropout3d` 的接口参数 p 置为0.

## 2 示例代码

### 2.1 快速上手

这个示例定义了一个 nn.Module 类型的简单网络，在进行数据采集时使用原型函数 PrecisionDebugger 传入 config_path 参数和 model 参数。

```python
# 根据需要import包
import torch
import torch.nn as nn
import torch_npu  # 需安装 torch_npu
import torch.nn.functional as F
from msprobe.pytorch import PrecisionDebugger, seed_all
# 在模型训练开始前固定随机性
seed_all()

torch.npu.set_device("npu:0")
# 定义网络
class ModuleOP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear_1 = nn.Linear(in_features=8,out_features=4)
        self.linear_2 = nn.Linear(in_features=4,out_features=2)

    def forward(self,x):
        x1 = self.linear_1(x)
        x2 = self.linear_2(x1)
        r1 = F.relu(x2)
        return r1

if __name__ == "__main__":
	module = ModuleOP()
	# 注册工具
	debugger = PrecisionDebugger('./config.json', model=module)
	debugger.start()
	x = torch.randn(10,8)
	out = module(x)
	loss = out.sum()
	loss.backward()
	debugger.stop()

```
### 2.2 采集完整的前反向数据

```Python
from msprobe.pytorch import PrecisionDebugger, seed_all
# 在模型训练开始前固定随机性
seed_all()
# 请勿将PrecisionDebugger的初始化流程插入到循环代码中
debugger = PrecisionDebugger(config_path="./config.json", dump_path="./dump_path")
# 模型、损失函数的定义及初始化等操作
# ...
# 数据集迭代的位置一般为模型训练开始的位置
for data, label in data_loader:
	debugger.start() # 开启数据dump
	# 如下是模型每个step执行的逻辑
    output = model(data)
    #...
    loss.backward()
	debugger.stop() # 关闭数据dump
	debugger.step() # 结束一个step的dump
```

### 2.3 采集指定代码块的前反向数据

```Python
from msprobe.pytorch import PrecisionDebugger, seed_all
# 在模型训练开始前固定随机性
seed_all()
# 请勿将PrecisionDebugger的初始化流程插入到循环代码中
debugger = PrecisionDebugger(config_path="./config.json", dump_path="./dump_path")

# 模型、损失函数的定义及初始化等操作
# ...
# 数据集迭代的位置一般为模型训练开始的位置
for data, label in data_loader:
	debugger.start() # 开启数据dump
	# 如下是模型每个step执行的逻辑
    output = model(data)
    debugger.forward_backward_dump_end() # 插入该函数到start函数之后，只dump start函数到该函数之间代码的前反向数据，本函数到stop函数之间的数据则不dump
    #...
    loss.backward()
	debugger.stop() # 关闭数据dump
	debugger.step() # 结束一个step的dump
```

### 2.4 采集函数模块化数据

```Python
# 根据需要import包
import torch
import torch.nn as nn
import torch_npu  # 需安装 torch_npu
import torch.nn.functional as F
from msprobe.pytorch import PrecisionDebugger, module_dump, module_dump_end

torch.npu.set_device("npu:0")
# 定义网络
class ModuleOP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear_1 = nn.Linear(in_features=8, out_features=4)
        self.linear_2 = nn.Linear(in_features=4, out_features=2)

    def forward(self, x):
        x1 = self.linear_1(x)
        x2 = self.linear_2(x1)
        r1 = F.relu(x2)
        return r1

if __name__ == "__main__":
    module = ModuleOP()
    # 注册工具
    debugger = PrecisionDebugger(config_path='./config.json')
    debugger.start()                   # 开启数据dump
    
    x = torch.randn(10, 8)             
    # ...                              # start和module_dump接口之间的数据正常dump
    module_dump(module, "MyModuleOP")  # 开启模块级精度数据dump
    out = module(x)                    # module内部的child modules或API不会被dump
    module_dump_end()                  # 关闭模块级精度数据dump
    loss = out.sum()                   # module_dump_end和stop接口之间的数据正常dump
    loss.backward()          

    debugger.stop()                    # 关闭数据dump
```

## 3 dump 结果文件介绍

训练结束后，工具将 dump 的数据保存在 dump_path 参数指定的目录下。目录结构示例如下：

```Python
├── dump_path
│   ├── step0
│   |   ├── rank0
│   |   │   ├── dump_tensor_data
|   |   |   |    ├── Tensor.permute.1.forward.pt
|   |   |   |    ├── MyModule.0.forward.input.pt        # 开启模块级精度数据dump时存在模块级的dump数据文件
|   |   |   |    ...
|   |   |   |    └── Fcuntion.linear.5.backward.output.pt
│   |   |   ├── dump.json        # 保存前反向算子、算子的统计量信息或溢出算子信息。包含dump数据的API名称（命名格式为：`{api_type}_{api_name}_{API调用次数}_{前向反向}_{input/output}.{参数序号}`）、dtype、 shape、各数据的max、min、mean、L2norm统计信息以及当配置summary_mode="md5"时的CRC-32数据。其中，“参数序号”表示该API下的第n个参数，例如1，则为第一个参数，若该参数为list格式，则根据list继续排序，例如1.1，表示该API的第1个参数的第1个子参数；L2norm表示L2范数（平方根）
│   |   |   ├── stack.json        # 算子调用栈信息
│   |   |   └── construct.json        # 分层分级结构
│   |   ├── rank1
|   |   |   ├── dump_tensor_data
|   |   |   |   └── ...
│   |   |   ├── dump.json
│   |   |   ├── stack.json
|   |   |   └── construct.json
│   |   ├── ...
│   |   |
|   |   └── rank7
│   ├── step1
│   |   ├── ...
│   ├── step2
```

dump 过程中，pt 文件在对应算子或者模块被执行后就会落盘，而 json 文件则需要在正常执行 PrecisionDebugger.stop() 后才会写入完整数据，异常的程序终止会保存终止前被执行算子的相关 npy 文件，可能会导致 json 文件中数据丢失。

其中 rank 为设备上各卡的 ID，每张卡上 dump 的数据会生成对应 dump 目录。非分布式场景下没有 rank ID，目录名称为 rank。

pt 文件保存的前缀和 PyTorch 对应关系如下：

| 前缀        | Torch模块           |
| ----------- | ------------------- |
| Tensor      | torch.Tensor        |
| Torch       | torch               |
| Functional  | torch.nn.functional |
| NPU         | NPU 亲和算子         |
| VF          | torch._VF           |
| Aten        | torch.ops.aten      |
| Distributed | torch.distributed   |



