Function get_opt_state_spec_and_shape

tools/train/train.py:946–1000

()

    # get PartitionSpec for optimizer state
    def get_opt_state_spec_and_shape():
        # get opt_state shape without actual init
        opt_state_shape = {}
        for k, p in split_params(params_shape).items():
            if "scanned" not in k:
                opt_state_shape[k] = jax.eval_shape(optimizer[k].init, p)
            else:
                opt_state_shape[k] = jax.eval_shape(jax.vmap(optimizer[k].init), p)

        if training_args.optim == "adafactor":
            # factorized state must be replicated (rank different than params)
            opt_state_spec = {k: None for k in split_params(params_shape)}

        elif training_args.optim in ["adam", "distributed_shampoo"]:

            def _opt_state_spec_per_leaf(x, spec):
                if isinstance(x, FrozenDict):
                    # variables with same structure as params
                    return spec
                else:
                    # other variables such as count
                    return None

            split_spec = split_params(set_partitions(params_shape, False))
            opt_state_spec = {}
            for k, p in split_params(params_shape).items():
                if "scanned" in k:
                    p = jax.eval_shape(lambda x: jax.tree_map(lambda y: y[0], x), p)
                if training_args.optim == "adam":
                    opt_state_spec[k] = jax.tree_map(
                        _opt_state_spec_per_leaf,
                        opt_state_shape[k],
                        split_spec[k],
                        # return None spec for empty elements
                        is_leaf=lambda x: isinstance(x, (FrozenDict, optax.EmptyState)),
                    )
                elif training_args.optim == "distributed_shampoo":
                    opt_state_spec[k] = opt_fn[k].pspec_fn(
                        p,
                        split_spec[k],
                        statistics_partition_spec,
                    )
                # add dimension for scanned params
                if "scanned" in k:
                    opt_state_spec[k] = jax.tree_map(
                        lambda x: PartitionSpec(*(None,) + x)
                        if x is not None
                        else None,
                        opt_state_spec[k],
                        is_leaf=lambda x: isinstance(x, PartitionSpec),
                    )

        else:
            raise NotImplementedError
        return freeze(opt_state_spec), freeze(opt_state_shape)

    opt_state_spec, opt_state_shape = get_opt_state_spec_and_shape()