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

Code before:
    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:
        raise ValueError(
            f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
        )

    jax.config.update("jax_threefry_partitionable", True)  # noqa: FBT003
    jax.config.update("jax_compilation_cache_dir", str(epath.Path("~/.cache/jax").expanduser()))

    rng = jax.random.key(config.seed)
    train_rng, init_rng = jax.random.split(rng)

    # TODO: replace with jax.make_mesh when available
    mesh_shape = (jax.device_count(), 1) if not config.fsdp else (1, jax.device_count())
    mesh = jax.sharding.Mesh(mesh_utils.create_device_mesh(mesh_shape), ("batch", "model"))
    data_parallel_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(("batch", "model")))
    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())

    checkpoint_manager, resuming = _checkpoints.initialize_checkpoint(
        config.checkpoint_dir,
        keep_interval=config.keep_interval,
        overwrite=config.overwrite,
        resume=config.resume,
    )
    init_wandb(config, resuming=resuming, enabled=config.wandb_enabled)

    model = config.create_model()

    data_loader = _data_loader.create_data_loader(
        config, model, sharding=data_parallel_sharding, num_workers=config.num_workers
    )
    data_iter = iter(data_loader)
    batch = next(data_iter)
    logging.info(f"Data loader initialized: {training_utils.to_tree_info(batch)}")

    train_state, train_state_sharding = init_train_state(config, model, init_rng, batch, mesh, resume=resuming)
    jax.block_until_ready(train_state)
    logging.info(f"Initialized train state:\n{training_utils.to_tree_info(train_state.params)}")


Review comment (suggestion):
I am wondering if we should make the ratio between data and model gpus configurable. The default would still be data parallel but people would be able to allocate more GPUs for the model using a command line flag.

A possible way to define it is through a ratio:
```
assert devices % flag == 0
data = devices // flag
model = devices // data
```

when flag==1, we get the default behavior.

Code after:
    mesh = jax.make_mesh(mesh_shape, ("batch", "model"))
    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(("batch", "model")))
    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())

    checkpoint_manager, resuming = _checkpoints.initialize_checkpoint(
        config.checkpoint_dir,
        keep_interval=config.keep_interval,
        overwrite=config.overwrite,
        resume=config.resume,
    )
    init_wandb(config, resuming=resuming, enabled=config.wandb_enabled)

    model = config.create_model()

    data_loader = _data_loader.create_data_loader(
        config,
        model,
        sharding=data_sharding,
        num_workers=config.num_workers,
        shuffle=True,
    )
    data_iter = iter(data_loader)
    batch = next(data_iter)
    logging.info(f"Data loader initialized: {training_utils.to_tree_info(batch)}")

    train_state, train_state_sharding = init_train_state(
        config, model, init_rng, batch, mesh, data_sharding, resume=resuming
    )
    jax.block_until_ready(train_state)
    logging.info(f"Initialized train state:\n{training_utils.to_tree_info(train_state.params)}")

    if resuming:
        train_state = _checkpoints.restore_state(checkpoint_manager, train_state, data_loader)

    ptrain_step = jax.jit(
        train_step,
        in_shardings=(replicated_sharding, train_state_sharding, None, data_sharding),
        out_shardings=(train_state_sharding, replicated_sharding),
        donate_argnums=(1,),
    )

    start_step = int(train_state.step)
    pbar = tqdm.tqdm(
        range(start_step, config.num_train_steps),
        initial=start_step,
        total=config.num_train_steps,
        dynamic_ncols=True,
    )

    infos = []
    for step in pbar: