| 55 | |
| 56 | @total_ordering |
| 57 | class DSZeRO(Enum): |
| 58 | NONE = 0 |
| 59 | """No DeepSPeed ZeRO; requires the most amount of memory.""" |
| 60 | |
| 61 | STAGE_1 = 1 |
| 62 | """ZeRO stage 1 shards the optimizer states across the data parallel group.""" |
| 63 | |
| 64 | STAGE_2 = 2 |
| 65 | """ZeRO stage 2 shards the optimizer states and gradients across the data parallel |
| 66 | group.""" |
| 67 | |
| 68 | STAGE_3 = 3 |
| 69 | """ZeRO stage 3 shards the optimizer states, gradients, and model weights across the |
| 70 | data parallel group.""" |
| 71 | |
| 72 | def __lt__(self, other): |
| 73 | if self.__class__ is other.__class__: |
| 74 | return self.value < other.value |
| 75 | return NotImplemented |
| 76 | |
| 77 | |
| 78 | class LLMAnalysis: |