Skip to content

llm.engine.base

/home/runner/work/llm-pytorch/llm-pytorch/llm/engine/base.py

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: torch.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()