Function load_model

axlearn/cloud/gcp/examples/colocated_python_benchmark.py:200–227

Load model from checkpoint. Args: tensorstore_specs: TensorStore specifications for each array. shardings: Target shardings for the restored arrays. global_shapes: Global shapes for each array. dtypes: Data types for each array. Returns: List of rest

(
    tensorstore_specs: Sequence[Dict[str, Any]],
    shardings: Sequence[jax.sharding.NamedSharding],
    global_shapes: Sequence[tuple],
    dtypes: Sequence[jnp.dtype],
)

def load_model(
    tensorstore_specs: Sequence[Dict[str, Any]],
    shardings: Sequence[jax.sharding.NamedSharding],
    global_shapes: Sequence[tuple],
    dtypes: Sequence[jnp.dtype],
):
    """Load model from checkpoint.

    Args:
        tensorstore_specs: TensorStore specifications for each array.
        shardings: Target shardings for the restored arrays.
        global_shapes: Global shapes for each array.
        dtypes: Data types for each array.

    Returns:
        List of restored JAX arrays.
    """
    manager = GlobalAsyncCheckpointManager()
    restored_values = manager.deserialize(
        shardings=shardings,
        tensorstore_specs=tensorstore_specs,
        global_shapes=global_shapes,
        dtypes=dtypes,
        concurrent_gb=400,
    )
    print(f"Loaded {len(restored_values)} arrays")

    return restored_values


def main():