mindformers.modules.transformer.TransformerRecomputeConfig

class mindformers.modules.transformer.TransformerRecomputeConfig(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.

参数

• 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

实际案例

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