SageMaker TF 2.3 distributed training

Question

Using SageMaker v2.29.2 and Tensorflow v2.3.2 I'm trying to implement distributed training as explained in the following blogpost:

https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-customize-training-script-tf.html#model-parallel-customize-training-script-tf-23

However I'm having difficulties importing the smdistributed script.

Here is my code:

import tensorflow as tf
import smdistributed.modelparallel.tensorflow as smp

Error:

Traceback (most recent call last):
  File "temp.py", line 2, in <module>
    import smdistributed.modelparallel.tensorflow as smp
ModuleNotFoundError: No module named 'smdistributed'

What am I missing?

Answer 1

smdistributed is only available on the SageMaker containers. It is supported for specific TensorFlow versions and you must add:

distribution={'smdistributed': {
            'dataparallel': {
                'enabled': True
            }
        }}

On the estimator code in order to enable it