Function train_step

src/MaxText/train.py:222–350

Args: model: A nn.Module state: A pytree of the current state of the model data: Batch of data to apply to the model dropout_rng: A key to use to generate rng for dropout Returns: new_state: Same format as state. metrics: Dictionary of model metrics such as loss, train

(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng)

def train_step(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng):
  """

  Args:
    model: A nn.Module
    state: A pytree of the current state of the model
    data: Batch of data to apply to the model
    dropout_rng: A key to use to generate rng for dropout

  Returns:
    new_state: Same format as state.
    metrics: Dictionary of model metrics such as loss, training rate, etc.
    rng2: A new rng key that can be used in future calls.

  """
  reference_params, reference_params_sharding, extra_dpo_args, _loss_fn = (
      [],
      [],
      [],
      loss_fn,
  )
  if config.use_dpo:
    state, reference_params = _split_dpo_state(state)
    state_mesh_shardings, reference_params_sharding = _split_dpo_state(state_mesh_shardings)
    extra_dpo_args = [reference_params]
    _loss_fn = dpo_loss_fn

  params = state.params

  if config.gradient_accumulation_steps > 1:
    loss, aux, raw_grads = gradient_accumulation_loss_and_grad(
        _loss_fn,
        config,
        model,
        params,
        params_shardings,
        data,
        dropout_rng,
        extra_dpo_args,
    )
  else:
    if config.optimizer_memory_host_offload:
      if config.use_dpo:
        reference_params = jax.device_put(
            reference_params,
            max_utils.with_memory_kind(reference_params_sharding, "device"),
        )
        extra_dpo_args = [reference_params]
    if config.shard_optimizer_over_data:
      params = jax.tree.map(
          functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode),
          params,
          params_shardings,
      )
    grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
    (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, params, *extra_dpo_args, is_train=True)

  raw_grads = jax.tree_util.tree_map(