DtypeOverrider#

class DtypeOverrider(patterns=None, class_patterns=None, dtype='float32', methods_to_override=None, override_parameters=False, verbose=False, check_interval=10)[source]#

Bases: HookBase

Force matched layers to compute (and optionally store params) in a dtype.

Runs once in before_train: walks the model and, for every submodule whose dotted name matches patterns or whose class name matches class_patterns, wraps each method in methods_to_override so its tensor input is cast to dtype and its output cast back to the original dtype. With override_parameters=True the matched module’s own parameters are also cast to dtype in place. Useful for keeping precision-sensitive layers (e.g. LayerNorm) in fp32 under mixed-precision training. Registered as DtypeOverrider.

Parameters:

  • patterns (list, optional) – Regex patterns matched against dotted module names. Defaults to None (treated as []).

  • class_patterns (list, optional) – Regex patterns matched against module class names. Defaults to None (treated as []).

  • dtype (str) – Target dtype name (key of the module’s dtype table, e.g. "float32", "bfloat16"), resolved to a torch.dtype. Defaults to "float32".

  • methods_to_override (list, optional) – Method names on matched modules to wrap with the cast. Defaults to None (treated as ["forward"]).

  • override_parameters (bool) – If True, also cast matched modules’ parameters to dtype in place. Defaults to False.

  • verbose (bool) – If True, log each wrapped method/parameter and a summary. Defaults to False.

  • check_interval (int) – Reserved cadence for periodic parameter-dtype re-checks. Defaults to 10.

Note

Only the first positional tensor argument of a wrapped method is cast; methods whose first argument is not a tensor are called unchanged. The DDP wrapper is handled by walking the underlying model.

Example

Add to cfg.hooks; once in before_train it wraps matched modules so they compute in dtype regardless of AMP autocast:

hooks = [
    dict(type="DtypeOverrider",
         patterns=["layer_norm", "LayerNorm", "norm"],
         dtype="float32",
         override_parameters=True,
         verbose=True),
]
# → each matched module's forward() now casts its input to float32 and
#   casts the output back to the original dtype; with
#   override_parameters=True the module's own params are converted to
#   float32 in place (each wrap/conversion logged via verbose)
before_train()[source]#

Apply dtype overriding before training starts.