This document is relevant for: Inf2, Trn1, Trn2

Developer guide for Standard Mixed Precision

This document will introduce the concept of Standard Mixed Precision in NxD. It’s newly introduced in neuron release 2.20. It is recommended to use this setting for training large models using NxD. When enabled, the optimizer will maintain a copy of weights and their grads in FP32 data type.

Note

Using this can increase memory pressure as we are using master weights and also performing optimiizer updates in higher precision. This can result in increased memory pressure and a slighly lower throughpout

Standard Mixed Precision offers few config settings that can be tuned by users

Compared to legacy mixed precision setting (i.e. before this feature’s addition), Standard Mixed Precision includes these components:

• Use FP32 for precision sensitive operators

• Use FP32 master weights and optimizer states for ZeRO-1 optimizer

• Use FP32 in local gradients accumulation

• Turn off stochastic rounding

Note

The feature is tightly integrated with the NeuronZero1Optimizer, to make Standard Mixed Precision take effect, ZeRO-1 optimizer needs to be enabled.

NxD Config Update

Newly introduced NxD config is as below:

mixed_precision_config = {
    "use_master_weights": True,
    "use_fp32_grad_acc": True,
    "use_master_weights_in_ckpt": False,
}

config = {
    ...
    "mixed_precision_config": mixed_precision_config,
}

In NxD training config, a new field mixed_precision_config (default value is None, see details in the following sections) is added. It contains three sub-fields: use_master_weights, use_fp32_grad_acc, and use_master_weights_in_ckpt. Default value of use_master_weights and use_fp32_grad_acc is whether ZeRO-1 optimizer is enabled. Field use_master_weights controls whether to use FP32 master weights. Field use_fp32_grad_acc controls whether to enable FP32 gradient accumulation buffer. Default value of use_master_weights_in_ckpt is False. This field controls whether to save master weights in checkpoints.

# same as `mixed_precision_config = None`
mixed_precision_config = {
    "use_master_weights": optimizer_config["zero_one_enabled"],
    "use_fp32_grad_acc": optimizer_config["zero_one_enabled"],
    "use_master_weights_in_ckpt": False,
}

config = {
    ...
    "mixed_precision_config": mixed_precision_config,
}

Note that only when ZeRO-1 optimizer is enabled, Standard Mixed Precision will take effect.

To disable this Standard Mixed Precision setting, just change NxD config:

mixed_precision_config = {
    "use_master_weights": False,
    "use_fp32_grad_acc": False,
    "use_master_weights_in_ckpt": False,
}

config = {
    ...
    "mixed_precision_config": mixed_precision_config,
}

This document is relevant for: Inf2, Trn1, Trn2