Main Training loop.
(config, elastic_manager, recorder, state=None)
| 170 | |
| 171 | |
| 172 | def train_loop(config, elastic_manager, recorder, state=None): |
| 173 | """Main Training loop.""" |
| 174 | ( |
| 175 | init_rng, |
| 176 | checkpoint_manager, |
| 177 | state_mesh_shardings, |
| 178 | model, |
| 179 | mesh, |
| 180 | learning_rate_schedule, |
| 181 | data_iterator, |
| 182 | _, |
| 183 | _, |
| 184 | _, |
| 185 | state, |
| 186 | ) = setup_train_loop(config, recorder) |
| 187 | |
| 188 | p_train_step, _ = train_utils.jit_train_and_eval_step(config, model, mesh, state, state_mesh_shardings, train_step) |
| 189 | with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): |
| 190 | shaped_batch = maxtext_utils.get_shaped_batch(config) |
| 191 | compiled = p_train_step.lower(state, shaped_batch, init_rng).compile() |
| 192 | compiled_stats = compiled.memory_analysis() |
| 193 | max_utils.print_compiled_memory_stats(compiled_stats) |
| 194 | |
| 195 | start_step = get_first_step(state) # this is the start_step for training |
| 196 | prof = profiler.Profiler(config, offset_step=start_step) |
| 197 | |
| 198 | step = start_step |
| 199 | |
| 200 | elastic_manager.maybe_snapshot( |
| 201 | step, |
| 202 | snapshot_jax_arrays={ |
| 203 | "params": state.params, |
| 204 | "opt_state": state.opt_state, |
| 205 | }, |
| 206 | force=True, |
| 207 | block=True, |
| 208 | ) |
| 209 | |
| 210 | data_loader = DataLoader(config, mesh, data_iterator, recorder) |
| 211 | metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) |
| 212 | |
| 213 | # Write train config params, num model params, and XLA flags to tensorboard |
| 214 | metric_logger.write_setup_info_to_tensorboard(state.params) |
| 215 | |
| 216 | last_step_completion = datetime.datetime.now() |
| 217 | |
| 218 | # Using while loop instead of a for loop because with elasticity |
| 219 | # the step is restored back to the latest snapshot when a slice is lost |
| 220 | while step < config.steps: |
| 221 | try: |
| 222 | prof.maybe_activate_profiler(step, state) |
| 223 | |
| 224 | max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}") |
| 225 | with ( |
| 226 | mesh, |
| 227 | nn_partitioning.axis_rules(config.logical_axis_rules), |
| 228 | jax.default_device(elastic_manager.default_device), |
| 229 | ): |
no test coverage detected