Repository: google-research/text-to-text-transfer-transformer

Issue:
Creating MtfModel with GPU
We are trying to create a task, a model, and fine tune on a GPU:

```
[Some imports]
[Some constants defined here]

def create_task(name, dataset_fn):
  t5.data.TaskRegistry.add(
    name,
    # Supply a function which returns a tf.data.Dataset.
    dataset_fn=dataset_fn,
    splits=["train", "validation"],
    # Supply a function which preprocesses text from the tf.data.Dataset.
    text_preprocessor=None,
    # Use the same vocabulary that we used for pre-training.
    sentencepiece_model_path=t5.data.DEFAULT_SPM_PATH,
    # Lowercase targets before computing metrics.
    postprocess_fn=t5.data.postprocessors.lower_text,
    # We'll use accuracy as our evaluation metric.
    metric_fns=[t5.evaluation.metrics.accuracy]
  )

def create_model():
  return t5.models.MtfModel(
    model_dir=MODEL_DIR,
    tpu=None,
    model_parallelism=1,
    batch_size=1,
    sequence_length=None,  # {"inputs": 128, "targets": 32},
    learning_rate_schedule=0.003,
    save_checkpoints_steps=50,
    keep_checkpoint_max=None,
    iterations_per_loop=100,
  )

gin.parse_config_file(os.path.join(PRETRAINED_DIR, 'operative_config.gin'))
create_task("s_test",s_dataset_fn)
test_model = create_model()

test_model.finetune(
    mixture_or_task_name="s_test",
    pretrained_model_dir=model.PRETRAINED_DIR,
    finetune_steps=FINETUNE_STEPS
)
```

Gin file includes the following (although given the comment below I think this is either the wrong way to do it or is currently being ignored):
```
run.mesh_shape = 'model:1,batch:1'
run.mesh_devices = ['gpu:0']
```

However, we are running into some issues with GPU support.
1. It is trying to run on our CPU and not our GPU
2. The "mesh shape" is being stuck on 8 which does work when we want the batch size to only be one.

We were looking at `mtf_model.py` to understand how to properly set these and saw on the todo:
```
# TODO(adarob): Allow setting mesh_shape and mesh_devices directly for GPU.
```

Does T5 currently not support GPU for the above approach? If it does, is there something that needs to be changed with this configuration to make sure that T5 is properly using our desired GPU and mesh shape?

Thank you!


Discussion:
The model API (`t5.models.MtfModel`) hardcodes a TPU mesh shape, so you can't use the model API with a GPU currently (that's what @adarob 's note is about). You need to call `utils.run` from the Mesh TF library directly, as is done here:
https://github.com/google-research/text-to-text-transfer-transformer/blob/master/t5/models/mesh_transformer_main.py#L203
That should respect the gin config values you have in your .gin file.
It's also quite easy to add this feature to MtfModel, although I can't do
it until next week.

On Wed, Feb 26, 2020, 7:34 AM Colin Raffel <notifications@github.com> wrote:

> The model API (t5.models.MtfModel) hardcodes a TPU mesh shape, so you
> can't use the model API with a GPU currently (that's what @adarob
> <https://github.com/adarob> 's note is about). You need to call utils.run
> from the Mesh TF library directly, as is done here:
>
> https://github.com/google-research/text-to-text-transfer-transformer/blob/master/t5/models/mesh_transformer_main.py#L203
> That should respect the gin config values you have in your .gin file.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/google-research/text-to-text-transfer-transformer/issues/107?email_source=notifications&email_token=AAIJV2AEQGY5YEHZPCEFBD3RE2R23A5CNFSM4K4KBPN2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOENBFDGA#issuecomment-591548824>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAIJV2DSZOB6MPR7QDZLK3TRE2R23ANCNFSM4K4KBPNQ>
> .
>

Thank you. We will use `utils.run` for now and will look out for any changes to MtfModel.
Thank you both for quick responses and help.

We are currently using `utils.run` for now. It is now respecting our gin config values and using the GPU.

There is still however a problem - GPU usage never goes above 4% even during training and it is excruciatingly slow (1 step per minute). (It is however recognizing the card and allocating all available GPU memory).

Is it possible that the code is still not properly utilizing the GPU?

```
tf.compat.v1.enable_eager_execution()
gin.parse_config_file(os.path.join(model.PRETRAINED_DIR, 'operative_config.gin'))

ON_CLOUD = False
FINETUNE_STEPS = 25000
mixture_or_task_name = 's_test'
create_task.create_task(mixture_or_task_name, create_task.s_dataset_fn2)

# Using Mesh TF library
ckpt = tf.train.latest_checkpoint(model.PRETRAINED_DIR)
if ckpt is None:
    raise ValueError("No checkpoints found in model directory: %s" % model_dir)

checkpoint_step = int(re.sub(".*ckpt-", "", ckpt))
model_ckpt = "model.ckpt-" + str(checkpoint_step)
init_checkpoint = os.path.join(model.PRETRAINED_DIR, model_ckpt)

vocabulary = t5.data.get_mixture_or_task(mixture_or_task_name).get_vocabulary()
dataset_fn = functools.partial(
        t5.models.mesh_transformer.mesh_train_dataset_fn, mixture_or_task_name=mixture_or_task_name)

utils.run(
        tpu_job_name=None,
        tpu=None,
        gcp_project=None,
        tpu_zone=None,
        model_dir=model.MODEL_DIR,
        init_checkpoint=init_checkpoint,
        vocabulary=vocabulary,
        train_dataset_fn=dataset_fn,
        train_steps=(checkpoint_step + FINETUNE_STEPS)
)
```

(We will also test with the new update once it is merged and deployed)

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='t5',
-    version='0.2.1',
+    version='0.3.0',
     description='Text-to-text transfer transformer',
     long_description=_LONG_DESCRIPTION,
     long_description_content_type='text/markdown',
@@ -39,7 +39,7 @@
         'absl-py',
         'babel',
         'gin-config',
-        'mesh-tensorflow[transformer]>=0.1.10',
+        'mesh-tensorflow[transformer]>=0.1.12',
         'nltk',
         'numpy',
         'pandas',
@@ -49,7 +49,7 @@
         'scipy',
         'sentencepiece',
         'six>=1.14',  # TODO(adarob): Remove once rouge-score is updated.
-        'tensorflow-text==1.15.0rc0',
+        'tensorflow-text==1.15',
         'tfds-nightly',
     ],
     extras_require={
diff --git a/t5/models/mtf_model.py b/t5/models/mtf_model.py
--- a/t5/models/mtf_model.py
+++ b/t5/models/mtf_model.py
@@ -75,6 +75,8 @@ def __init__(
       sequence_length=None,
       model_type="bitransformer",
       layout_rules="ensemble:ensemble,batch:batch,d_ff:model,heads:model,vocab:model,experts:batch",
+      mesh_shape=None,
+      mesh_devices=None,
       autostack=True,
       learning_rate_schedule=None,
       keep_checkpoint_max=None,
@@ -101,6 +103,12 @@ def __init__(
         the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
       model_type: str, a model type from mesh tf models.
       layout_rules: an input to mtf.convert_to_layout_rules()
+      mesh_shape: an mtf.Shape or string (e.g., "model:2,batch:4") specifying
+        how the data/model should be parallelized. If None (default), the mesh
+        shape will be constructed using the supplied `tpu_topology` and
+        `model_parallelism` arguments.
+      mesh_devices: a list of strings, the device names to use for each mesh
+        slice. Only required for GPU.
       autostack: boolean, internally combine variables.
       learning_rate_schedule: an optional function taking the scalar name
         argument `step` and the numeric argument `total_train_steps` and return
@@ -121,9 +129,7 @@ def __init__(
       ensemble_inputs: an integer, see `train_model` docstring for details.
       iterations_per_loop: integer, steps per train loop
     """
-
-    # TODO(adarob): Allow setting mesh_shape and mesh_devices directly for GPU.
-    mesh_shape = (
+    mesh_shape = mesh_shape or (
         utils.tpu_mesh_shape(tpu_topology, model_parallelism) if tpu else "")
 
     sequence_length = sequence_length or {"inputs": 512, "targets": 512}
@@ -144,6 +150,7 @@ def __init__(
 
     self._layout_rules = mtf.convert_to_layout_rules(layout_rules)
     self._mesh_shape = mtf.convert_to_shape(mesh_shape)
+    self._mesh_devices = mesh_devices
 
     self._autostack = autostack
     self._keep_checkpoint_max = keep_checkpoint_max
@@ -182,6 +189,7 @@ def estimator(self, vocabulary, init_checkpoint=None, disable_tpu=False):
         vocabulary=vocabulary,
         layout_rules=self._layout_rules,
         mesh_shape=mtf.Shape([]) if disable_tpu else self._mesh_shape,
+        mesh_devices=self._mesh_devices,
         model_dir=self._model_dir,
         batch_size=self.batch_size,
         sequence_length=self._sequence_length,


Test patch: