mindformers.modules.transformer.TransformerRecomputeConfig

class mindformers.modules.transformer.TransformerRecomputeConfig(recompute=False, select_recompute=False, parallel_optimizer_comm_recompute=False, mp_comm_recompute=True, recompute_slice_activation=False)

TransformerRecomputeConfig for the setting recompute attributes for encoder/decoder layers.

Args:

recompute (bool): Enable recomputation of the transformer block or not. Default: False. parallel_optimizer_comm_recompute (bool): Specifies whether the communication operator allgathers

introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode. Default: False.

mp_comm_recompute (bool): Specifies whether the model parallel communication operators

in the cell are recomputed in auto parallel or semi auto parallel mode. Default: True.

recompute_slice_activation (bool): Slice the cell output which would remains in memory. Default: False.

Supported Platforms:

Ascend GPU

Examples:

>>> from mindformers.modules.transformer import TransformerRecomputeConfig
>>> config=TransformerRecomputeConfig(recompute=True, parallel_optimizer_comm_recompute=True, \
...                                   mp_comm_recompute=True, recompute_slice_activation=True)