Maximising multi-instance and multi-GPU utilisation

Hi,

I've been trying to distribute the MNIST example across instances and GPUs with SageMaker Tensorflow, but I'm not seeing the kind of benefit that I was hoping for. I'm not sure whether I'm just setting the job up incorrectly or whether this example is just not suited to distribution and was wondering if anyone has any ideas which it might be please?

I'm using TensorFlow 1.9 because I think MPI/Horovod doesn't work with 1.12 which I was using originally?

In my test I get these results with a batch size of 512:
1 Instance 1 GPU 13.5 global_step/sec
2 Instance 1 GPU 7.3 global_step/sec
1 Instance 8 GPU 18.8 global_step/sec

If I reduce the batch size I get lower overall throughput, and see little benefit when increasing beyond 512.

My job specification looks as below and mnist.py is the file which comes with the examples, though I changed batch_size=100 to batch_size=512 in the script:

estimator = TensorFlow(entry_point='mnist.py',
                  role=role,
                  framework_version='1.9.0',
                  training_steps=1250, 
                  evaluation_steps=10,
                  train_max_run=5*60,    
                  output_path=output_location,
                  checkpoint_path=output_location,
                  code_location=output_location,
                  model_dir=output_location,
                  train_instance_count=1,
                  train_instance_type='ml.p3.16xlarge',
                  base_job_name='PerformanceTest-p3-16xlarge-1-instance',
                  distributions={
                    'mpi': {
                      'enabled': True,
                      'processes_per_host': 8,
                      'custom_mpi_options': '--NCCL_DEBUG INFO'
                    }
                  })

What I was hoping to see was a single 8-GPU instance hitting global_step/sec of 70.2-97.2. Based on 8x the global_step/sec of a single instance, scaled with 60-90% efficiency. Any help or clarification on this would be greatly appreciated!

Thanks,

Carl