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

Issue:
Unable to export a model finetuned on multiple-GPUs
I've successfully finetuned a model using 8 GPUs. However, when I tried to export the model using the function below:

```
saved_model_path = model.export(
    export_dir,
    checkpoint_step=-1,  # use most recent
    beam_size=1,  # no beam search
    temperature=0,  # sample according to predicted distribution
)
```

I got an error:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-43-0f172c37d968> in <module>
      4     checkpoint_step=-1,  # use most recent
      5     beam_size=1,  # no beam search
----> 6     temperature=0,  # sample according to predicted distribution
      7 )
      8 print("Model saved to:", saved_model_path)

~/envs/t5/lib64/python3.7/site-packages/t5/models/mtf_model.py in export(self, export_dir, checkpoint_step, beam_size, temperature, vocabulary)
    405         self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
    406         self._sequence_length, batch_size=self.batch_size,
--> 407         checkpoint_path=os.path.join(self._model_dir, model_ckpt))

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/transformer/utils.py in export_model(estimator, export_dir, vocabulary, sequence_length, batch_size, checkpoint_path)
   1761 
   1762   return estimator.export_saved_model(
-> 1763       export_dir, serving_input_fn, checkpoint_path=checkpoint_path)
   1764 
   1765 

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py in export_saved_model(self, export_dir_base, serving_input_receiver_fn, assets_extra, as_text, checkpoint_path, experimental_mode)
    726         as_text=as_text,
    727         checkpoint_path=checkpoint_path,
--> 728         strip_default_attrs=True)
    729 
    730   def experimental_export_all_saved_models(self,

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py in _export_all_saved_models(self, export_dir_base, input_receiver_fn_map, assets_extra, as_text, checkpoint_path, strip_default_attrs)
    864             save_variables,
    865             mode=ModeKeys.PREDICT,
--> 866             strip_default_attrs=strip_default_attrs)
    867         save_variables = False
    868 

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py in _add_meta_graph_for_mode(self, builder, input_receiver_fn_map, checkpoint_path, save_variables, mode, export_tags, check_variables, strip_default_attrs)
   2878           export_tags=export_tags,
   2879           check_variables=check_variables,
-> 2880           strip_default_attrs=strip_default_attrs))
   2881 
   2882     if self._export_to_tpu and mode == model_fn_lib.ModeKeys.PREDICT:

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py in _add_meta_graph_for_mode(self, builder, input_receiver_fn_map, checkpoint_path, save_variables, mode, export_tags, check_variables, strip_default_attrs)
    943           labels=getattr(input_receiver, 'labels', None),
    944           mode=mode,
--> 945           config=self.config)
    946 
    947       export_outputs = export_lib.export_outputs_for_mode(

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py in _call_model_fn(self, features, labels, mode, config)
   2918     else:
   2919       return super(TPUEstimator, self)._call_model_fn(features, labels, mode,
-> 2920                                                       config)
   2921 
   2922   def _call_model_fn_for_inference(self, features, labels, mode, config):

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py in _call_model_fn(self, features, labels, mode, config)
   1174 
   1175     logging.info('Calling model_fn.')
-> 1176     model_fn_results = self._model_fn(features=features, **kwargs)
   1177     logging.info('Done calling model_fn.')
   1178 

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py in _model_fn(features, labels, mode, config, params)
   3178           tf.compat.v1.logging.info('Running %s on CPU/GPU', mode)
   3179           estimator_spec = model_fn_wrapper.call_without_tpu(
-> 3180               features, labels, is_export_mode=is_export_mode)
   3181           if (self._log_every_n_steps is not None or
   3182               self._log_every_n_secs is not None):

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py in call_without_tpu(self, features, labels, is_export_mode)
   1698 
   1699   def call_without_tpu(self, features, labels, is_export_mode):
-> 1700     return self._call_model_fn(features, labels, is_export_mode=is_export_mode)
   1701 
   1702   def _add_embedding_features(self, features, hook_dummy_table_variables):

~/envs/t5/lib64/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py in _call_model_fn(self, features, labels, is_export_mode)
   2029       _add_item_to_params(params, _CTX_KEY, user_context)
   2030 
-> 2031     estimator_spec = self._model_fn(features=features, **kwargs)
   2032     print(f'printing estimator_spec', estimator_spec)
   2033     print(f'Debugging: estimator_spec = self._model_fn(features=features, **kwargs)')

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/transformer/utils.py in my_model_fn(***failed resolving arguments***)
    517       mtf_samples = mtf.anonymize(mtf_samples)
    518       inputs = mtf.anonymize(inputs)
--> 519       lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)
    520       inputs = clean_decodes(lowering.export_to_tf_tensor(inputs))
    521       outputs = clean_decodes(lowering.export_to_tf_tensor(mtf_samples))

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/ops.py in __init__(self, graph, mesh_to_impl, autostack, log_file)
    724       # tf.logging.info("Lowering operation %s" % op.to_string)
    725       with tf.name_scope(op.name):
--> 726         op.lower(self)
    727       for out in op.outputs:
    728         self.add_counter(

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/ops.py in lower(self, lowering)
   4568     def reshape_fn(x):
   4569       return tf.reshape(x, new_slice_shape)
-> 4570     slices = mesh_impl.slicewise_delay_allreduce(reshape_fn, slices)
   4571     for mesh_axis, tensor_axis in allsplit_after_reshape:
   4572       slices = mesh_impl.allsplit(slices, mesh_axis, tensor_axis)

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/ops.py in slicewise_delay_allreduce(self, fn, *inputs)
   1047           add_counter_fn=inputs[0].add_counter_fn)
   1048     else:
-> 1049       return self.slicewise(fn, *inputs)
   1050 
   1051   def Print(self, x, data, message, **kwargs):  # pylint: disable=invalid-name

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/placement_mesh_impl.py in slicewise(self, fn, *inputs)
    171     inputs = [x.tensor_list if isinstance(x, self.LaidOutTensor)
    172               else [x] * len(self.devices) for x in inputs]
--> 173     ret = mtf.parallel(self.devices, fn, *inputs)
    174     if isinstance(ret[0], tuple):
    175       ret = mtf.transpose_list_of_lists(ret)

~/envs/t5/lib64/python3.7/site-packages/mesh_tensorflow/ops.py in parallel(devices, fn, *args, **kwargs)
   5649       raise ValueError(
   5650           "Argument not a list with same length as devices "
-> 5651           "arg=%s devices=%s" % (x, devices))
   5652   ret = []
   5653   for i, device in enumerate(devices):

ValueError: Argument not a list with same length as devices arg=[<tf.Tensor 'Print:0' shape=(10, 128) dtype=int32>] devices=['gpu:0', 'gpu:1', 'gpu:2', 'gpu:3', 'gpu:4', 'gpu:5', 'gpu:6', 'gpu:7']
```

Can someone please shed some light?


Discussion:
We don't yet have a way to use multiple accelerators with an exported saved model, unfortunately. See #138.
hi @adarob , thanks for getting this back. Just want to clarify what I wanted to do. I trained the model on 8 GPUs but I don't plan to use 8 GPUs for serving. I just wanted to be able to export the model so it can be served using a single GPU/CPU instance. I.e. I don't need to use multiple accelerators. Can this be done at the moment?
Sorry about that, I misread. This should certainly be possible. Can you share your operative config?
No worries @adarob . I'm not sure what you mean by operative config but below are how I configured and finetuned the model:

```
MODEL_SIZE = "small" #@param["small", "base", "large", "3B", "11B"]
# Public GCS path for T5 pre-trained model checkpoints
BASE_PRETRAINED_DIR = "gs://t5-data/pretrained_models"
PRETRAINED_DIR = os.path.join(BASE_PRETRAINED_DIR, MODEL_SIZE)
MODEL_DIR = os.path.join(MODELS_DIR, MODEL_SIZE)

# Set parallelism and batch size to fit on v2-8 TPU (if possible).
# Limit number of checkpoints to fit within 5GB (if possible).
model_parallelism, train_batch_size, keep_checkpoint_max = {
    "small": (2, 128, 4),
    "base": (2, 128, 8),
    "large": (8, 64, 4),
    "3B": (8, 16, 1),
    "11B": (8, 16, 1)}[MODEL_SIZE]

tf.io.gfile.makedirs(MODEL_DIR)

# Init a model
model = t5.models.MtfModel(
    model_dir=MODEL_DIR,
    tpu=None,
    mesh_shape='model:2,batch:4',
    mesh_devices=['gpu:0','gpu:1','gpu:2','gpu:3','gpu:4','gpu:5','gpu:6','gpu:7'],
    #tpu_topology=TPU_TOPOLOGY,
    model_parallelism=model_parallelism,
    batch_size=train_batch_size,
    sequence_length={"inputs": 128, "targets": 32},
    learning_rate_schedule=0.003,
    save_checkpoints_steps=5000,
    keep_checkpoint_max=8,
    iterations_per_loop=100,
)

# finetune the model
FINETUNE_STEPS = 35000
model.finetune(
    mixture_or_task_name="question_all",
    pretrained_model_dir=PRETRAINED_DIR,
    finetune_steps=FINETUNE_STEPS
)

export_dir = os.path.join(MODEL_DIR, "export")

model.batch_size = 10 # make one prediction per call
saved_model_path = model.export(
    export_dir,
    checkpoint_step=-1,  # use most recent
    beam_size=1,  # no beam search
    temperature=0,  # sample according to predicted distribution
)
print("Model saved to:", saved_model_path)
```

The error happened at the last step when I tried to run the model.export function to export the model.
can you try setting `model._mesh_devices=None` before calling export?

Resolving patch:
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
@@ -195,7 +195,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,
+        mesh_devices=None if disable_tpu else self._mesh_devices,
         model_dir=self._model_dir,
         batch_size=self.batch_size,
         sequence_length=self._sequence_length,


Test patch: