Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds test cases for t5x finetuning+eval with and without TE #359

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Adds test cases for t5x finetuning+eval with and without TE #359

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
68bf6ec
Updates test-t5x.sh to support finetuning+eval code paths
terrykong Nov 6, 2023
a845ddb
T5x: two new test cases for finetuning w/ and w/o TE
terrykong Nov 6, 2023
0187e4b
typo
terrykong Nov 6, 2023
4c4f42b
Merge branch 'main' into t5x-finetuning
yhtang Dec 4, 2023
ce3b34c
Merge branch 'main' into t5x-finetuning
terrykong Dec 12, 2023
e5ca9b2
casing of squad
terrykong Dec 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 96 additions & 14 deletions .github/container/test-t5x.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ print_var() {
}

usage() {
echo "Test T5X throughput on a fake-data Wikipedia benchmark."
echo "Test T5X throughput on a fake-data Wikipedia benchmark or real SQuAD."
echo ""
echo "Usage: $0 [OPTIONS]"
echo ""
Expand All @@ -18,6 +18,7 @@ usage() {
echo " -d, --dtype Data type, defaults to bfloat16."
echo " --enable-te {0,1} 1 to enable, 0 to disable; defaults to ENABLE_TE in env or 0 if unset"
echo " -e, --epochs Number of epochs to run, defaults to 7."
echo " -f, --finetune Uses squad1 finetuning configs (w/ dropout). By default will only use pre-training configs (no dropout)."
echo " --multiprocess Enable the multiprocess GPU mode."
echo " -o, --output NAME Name for the output folder, a temporary folder will be created if none specified."
echo " --seed INT Random seed for deterministim. Defaults to 42."
Expand All @@ -26,7 +27,7 @@ usage() {
exit $1
}

args=$(getopt -o a:b:cd:e:ho:s: --long additional-args:,batch-size:,use-contrib-configs,dtype:,enable-te:,epochs:,help,multiprocess,output:,seed:,steps-per-epoch: -- "$@")
args=$(getopt -o a:b:cd:e:fho:s: --long additional-args:,batch-size:,use-contrib-configs,dtype:,enable-te:,epochs:,finetune,help,multiprocess,output:,seed:,steps-per-epoch: -- "$@")
if [[ $? -ne 0 ]]; then
exit 1
fi
Expand All @@ -38,6 +39,7 @@ BATCH_SIZE=0
USE_CONTRIB_CONFIGS=0
DTYPE=bfloat16
EPOCHS=7
FINETUNE=0
MULTIPROCESS=0
OUTPUT=$(mktemp -d)
SEED=42
Expand Down Expand Up @@ -71,6 +73,10 @@ while [ : ]; do
EPOCHS="$2"
shift 2
;;
-f | --finetune)
FINETUNE=1
shift 1
;;
-h | --help)
usage 1
;;
Expand Down Expand Up @@ -115,6 +121,7 @@ print_var USE_CONTRIB_CONFIGS
print_var DTYPE
print_var ENABLE_TE
print_var EPOCHS
print_var FINETUNE
print_var OUTPUT
print_var MULTIPROCESS
print_var STEPS_PER_EPOCH
Expand Down Expand Up @@ -158,41 +165,115 @@ seqio.TaskRegistry.add(
EOF

## Create GIN file
cat > benchmark.gin <<EOF
cat > pretrain.gin <<EOF
from __gin__ import dynamic_registration
from t5x import partitioning

# Register Dummy Wikipedia Seqio Task for benchmarking
import dummy_wikipedia
$(
if [[ "$USE_CONTRIB_CONFIGS" -eq 0 ]]; then
if [[ $USE_CONTRIB_CONFIGS -eq 1 ]]; then
echo "from t5x.contrib.gpu.t5 import network"
echo "include 't5x/contrib/gpu/t5/t5_1_1/small.gin'"
echo "include 't5x/contrib/gpu/t5/configs/runs/pretrain.gin'"
echo "include 't5x/contrib/gpu/t5/t5_1_1/adamw_opt.gin'"
else
echo "from t5x.examples.t5 import network"
echo "include 't5x/examples/t5/t5_1_1/small.gin'"
echo "include 't5x/configs/runs/pretrain.gin'"
else
fi
)
MIXTURE_OR_TASK_NAME = "dummy_wikipedia"
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 114}

DROPOUT_RATE = 0.0
USE_CACHED_TASKS = False
TRAIN_STEPS = %gin.REQUIRED
BATCH_SIZE = %gin.REQUIRED

partitioning.PjitPartitioner:
num_partitions = 1
EOF

cat > finetune.gin <<EOF
from __gin__ import dynamic_registration
# Need t5.data.mixtures for squad
from t5.data import mixtures
from t5x import partitioning
from t5x import utils

# Based on t5x/contrib/gpu/t5/t5_1_1/examples/large_squad1_finetune_adam.gin
import t5x.contrib.gpu.scripts_gpu.seqio_tasks
$(
if [[ $USE_CONTRIB_CONFIGS -eq 1 ]]; then
echo "from t5x.contrib.gpu.t5 import network"
echo "include 't5x/contrib/gpu/t5/t5_1_1/small.gin'"
echo "include 't5x/contrib/gpu/t5/configs/runs/pretrain.gin'"
echo "include 't5x/contrib/gpu/t5/t5_1_1/adamw_opt.gin'"
fi
else
echo "from t5x.examples.t5 import network"
echo "include 't5x/examples/t5/t5_1_1/small.gin'"
fi
)
# Using contrib config since it helps avoid repeating several gin overrides (e.g., DatasetConfig.split)
include 't5x/contrib/gpu/t5/configs/runs/finetune_squad1.gin'

# Register Dummy Wikipedia Seqio Task for benchmarking
import dummy_wikipedia
MIXTURE_OR_TASK_NAME = "squad_v010_allanswers"
# Default is {"inputs": 956, "targets": 256}
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 256}
INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/t5_1_1_large/checkpoint_1000000"
# Setting all splits to "validation" to overfit
train/utils.DatasetConfig.split = 'validation'
train_eval/utils.DatasetConfig.split = 'validation'
infer_eval/utils.DatasetConfig.split = 'validation'

MIXTURE_OR_TASK_NAME = "dummy_wikipedia"
TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 114}
DROPOUT_RATE = 0.0
DROPOUT_RATE = 0.1
USE_CACHED_TASKS = False
TRAIN_STEPS = %gin.REQUIRED
BATCH_SIZE = %gin.REQUIRED

# Default is a linear_warmup (base_learning_rate=0.001)
utils.create_learning_rate_scheduler.base_learning_rate = 0.01
utils.create_learning_rate_scheduler.factors = 'constant'

partitioning.PjitPartitioner:
num_partitions = 1
EOF

## Launch
set -exou pipefail

ENABLE_TE=$ENABLE_TE python -m t5x.train \
--gin_file benchmark.gin \
export ENABLE_TE=$ENABLE_TE
if [[ $FINETUNE -eq 1 ]]; then
# Train for a single step
python -m t5x.train \
--gin_file pretrain.gin \
--gin.MODEL_DIR=\"${OUTPUT}\" \
--gin.network.T5Config.dtype=\"${DTYPE}\" \
--gin.TRAIN_STEPS=0 \
--gin.BATCH_SIZE=${BATCH_SIZE} \
--gin.train.eval_steps=0 \
--gin.train.eval_period=${STEPS_PER_EPOCH} \
--gin.train/utils.DatasetConfig.seed=${SEED} \
--gin.train_eval/utils.DatasetConfig.seed=${SEED} \
--gin.train.random_seed=${SEED} \
$ADDITIONAL_ARGS \
$([[ $MULTIPROCESS != 0 ]] && echo --multiprocess_gpu)
python -m t5x.train \
--gin_file finetune.gin \
--gin.MODEL_DIR=\"${OUTPUT}\" \
--gin.network.T5Config.dtype=\"${DTYPE}\" \
--gin.TRAIN_STEPS=${TRAIN_STEPS} \
--gin.BATCH_SIZE=${BATCH_SIZE} \
--gin.train.eval_period=${STEPS_PER_EPOCH} \
--gin.CheckpointConfig.save=None \
--gin.train/utils.DatasetConfig.seed=${SEED} \
--gin.train_eval/utils.DatasetConfig.seed=${SEED} \
--gin.train.random_seed=${SEED} \
$ADDITIONAL_ARGS \
$([[ $MULTIPROCESS != 0 ]] && echo --multiprocess_gpu)
else
python -m t5x.train \
--gin_file pretrain.gin \
--gin.MODEL_DIR=\"${OUTPUT}\" \
--gin.network.T5Config.dtype=\"${DTYPE}\" \
--gin.TRAIN_STEPS=${TRAIN_STEPS} \
Expand All @@ -205,4 +286,5 @@ ENABLE_TE=$ENABLE_TE python -m t5x.train \
--gin.train.random_seed=${SEED} \
$ADDITIONAL_ARGS \
$([[ $MULTIPROCESS != 0 ]] && echo --multiprocess_gpu)
fi
echo "Output at ${OUTPUT}"
10 changes: 10 additions & 0 deletions .github/workflows/_test_t5x_rosetta.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,16 @@ jobs:
N_NODE: 2
ADDITIONAL_ARGS: "--enable-te 0"
EXTRA_GIN_ARGS: ""
- TEST_NAME: "1N2G-ft_te-1"
N_GPU: 2
N_NODE: 1
ADDITIONAL_ARGS: "--finetune"
EXTRA_GIN_ARGS: ""
- TEST_NAME: "1N2G-ft_te-0"
N_GPU: 2
N_NODE: 1
ADDITIONAL_ARGS: "--enable-te 0 --finetune"
EXTRA_GIN_ARGS: ""
fail-fast: false

runs-on: ubuntu-22.04
Expand Down
Loading