Spaces:
Build error
Build error
File size: 1,059 Bytes
b100e1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# Evaluate a Mixture of Experts model.
#
#
# You must also include a binding for MODEL.
#
# Required to be set:
#
# - NUM_EXPERTS
# - NUM_MODEL_PARTITIONS (1 if no model parallelism)
# - MIXTURE_OR_TASK_NAME
# - CHECKPOINT_PATH
# - EVAL_OUTPUT_DIR
#
# Commonly overridden options (see also t5x/configs/runs/eval.gin):
#
# - DROPOUT_RATE
# - BATCH_SIZE
from __gin__ import dynamic_registration
import __main__ as eval_script
from t5x.contrib.moe import partitioning as moe_partitioning
from t5x import utils
include 't5x/configs/runs/eval.gin'
NUM_EXPERTS = %gin.REQUIRED
NUM_MODEL_PARTITIONS = %gin.REQUIRED
# We use the MoE partitioner.
eval_script.evaluate.partitioner = @moe_partitioning.MoePjitPartitioner()
moe_partitioning.MoePjitPartitioner:
num_experts = %NUM_EXPERTS
num_partitions = %NUM_MODEL_PARTITIONS
logical_axis_rules = @moe_partitioning.standard_logical_axis_rules()
moe_partitioning.standard_logical_axis_rules:
num_experts = %NUM_EXPERTS
num_partitions = %NUM_MODEL_PARTITIONS
utils.DatasetConfig.batch_size = %BATCH_SIZE
|