Creates the init_rng, optimizer, learning rate schedule, and checkpoint manager.
(config, model, mesh)
| 33 | |
| 34 | |
| 35 | def create_training_tools(config, model, mesh): |
| 36 | """Creates the init_rng, optimizer, learning rate schedule, and checkpoint manager.""" |
| 37 | init_rng = jax.random.PRNGKey(config.init_weights_seed) |
| 38 | learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) |
| 39 | # pass in model for muon |
| 40 | tx = optimizers.get_optimizer(config, learning_rate_schedule, model) |
| 41 | logger = checkpointing.setup_checkpoint_logger(config) |
| 42 | if config.enable_multi_tier_checkpointing: |
| 43 | checkpoint_manager = checkpointing.create_orbax_emergency_replicator_checkpoint_manager( |
| 44 | config.local_checkpoint_directory, |
| 45 | config.local_checkpoint_period, |
| 46 | mesh, |
| 47 | ) |
| 48 | elif config.enable_emergency_checkpoint: |
| 49 | abstract_state, _, _ = maxtext_utils.get_abstract_state(model, tx, config, init_rng, mesh, is_training=True) |
| 50 | checkpoint_manager = checkpointing.create_orbax_emergency_checkpoint_manager( |
| 51 | config.local_checkpoint_directory, |
| 52 | config.checkpoint_dir, |
| 53 | mesh, |
| 54 | abstract_state, |
| 55 | config.local_checkpoint_period, |
| 56 | config.checkpoint_period, |
| 57 | logger, |
| 58 | ) |
| 59 | else: |
| 60 | # TODO(b/368121306): Remove this once zarr3 support is plumbed on the backend |
| 61 | use_ocdbt = config.checkpoint_storage_use_ocdbt |
| 62 | use_zarr3 = config.checkpoint_storage_use_zarr3 |
| 63 | if config.enable_single_controller: |
| 64 | use_ocdbt, use_zarr3 = False, False |
| 65 | |
| 66 | checkpoint_dir = "" |
| 67 | if config.enable_checkpointing: |
| 68 | checkpoint_dir = config.checkpoint_dir |
| 69 | checkpoint_manager = checkpointing.create_orbax_checkpoint_manager( |
| 70 | checkpoint_dir, |
| 71 | config.enable_checkpointing, |
| 72 | config.async_checkpointing, |
| 73 | config.checkpoint_period, |
| 74 | config.dataset_type, |
| 75 | logger, |
| 76 | use_ocdbt, |
| 77 | use_zarr3, |
| 78 | config.enable_continuous_checkpointing, |
| 79 | config.max_num_checkpoints_to_keep, |
| 80 | ) |
| 81 | |
| 82 | return init_rng, checkpoint_manager, learning_rate_schedule, tx |
| 83 | |
| 84 | |
| 85 | def jit_train_step(config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings): |