Set up prerequisites for the training loop - checkpoint_manager, PRNG keys, Mesh, Model and optimizer. Set up data iterator and tokenizer, initialize the model. Args: config recorder Returns: init_rng: checkpoint_manager: Orbax checkpointer state_mesh_annotations: the
(config, recorder, devices=None)
| 158 | |
| 159 | |
| 160 | def setup_train_loop(config, recorder, devices=None): |
| 161 | """Set up prerequisites for the training loop - |
| 162 | |
| 163 | checkpoint_manager, PRNG keys, Mesh, Model and optimizer. |
| 164 | Set up data iterator and tokenizer, initialize the model. |
| 165 | |
| 166 | Args: config recorder |
| 167 | |
| 168 | Returns: |
| 169 | init_rng: |
| 170 | checkpoint_manager: Orbax checkpointer |
| 171 | state_mesh_annotations: the mesh annotations for the train state |
| 172 | model: |
| 173 | mesh: |
| 174 | learning_rate_schedule: |
| 175 | data_iterator: |
| 176 | data_loader: |
| 177 | rampup_manager: the class managing rampup batch sizes |
| 178 | state: the initialized train state |
| 179 | """ |
| 180 | |
| 181 | with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT): |
| 182 | model = model_creation_utils.from_config(config, devices) |
| 183 | mesh = model.mesh |
| 184 | init_rng, checkpoint_manager, learning_rate_schedule, tx = create_training_tools(config, model, mesh) |
| 185 | |
| 186 | with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION): |
| 187 | data_iterator, eval_data_iterator = create_data_iterator(config, mesh) |
| 188 | rampup_manager = create_rampup_manager(config, checkpoint_manager) |
| 189 | data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager) |
| 190 | context_parallel_size = mesh.shape["context"] |
| 191 | # Check if context parallelism is being used with sequence packing |
| 192 | if context_parallel_size > 1 and config.packing and config.dataset_type != "synthetic": |
| 193 | raise ValueError( |
| 194 | "Context parallelism cannot be used with sequence packing. " |
| 195 | "Disable sequence packing (set packing=False). " |
| 196 | "Context parallelism with packing support will be added soon." |
| 197 | ) |
| 198 | |
| 199 | # Apply reordering wrapper to data iterators if context parallelism is enabled |
| 200 | with jax.set_mesh(mesh): |
| 201 | if context_parallel_size > 1 and config.context_parallel_load_balance: |
| 202 | data_iterator = map(maxtext_utils.get_reorder_callable(context_parallel_size, config.shard_mode), data_iterator) |
| 203 | if eval_data_iterator: |
| 204 | eval_data_iterator = map( |
| 205 | maxtext_utils.get_reorder_callable(context_parallel_size, config.shard_mode), |
| 206 | eval_data_iterator, |
| 207 | ) |
| 208 | |
| 209 | state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state( |
| 210 | model, data_iterator, tx, config, init_rng, mesh, checkpoint_manager |
| 211 | ) |
| 212 | |
| 213 | # TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal |
| 214 | if not config.using_pipeline_parallelism and not config.use_multimodal: |
| 215 | # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage |
| 216 | sharding.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance) |
| 217 |