Main Training loop.
(config, recorder, state=None)
| 387 | |
| 388 | |
| 389 | def train_loop(config, recorder, state=None): |
| 390 | """Main Training loop.""" |
| 391 | ( |
| 392 | init_rng, |
| 393 | checkpoint_manager, |
| 394 | state_mesh_shardings, |
| 395 | model, |
| 396 | mesh, |
| 397 | learning_rate_schedule, |
| 398 | data_iterator, |
| 399 | data_loader, |
| 400 | rampup_manager, |
| 401 | eval_data_iterator, |
| 402 | state, |
| 403 | ) = train_utils.setup_train_loop(config, recorder) |
| 404 | |
| 405 | if config.use_dpo: |
| 406 | if "reference_params" not in state.params: |
| 407 | reference_params = jax.tree.map(jnp.copy, state.params["params"]) |
| 408 | state = _merge_dpo_state(state, reference_params) |
| 409 | state_mesh_shardings = _merge_dpo_state(state_mesh_shardings, state_mesh_shardings.params["params"]) |
| 410 | |
| 411 | params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) |
| 412 | |
| 413 | p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( |
| 414 | config, |
| 415 | model, |
| 416 | mesh, |
| 417 | state, |
| 418 | state_mesh_shardings, |
| 419 | train_step, |
| 420 | eval_step, |
| 421 | eval_data_iterator, |
| 422 | params_shardings, |
| 423 | ) |
| 424 | |
| 425 | with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): |
| 426 | shaped_batch = maxtext_utils.get_shaped_batch(config) |
| 427 | if config.shard_optimizer_over_data: |
| 428 | state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode) |
| 429 | if config.compiled_trainstep_file == "": # compile only when there is no pre-compiled file loaded |
| 430 | compiled = p_train_step.lower(state, shaped_batch, init_rng).compile() |
| 431 | compiled_stats = compiled.memory_analysis() |
| 432 | max_utils.print_compiled_memory_stats(compiled_stats) |
| 433 | |
| 434 | start_step = get_first_step(state) # this is the start_step for training |
| 435 | prof = profiler.Profiler(config, offset_step=start_step) |
| 436 | metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) |
| 437 | |
| 438 | # Write train config params, num model params, and XLA flags to tensorboard |
| 439 | metric_logger.write_setup_info_to_tensorboard(state.params) |
| 440 | |
| 441 | try: |
| 442 | last_step_completion = datetime.datetime.now() |
| 443 | for step in np.arange(start_step, config.steps): |
| 444 | prof.maybe_activate_profiler(step, state) |
| 445 | |
| 446 | with jax.profiler.StepTraceAnnotation("train", step_num=step): |
no test coverage detected