colossalai.amp
- colossalai.amp.convert_to_amp(model, optimizer, criterion, mode, amp_config=None)[source]
A helper function to wrap training components with Torch AMP modules.
- Parameters
model (param) – your model object.
optimizer (
torch.optim.Optimizer
) – your optimizer object.criterion (
torch.nn.modules.loss._Loss
) – your loss function object.mode (
colossalai.amp.AMP_TYPE
) – amp mode.amp_config (Union[
colossalai.context.Config
, dict]) – configuration for different amp modes.
- Returns
A tuple (model, optimizer, criterion).
Note
amp_config
may vary from different mode you choose. You should check the corresponding amp mode for more details aboutamp_config
. Forapex_amp
, please check apex_amp config. Fornaive_amp
, please check naive_amp config. Fortorch_amp
, please check torch_amp config.
- colossalai.amp.convert_to_naive_amp(model, optimizer, amp_config)[source]
A helper function to wrap training components with naive AMP modules. In this mode, we forcibly cast the model weights and inputs to FP16, and cast the model outputs to FP32 to calculate loss, which is equivalent to Apex O3.
- Parameters
model (
torch.nn.Module
) – your model objectoptimizer (
torch.optim.Optimizer
) – your optimizer objectamp_config (
colossalai.context.Config
or dict) – configuration for naive mode amp.
- Returns
A tuple (model, optimizer)
- Return type
Tuple
The
amp_config
should contain parameters below:verbose (bool, optional): if set to `True`, will print debug info (Default: False). clip_grad_norm (float, optional): clip gradients with this global L2 norm (Default 0). Note that clipping is ignored if clip_grad == 0. dynamic_grad_scale (bool): whether to use dynamic grad scaler.
- colossalai.amp.convert_to_apex_amp(model, optimizer, amp_config)[source]
A helper function to wrap training components with Apex AMP modules
- Parameters
model (
torch.nn.Module
) – your model object.optimizer (
torch.optim.Optimizer
) – your optimizer object.amp_config (Union[
colossalai.context.Config
, dict]) – configuration for initializing apex_amp.
- Returns
A tuple (model, optimizer).
- Return type
Tuple
The
amp_config
should include parameters below:enabled (bool, optional, default=True) opt_level (str, optional, default="O1") cast_model_type (``torch.dtype``, optional, default=None) patch_torch_functions (bool, optional, default=None) keep_batchnorm_fp32 (bool or str, optional, default=None master_weights (bool, optional, default=None) loss_scale (float or str, optional, default=None) cast_model_outputs (torch.dtype, optional, default=None) num_losses (int, optional, default=1) verbosity (int, default=1) min_loss_scale (float, default=None) max_loss_scale (float, default=2.**24)
More details about
amp_config
refer to amp_config.
- colossalai.amp.convert_to_torch_amp(model, optimizer, criterion=None, amp_config=None)[source]
A helper function to wrap training components with Pytorch AMP modules
- Parameters
model (
torch.nn.Module
) – your model object.optimizer (
torch.optim.Optimizer
) – your optimizer objectcriterion (
torch.nn.modules.loss._Loss
, optional) – your loss function objectamp_config (
colossalai.context.Config
or dict, optional) – configuration for Pytorch AMP.
The
amp_config
should include parameters below:init_scale (float, optional, default=2.**16) growth_factor (float, optional, default=2.0) backoff_factor (float, optional, default=0.5) growth_interval (int, optional, default=2000) enabled (bool, optional, default=True)
- Returns
A tuple (model, optimizer, criterion)