fused_layernorm#

Module fused_layernorm.

Functions

`apply_fused_layernorm`	Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm.
`check_if_apex_installed`	Module `check_if_apex_installed`.
`from_LayerNorm`	Defines a replacement policy from a torch.nn.LayerNorm to a apex.normalization.fused_layer_norm

Classes

`Algorithm`	Base class for algorithms.
`Event`	Enum to represent training loop events.
`FusedLayerNorm`	Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm.
`Logger`	An interface to record training data.
`State`	The state of the trainer.

Exceptions

NoEffectWarning

Warns when an algorithm did not have an effect.

Attributes

APEX_INSTALLED
Dict
Optional
Sequence
Type
Union
annotations
log

class composer.algorithms.fused_layernorm.fused_layernorm.FusedLayerNorm[source]#

Bases: composer.core.algorithm.Algorithm

Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm.

By fusing multiple kernel launches into one, this usually improves GPU utilization.

Runs on Event.INIT, so it can replace all instances of torch.nn.LayerNorm before the model is DDP wrapped. Has no hyperparameters.

Example

from composer.algorithms import FusedLayerNorm

algorithm = FusedLayerNorm()
trainer = Trainer(
    model=model,
    train_dataloader=train_dataloader,
    max_duration="1ep",
    algorithms=[algorithm],
    optimizers=[optimizer]
)

composer.algorithms.fused_layernorm.fused_layernorm.apply_fused_layernorm(model, optimizers)[source]#

Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm.

By fusing multiple kernel launches into one, this usually improves GPU utilization.

composer.algorithms.fused_layernorm.fused_layernorm.from_LayerNorm(layer, module_index)[source]#: Defines a replacement policy from a torch.nn.LayerNorm to a apex.normalization.fused_layer_norm