| 66 | |
| 67 | |
| 68 | class FSDPConfig: |
| 69 | def __init__( |
| 70 | self, |
| 71 | sharding_strategy, |
| 72 | backward_prefetch, |
| 73 | cpu_offload, |
| 74 | num_replicate, |
| 75 | num_shard=8, |
| 76 | ): |
| 77 | self.sharding_strategy = sharding_strategy |
| 78 | self.backward_prefetch = backward_prefetch |
| 79 | self.cpu_offload = cpu_offload |
| 80 | self.num_replicate = num_replicate |
| 81 | self.num_shard = num_shard |
| 82 | |
| 83 | |
| 84 | def fsdp_wrapper(original_model, fsdp_config, ignored_modules=[]): |