mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.1__py3-none-any.whl

msprobe/docs/19.monitor.md

@@ -0,0 +1,468 @@
+# 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。|

msprobe/docs/20.monitor_performance_baseline.md

@@ -0,0 +1,52 @@
+# 训练状态监控工具标准性能基线报告
+
+## 环境信息
+
+NPU：Atlas A2 训练系列产品
+
+CPU：
+
+![输入图片说明](img/monitor/cpu_info.png)
+
+Torch：2.1.0
+
+CANN：8.0.RC2
+
+除上述环境信息影响性能外，被检控的模块的数量和结构会对性能产生影响，因此本次选取典型网络进行测试，并且选取耗时稳定后的步数进行测试。工具输出键小，对内存无要求。
+
+## 模型信息和性能基线
+
+以下场景的性能基线测试数据均为多次测试后取平均值，因此实际运行时性能数据可能会根据环境状态稍有浮动。
+
+### LLAMA2-13B
+
+主要数据类型：BFLOAT16
+
+模型层数：40
+
+配置文件（采了10层）：
+```
+{  
+    "targets": {  
+        "language_model.encoder.layers.0": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.1": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.2": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.3": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.4": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.5": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.6": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.7": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.8": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}，
+        "language_model.encoder.layers.9": {"input": "tuple[2]:0", "output": "tensor", "input_grad":"tuple[2]:0", "output_grad":"tuple[1]:0"}
+    },
+    "module_ranks": "0"
+}  
+```
+
+启动命令参数：python3 -u pretrain_gpt.py --local-rank=1 --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --sequence-parallel --num-layers 40 --hidden-size 5120 --ffn-hidden-size 13824 --num-attention-heads 40 --tokenizer-type Llama2Tokenizer --tokenizer-model /new_data/LLM/checkpoint_origin/llama2-13b-hf/tokenizer.model --seq-length 4096 --max-position-embeddings 4096 --micro-batch-size 2 --global-batch-size 16 --make-vocab-size-divisible-by 1 --lr 1e-6 --train-iters 5000 --lr-decay-style cosine --untie-embeddings-and-output-weights --disable-bias-linear --attention-dropout 0.0 --init-method-std 0.01 --hidden-dropout 0.0 --position-embedding-type rope --normalization RMSNorm --use-fused-rmsnorm --swiglu --use-flash-attn --no-masked-softmax-fusion --attention-softmax-in-fp32 --min-lr 1e-8 --weight-decay 1e-1 --lr-warmup-fraction 0.01 --clip-grad 1.0 --adam-beta1 0.9 --initial-loss-scale 4096 --adam-beta2 0.95 --no-gradient-accumulation-fusion --load /data/LLM/checkpoint_magatron/llama2_13b_tp1_pp8 --no-load-optim --no-load-rng --use-fused-swiglu --use-fused-rotary-pos-emb --use-mc2 --bf16 --data-path /data/LLM/data_modellink/llama2_13b/alpaca_text_document --split 949,50,1 --log-interval 1 --save-interval 10000 --eval-interval 1000 --eval-iters 10 --distributed-backend nccl --save ./ckpt
+
+不加工具原始耗时：**4s**
+
+加工具后单卡耗时：**4.25s**
+
+加工具后多卡耗时：**4.35s**