gm.ckpts.load_params

gm.ckpts.load_params#

gemma.gm.ckpts.load_params(
path: str | os.PathLike,
*,
params: collections.abc.Mapping[str, Any] | None = None,
donate: bool = True,
text_only: bool = False,
sharding: kauldron.ktyping.pytree.PyTree[None | Sharding | Callable[list, str]] | None = None,
quantize: bool = False,
) → collections.abc.Mapping[str, Any][source]

Restore the params from a checkpoint.

Parameters:

  • path – The path to the orbax checkpoint.

  • params – The state matching the checkpoint structure. Is used to restore the params with the correct sharding.

  • donate – If True and params is provided, the memory from params will be released.

  • text_only – If True, only the text params are restored, and the multimodal params are ignored.

  • sharding – If provided, the params will be restored with this sharding. This is mutually exclusive with params.

  • quantize – If True, the params will be mapped to enable quantization aware training.

Returns:

The restored params.