Get shaped abstractions of inputs to train_step: state, batch and rng
(topology_mesh, config)
| 86 | |
| 87 | |
| 88 | def get_shaped_inputs(topology_mesh, config): |
| 89 | """Get shaped abstractions of inputs to train_step: state, batch and rng""" |
| 90 | # Construct the model and optimizer to get shaped versions of the state |
| 91 | quant = quantizations.configure_quantization(config) |
| 92 | model = Transformer(config, topology_mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) |
| 93 | # The learning_rate_schedule is baked into the compiled object. |
| 94 | learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) |
| 95 | # pass in model for muon |
| 96 | tx = optimizers.get_optimizer(config, learning_rate_schedule, model) |
| 97 | |
| 98 | # Shaped RNG keys |
| 99 | _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2) |
| 100 | shaped_rng = jax.ShapeDtypeStruct(example_rng.shape, example_rng.dtype) |
| 101 | |
| 102 | # Shaped state |
| 103 | abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state( |
| 104 | model, tx, config, example_rng, topology_mesh |
| 105 | ) |
| 106 | |
| 107 | # Shaped batch |
| 108 | shaped_batch = maxtext_utils.get_shaped_batch(config) |
| 109 | |
| 110 | shaped_train_args = (abstract_state, shaped_batch, shaped_rng) |
| 111 | shaped_train_kwargs = {} |
| 112 | return shaped_train_args, shaped_train_kwargs, state_mesh_shardings, model |
| 113 | |
| 114 | |
| 115 | def jit_and_compile( |