DeepSpeedPrecisionPlugin¶
- class lightning.pytorch.plugins.precision.DeepSpeedPrecisionPlugin(precision)[source]¶
Bases:
lightning.pytorch.plugins.precision.precision_plugin.PrecisionPlugin
Precision plugin for DeepSpeed integration.
Warning
This is an experimental feature.
- Parameters
precision¶ (
Literal
[‘32-true’, ‘16-mixed’, ‘bf16-mixed’]) – Full precision (32), half precision (16) or bfloat16 precision (bf16).- Raises
ValueError – If unsupported
precision
is provided.
- backward(tensor, model, optimizer, *args, **kwargs)[source]¶
Performs back-propagation using DeepSpeed’s engine.
- Parameters
model¶ (
LightningModule
) – the model to be optimized*args¶ – additional positional arguments for the
deepspeed.DeepSpeedEngine.backward()
call**kwargs¶ – additional keyword arguments for the
deepspeed.DeepSpeedEngine.backward()
call
- Return type