| 973 | """Get a shaped abstraction of the state (including optimizer)""" |
| 974 | |
| 975 | def init_kv_cache(model, config): |
| 976 | input_shape = ( |
| 977 | config.micro_batch_size_to_train_on, |
| 978 | config.max_prefill_predict_length, |
| 979 | ) |
| 980 | image_shape = multimodal_utils.get_dummy_image_shape_for_init( |
| 981 | config.model_name, batch_size=config.micro_batch_size_to_train_on |
| 982 | ) |
| 983 | |
| 984 | model_vars = model.init( |
| 985 | {"params": rng, "dropout": rng, "aqt": rng}, |
| 986 | jnp.ones(input_shape), |
| 987 | jnp.ones(input_shape), |
| 988 | encoder_images=jnp.ones(image_shape) if config.use_multimodal else None, |
| 989 | model_mode=MODEL_MODE_PREFILL, |
| 990 | slot=0, |
| 991 | page_state=page_state, |
| 992 | ) |
| 993 | return model_vars["cache"] |
| 994 | |
| 995 | with nn_partitioning.axis_rules(config.logical_axis_rules): |
| 996 | init_kv_cache_partial = functools.partial(init_kv_cache, model, config) |