Main Checkpointing loop. Saves checkpoints. Args: config: state: ckpt_path: Returns:
(config, state=None)
| 45 | |
| 46 | |
| 47 | def checkpoint_loop(config, state=None): |
| 48 | """Main Checkpointing loop. |
| 49 | Saves checkpoints. |
| 50 | Args: |
| 51 | config: |
| 52 | state: |
| 53 | ckpt_path: |
| 54 | Returns: |
| 55 | """ |
| 56 | model = mt.from_config(config) |
| 57 | mesh = model.mesh |
| 58 | init_rng, checkpoint_manager, _, tx = train_utils.create_training_tools(config, model, mesh) |
| 59 | |
| 60 | unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(model, tx, config, init_rng, mesh, is_training=True) |
| 61 | # A barrier to sync all hosts before starting to restore checkpoint |
| 62 | jax.experimental.multihost_utils.sync_global_devices("Barrier before load") |
| 63 | checkpoint_load_start = datetime.datetime.now() |
| 64 | with nn_partitioning.axis_rules(config.logical_axis_rules): |
| 65 | state, _ = checkpointing.load_state_if_possible( |
| 66 | checkpoint_manager, |
| 67 | None, |
| 68 | config.load_parameters_path, |
| 69 | config.load_full_state_path, |
| 70 | config.checkpoint_storage_concurrent_gb, |
| 71 | unboxed_abstract_state, |
| 72 | use_ocdbt=config.checkpoint_storage_use_ocdbt, |
| 73 | use_zarr3=config.checkpoint_storage_use_zarr3, |
| 74 | ) |
| 75 | if state: |
| 76 | state = state["items"] |
| 77 | |
| 78 | jax.block_until_ready(state) |
| 79 | checkpoint_load_end = datetime.datetime.now() |
| 80 | if state is not None: # Checkpoint was available for restore |
| 81 | if jax.process_index() == 0: |
| 82 | max_logging.log(f"STANDALONE CHECKPOINTER : Checkpoint restored in : {checkpoint_load_end - checkpoint_load_start}") |
| 83 | else: # Checkpoint was unavailable, state needs to be initialized |
| 84 | state, _, _, _ = maxtext_utils.setup_training_state(model, None, tx, config, init_rng, mesh, checkpoint_manager) |
| 85 | state = add_entropy_to_checkpoint(state) |
| 86 | |
| 87 | start_step = get_first_step(state) # this is the start_step for training |
| 88 | for step in np.arange(start_step, config.steps): |
| 89 | if checkpoint_manager is not None: |
| 90 | start_time = datetime.datetime.now() |
| 91 | # A barrier to sync all hosts before starting to save checkpoint |
| 92 | jax.experimental.multihost_utils.sync_global_devices("Barrier before save") |
| 93 | if checkpointing.save_checkpoint(checkpoint_manager, int(step), state): |
| 94 | checkpoint_manager.wait_until_finished() |
| 95 | end_time = datetime.datetime.now() |
| 96 | if jax.process_index() == 0: |
| 97 | max_logging.log( |
| 98 | f"STANDALONE CHECKPOINTER : Checkpoint saved in {end_time - start_time} ,step {step}, on host 0" |
| 99 | ) |
| 100 | |
| 101 | return state |
| 102 | |
| 103 | |
| 104 | def add_entropy_to_checkpoint(state): |
no test coverage detected