[content] train_lam.py (python)
"""
Helper functions for building components and using them both from CLI and notebooks
without duplicating logic in the training loop.
"""


def build_model(a: Args) -> LatentActionModel:
    rng = jax.random.key(a.seed)
    _, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    return LatentActionModel(
        in_dim=a.image_channels,
        model_dim=a.model_dim,
        ffn_dim=a.ffn_dim,
        latent_dim=a.latent_dim,
        num_latents=a.num_latents,
        patch_size=a.patch_size,
        num_blocks=a.num_blocks,
        num_heads=a.num_heads,
        dropout=a.dropout,
        codebook_dropout=a.codebook_dropout,
        param_dtype=a.param_dtype,
        dtype=a.dtype,
        use_flash_attention=a.use_flash_attention,
        rngs=rngs,
    )


def build_optimizer(model: LatentActionModel, a: Args):
    lr_schedule = get_lr_schedule(
        a.lr_schedule,
        a.init_lr,
        a.max_lr,
        a.decay_end,
        a.num_steps,
        a.warmup_steps,
        a.wsd_decay_steps,
    )
    tx = optax.adamw(
        learning_rate=lr_schedule,
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
        mu_dtype=a.param_dtype,
    )
    optimizer = nnx.Optimizer(model, tx)
    return optimizer, lr_schedule


def build_mesh_and_sharding(num_devices: int):
    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
    return mesh, replicated_sharding, videos_sharding


def shard_optimizer_states(optimizer: nnx.Optimizer, replicated_sharding: NamedSharding) -> None:
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.lax.with_sharding_constraint(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)


def build_dataloader(a: Args):
    image_shape = (a.image_height, a.image_width, a.image_channels)
    array_record_files = [
        os.path.join(a.data_dir, x)
        for x in os.listdir(a.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        a.seq_len,
        a.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=a.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
    return grain_dataloader, grain_iterator


def build_checkpoint_manager(a: Args):
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add("model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler)
    handler_registry.add("model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler)
    handler_registry.add(
        "dataloader_state",
        grain.checkpoint.CheckpointSave,
        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
    )
    handler_registry.add(
        "dataloader_state",
        grain.checkpoint.CheckpointRestore,
        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
    )
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=a.log_checkpoint_interval,
        max_to_keep=3,
        keep_period=a.log_checkpoint_keep_period,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    checkpoint_manager = ocp.CheckpointManager(
        a.ckpt_dir,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
    return checkpoint_manager


def restore_checkpoint_if_needed(a: Args, checkpoint_manager, optimizer, grain_iterator):
    step = 0
    if a.restore_ckpt:
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),  # type: ignore
            ),
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
    return step, optimizer, grain_iterator


def enable_sowing(lam: LatentActionModel) -> None:
    for model in [lam.encoder, lam.decoder]:
        setattr(model, "sow_logits", True)
        for blk in getattr(model, "blocks", []):
            setattr(blk, "sow_weights", True)
            setattr(blk, "sow_activations", True)