Training Metric Monitoring and Profiling

The MindSpore Transformers dynamic graph (PyNative) provides two types of observability capabilities in the training loop:

Training metric monitoring (monitor): collects fine-grained metrics such as gradient/norm and loss by micro-step (including the reserved configuration for MoE expert load monitoring). These metrics are currently output only through training logs (logger.info).
Performance data collection (profiler): collects performance data such as operators, communication, memory, and call stacks within a specified step range. The results are flushed to the disk as MindSpore Profiler products for performance analysis and optimization.

Both capabilities are enabled through the training YAML configuration and are independent of each other. They can be used separately. If you have not performed a training task, you are advised to read Quick Start first. For details about the overall structure of the configuration file, see Configuration File Description.

Output Description

Currently, dynamic graph monitoring metrics (in the monitor section) are written only to logs. Each record is printed by Monitor._flush_logger in { key: value, ... } format to the training log. To visualize the norm/loss curve, you need to parse the logs.

1 Training Metric Monitoring

The monitor section has two sub-configurations: train_state and moe_monitor.

The monitoring is instantiated by MonitorGroup during training initialization based on the configuration. Only the activated sub-monitoring is created: train_state is activated when any norm or loss switch is turned on; moe_monitor is activated when save_tokens_per_expert_interval is a positive integer. Metrics are flushed to logs by MonitorCallback (on_step_end) at the end of each training step.

1.1 train_state (Monitoring norm and loss)

The gradient norms and losses at the local and device levels are collected by micro-step. The four switches are independent of each other and can be combined in any way.

Parameter	Data Type	Required/Optional	Default Value	Description
`local_norm`	Boolean/String/list[str]	Optional	`""`	Specifies whether to monitor the local gradient norm. `True` indicates that all trainable parameters are monitored. A string or multiple strings in a list are used to filter parameters by parameter name substring (`"layers.0"` matches parameters whose names contain the substring). `""` or `False` indicates that this function is disabled.
`local_loss`	Boolean	Optional	`False`	Specifies whether to monitor the local loss of each micro-batch.
`device_norm`	Boolean/String/list[str]	Optional	`False`	Specifies whether to monitor the device-level (accumulated) gradient norm. The value meaning is the same as that of `local_norm`.
`device_loss`	Boolean	Optional	`False`	Specifies whether to monitor the device-level loss (loss after gradient accumulation/all-reduce).

device_norm forcibly collects local_norm.

In the source code, the accumulated norm of device_norm comes from the prev_local_norms cached during local_norm collection. Therefore, as long as device_norm is enabled, the framework computes the norm of the required parameters during the local_norm collection phase, even if local_norm itself is disabled. Both of them independently filter the parameters to be printed based on their parameter name filters.

Output Interpretation

local_norm is the incremental norm of this step. The framework maintains the accumulated norm accumulated for each parameter and performs differentiation by micro-step:
- micro-step 0: actual_norm = accumulated;
- Subsequent steps: actual_norm = accumulated – prev (prev is the value of accumulated in the previous micro-step). In this way, each record reflects the incremental contribution of the current micro-batch to the gradient norm, rather than the total norm up to the current point.
device_norm is an accumulated value and directly prints the accumulated norm cached in prev_local_norms.
Logs are printed by parameter, for example, { "local_norm": decoder.layers.0.self_attention.q_layernorm.weight: 0.000381 }. Before multiple norm records of the same micro-step are printed, a record in the format of { "step": 4, "micro_step": 2 } is printed as a group marker.
The local_loss record is in the format of { "micro_step": 2, "local_loss": 0.123 }. device_loss is recorded after the optimizer is updated.

Applicable Scenarios

Scenario A (Simplest): Viewing Only the Global Loss Curve

If you are only concerned about the overall convergence and do not need to compute the norm of each parameter, you do not need to enable any norm of train_state. The step-level loss is directly printed to logs by LossCallback of the framework. The monitoring segment can be left empty or only device_loss can be enabled.

monitor:
  train_state:
    device_loss: True       # Print the reduced loss after each step optimizer is updated.

Scenario B: Checking for Gradient Explosion or Norm Exceptions at a Layer

When a loss spike or NaN occurs, or you suspect that the gradient at some layers is abnormal, use local_norm to filter and locate the fault layer by parameter name. The incremental norm semantics (see "Output Interpretation" above) allow you to check which micro-batch and which parameter are used to contribute to the abnormal increment.

monitor:
  train_state:
    # Only the incremental gradient norm of these types of parameters is monitored to avoid screen flooding by parameter.
    local_norm: ["attention", "layernorm"]
    local_loss: True        # Observe the local loss of each micro-batch to locate abnormal micro-batches.

Note: To monitor all parameters, use local_norm: True. To monitor a single parameter, use a string, for example, local_norm: "embedding".

Scenario C: Observing Accumulation by Using device_norm and local_norm Together

If you need to view both the incremental value per step and the total accumulated value, enable both of them. Note that device_norm reuses the collection result of local_norm (see the preceding note block).

monitor:
  train_state:
    local_norm: True        # Incremental norm of each micro-step.
    device_norm: True       # Accumulated norm after gradient accumulation (local_norm collection is forcibly triggered).
    device_loss: True

1.2 moe_monitor (MoE Expert Load)

This function is designed to monitor the number of tokens (tokens-per-expert) allocated to each expert at the MoE decoding layer and is used to diagnose uneven expert load. If a few experts are saturated for a long time and the rest are idle, the routing is unbalanced. In this case, you need to adjust the load balancing loss or routing strategy.

Parameter	Data Type	Required/Optional	Default Value	Description
`save_tokens_per_expert_interval`	int / None	Optional	`None`	Prints the micro-step interval recorded by tokens-per-expert. You can set this parameter to `None` (indicating that the monitoring is disabled) or a positive integer. If you set this parameter to a Boolean value, an error will be reported.
`target_layers`	Integer/list[int]/None	Optional	`None`	Specifies the target layer to be monitored. `int` indicates the first N layers of `range(N)`, matching only the decoder layer. `list[int]` indicates the specified layer ID. `None` indicates that all MoE modules containing `tokens_per_expert` in the model are automatically discovered.

Semantic differences between int and list of target_layers

int (for example, 3) is expanded to range(3), that is, layers 0, 1, and 2 of the model.

list[int] (for example, [0, 2, 4]) is used to match these specified layer IDs.

If the layer with the specified ID is not an MoE layer (without tokens_per_expert), the alarm decoder.layers.X is not a MoE layer is reported in the framework.

Output Format

MoE monitoring information is printed to logs in the tokens-per-expert JSON format of Megatron (one JSON record per line). The fields are described as follows.

Field	Description
`iter`	Global micro-step count (`global_micro_step`).
`step`	Training step ID.
`micro_step`	Sequence number of micro-batch in a step.
`block`	`"decoder"` or `"mtp"`.
`layer`	Layer ID.
`mtp_idx`	MTP sublayer, which is displayed only for `block == "mtp"`.
`tpe`	List of tokens allocated to each expert in the micro-batch (increment obtained based on expert differentiation).

Configuration Example: MoE Load Imbalance Diagnosis

The configuration of the monitoring is as follows:

monitor:
  moe_monitor:
    save_tokens_per_expert_interval: 10   # Record every 10 micro-steps (design semantics).
    target_layers: [0, 5, 10]             # Sample the first, middle, and last layers to check for layer imbalance.

You can also use target_layers: 4 to monitor the first four layers, or omit this field to enable the framework to automatically discover all MoE layers (resulting in a larger data volume). The value of save_tokens_per_expert_interval must be a positive integer. If True or False is used, TypeError will be reported during configuration verification.

1.3 MaxLogits Health Monitoring (callbacks: MaxLogitsMonitor)

During large-scale long-term training, attention logits value overflow and gradient spikes can undermine training stability. MaxLogitsMonitor is a health monitoring callback for the actual values consumed in the dynamic graph. For the source code, see mindformers/pynative/callback/max_logits_monitor.py. Unlike the sub-configurations under the monitor section, it is registered through the callbacks section.

MaxLogitsMonitor collects the maximum attention logit of each layer of the model at each training step, prints the logit by layer to the log, and outputs the mean and maximum values of all layers. Then, it resets the accumulated value within each layer at the end of each step, facilitating the observation of whether the attention value increases abnormally step by step.

Output format (_dump and max_logits_monitor.py:88-108):

One line per layer. The log tag is max_attention_logit/<param_name>, and the value is a list of items at the layer (four significant digits are retained).
Two summary lines: max_attention_logit/mean and max_attention_logit/max.
The prefix of each line is step:[Current step/Total steps], which is the same as the printing format of TrainingStateMonitor.

Application Scenarios

If you suspect that the loss jitter or NaN is caused by attention logits overflow or numerical explosion, use this function to locate the layer and step where the abnormal increase starts.
When the Muon optimizer is used and the QK clip is enabled (in this case, the framework automatically enables tracing, as described in the following scenario), use this function to verify whether the clip takes effect.

Parameter Configuration

Register the callback through the callbacks section.

callbacks:
  - type: MaxLogitsMonitor
    step_interval: 1   # Number of steps between outputs. The value must be a positive integer.

Parameter	Data Type	Required/Optional	Default Value	Description
`step_interval`	int	Optional	`1`	Number of steps between outputs. The value must be a positive integer. Otherwise, `ValueError` (`max_logits_monitor.py:51-54`) is thrown during construction.

Behavior details (on_step_end and max_logits_monitor.py:58-78):

When step_interval > 1 and the current step is not an output step, only the accumulated value is reset and is not printed.
In the output step, get_max_attention_logit() of each layer is collected first. If no data is available, only the value is reset.
Regardless of whether the value is printed, reset_max_attention_logit() is called at the end of each step to clear the value, ensuring that the next step starts from a clean state.

Scenario: Enabling the Automatic Tracking with Muon and QK Clip

model.track_max_attention_logit is used to specify whether to collect the max attention logit on the model side. The framework automatically specifies whether to enable this function in configure_max_logits_tracking (max_logits_monitor.py:120-132).

The optimizer is Muon and qk_clip_enabled is True.
MaxLogitsMonitor has been configured in callbacks.

When either of the preceding condition is met, track_max_attention_logit is set to True. This assembly is triggered during Trainer initialization (pynative/trainer/trainer.py:153), and ensures that a MaxLogitsMonitor is present in ensure_max_logits_reset_callback to handle the reset (trainer.py:536-538).

Therefore, in the Muon + QK clip scenario, even if no MaxLogitsMonitor is explicitly written, the framework automatically adds one for resetting. To view layer-by-layer logs, it is still recommended to explicitly write one and set step_interval.

# Scenario: Explicitly enabling layer-by-layer max logit observation in Muon + QK clip
optimizer:
  type: Muon
  qk_clip_enabled: True       # Enable QK clip. The framework automatically enables track_max_attention_logit based on this setting.

callbacks:
  - type: MaxLogitsMonitor
    step_interval: 50         # Print the layer-by-layer max logit and mean/max every 50 steps.

Tracking (instead of printing) is automatically enabled.

For Muon + QK clip, only model-side tracking is automatically enabled, and a callback for resetting is added. To view the layer-by-layer numerical values in the logs, you still need to explicitly configure MaxLogitsMonitor (or accept the step-by-step output of the default step_interval=1).

2 Profile Data Collection

The profiler section controls profile data collection. The collection takes effect only when enable_profiling: True and the current rank hits profiler_rank. The framework uses MindSpore schedule to collect data within the [start_step,end_step] range. (Training outside the range is normal and no data is collected.)

Parameter	Data Type	Required/Optional	Default Value	Description
`enable_profiling`	Boolean	Optional	`False`	Specifies whether to enable profile data collection.
`start_step`	int	Optional	`1`	Specifies the step at which collection starts (corresponding to `skip_first` of `schedule`). The value must be a positive integer.
`end_step`	int	Optional	`1`	Specifies the step at which collection ends. The value must be greater than or equal to `start_step` (the number of collected steps = `end_step – start_step + 1`). If a non-default `start_step` is set, `end_step` must also be set. Otherwise, the default `1` may cause an invalid range.
`output_path`	String/None	Optional	`None`	Specifies the path for saving the result. The result is saved in `output_path/rank_x` by rank. If this parameter is set to `None`, the result is rolled back to save in `Current working directory/profile/rank_x`.
`profiler_rank`	list[int]/None	Optional	`None`	Specifies the list of rank IDs for which collection is enabled. `None` indicates that all ranks are collected.
`profiler_level`	int	Optional	`0`	Specifies the collection level, which can be `0`, `1`, or `2`. A higher level indicates more detailed information. (The value is mapped to `ProfilerLevel.Level0/1/2`. If the value is invalid, `LevelNone` is used.)
`mstx`	Boolean	Optional	`False`	Specifies whether to enable lightweight mstx dotting (transparently transmitted to `_ExperimentalConfig`) to insert lightweight markers on the timeline.
`profile_memory`	Boolean	Optional	`False`	Specifies whether to collect tensor memory data.
`profile_cpu`	Boolean	Optional	`True`	Specifies whether to collect CPU profiling activities.
`with_stack`	Boolean	Optional	`True`	Specifies whether to collect call stack data on the Python side.

Avoiding the first step for warm-up during range selection

The first step usually includes one-off overheads such as compilation and initialization, and the collected data does not represent the steady state. You are advised to set start_step to several steps later (for example, start_step: 5) and collect data for only a few steps (for example, 2 to 3 steps) to avoid warm-up and control the product size and overhead.

Overhead trade-off

profile_memory: True and with_stack: True significantly increase the collection overhead and product size, which may slow down the collected steps. Enable this option only when checking memory or call stack issues. It is recommended that profile_memory be disabled when analyzing the time consumption of pure operators.

2.1 Rank Selection in Multi-Device Scenarios

In multi-device training, you do not need to collect data from all devices because the data volume is large and the data is redundant. You can use profiler_rank to specify the rank to be collected.

In single-device tuning, only [0] is collected.
When checking for communication imbalance, you can collect data from different ranks in the same parallel group for comparison (for example, [0, 1]).
For details about the mapping between ranks and parallel dimensions (dp/tp/pp/cp), see Distributed Parallel Training. Select representative ranks based on the mapping.

2.2 Output Directory and Viewing Method

Directory structure: output_path/rank_x/ if output_path is set; <cwd>/profile/rank_x/ if None is used.

2.3 Scenario-based Configuration (Three Levels Based on profiler_level)

Scenario A: Lightweight Operator-Level Analysis (Level 0)

Only the operator time consumption distribution is viewed, and hotspot kernels are located, resulting in the minimum overhead. Disable the memory and call stack.

profiler:
  enable_profiling: True
  start_step: 5            # Skip the first four warm-up steps.
  end_step: 7              # Perform three-step collection.
  profiler_rank: [0]
  profiler_level: 0
  profile_memory: False
  profile_cpu: False
  with_stack: False
  output_path: "./output/profile"

Scenario B: Communication Analysis (Level 1)

If you need to analyze the communication time and the overlap between compute and communication in distributed training, increase the level to Level 1 and perform comparison collection for multiple ranks as required.

profiler:
  enable_profiling: True
  start_step: 5
  end_step: 7
  profiler_rank: [0, 1]    # Comparison within the same group to observe communication imbalance.
  profiler_level: 1
  mstx: True               # Lightweight dotting, facilitating alignment of key phases on the timeline.
  profile_memory: False
  with_stack: False
  output_path: "./output/profile"

Scenario C: Including memory and call stack (Level 2)

To check the peak graphics memory usage or attribute time consumption to the Python call stack, enable Level 2 and enable memory and call stack collection (which has the highest overhead). In this case, you are advised to collect data for only a single device and a small number of steps):

profiler:
  enable_profiling: True
  start_step: 5
  end_step: 6              # Perform two-step collection to control the size of the product.
  profiler_rank: [0]
  profiler_level: 2
  profile_memory: True     # Collect tensor memory.
  with_stack: True         # Collect the Python call stack.
  output_path: "./output/profile"