diff --git a/src/pytorch_lightning/strategies/deepspeed.py b/src/pytorch_lightning/strategies/deepspeed.py index e7fbcf91967fc..2d32503dd406a 100644 --- a/src/pytorch_lightning/strategies/deepspeed.py +++ b/src/pytorch_lightning/strategies/deepspeed.py @@ -152,7 +152,8 @@ def __init__( Arguments: - zero_optimization: Enable ZeRO optimization. This is only compatible with precision=16. + zero_optimization: Enable ZeRO optimization. This is compatible with either `precision=16` or + `precision="bf16"`. stage: Different stages of the ZeRO Optimizer. 0 is disabled, 1 is optimizer state partitioning, 2 is optimizer+gradient state partitioning,