optimizers.rst 792 Bytes
Newer Older
aiss's avatar
aiss committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
Optimizers
===================

DeepSpeed offers high-performance implementations of ``Adam`` optimizer on CPU; ``FusedAdam``, ``FusedLamb``, ``OnebitAdam``, ``OnebitLamb`` optimizers on GPU.

Adam (CPU)
----------------------------
.. autoclass:: deepspeed.ops.adam.DeepSpeedCPUAdam

FusedAdam (GPU)
----------------------------
.. autoclass:: deepspeed.ops.adam.FusedAdam

FusedLamb (GPU)
----------------------------
.. autoclass:: deepspeed.ops.lamb.FusedLamb

OneBitAdam (GPU)
----------------------------
.. autoclass:: deepspeed.runtime.fp16.onebit.adam.OnebitAdam

ZeroOneAdam (GPU)
----------------------------
.. autoclass:: deepspeed.runtime.fp16.onebit.zoadam.ZeroOneAdam

OnebitLamb (GPU)
----------------------------
.. autoclass:: deepspeed.runtime.fp16.onebit.lamb.OnebitLamb