ColossalAIOptimWrapper

class mmengine._strategy.colossalai.ColossalAIOptimWrapper(optimizer, booster=None, accumulative_counts=1)

OptimWrapper for ColossalAI.

The available optimizers are:

• CPUAdam

• FusedAdam

• FusedLAMB

• FusedSGD

• HybridAdam

• Lamb

• Lars

You can find more details in the colossalai tutorial

Parameters:

• optimizer (dict or torch.optim.Optimizer) – The optimizer to be wrapped.

• accumulative_counts (int) – The number of iterations to accumulate gradients. The parameters will be updated per accumulative_counts.

• booster (None) –

backward(loss, **kwargs)

Perform gradient back propagation.

Provide unified backward interface compatible with automatic mixed precision training. Subclass can overload this method to implement the required logic. For example, torch.cuda.amp require some extra operation on GradScaler during backward process.

Note

If subclasses inherit from OptimWrapper override backward, _inner_count +=1 must be implemented.

Parameters:

• loss (torch.Tensor) – The loss of current iteration.

• kwargs – Keyword arguments passed to torch.Tensor.backward().

Return type:

None

optim_context(model)

A Context for gradient accumulation and automatic mix precision training.

If subclasses need to enable the context for mix precision training, e.g., :class:`AmpOptimWrapper, the corresponding context should be enabled in optim_context. Since OptimWrapper uses default fp32 training, optim_context will only enable the context for blocking the unnecessary gradient synchronization during gradient accumulation

If model is an instance with no_sync method (which means blocking the gradient synchronization) and self._accumulative_counts != 1. The model will not automatically synchronize gradients if cur_iter is divisible by self._accumulative_counts. Otherwise, this method will enable an empty context.

Parameters:

model (nn.Module) – The training model.