colossalai.amp

This module contains different amp models and optimizers.

• colossalai.amp.apex_amp
• colossalai.amp.naive_amp
• colossalai.amp.torch_amp

• colossalai.amp.amp_type

colossalai.amp.convert_to_amp(model, optimizer, criterion, mode, amp_config=None)

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.