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