colossalai.zero

colossalai.zero.convert_to_zero(model, optimizer, level, zero_config)

A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading

Parameters

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

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

  • level (int) – Optimizer level, can be 2 or 3

  • zero_config (dict) – Configuration for zero

Returns

(model, optimizer)

Return type

Tuple