Main training loop for SFT.
(config, recorder, state=None)
| 53 | |
| 54 | |
| 55 | def train_loop(config, recorder, state=None): |
| 56 | """Main training loop for SFT.""" |
| 57 | if not config.use_sft: |
| 58 | raise TypeError("Set use_sft to True to run Supervised Fine Tuning.") |
| 59 | |
| 60 | ( |
| 61 | init_rng, |
| 62 | checkpoint_manager, |
| 63 | state_mesh_shardings, |
| 64 | model, |
| 65 | mesh, |
| 66 | learning_rate_schedule, |
| 67 | data_iterator, |
| 68 | _, |
| 69 | _, |
| 70 | eval_data_iterator, |
| 71 | state, |
| 72 | ) = setup_train_loop(config, recorder) |
| 73 | |
| 74 | p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( |
| 75 | config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator |
| 76 | ) |
| 77 | |
| 78 | with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): |
| 79 | shaped_batch = maxtext_utils.get_shaped_batch(config) |
| 80 | compiled = p_train_step.lower(state, shaped_batch, init_rng).compile() |
| 81 | compiled_stats = compiled.memory_analysis() |
| 82 | max_utils.print_compiled_memory_stats(compiled_stats) |
| 83 | |
| 84 | start_step = get_first_step(state) # this is the start_step for training |
| 85 | prof = profiler.Profiler(config, offset_step=start_step) |
| 86 | data_loader = DataLoader(config, mesh, data_iterator, recorder) |
| 87 | metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) |
| 88 | |
| 89 | # Write train config params, num model params, and XLA flags to tensorboard |
| 90 | metric_logger.write_setup_info_to_tensorboard(state.params) |
| 91 | |
| 92 | try: |
| 93 | last_step_completion = datetime.datetime.now() |
| 94 | for step in np.arange(start_step, config.steps): |
| 95 | prof.maybe_activate_profiler(step, state) |
| 96 | |
| 97 | with jax.profiler.StepTraceAnnotation("train", step_num=step): |
| 98 | example_batch = data_loader.load_next_batch() |
| 99 | # pylint: disable=not-callable |
| 100 | nextrng = jax.jit(jax.random.fold_in)(init_rng, step) |
| 101 | with maybe_record_goodput(recorder, GoodputEvent.STEP, step): |
| 102 | with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): |
| 103 | state, metrics = p_train_step(state, example_batch, nextrng) |
| 104 | |
| 105 | step_time_delta = datetime.datetime.now() - last_step_completion |
| 106 | last_step_completion = datetime.datetime.now() |
| 107 | |
| 108 | checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step) |
| 109 | |
| 110 | if config.dump_hlo and step == start_step: |
| 111 | jax.block_until_ready(state) # Ensure compilation has finished. |
| 112 | gcs_utils.upload_dump( |
no test coverage detected