kubeflow / kubeflow/mpi-operator
TF2 Jobs latches on to CPUs if both CPU and GPU are provided in the container resource requests/limit section
- Dominant language
- Go
- Stars
- 535
- Forks
- 238
- Avg merge
- 2d 3h
- Merged PRs (30d)
- 10
Description
I have applied the following MPI Job Yaml. I observe that when I run the workers with only the GPU specified in the resources section the TF2 Job proceeds very fast with `3s` per epoch. The TF2 Job is also after the MPI Yaml. However, if I provide both cpu:1, gpu:1 and memory:8g in the worker then the TF2 job starts using the CPU instead of the GPU and it takes `15s` per epoch.
```
apiVersion: kubeflow.org/v1alpha2
kind: MPIJob
metadata:
name: horovod-asaha-t256jhci2pi
spec:
slotsPerWorker: 1
cleanPodPolicy: Running
mpiReplicaSpecs:
Launcher:
replicas: 1
template:
metadata:
annotations:
iam.amazonaws.com/role: arn:aws:iam:::role/lyftlearn-production-iad
lyft.net/secret-inject: required
labels:
lyft.com/ml-platform: ''
environment: production
secretsIam: lyftlearn-production-iad
version: dummy
spec:
containers:
- name: horovod-asaha-t256jhci2pi-launcher
image:/pythonlyftdistributed:lyftlearn.2827e895dcfca029efd72fe35e18fa7d6b18a311
command:
- mpirun
args:
- '-np'
- '2'
- '--allow-run-as-root'
- '-bind-to'
- none
- '-map-by'
- slot
- '-x'
- NCCL_DEBUG=INFO
- '-x'
- LD_LIBRARY_PATH
- '-x'
- PATH
- '-x'
- NCCL_SOCKET_IFNAME=eth0
- '-mca'
- pml
- ob1
- '-mca'
- btl
- ^openib
- python
- /mnt/user-home/distributed-training-exploration/tf2_keras_horovod_mnist.py
resources:
limits:
cpu: 1
memory: 2Gi
volumeMounts:
- mountPath: /mnt/user-home
name: nfs
volumes:
- name: nfs
persistentVolumeClaim:
claimName: asaha
Worker:
replicas: 2
template:
metadata:
annotations:
iam.amazonaws.com/role: arn:aws:iam::173840052742:role/lyftlearn-production-iad
lyft.com/user-job-name: NOTEBOOK-1611821340742
lyft.net/secret-inject: required
labels:
lyft.com/ml-platform: ''
environment: production
secretsIam: lyftlearn-production-iad
version: dummy
spec:
containers:
- name: horovod-asaha-t256jhci2pi-worker
image: /pythonlyftdistributed:lyftlearn.2827e895dcfca029efd72fe35e18fa7d6b18a311
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- mountPath: /mnt/user-home
name: nfs
volumes:
- name: nfs
persistentVolumeClaim:
claimName: asaha
tolerations:
- key: lyft.net/gpu
operator: Equal
value: dedicated
effect: NoSchedule
```
The applied TF2 Job from the Horovod Repo
```
# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import tensorflow as tf
import horovod.tensorflow.keras as hvd
# Horovod: initialize Horovod.
hvd.init()
# Horovod: pin GPU to be used to process local rank (one GPU per process)
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")
(mnist_images, mnist_labels), _ = tf.keras.datasets.mnist.load_data(
path="mnist-%d.npz" % hvd.rank()
)
dataset = tf.data.Dataset.from_tensor_slices(
(
tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32),
tf.cast(mnist_labels, tf.int64),
)
)
dataset = dataset.repeat().shuffle(10000).batch(128)
mnist_model = tf.keras.Sequential(
[
tf.keras.layers.Conv2D(32, [3, 3], activation="relu"),
tf.keras.layers.Conv2D(64, [3, 3], activation="relu"),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Dropout(0.25),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation="relu"),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation="softmax"),
]
)
# Horovod: adjust learning rate based on number of GPUs.
scaled_lr = 0.001 * hvd.size()
opt = tf.optimizers.Adam(scaled_lr)
# Horovod: add Horovod DistributedOptimizer.
opt = hvd.DistributedOptimizer(
opt, backward_passes_per_step=1, average_aggregated_gradients=True
)
# Horovod: Specify `experimental_run_tf_function=False` to ensure TensorFlow
# uses hvd.DistributedOptimizer() to compute gradients.
mnist_model.compile(
loss=tf.losses.SparseCategoricalCrossentropy(),
optimizer=opt,
metrics=["accuracy"],
experimental_run_tf_function=False,
)
callbacks = [
# Horovod: broadcast initial variable states from rank 0 to all other processes.
# This is necessary to ensure consistent initialization of all workers when
# training is started with random weights or restored from a checkpoint.
hvd.callbacks.BroadcastGlobalVariablesCallback(0),
# Horovod: average metrics among workers at the end of every epoch.
#
# Note: This callback must be in the list before the ReduceLROnPlateau,
# TensorBoard or other metrics-based callbacks.
hvd.callbacks.MetricAverageCallback(),
# Horovod: using `lr = 1.0 * hvd.size()` from the very beginning leads to worse final
# accuracy. Scale the learning rate `lr = 1.0` ---> `lr = 1.0 * hvd.size()` during
# the first three epochs. See https://arxiv.org/abs/1706.02677 for details.
hvd.callbacks.LearningRateWarmupCallback(
initial_lr=scaled_lr, warmup_epochs=3, verbose=1
),
]
# Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them.
if hvd.rank() == 0:
callbacks.append(tf.keras.callbacks.ModelCheckpoint("./checkpoint-{epoch}.h5"))
# Horovod: write logs on worker 0.
verbose = 1 if hvd.rank() == 0 else 0
# Train the model.
# Horovod: adjust number of steps based on number of GPUs.
mnist_model.fit(
dataset,
steps_per_epoch=500 // hvd.size(),
callbacks=callbacks,
epochs=18,
verbose=verbose,
)
```
Contributor guide
Research direction
Start with the MPIJob worker resource configuration and compare runs that specify only nvidia.com/gpu with runs that also specify CPU and memory. Then inspect the TensorFlow GPU-detection code in the supplied training script and determine whether the operator or container setup changes GPU visibility; done means the cause and a reproducible configuration or fix are identified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- kubernetes, python, tensorflow
- Domain
- distributed-systems, infrastructure, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100