Skip to content

llm.engine.amp

Utilities for easy automatic mixed precision training.

AMPCriterion 

AMPCriterion(criterion: Module, autocast: autocast)

Bases: Module

Wrap a loss function for AMP training.

Parameters:

  • criterion (Module) –

    Loss function to wrap.

  • autocast (autocast) –

    Autocast context manager to compute loss inside.

Source code in llm/engine/amp.py 
23
24
25
26
27
28
29
30
def __init__(
    self,
    criterion: torch.nn.Module,
    autocast: torch.autocast,
) -> None:
    super().__init__()
    self._criterion = criterion
    self._autocast = autocast

forward 

forward(*args: Any, **kwargs: Any) -> Any

Compute the loss inside the autocast.

Source code in llm/engine/amp.py 
32
33
34
35
def forward(self, *args: Any, **kwargs: Any) -> Any:
    """Compute the loss inside the autocast."""
    with self._autocast:
        return self._criterion(*args, **kwargs)

AMPModel 

AMPModel(model: Module, autocast: autocast)

Bases: Module

Wrap a model for AMP training.

Parameters:

  • model (Module) –

    Model to wrap.

  • autocast (autocast) –

    Autocast context manager to compute loss inside.

Source code in llm/engine/amp.py 
46
47
48
49
50
51
52
53
def __init__(
    self,
    model: torch.nn.Module,
    autocast: torch.autocast,
) -> None:
    super().__init__()
    self._model = model
    self._autocast = autocast

forward 

forward(*args: Any, **kwargs: Any) -> Any

Perform a forward pass inside the autocast.

Source code in llm/engine/amp.py 
55
56
57
58
def forward(self, *args: Any, **kwargs: Any) -> Any:
    """Perform a forward pass inside the autocast."""
    with self._autocast:
        return self._model(*args, **kwargs)

AMPOptimizer 

AMPOptimizer(
    model: Module,
    optimizer: Optimizer,
    scaler: GradScaler,
    max_norm: float | None = None,
)

Bases: BaseOptimizer

Wrap an optimizer for AMP training.

Parameters:

  • model (Module) –

    Model being optimized.

  • optimizer (Optimizer) –

    Optimizer to wrap.

  • scaler (GradScaler) –

    Gradient scaler.

  • max_norm (float | None, default: None ) –

    Optionally clip gradient norm.

Source code in llm/engine/amp.py 
71
72
73
74
75
76
77
78
79
80
81
def __init__(
    self,
    model: torch.nn.Module,
    optimizer: Optimizer,
    scaler: GradScaler,
    max_norm: float | None = None,
) -> None:
    super().__init__(optimizer)
    self._model = model
    self._scaler = scaler
    self._max_norm = max_norm

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)

state_dict 

state_dict() -> dict[str, Any]

Dictionary containing references to the whole state of the module.

Includes the state of the grad_scaler.

Source code in llm/engine/amp.py 
83
84
85
86
87
88
89
90
91
def state_dict(self) -> dict[str, Any]:
    """Dictionary containing references to the whole state of the module.

    Includes the state of the `grad_scaler`.
    """
    state_dict = self._optimizer.state_dict()
    assert 'grad_scaler' not in state_dict
    state_dict['grad_scaler'] = self._scaler.state_dict()
    return state_dict

load_state_dict 

load_state_dict(state_dict: dict[str, Any]) -> None

Copy the state into this module.

Source code in llm/engine/amp.py 
93
94
95
96
97
98
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Copy the state into this module."""
    scaler_state_dict = state_dict.pop('grad_scaler', None)
    if scaler_state_dict is not None:  # pragma: no branch
        self._scaler.load_state_dict(scaler_state_dict)
    self._optimizer.load_state_dict(state_dict)

backward 

backward(loss: Tensor) -> None

Perform a backward pass and correctly scale the loss.

Source code in llm/engine/amp.py 
100
101
102
def backward(self, loss: torch.Tensor) -> None:
    """Perform a backward pass and correctly scale the loss."""
    self._scaler.scale(loss).backward()

step 

step(closure: Callable[[], float] | None = None) -> None

Perform an optimization using the gradient scaler.

Source code in llm/engine/amp.py 
104
105
106
107
108
109
110
111
112
113
def step(self, closure: Callable[[], float] | None = None) -> None:
    """Perform an optimization using the gradient scaler."""
    if self._max_norm is not None:
        self._scaler.unscale_(self._optimizer)
        torch.nn.utils.clip_grad_norm_(
            self._model.parameters(),
            self._max_norm,
        )
    self._scaler.step(self._optimizer)
    self._scaler.update()

initialize 

initialize(
    model: Module,
    optimizer: Optimizer,
    criterion: Module,
    dtype: dtype = torch.float16,
    max_norm: float | None = None,
    **kwargs: Any
) -> tuple[AMPModel, AMPOptimizer, AMPCriterion]

Initialize AMP training.

Parameters:

  • model (Module) –

    Model being optimized.

  • optimizer (Optimizer) –

    Optimizer to wrap.

  • criterion (Module) –

    Loss function to wrap.

  • dtype (dtype, default: float16 ) –

    Data type to perform mixed precision in. Typically torch.float16 or torch.bfloat16.

  • max_norm (float | None, default: None ) –

    Optionally clip gradient norm.

  • kwargs (Any, default: {} ) –

    Additional keyword arguments to pass to the GradScaler.

Returns:

Source code in llm/engine/amp.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def initialize(
    model: torch.nn.Module,
    optimizer: Optimizer,
    criterion: torch.nn.Module,
    dtype: torch.dtype = torch.float16,
    max_norm: float | None = None,
    **kwargs: Any,
) -> tuple[AMPModel, AMPOptimizer, AMPCriterion]:
    """Initialize AMP training.

    Args:
        model: Model being optimized.
        optimizer: Optimizer to wrap.
        criterion: Loss function to wrap.
        dtype: Data type to perform mixed precision in. Typically
            `torch.float16` or `torch.bfloat16`.
        max_norm: Optionally clip gradient norm.
        kwargs: Additional keyword arguments to pass to the
            [`GradScaler`][torch.cuda.amp.GradScaler].

    Returns:
        A tuple of the wrapped model, optimizer, and loss.
    """
    device = 'cuda' if next(model.parameters()).is_cuda else 'cpu'
    autocast = torch.autocast(device, dtype=dtype)

    # GradScaler only works on CUDA tensors so we disable on CPU
    scaler = GradScaler(**kwargs, enabled=device == 'cuda')

    return (
        AMPModel(model, autocast),
        AMPOptimizer(model, optimizer, scaler, max_norm),
        AMPCriterion(criterion, autocast),
    )