| 238 | |
| 239 | @contextmanager |
| 240 | def processor_init_context(accelerator: Accelerator): |
| 241 | if accelerator.distributed_type != DistributedType.DEEPSPEED: |
| 242 | yield |
| 243 | return |
| 244 | |
| 245 | plugin = accelerator.state.deepspeed_plugin |
| 246 | if plugin is None or not plugin.is_zero3_init_enabled(): |
| 247 | yield |
| 248 | return |
| 249 | |
| 250 | import deepspeed |
| 251 | |
| 252 | with plugin.zero3_init_context_manager(enable=False): |
| 253 | deepspeed.zero.partition_parameters.shutdown_init_context() |
| 254 | try: |
| 255 | yield |
| 256 | finally: |
| 257 | deepspeed.zero.partition_parameters.restore_init_context() |
| 258 | |
| 259 | |
| 260 | def model_init_context(accelerator: Accelerator): |