Function elastic_handler

src/MaxText/elastic_train.py:91–169

Reconfigures the workload onto the currently available slices. This is called by the elastic manager's maybe_reshard_up/down functions and is responsible for creating a new mesh, reinitializing the state and any objects that depend on the mesh. It returns all of the reinitialized objects.

(
    config: pyconfig.HyperParameters,
    elastic_manager,
    checkpoint_manager,
    recorder,
)

@timing.timeit
def elastic_handler(
    config: pyconfig.HyperParameters,
    elastic_manager,
    checkpoint_manager,
    recorder,
):
  """Reconfigures the workload onto the currently available slices.

  This is called by the elastic manager's maybe_reshard_up/down
  functions and is responsible for creating a new mesh,
  reinitializing the state and any objects that depend on the mesh.

  It returns all of the reinitialized objects.

  maybe_reshard_up/down take this function and its arguments and if
  there is an elastic event, those functions will call this function
  and return its returns.
  """
  # We use train_utils.create_training_tools because it contains most of the
  # reconfiguration. Depending on the configuration, the checkpoint
  # manager depends on the mesh and must be recreated. Therefore, we
  # close the previous checkpoint manager and get a new checkpoint
  # manager from create_training_tools.
  if checkpoint_manager is not None:
    checkpoint_manager.close()

  with jax.default_device(elastic_manager.default_device):
    (
        init_rng,
        checkpoint_manager,
        state_mesh_shardings,
        model,
        mesh,
        learning_rate_schedule,
        data_iterator,
        _,
        _,
        _,
        state,
    ) = setup_train_loop(config, recorder, elastic_manager.good_devices)

    p_train_step, _ = train_utils.jit_train_and_eval_step(config, model, mesh, state, state_mesh_shardings, train_step)

    step, snapshot_jax_arrays, _ = elastic_manager.get_resharded_snapshot(mesh)
    state = state.replace(**snapshot_jax_arrays)
    state = state.replace(step=state.step.at[None].set(step))
    jax.block_until_ready(state)

    # We do not want to restore from the previous checkpoint but instead
    # restore from the host offloaded snapshot.
    if checkpoint_manager is not None:
      latest_step = checkpoint_manager.latest_step()

      # If we checkpointed after the latest snapshot, the checkpoint manager
      # will try to take another checkpoint and fail because it already
      # exists. Therefore, we delete the checkpoint and let the checkpoint
      # manager re-take the checkpoint.
      if latest_step is not None and latest_step >= step: