Distributed training on Kubernetes (Azure AKS) cannot utilize GPUs properly

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not already been filed.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/blob/master/research/object_detection/dockerfiles/tf1/Dockerfile

2. Describe the bug

When training the SSD_inception_V2 model (from the TF1 model zoo) on azure Kubernetes with multiple GPUs, the error tensorflow/stream_executor/cuda/cuda_driver.cc:175] Check failed: err == cudaSuccess || err == cudaErrorInvalidValue Unexpected CUDA error: out of memory occurs sometimes. (The training works fine under certain settings).

3. Steps to reproduce

• Build the image from the tensorflow object detection api tf1 docker image here.
• Have an AKS cluster with two nodepools: 1) 4 Standard_DS2_VS2 2) 8 Standard_NC_6. (Standard_NC_6 have 1 K80 GPUs each.)
• Deploy GPU daemon on cluster.
• Deploy Kubeflow to get TFJob custom resource.
• Deploy PVC (Model data and checkpoints will be read from/written to here)
• Configure pipeline.config for ssd_inception_v2 with a batch size of 1
• Configure tfjob yaml with 1 Master (GPU), 7 Workers (GPU), 3 Parameter Servers (CPU).
• Deploy Job

4. Expected behavior

Job runs and finishes successfully.

5. Additional context

The initial test was done with a small amount of data, about 150 Mb. With 1 M, 7 W and 3 PS, the job failed with the following error:
tensorflow/stream_executor/cuda/cuda_driver.cc:175] Check failed: err == cudaSuccess || err == cudaErrorInvalidValue Unexpected CUDA error: out of memory

Reducing the configuration to 1 M, 3 W, and 3 PS made the job run without error. It finished successfully (although it was a bit slow).

The second time I tried training I had a larger amount of data, about 12 GB. Using my previously working configuration of 1 M, 3 W and 3 PS, the job failed with the same error as earlier.

Again, reducing the workers to 1, ie, Using 1 M, 1 W, the job again ran without error (albeit even slower).

Previously I used to train object detection api models on GCP with gcloud ai-platform jobs submit training with the cluster specs of 1 M, 7W and 3 PS, all of type n1-highmem-8, and master and workers each had a single NVIDIA K-80 GPU attached. I expected to be able to run training jobs with similar specs in Azure Kubernetes.

I have posted an issue on the Kubeflow github repository in prior, consulted an Azure consultant, and asked a stack overflow question tagged tensorflow and object-detection as suggested in the Getting Help section, but haven't found satisfactory answers yet.

I have attached the pipeline.config and job yaml here for additional reference.

6. System information

• Linux Ubuntu 18.04
• Tf1 Object Detection API docker image built from the repository
• GPU model and memory: 8 x NVIDIA Tesla K80 nodes, each having 12 GB GPU memory