Repository: Physical-Intelligence/openpi
File: scripts/train.py (Python)

Code before:
    mesh: jax.sharding.Mesh,
    *,
    resume: bool,
) -> tuple[training_utils.TrainState, Any]:
    weight_decay_mask = None
    freeze_mask = None
    tx = _optimizer.create_optimizer(config.optimizer, config.lr_schedule, weight_decay_mask, freeze_mask)

    def init(rng: at.KeyArrayLike, data: tuple[_common.Observation, _common.Actions]) -> training_utils.TrainState:
        rng, model_rng = jax.random.split(rng)
        observation, actions = data
        params = model.init_params(model_rng, observation, actions)
        params = jax.experimental.io_callback(
            partial(_load_weights_and_validate, config.weight_loader), params, params, ordered=True
        )
        return training_utils.TrainState(
            step=0,
            params=params,
            opt_state=tx.init(params),
            tx=tx,
            ema_decay=config.ema_decay,
            ema_params=None if config.ema_decay is None else params,
        )

    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
    data_parallel_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(("batch", "model")))

    train_state_shape = jax.eval_shape(init, init_rng, batch)
    # This is where we may want to shard the train state (e.g., FSDP).
    state_sharding = sharding(train_state_shape, mesh, log=True)

    if resume:
        return train_state_shape, state_sharding

    train_state = jax.jit(
        init, in_shardings=(replicated_sharding, data_parallel_sharding), out_shardings=state_sharding
    )(init_rng, batch)
    return train_state, state_sharding


@at.typecheck
def train_step(
    rng: at.KeyArrayLike,
    state: training_utils.TrainState,
    model: _model.Model,
    batch: tuple[_common.Observation, _common.Actions],
) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
    def loss_fn(params: at.Params, rng: at.KeyArrayLike, observation: _common.Observation, actions: _common.Actions):
        chunked_loss = model.compute_loss(rng, observation, actions, params=params, train=True)
        return jnp.mean(chunked_loss)


Review comment (refactor):
Also, what do you think about providing the data sharding externally, or at the very least putting the logic inside a helper function. The problem is that we create it in two places now.

Another option not to specify it explicitly at all, since the input data should be already sharded. The data loader takes care of that.

Code after:
        return train_state_shape, state_sharding

    train_state = jax.jit(
        init,
        in_shardings=(replicated_sharding, data_sharding),
        out_shardings=state_sharding,
        static_argnums=(2,),
    )(init_rng, batch, replicated_sharding)
    return train_state, state_sharding


@at.typecheck
def train_step(
    rng: at.KeyArrayLike,
    state: training_utils.TrainState,
    model: _model.Model,
    batch: tuple[_common.Observation, _common.Actions],
) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
    def loss_fn(params: at.Params, rng: at.KeyArrayLike, observation: _common.Observation, actions: _common.Actions):
        chunked_loss = model.compute_loss(rng, observation, actions, params=params, train=True)
        return jnp.mean(chunked_loss)

    train_rng = jax.random.fold_in(rng, state.step)
    observation, actions = batch
    loss, grads = jax.value_and_grad(loss_fn)(state.params, train_rng, observation, actions)
    updates, new_opt_state = state.tx.update(grads, state.opt_state, state.params)
    new_params = optax.apply_updates(state.params, updates)

    new_state = state.replace(step=state.step + 1, params=new_params, opt_state=new_opt_state)
    if state.ema_decay is not None:
        new_state = new_state.replace(
            ema_params=jax.tree.map(
                lambda old, new: state.ema_decay * old + (1 - state.ema_decay) * new, state.ema_params, new_params
            )
        )

    kernel_mask = training_utils.mask_from_regex(r".*\['kernel'\]", state.params)
    kernel_params = jax.tree.map(lambda p, m: p if m else None, state.params, kernel_mask)
    info = {
        "loss": loss,
        "grad_norm": optax.global_norm(grads),  # TODO: do not compute norm for frozen params
        "param_norm": optax.global_norm(kernel_params),
    }
    return new_state, info


def main(config: _config.TrainConfig):
    init_logging()
    logging.info(f"Running on: {platform.node()}")

    if config.batch_size % jax.device_count() != 0: