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 about amp_config. For apex_amp, please check apex_amp config. For naive_amp, please check naive_amp config. For torch_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 object

  • optimizer (torch.optim.Optimizer) – your optimizer object

  • amp_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 object

  • criterion (torch.nn.modules.loss._Loss, optional) – your loss function object

  • amp_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)

class colossalai.amp.AMP_TYPE(value)[source]

An enumeration.