| 88 | tx = _optimizer.create_optimizer(config.optimizer, config.lr_schedule, weight_decay_mask=None) |
| 89 | |
| 90 | def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> training_utils.TrainState: |
| 91 | rng, model_rng = jax.random.split(rng) |
| 92 | # initialize the model (and its parameters). |
| 93 | model = config.model.create(model_rng) |
| 94 | |
| 95 | # Merge the partial params into the model. |
| 96 | if partial_params is not None: |
| 97 | graphdef, state = nnx.split(model) |
| 98 | # This will produce an error if the partial params are not a subset of the state. |
| 99 | state.replace_by_pure_dict(partial_params) |
| 100 | model = nnx.merge(graphdef, state) |
| 101 | |
| 102 | params = nnx.state(model) |
| 103 | # Convert frozen params to bfloat16. |
| 104 | params = nnx_utils.state_map(params, config.freeze_filter, lambda p: p.replace(p.value.astype(jnp.bfloat16))) |
| 105 | |
| 106 | return training_utils.TrainState( |
| 107 | step=0, |
| 108 | params=params, |
| 109 | model_def=nnx.graphdef(model), |
| 110 | tx=tx, |
| 111 | opt_state=tx.init(params.filter(config.trainable_filter)), |
| 112 | ema_decay=config.ema_decay, |
| 113 | ema_params=None if config.ema_decay is None else params, |
| 114 | ) |
| 115 | |
| 116 | train_state_shape = jax.eval_shape(init, init_rng) |
| 117 | state_sharding = sharding.fsdp_sharding(train_state_shape, mesh, log=True) |