colossalai.amp.apex_amp

colossalai.amp.apex_amp.convert_to_apex_amp(model, optimizer, amp_config)

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.

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)

Returns

A tuple (model, optimizer).

Return type

Tuples

More details about amp_config refer to amp_config.

class colossalai.amp.apex_amp.ApexAMPOptimizer(optim)

A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm methods

backward(loss)

Backward pass to get all gradients

Parameters

loss (torch.Tensor) – Loss computed by a loss function

clip_grad_norm(model, max_norm)

Clip gradients’ norm

Parameters

• model (torch.nn.Module) – Your model object

• max_norm (float) – The max norm value for gradient clipping