Initializes the elastic manager and pyconfig to support elastic training Args: devices: The devices used for training Returns: The initialized elastic manager
(devices: Sequence[jax.Device])
| 331 | |
| 332 | |
| 333 | def elastic_initialize(devices: Sequence[jax.Device]) -> manager.Manager: |
| 334 | """Initializes the elastic manager and pyconfig to support elastic training |
| 335 | |
| 336 | Args: |
| 337 | devices: The devices used for training |
| 338 | |
| 339 | Returns: |
| 340 | The initialized elastic manager |
| 341 | """ |
| 342 | elastic_manager = manager.Manager( |
| 343 | devices, |
| 344 | reshard_check_period=1, |
| 345 | snapshot_period=5, |
| 346 | max_elastic_down_event_count=100, |
| 347 | max_reshard_retry_count=3, |
| 348 | ) |
| 349 | |
| 350 | # Do not start training until all slices are available |
| 351 | # TODO: b/408455557 - Migrate to pathwaysutils and make configurable |
| 352 | wait_for_all_slices(elastic_manager) |
| 353 | |
| 354 | pyconfig.HyperParameters.global_batch_size_to_train_on = property( |
| 355 | lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_train_on"]) |
| 356 | ) |
| 357 | pyconfig.HyperParameters.global_batch_size_to_load = property( |
| 358 | lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_load"]) |
| 359 | ) |
| 360 | pyconfig.HyperParameters.micro_batch_size_to_train_on = property( |
| 361 | lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["micro_batch_size_to_train_on"]) |
| 362 | ) |
| 363 | pyconfig.HyperParameters.num_slices = property(lambda self: elastic_manager.good_slice_count) |
| 364 | |
| 365 | return elastic_manager |
| 366 | |
| 367 | |
| 368 | def main(argv: Sequence[str]) -> None: |
no test coverage detected