| 135 | |
| 136 | @at.typecheck |
| 137 | def train_step( |
| 138 | config: _config.TrainConfig, |
| 139 | rng: at.KeyArrayLike, |
| 140 | state: training_utils.TrainState, |
| 141 | batch: tuple[_model.Observation, _model.Actions], |
| 142 | ) -> tuple[training_utils.TrainState, dict[str, at.Array]]: |
| 143 | model = nnx.merge(state.model_def, state.params) |
| 144 | model.train() |
| 145 | |
| 146 | @at.typecheck |
| 147 | def loss_fn( |
| 148 | model: _model.BaseModel, rng: at.KeyArrayLike, observation: _model.Observation, actions: _model.Actions |
| 149 | ): |
| 150 | chunked_loss = model.compute_loss(rng, observation, actions, train=True) |
| 151 | return jnp.mean(chunked_loss) |
| 152 | |
| 153 | train_rng = jax.random.fold_in(rng, state.step) |
| 154 | observation, actions = batch |
| 155 | |
| 156 | # Filter out frozen params. |
| 157 | diff_state = nnx.DiffState(0, config.trainable_filter) |
| 158 | loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions) |
| 159 | |
| 160 | params = state.params.filter(config.trainable_filter) |
| 161 | updates, new_opt_state = state.tx.update(grads, state.opt_state, params) |
| 162 | new_params = optax.apply_updates(params, updates) |
| 163 | |
| 164 | # Update the model in place and return the new full state. |
| 165 | nnx.update(model, new_params) |
| 166 | new_params = nnx.state(model) |
| 167 | |
| 168 | new_state = dataclasses.replace(state, step=state.step + 1, params=new_params, opt_state=new_opt_state) |
| 169 | if state.ema_decay is not None: |
| 170 | new_state = dataclasses.replace( |
| 171 | new_state, |
| 172 | ema_params=jax.tree.map( |
| 173 | lambda old, new: state.ema_decay * old + (1 - state.ema_decay) * new, state.ema_params, new_params |
| 174 | ), |
| 175 | ) |
| 176 | |
| 177 | # Filter out params that aren't kernels. |
| 178 | kernel_params = nnx.state( |
| 179 | model, |
| 180 | nnx.All( |
| 181 | nnx.Param, |
| 182 | nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")), |
| 183 | lambda _, x: x.value.ndim > 1, |
| 184 | ), |
| 185 | ) |
| 186 | info = { |
| 187 | "loss": loss, |
| 188 | "grad_norm": optax.global_norm(grads), |
| 189 | "param_norm": optax.global_norm(kernel_params), |
| 190 | } |
| 191 | return new_state, info |
| 192 | |
| 193 | |
| 194 | def main(config: _config.TrainConfig): |