composer.algorithms.fused_layernorm.fused_layernorm#
composer.algorithms.fused_layernorm.fused_layernorm
Functions
Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm. |
|
|
composer.algorithms.fused_layernorm.fused_layernorm.check_if_apex_installed |
Defines a replacement policy from a torch.nn.LayerNorm to a apex.normalization.fused_layer_norm |
Classes
Base class for algorithms. |
|
Enum to represent training loop events. |
|
Replaces all instances of torch.nn.LayerNorm with a apex.normalization.fused_layer_norm.FusedLayerNorm. |
|
An interface to record training data. |
|
The state of the trainer. |
Exceptions
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.