# Monitor 训练状态轻量化监控工具

## 简介

训练状态轻量化监控工具，能够在较低性能损耗下收集和记录模型训练过程中的激活值、权重梯度、优化器状态和通信算子的中间值，实时呈现训练状态。

- [快速上手](#快速上手)
    - [权重监控](#权重监控)
    - [权重梯度监控](#权重梯度监控)
    - [激活值监控](#激活值监控)
    - [优化器状态监控](#优化器状态监控)
- [详细配置](#详细配置)

## 安装
参见[msprobe安装](./01.installation.md)

## 快速上手
根据需求监控相应对象。比如在loss上扬，grad norm正常的异常训练过程中，优先考虑监控模型前向过程；在grad norm异常的训练过程中，监控权重和激活值的梯度。
推荐使用方式：权重梯度的监控性能损耗小（20B dense模型全量权重梯度监控，时间增加<1%，内存增加<1%），可以长期开启。激活值监控性能损耗大，在必要时开启或者仅监控部分。  

### 工具使能
在训练脚本中使能工具，在配置文件（json）中控制工具行为。
```python
# megatorn中构建初始化模型和优化器。在实际训练代码中找到模型、优化器（optional）定义的位置
# Megatron-LM(core_r0.6.0)  megatron/training.py, def pretrain:
model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
        model_provider, model_type) 

# 使能工具
from msprobe.pytorch import TrainerMon
# 监控工具初始化
monitor = TrainerMon(
    config_file_path="./monitor_config.json",
    process_group=None,
    params_have_main_grad=True,  # 权重是否使用main_grad，通常megatron为True，deepspeed为False。默认为True。
    opt_ty=None  # 优化器类型，默认为None，具体取值参考公开接口
) 
monitor.set_wrapped_optimizer(optimizer)
# 挂载监控对象
monitor.monitor_gnorm_with_ad(
    model,
    grad_acc_steps=args.global_batch_size//args.data_parallel_size//args.micro_batch_size,
    optimizer=None,
    dp_group=None,
    tp_group=None) 


# optional
# 可以在任意位置获取当前的参数梯度统计量
reduced, unreduced = monitor.generate_wgrad_metrics()
# 可以在任意位置获取当前的激活值、激活值梯度统计量
actv, actv_grad = monitor.generate_xy_metrics()
```

补充deepspeed下常用框架的使能位置，提供参考

- accelerate

```python
model, optimizer, trainloader, evalloader, schedular = accelerator.prepare(...)

monitor = TrainerMon(...)
monitor.set_wrapped_optimizer(optimizer.optimizer)  # optimizer.optimizer为DeepSpeedZeroOptimizer
monitor.monitor_gnorm_with_ad(....optimizer=optimizer.optimizer)
```

- transformers

```python
# src/transformers/trainer.py
class Trainer:
    def _inner_training_loop:
        ...
        monitor = TrainerMon(...)
        monitor.set_wrapped_optimizer(self.optimizer.optimizer)
        monitor.monitor_gnorm_with_ad(....optimizer=self.optimizer.optimizer)

        for epoch in range(epochs_trained, num_train_epochs):
            ...
```

### 指定监控对象

工具支持对nn.Module（**激活值监控**）和nn.Parameter（**梯度监控、优化器监控**）对象实现相应的监控行为，在配置文件的"targets"（dict）字段指定，targets格式为{module_name/param_name: {filed: format}}。

- 打印模型结构
工具提供可选项"print_struct"打印模型结构，帮助配置targets。工具会在在第一个step后打印结构并停止训练进程。  
```json
{
    "print_struct": true
}
```

输出样例:
"config"字段用于配置文件中指定module target。其余为各个元素的shape和dtype。   

```json
"0:63.mlp.linear_fc2": {
    "input": {
        "config": "tuple[1]",
        "0": "size=(4096, 4, 1024), dtype=torch.bfloat16"
    },
    "output": {
        "config": "tuple[2]",
        "0": "size=(2048, 4, 512), dtype=torch.bfloat16",
        "1": "size=(512,), dtype=torch.bfloat16"
    },
    "input_grad": {
        "config": "tuple[1]",
        "0": "size=(4096, 4, 1024), dtype=torch.bfloat16"
    },
    "output_grad": {
        "config": "tuple[2]",
        "0": "size=(2048, 4, 512), dtype=torch.bfloat16",
        "1": "size=(512,), dtype=torch.bfloat16"
    }
},
```

- Module
  对于module对象，通常关心其前向的输入(input)输出(output)和反向的输入--前向输出的梯度(output_grad)和输出--前向输入的梯度（input_grad）。同时需要声明这些对象的类型，通常为"tensor"或"tuple\[length]"。

  "tensor"可以直接用来计算统计量，"tuple"需要进一步指定监控的索引。如"tuple[2]:0",表示该对象为长度2的tuple，对第0元素进行监控；不指定索引时，默认对第0元素进行监控。

  module_name可以通过nn.Module的接口`named_modules()`获取。  
```json
// 示例：对一个名为"module.encoder.layers.0.mlp"的module，监控其前向输入第0元素和输出。
{
    "targets": {
        "module.encoder.layers.0.mlp": {
            "input": "tuple[2]:0", 
            "output": "tensor"
        }
    }
}
```
**Module全量监控**：工具提供简便的全量module监控方式。

```json
{
    "targets": {},
    "all_xy": true
}
```


- Parameter
  对于parameter对象，通常会关注其在一个训练迭代中的梯度（weight grad）、adam类优化器中的动量（1st moment, 2nd moment）。
  parameter归属于某一module，也可以通过指定module_name来监控包含在这一module中的**所有**parameter。

  param_name可以通过nn.Module的接口`named_parameters()`获取。
```json
// 示例：监控"module.encoder.layers.0.mlp"的所有参数和"module.embedding.word_embedding.weight"这一参数
{
    "targets": {
        "module.encoder.layers.0.mlp": {},
        "module.embedding.word_embedding.weight": {}
    }
}
```

**Parameter全量监控**：工具提供简便的全量parameter监控方式。

```json
{
    "targets": {"":{}}
}
```

### 输出格式和统计量
工具配置示例：
```json
{
    "format": "csv",
    "ops": ["norm", "min", "max", "mean", "nans", "zeros"],
    "ndigits": 12
}
```

- 输出路径 
通过环境变量`MONITOR_OUTPUT_DIR`设置，默认为"monitor_output"。
```shell
export MONITOR_OUTPUT_DIR=/xxx/output_dir
```

- 输出格式 
  通过可选配置项`format`指定。可以是\["tensorboard"（缺省值）, "csv", "api"\]。

    - format: tensorboard 
      监控结果写入tensorboard的event文件，启动tensorboard查看
    ```shell
    tensorboard --logdir=$MONITOR_OUTPUT_DIR
    ```
    之后，运行以下SSH命令来建立端口转发，可以在本地通过http://localhost:6006访问tensorboard：
    ```shell
    ssh -N -L localhost:6006:localhost:6006 your_username@remote_server_address
    ```

    - format: csv 
      监控结果写入csv文件中，可以通过`ndigits`字段设置小数位数.

    - format: api 
      监控结果不落盘，在训练过程中可以通过`generate_wgrad_metrics`、`generate_xy_metrics`等接口获取。

- 统计量 
通过必选配置项"ops"指定。可以是["norm", "min", "max", "mean", "nans"，"zeros"]。其中"nans"统计tensor中nan的数量，"zeros"统计tensor中数值小于"eps"的比例。  

### 权重监控
- 工具配置示例：
```json
{  
    "targets": {
        "": {}
    },
    "param_distribution": true,
    "format": "csv",
    "ops": ["norm", "min", "max", "nans"]
}  
```
"targets"中指定module包含的所有权重都会被监控。整个model的name为空字符串可以覆盖全量权重。
设置"param_distribution"开启权重监控功能。

使用deepspeed的zero优化器时，需要在工具中指定优化器类型并传入优化器，获取梯度切分行为已还原参数梯度。
```python
from msprobe.pytorch import TrainerMon
# 以zero1优化器举例，opt_ty取值DeepSpeedZeroOptimizer_Stage1_or_2
# 示例为deepspeed，params_have_main_grad取值False
monitor = TrainerMon("./monitor_config.json", params_have_main_grad=False, opt_ty="DeepSpeedZeroOptimizer_Stage1_or_2")
monitor.set_wrapped_optimizer(optimizer) # optimzier为训练框架自定义的优化器
monitor.monitor_gnorm_with_ad(
        model, grad_acc_steps=model.grad_acc_steps, optimizer=optimizer)
```


### 权重梯度监控
- 工具配置示例：
```json
{  
    "targets": {
        "": {}
    },
    "wg_distribution": true,
    "format": "csv",
    "ops": ["norm", "min", "max", "nans"]
}  
```
"targets"中指定module包含的所有权重都会被监控。整个model的name为空字符串可以覆盖全量梯度。
设置"wg_distribution"(weight grad, noted as `wg`)开启梯度监控功能。

使用deepspeed的zero优化器时，需要在工具中指定优化器类型并传入优化器，获取梯度切分行为已还原参数梯度。
```python
from msprobe.pytorch import TrainerMon
# 以zero1优化器举例，opt_ty取值DeepSpeedZeroOptimizer_Stage1_or_2
# 示例为deepspeed，params_have_main_grad取值False
monitor = TrainerMon("./monitor_config.json", params_have_main_grad=False, opt_ty="DeepSpeedZeroOptimizer_Stage1_or_2")
monitor.set_wrapped_optimizer(optimizer) # optimzier为训练框架自定义的优化器
monitor.monitor_gnorm_with_ad(
        model, grad_acc_steps=model.grad_acc_steps, optimizer=optimizer)
```


### 梯度异常时序判断
1. 训练前配置相关参数

工具支持自动判断训练过程中的梯度异常，需要在配置文件中设置alert相关字段。"AnomalyTurbulence"会将当前数值与历史均值比较，如果相对偏差超过阈值，会在打屏信息中提示用户。如果打开"dump"选项，则会将异常梯度相关信息落盘到目录`monitor_output/anomaly_detected`，用于后续时序判断。
```json
    "alert": {
        "rules": [{"rule_name": "AnomalyTurbulence", "args": {"threshold": 0.5}}],
        "dump": true
    },
```
2. 实例化工具时传入流水线并行group
```python
monitor = TrainerMon("./monitor_config.json", process_group=mpu.get_pipeline_model_parallel_group(), params_have_main_grad=True)
```
照常开始训练

3. 启动异常分析脚本：

```shell
python3 -m msprobe.pytorch.monitor.anomaly_analyse -d $MONITOR_OUTPUT_DIR/anomaly_detected
```
支持以下参数配置
| 字段名                                                       | 解释    | 是否必选   | 
| ------ | -------- | -------- | 
|-d 或 --data_path| 指定梯度异常落盘文件夹，梯度监控功能输出，一般为$MONITOR_OUTPUT_DIR/anomaly_detected。|是 |
|-o 或 --out_path| 排序后的异常落盘文件地址，默认在--data_path路径下落盘一个anomaly_analyse.json文件| 否 |
|-k 或 --topk| 指定保留前topk个异常，默认为8| 否 |
|-s 或 --step_list| 指定分析的step范围，默认为[]| 否 |



### 激活值监控

- 工具配置
```json
{  
    "targets": {
        "module.module.language_model.encoder.layers.0": {
            "input": "tuple[2]", 
            "output": "tensor"
        }
    },
    "print_struct": false,
    "xy_distribution": true,
    "forward_only": true,
    "backward_only": false,
    "all_xy": true,
    "format": "csv",
    "ops": ["norm", "min", "max", "nans"]
}  
```
设置"xy_distribution"为true表示开启激活值监控功能，"all_xy"为true表示监控全量module激活值。

forward_only和backward_only均为true时，触发warning，前反向均不采集；均为false时，前反向均采集。
```python
from msprobe.pytorch import TrainerMon
# 以zero1优化器举例，opt_ty取值DeepSpeedZeroOptimizer_Stage1_or_2
# 示例为deepspeed，params_have_main_grad取值False
monitor = TrainerMon("./monitor_config.json", params_have_main_grad=False, opt_ty="DeepSpeedZeroOptimizer_Stage1_or_2")
monitor.set_wrapped_optimizer(optimizer) # optimzier为训练框架自定义的优化器
monitor.monitor_gnorm_with_ad(
        model, grad_acc_steps=model.grad_acc_steps, optimizer=optimizer)
```



### 功能重载
- 统计量
可以在训练过程中修改`TrainerMon`实例的`ops`属性, 调整监控的统计量。
```python
if {some condition}:
    monitor.ops = ["min", "max"]
```

- 训练过程中开关激活值监控
激活值监控的性能损耗较大, 推荐仅在必要时开启, 比如发现loss出现尖刺, 根据loss的异常开启激活值监控.
```python
if {some condition}:
    monitor.reload_xy(xy_distribution=True)
```

### 优化器状态监控
- 工具配置示例：
```json
{  
    "targets": {
        "module.encoder.layers.0": {},
        "module.embedding.word_embedding.weight": {}
    },
    "mv_distribution": true,
    "format": "csv",
    "ops": ["norm", "min", "max", "nans"]
}  
```
"targets"中指定module包含的所有权重都会被监控。
设置"mv_distribution"表示开启优化监控功能（1st moment noted as `m`, 2nd moment noted as `v`）。[什么是mv](https://arxiv.org/pdf/1412.6980)

本工具针对分布式计算框架megatron和deepspeed框架做了适配，暂不支持其他框架。

```python
from msprobe.pytorch import TrainerMon
# 以zero1优化器举例，opt_ty取值DeepSpeedZeroOptimizer_Stage1_or_2
# 示例为deepspeed，params_have_main_grad取值False
monitor = TrainerMon("./monitor_config.json", params_have_main_grad=False, opt_ty="DeepSpeedZeroOptimizer_Stage1_or_2")
monitor.set_wrapped_optimizer(optimizer) # optimzier为训练框架自定义的优化器
monitor.monitor_gnorm_with_ad(
        model, grad_acc_steps=model.grad_acc_steps, optimizer=optimizer)
```

## 公开接口

```python
TrainerMon.__init__(config_file_path, process_group=None, params_have_main_grad=True, opt_ty=None) -> None
```

| 参数  | 说明                  | 是否必选 |
| ----- | -------------------- | -------- |
| config_file_path |json配置文件路径。 | 是       |
| process_group | 传入ProcessGroup对象，用以确定pipeline并行不同rank异常间时序，megatron下通过core.parallel_state.get_pipeline_model_parallel_group()获得。 | 否       |
| params_have_main_grad |权重是否使用main_grad，通常megatron为True，deepspeed为False。默认为True。 | 否       |
| opt_ty |优化器类型，默认为None。<br>-Megatron_DistributedOptimizer：megatron分布式优化器；<br/>-Megatron_Float16OptimizerWithFloat16Params：megatron混合精度优化器；<br/>-DeepSpeedZeroOptimizer_Stage1_or_2：DeepSpeed Zero1和Zero2；<br/>-DeepSpeedZeroOptimizer_Stage3：DeepSpeed Zero3。 | 否      |

```python
TrainerMon.monitor_gnorm_with_ad(model, grad_acc_steps, optimizer, dp_group, tp_group) -> None
```
| 参数  | 说明                  | 是否必选 |
| ----- | -------------------- | -------- |
| model |需要监控的模型，需要是一个torch.nn.Module。 | 是       |
| grad_acc_steps | 梯度累积步数。 | 是      |
| optimizer | 需要patch的优化器 | 否      |
| dp_group | 数据并行的通信组。<br>dp域通信后，且没有使用分布式优化器时，group内所有rank的梯度相同，落盘数据冗余。<br>提供dp_group后，工具仅保留每个dp_group的第一个rank的梯度。 | 否      |
| tp_group | 张量并行的通信组。<br/>tp域通信后，group内部分参数所有rank的梯度相同，落盘数据冗余。<br/>提供tp_group后，工具仅保留每个tp_group中冗余参数在第一个rank的梯度。<br/>当前适配Megatron core_v0.6.0, 通过权重属性"tensor_model_parallel"判断是否冗余。 | 否      |


```python
TrainerMon.set_wrapped_optimizer(_wrapped_optimizer) -> None
```

| 参数  | 说明                  | 是否必选 |
| ----- | -------------------- | -------- |
| _wrapped_optimizer |megatron、deepspeed创建好的混合精度优化器。 | 是       |

```python
TrainerMon.generate_wgrad_metrics() -> tuple[dict[dict]]
```

```python
TrainerMon.generate_xy_metrics() -> tuple[dict[dict]]
```


##  详细配置

```json
{  
    "targets": {  
        "language_model.encoder.layers.0": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}  
    },
    "print_struct": false,
    "module_ranks": [0,1,2,3],
    "ur_distribution": true,
    "xy_distribution": true,
    "all_xy": true,
    "forward_only": false,
    "backward_only": false,
    "mv_distribution": true,
    "param_distribution": true,
    "wg_distribution": true,
    "cc_distribution": {"enable":true, "cc_codeline":[]},
    "alert": {
        "rules": [{"rule_name": "AnomalyTurbulence", "args": {"threshold": 0.5}}],
        "dump": false
    },
    "format": "tensorboard",
    "ops": ["min", "max", "norm", "zeros", "nans", "mean"],
    "eps": 1e-8,
    "ndigits": 12,
    "step_count_per_record": 1
}  
```

下面详细解释各个字段：

| 字段名字                                                       | 是否必选    | 解释   |
| ------------------------------------------------------------ | -------- | -------- |
|"targets"| 必选 |指定需要监控的模型层， 例如transformer的第0层language_model.encoder.layers.0。如果不清楚模型结构， 可以将 "print_struct" 字段设置为 true， 监控工具会打印模型中torch module的名字和详细结构，并在第1个step后退出。|
|"input"| 可选 |"tuple[2]:0"的意思是目标module的前向input参数为长度为2的tuple， 我们关心的是tuple第0个元素。|
|"output"| 必选 |"tensor"的意思是目标module的前向output参数类型为tensor|
|"input_grad"| 可选 |"tuple[2]:0"的意思是目标module的后向input_grad参数是长度为2的tuple， 我们关心的是tuple的第0个元素。|
|"output_grad"| 必选 |"tuple[1]:0"的意思是目标module的后向input_grad参数是长度为1的tuple， 我们关心的是tuple的第0个元素。|
|"print_struct"| 可选 |设置为true后监控工具会打印模型中torch module的名字和详细结构，并在第1个step后退出。不填默认为false。|
|"module_ranks"| 可选 |用于在分布式训练场景中希望控制在哪些rank开启module监控。如果不填，则默认在所有rank开启。|
|"ur_distribution"| 可选 |若为true则会统计adam优化器指定模块（targets中指定）参数的update和ratio向量的数值分布，并展示在heatmap里，默认为false，同时format字段必须设置为tensorboard。<br/>依赖histc算子， 需要CANN8.0.rc2以上版本， 否则会有严重的性能问题。 |
|"xy_distribution"| 可选 |若为true则会监控指定module（targets中指定）的输入输出张量。 默认为false。|
|"all_xy"| 可选 |开启xy_distribution后生效，若为true，监控所有module。默认为false。|
|"forward_only"| 可选 |开启xy_distribution后生效，若为true，仅监控指定module的前向，targets中的input_grad、output_grad不生效。默认为false。|
|"backward_only"| 可选 |开启xy_distribution后生效，若为true，仅监控指定module的反向，targets中的input、output不生效。默认为false。|
|"mv_distribution"| 可选 |若为true则会监控指定模块中的参数的优化器状态， 默认为false。需要在TrainerMon构造函数正确指定opt_ty。 目前支持megatron和Deepspeed的分布式优化器。<br/>-Megatron_DistributedOptimizer：megatron分布式优化器；<br/>-Megatron_Float16OptimizerWithFloat16Params：megatron混合精度优化器；<br/>-DeepSpeedZeroOptimizer_Stage0：DeepSpeed Zero0<br/>-DeepSpeedZeroOptimizer_Stage1_or_2：DeepSpeed Zero1和Zero2；<br/>-DeepSpeedZeroOptimizer_Stage3：DeepSpeed Zero3。 |
|"wg_distribution"| 可选 |若为true则会监控指定模块的参数梯度， 默认为false。 |
|"param_distribution"| 可选 |若为true则会监控指定模块的参数， 默认为false。 |
|"alert"| 可选 | "rules": 指定自动报警的异常检测机制及其相应的阈值。目前实现的异常检测是AnomalyTurbulence， 如果统计标量超出历史均值的指定浮动范围（threshold 0.5意味着上浮或者下浮50%）则在控制台打印报警信息。当"dump"字段配置为true表示异常事件写入文件，默认为false。 |
|"cc_distribution"|  可选 |其中"enable"字段控制通信监控模块的开关；需要监控通信算子时，务必尽量早地实例化`TrainerMon`, 因为监控通过劫持原始func后挂hook实现，部分加速库初始化时会保存原始function，避免监控失效。"cc_codeline"字段指定监控的代码行，如:`train.py\\[23\\]`，默认为空列表，不特别指定；"cc_pre_hook"字段控制是否监控通信前的数据； 模块会在第二个optimize.step之前打印通信日志，包括通信api的调用栈、输入dtype、通信group。 "cc_log_only"为true时，仅打印日志，不监控通信的输入输出，并在打印后中断训练。可以根据通信日志设置"cc_codeline"，规避与训练过程不相关的通信，比如一些时间、metrics的同步。|
|"format"| 可选 | 数据落盘格式，默认为tensorboard，可选 \["tensorboard", "csv", "api"\]。 |
|"ops"| 必选 |类型为list，与ur_distribution、xy_distribution、mv_distribution、wg_distribution、mg_direction、cc_distribution配合，监控所选张量的统计指标，目前支持"min"、"max"、"norm"、"mean"、"zeros"、"nans"。其中，zeros代表监控所选张量的元素小于eps的比例，nans代表张量中nan的数量。|
|"eps"|  可选 |若ops里包含"zeros"则需要配置，默认为1e-8。|
|"ndigits"| 可选 |"format"为"api"时，设置落盘文件中的小数位数，默认为6。|
|"step_count_per_record"| 可选 | "format"为"csv"时生效，每个csv记录多少个step的数据，默认为1。|
