colossalai.nn.optimizer.lars

Adapted from https://github.com/NUS-HPC-AI-Lab/LARS-ImageNet-PyTorch/blob/main/lars.py

class colossalai.nn.optimizer.lars.Lars(params, lr=0.001, momentum=0, eeta=0.001, weight_decay=0, epsilon=0.0)

Implements the LARS optimizer from “Large batch training of convolutional networks”.

Parameters

  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups

  • lr (float, optional) – learning rate (default: 1e-3)

  • momentum (float, optional) – momentum factor (default: 0)

  • eeta (float, optional) – LARS coefficient as used in the paper (default: 1e-3)

  • weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

step(closure=None)

Performs a single optimization step.

Parameters

closure (callable, optional) – A closure that reevaluates the model and returns the loss.