Documentation

Automatic AOE Tuning

The AOE tool continuously iterates tiling policies through a closed-loop feedback mechanism of policy generation, compilation, and verification in the operating environment, and finally obtains the optimal one. This helps fully utilize hardware resources, improve network performance, and achieve the optimal effect. During model training, enable the AOE tool to tune subgraphs, operators, and gradient segmentation. After the tuning is complete, the optimal scheduling policy is fixed in the repository. When the model is trained again, you can obtain performance gains using the repository without enabling tuning.

You are advised to use the AOE tool to perform tuning in the following sequence:

You can enable AOE tuning in training scenarios using either of the following methods:
  • Set the environment variable.
     # 1: subgraph tuning; 2: operator tuning; 4: gradient tuning
export AOE_MODE=2
  • Modify the training script and use aoe_mode to specify the tuning mode, for example:
    • In sess.run mode, modify the training script as follows:
      1
      custom_op.parameter_map["aoe_mode"].s = tf.compat.as_bytes("2")
    • In Estimator mode, modify the training script as follows:
      1
2
3
      config = NPURunConfig(
    session_config=session_config, 
    aoe_mode=2)
    • In keras mode, modify the training script as follows:
      1
      custom_op.parameter_map["aoe_mode"].s = tf.compat.as_bytes("2")

For details about the restrictions and functions of the AOE tool, see AOE Instructions.

Parent topic: Basic Tuning