| 285 | |
| 286 | @contextmanager |
| 287 | def processor_init_context(accelerator: Accelerator): |
| 288 | if accelerator.distributed_type != DistributedType.DEEPSPEED: |
| 289 | yield |
| 290 | return |
| 291 | |
| 292 | plugin = accelerator.state.deepspeed_plugin |
| 293 | if plugin is None or not plugin.is_zero3_init_enabled(): |
| 294 | yield |
| 295 | return |
| 296 | |
| 297 | import deepspeed |
| 298 | |
| 299 | with plugin.zero3_init_context_manager(enable=False): |
| 300 | deepspeed.zero.partition_parameters.shutdown_init_context() |
| 301 | try: |
| 302 | yield |
| 303 | finally: |
| 304 | deepspeed.zero.partition_parameters.restore_init_context() |
| 305 | |
| 306 | |
| 307 | def model_init_context(accelerator: Accelerator): |