def train(experiment_class, config, checkpointer, writer, periodic_actions=()):
"""Main training loop."""
logging.info("Training with config:\n%s", config)
is_chief = jax.host_id() == 0
rng = jax.random.PRNGKey(config.random_seed)
with utils.log_activity("experiment init"):
experiment = _initialize_experiment(experiment_class, "train", rng,
config.experiment_kwargs)
state = checkpointer.get_experiment_state("latest")
state.global_step = 0
state.experiment_module = experiment
state.train_step_rng = jnp.broadcast_to(rng, (jax.local_device_count(), ) +
rng.shape)
if checkpointer.can_be_restored("latest"):
with utils.log_activity("checkpoint restore"):
checkpointer.restore("latest")
periodic_actions += (utils.PeriodicAction(
_log_outputs,
interval_type=config.interval_type,
interval=config.log_tensors_interval), )
if config.train_checkpoint_all_hosts or is_chief:
if config.save_checkpoint_interval > 0:
periodic_actions += (utils.PeriodicAction(
lambda *_: checkpointer.save("latest"),
interval_type=config.interval_type,
interval=config.save_checkpoint_interval,
run_async=False), ) # run_async True would not be thread-safe.
if is_chief:
if writer is not None:
def write_scalars(global_step: int, scalar_values):
writer.write_scalars(global_step, scalar_values)
periodic_actions += (utils.PeriodicAction(
write_scalars,
interval_type=config.interval_type,
interval=config.log_train_data_interval,
log_all_data=config.log_all_train_data), )
for pa in periodic_actions:
pa.update_time(time.time(), state.global_step)
experiment.train_loop(config, state, periodic_actions, writer)
if is_chief:
with utils.log_activity("final checkpoint"):
checkpointer.save("latest")
# We occasionally see errors when the final checkpoint is being written if
# the other hosts exit. Here we force all hosts to participate in one final
# collective so the non-master hosts cannot exit before the master writes out
# the final checkpoint.
utils.rendezvous()
def test_log_success(self, mock_info):
"""Tests that logging an activity is successful."""
with utils.log_activity("for test"):
pass
mock_info.assert_any_call("[jaxline] %s starting...", "for test")
mock_info.assert_any_call("[jaxline] %s finished.", "for test")
def train_loop(
self,
config: config_dict.ConfigDict,
state,
periodic_actions: List[utils.PeriodicAction],
writer: Optional[utils.Writer] = None,
) -> None:
"""Default training loop implementation.
Can be overridden for advanced use cases that need a different training loop
logic, e.g. on device training loop with jax.lax.while_loop or to add custom
periodic actions.
Args:
config: The config of the experiment that is being run.
state: Checkpointed state of the experiment.
periodic_actions: List of actions that should be called after every
training step, for checkpointing and logging.
writer: An optional writer to pass to the experiment step function.
"""
@functools.partial(jax.pmap, axis_name="i")
def next_device_state(
global_step: jnp.ndarray,
rng: jnp.ndarray,
host_id: Optional[jnp.ndarray],
):
"""Updates device global step and rng in one pmap fn to reduce overhead."""
global_step += 1
step_rng, state_rng = tuple(jax.random.split(rng))
step_rng = utils.specialize_rng_host_device(
step_rng, host_id, axis_name="i", mode=config.random_mode_train)
return global_step, (step_rng, state_rng)
global_step_devices = np.broadcast_to(state.global_step,
[jax.local_device_count()])
host_id_devices = utils.host_id_devices_for_rng(config.random_mode_train)
step_key = state.train_step_rng
with utils.log_activity("training loop"):
while self.should_run_step(state.global_step, config):
with jax.profiler.StepTraceAnnotation(
"train", step_num=state.global_step):
scalar_outputs = self.step(
global_step=global_step_devices, rng=step_key, writer=writer)
t = time.time()
# Update state's (scalar) global step (for checkpointing).
# global_step_devices will be back in sync with this after the call
# to next_device_state below.
state.global_step += 1
global_step_devices, (step_key, state.train_step_rng) = (
next_device_state(global_step_devices,
state.train_step_rng,
host_id_devices))
for action in periodic_actions:
action(t, state.global_step, scalar_outputs)
def test_log_failure(self, mock_info, mock_exc):
"""Tests that an error thrown by an activity is correctly caught."""
with self.assertRaisesRegex(ValueError, "Intentional"):
with utils.log_activity("for test"):
raise ValueError("Intentional")
mock_info.assert_any_call("[jaxline] %s starting...", "for test")
mock_exc.assert_any_call("[jaxline] %s failed with error.", "for test")
def evaluate(
experiment_class,
config,
checkpointer: utils.Checkpointer,
writer: Optional[utils.Writer],
jaxline_mode: Optional[str] = None,
):
"""Main evaluation loop."""
if jaxline_mode is None:
jaxline_mode = FLAGS.jaxline_mode
logging.info("Evaluating with config:\n%s", config)
global_step = 0
eval_rng = jax.random.PRNGKey(config.random_seed)
experiment = _initialize_experiment(experiment_class, jaxline_mode,
eval_rng, config.experiment_kwargs)
if config.best_model_eval_metric and jax.host_id() == 0:
# Initialize best state.
best_state = checkpointer.get_experiment_state("best")
if config.best_model_eval_metric_higher_is_better:
best_state.best_eval_metric_value = float("-inf")
eval_metric_is_better_op = jnp.greater
eval_metric_comparison_str = ">"
else:
best_state.best_eval_metric_value = float("inf")
eval_metric_is_better_op = jnp.less
eval_metric_comparison_str = "<"
best_state.best_model_eval_metric = config.best_model_eval_metric
best_state.experiment_module = experiment
# Restore to preserve 'best_eval_metric_value' if evaluator was preempted.
if checkpointer.can_be_restored("best"):
with utils.log_activity("best checkpoint restore"):
checkpointer.restore("best")
# Will evaluate the latest checkpoint in the directory.
state = checkpointer.get_experiment_state("latest")
state.global_step = global_step
state.experiment_module = experiment
state.train_step_rng = None
eval_rng = jnp.broadcast_to(eval_rng,
(jax.local_device_count(), ) + eval_rng.shape)
host_id_devices = utils.host_id_devices_for_rng(config.random_mode_eval)
eval_rng = jax.pmap(functools.partial(utils.specialize_rng_host_device,
axis_name="i",
mode=config.random_mode_eval),
axis_name="i")(eval_rng, host_id_devices)
if config.one_off_evaluate:
checkpointer.restore("latest")
global_step_devices = utils.bcast_local_devices(
jnp.asarray(state.global_step))
scalar_values = utils.evaluate_should_return_dict(experiment.evaluate)(
global_step=global_step_devices, rng=eval_rng, writer=writer)
if writer is not None:
writer.write_scalars(state.global_step, scalar_values)
logging.info("Evaluated specific checkpoint, exiting.")
return
old_checkpoint_path = None
initial_weights_are_evaluated = False
while True:
checkpoint_path = checkpointer.restore_path("latest")
if (checkpoint_path is None and config.eval_initial_weights
and not initial_weights_are_evaluated):
# Skip restoring a checkpoint and directly call evaluate if
# `config.eval_initial_weights` but don"t do it more than once.
initial_weights_are_evaluated = True
else:
if checkpoint_path in (None, old_checkpoint_path):
logging.info(
"Checkpoint %s invalid or already evaluated, waiting.",
checkpoint_path)
time.sleep(10)
continue
checkpointer.restore("latest")
global_step_devices = utils.bcast_local_devices(
jnp.asarray(state.global_step))
scalar_values = utils.evaluate_should_return_dict(experiment.evaluate)(
global_step=global_step_devices, rng=eval_rng, writer=writer)
if writer is not None:
writer.write_scalars(state.global_step, scalar_values)
old_checkpoint_path = checkpoint_path
# Decide whether to save a "best checkpoint".
if config.best_model_eval_metric and jax.host_id() == 0:
if config.best_model_eval_metric not in scalar_values:
raise ValueError(
f"config.best_model_eval_metric has been specified "
f"as {config.best_model_eval_metric}, but this key "
f"was not returned by the evaluate method. Got: "
f"{scalar_values.keys()}")
current_eval_metric_value = scalar_values[
config.best_model_eval_metric]
old_eval_metric_value = best_state.best_eval_metric_value
if eval_metric_is_better_op(current_eval_metric_value,
old_eval_metric_value):
logging.info("%s: %s %s %s, saving new best checkpoint.",
config.best_model_eval_metric,
current_eval_metric_value,
eval_metric_comparison_str, old_eval_metric_value)
best_state.global_step = state.global_step
best_state.experiment_module = experiment
best_state.best_eval_metric_value = current_eval_metric_value
best_state.train_step_rng = state.train_step_rng
checkpointer.save("best")
if not experiment.should_run_step(state.global_step, config):
logging.info("Last checkpoint (iteration %d) evaluated, exiting.",
state.global_step)
break