Skip to content

llm.engine.base

BaseOptimizer 

BaseOptimizer(optimizer: Optimizer)

Bases: Optimizer

Base optimizer wrapper.

Compared to a standard PyTorch Optimizer, this optimizer adds the backward() that is used instead of directly calling loss.backward(). This is needed so the various optimizer wrappers can adjust how the loss is computed.

Otherwise, this optimizer behaves identically to the wrapped optimizer.

Parameters:

  • optimizer (Optimizer) –

    Standard PyTorch optimizer to wrap.

Source code in llm/engine/base.py 
24
25
def __init__(self, optimizer: Optimizer):
    self._optimizer = optimizer

step 

step(*args: Any, **kwargs: Any) -> None

Perform an optimization step.

Source code in llm/engine/base.py 
48
49
50
def step(self, *args: Any, **kwargs: Any) -> None:
    """Perform an optimization step."""
    self._optimizer.step(*args, **kwargs)

zero_grad 

zero_grad(*args: Any, **kwargs: Any) -> None

Zero the gradients of optimized parameters.

Source code in llm/engine/base.py 
52
53
54
def zero_grad(self, *args: Any, **kwargs: Any) -> None:
    """Zero the gradients of optimized parameters."""
    self._optimizer.zero_grad(*args, **kwargs)

backward 

backward(loss: Tensor) -> None

Perform a backward pass using the loss.

Source code in llm/engine/base.py 
56
57
58
def backward(self, loss: torch.Tensor) -> None:
    """Perform a backward pass using the loss."""
    loss.backward()