The objective of this notebook is to demonstrate splitting a log_prob and gradient computation across a number of GPU devices. For development purposes, this was prototyped in colab with a single GPU partitioned into multiple logical GPUs.
Note: Since it runs on a single GPU, performance is not representative of what can be achieved with multiple GPUs. Usage of tf.data can likely benefit from some tuning when deployed to multiple GPUs.
Needs a GPU: Edit > Notebook Settings: Hardware Accelerator => GPU
In [42]:
%tensorflow_version 2.x
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
tfb, tfd = tfp.bijectors, tfp.distributions
physical_gpus = tf.config.experimental.list_physical_devices('GPU')
print(physical_gpus)
tf.config.experimental.set_virtual_device_configuration(
physical_gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2000)] * 4)
gpus = tf.config.list_logical_devices('GPU')
print(gpus)
st = tf.distribute.MirroredStrategy(devices=tf.config.list_logical_devices('GPU'))
print(st.extended.worker_devices)
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
[LogicalDevice(name='/device:GPU:0', device_type='GPU'), LogicalDevice(name='/device:GPU:1', device_type='GPU'), LogicalDevice(name='/device:GPU:2', device_type='GPU'), LogicalDevice(name='/device:GPU:3', device_type='GPU')]
WARNING:tensorflow:NCCL is not supported when using virtual GPUs, fallingback to reduction to one device
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
In [43]:
# Draw samples from an MVN, then sort them. This way we can easily visually
# verify the correct partition ends up on the correct GPUs.
ndim = 3
def model():
Root = tfd.JointDistributionCoroutine.Root
loc = yield Root(tfb.Shift(.5)(tfd.MultivariateNormalDiag(loc=tf.zeros([ndim]))))
scale_tril = yield Root(tfb.FillScaleTriL()(tfd.MultivariateNormalDiag(loc=tf.zeros([ndim * (ndim + 1) // 2]))))
yield tfd.MultivariateNormalTriL(loc=loc, scale_tril=scale_tril)
dist = tfd.JointDistributionCoroutine(model)
tf.random.set_seed(1)
loc, scale_tril, _ = dist.sample(seed=2)
samples = dist.sample(value=([loc] * 1024, scale_tril, None), seed=3)[2]
samples = tf.round(samples * 1000) / 1000
for dim in reversed(range(ndim)):
samples = tf.gather(samples, tf.argsort(samples[:,dim]))
print(samples)
tf.Tensor( [[-4.534 -5.856 3.606] [-4.527 -5.875 1.671] [-4.269 -4.346 5.697] ... [ 2.158 5.71 -2.926] [ 2.302 6.658 -3.491] [ 2.632 5.67 -4.854]], shape=(1024, 3), dtype=float32)
In [44]:
print(loc)
print(scale_tril)
print(tf.reduce_mean(samples, 0))
tf.Tensor([-1.0574996 0.24829748 1.0737331 ], shape=(3,), dtype=float32) tf.Tensor( [[ 1.1475685 0. 0. ] [ 1.9094281 0.5724521 0. ] [-1.1899896 0.49813363 1.5088601 ]], shape=(3, 3), dtype=float32) tf.Tensor([-0.9953702 0.3626416 1.0675195], shape=(3,), dtype=float32)
Single batch of data resident on GPU.¶
In [45]:
%%time
def dataset_fn(ctx):
batch_size = ctx.get_per_replica_batch_size(len(samples))
d = tf.data.Dataset.from_tensor_slices(samples).batch(batch_size)
return d.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
ds = st.experimental_distribute_datasets_from_function(dataset_fn)
observations = next(iter(ds))
# print(observations)
@tf.function(autograph=False)
def log_prob_and_grad(loc, scale_tril, observations):
ctx = tf.distribute.get_replica_context()
with tf.GradientTape() as tape:
tape.watch((loc, scale_tril))
lp = tf.reduce_sum(dist.log_prob(loc, scale_tril, observations)) / len(samples)
grad = tape.gradient(lp, (loc, scale_tril))
return ctx.all_reduce('sum', lp), [ctx.all_reduce('sum', g) for g in grad]
@tf.function(autograph=False)
@tf.custom_gradient
def target_log_prob(loc, scale_tril):
lp, grads = st.run(log_prob_and_grad, (loc, scale_tril, observations))
return lp.values[0], lambda grad_lp: [grad_lp * g.values[0] for g in grads]
singleton_vals = tfp.math.value_and_gradient(target_log_prob, (loc, scale_tril))
kernel = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob, step_size=.35, num_leapfrog_steps=2)
kernel = tfp.mcmc.TransformedTransitionKernel(kernel, bijector=[tfb.Identity(), tfb.FillScaleTriL()])
@tf.function(autograph=False)
def sample_chain():
return tfp.mcmc.sample_chain(
num_results=200, num_burnin_steps=100,
current_state=[tf.ones_like(loc), tf.linalg.eye(scale_tril.shape[-1])],
kernel=kernel, trace_fn=lambda _, kr: kr.inner_results.is_accepted)
samps, is_accepted = sample_chain()
print(f'accept rate: {np.mean(is_accepted)}')
print(f'ess: {tfp.mcmc.effective_sample_size(samps)}')
print(tf.reduce_mean(samps[0], axis=0))
# print(tf.reduce_mean(samps[1], axis=0))
import matplotlib.pyplot as plt
for dim in range(ndim):
plt.figure(figsize=(10,1))
plt.hist(samps[0][:,dim], bins=50)
plt.title(f'loc[{dim}]: prior mean = 0.5, observation = {loc[dim]}')
plt.show()
accept rate: 0.8
ess: [<tf.Tensor: shape=(3,), dtype=float32, numpy=array([42.702564, 56.667793, 43.04328 ], dtype=float32)>, <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[57.469185, nan, nan],
[34.60636 , 43.471287, nan],
[13.848625, 30.807411, 72.34708 ]], dtype=float32)>]
tf.Tensor([-0.18598816 0.7396643 0.4074543 ], shape=(3,), dtype=float32)
CPU times: user 14.3 s, sys: 1.57 s, total: 15.8 s Wall time: 12.9 s
Two batches of data per log-prob eval (2x slower).¶
In [46]:
%%time
batches_per_eval = 2
def dataset_fn(ctx):
batch_size = ctx.get_per_replica_batch_size(len(samples))
d = tf.data.Dataset.from_tensor_slices(samples).batch(batch_size // batches_per_eval)
return d.shard(ctx.num_input_pipelines, ctx.input_pipeline_id).prefetch(2)
ds = st.experimental_distribute_datasets_from_function(dataset_fn)
@tf.function(autograph=False)
def log_prob_and_grad(loc, scale_tril, observations, prev_sum_lp, prev_sum_grads):
with tf.GradientTape() as tape:
tape.watch((loc, scale_tril))
lp = tf.reduce_sum(dist.log_prob(loc, scale_tril, observations)) / len(samples)
grad = tape.gradient(lp, (loc, scale_tril))
return lp + prev_sum_lp, [g + pg for (g, pg) in zip(grad, prev_sum_grads)]
@tf.function(autograph=False)
@tf.custom_gradient
def target_log_prob(loc, scale_tril):
sum_lp = tf.zeros([])
sum_grads = [tf.zeros_like(x) for x in (loc, scale_tril)]
sum_lp, sum_grads = st.run(
lambda *x: tf.nest.map_structure(tf.identity, x), (sum_lp, sum_grads))
def reduce_fn(state, observations):
sum_lp, sum_grads = state
return st.run(
log_prob_and_grad, (loc, scale_tril, observations, sum_lp, sum_grads))
sum_lp, sum_grads = ds.reduce((sum_lp, sum_grads), reduce_fn)
sum_lp = st.reduce('sum', sum_lp, None)
sum_grads = [st.reduce('sum', sg, None) for sg in sum_grads]
return sum_lp, lambda grad_lp: [grad_lp * sg for sg in sum_grads]
multibatch_vals = tfp.math.value_and_gradient(target_log_prob, (loc, scale_tril))
kernel = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob, step_size=.35, num_leapfrog_steps=2)
kernel = tfp.mcmc.TransformedTransitionKernel(kernel, bijector=[tfb.Identity(), tfb.FillScaleTriL()])
@tf.function(autograph=False)
def sample_chain():
return tfp.mcmc.sample_chain(
num_results=200, num_burnin_steps=100,
current_state=[tf.ones_like(loc), tf.linalg.eye(scale_tril.shape[-1])],
kernel=kernel, trace_fn=lambda _, kr: kr.inner_results.is_accepted)
samps, is_accepted = sample_chain()
print(f'accept rate: {np.mean(is_accepted)}')
print(f'ess: {tfp.mcmc.effective_sample_size(samps)}')
print(tf.reduce_mean(samps[0], axis=0))
# print(tf.reduce_mean(samps[1], axis=0))
import matplotlib.pyplot as plt
for dim in range(ndim):
plt.figure(figsize=(10,1))
plt.hist(samps[0][:,dim], bins=50)
plt.title(f'loc[{dim}]: prior mean = 0.5, observation = {loc[dim]}')
plt.show()
accept rate: 0.8
ess: [<tf.Tensor: shape=(3,), dtype=float32, numpy=array([42.702564, 56.667793, 43.043278], dtype=float32)>, <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[57.469193, nan, nan],
[34.606365, 43.471287, nan],
[13.848625, 30.807411, 72.34709 ]], dtype=float32)>]
tf.Tensor([-0.18598816 0.7396643 0.4074543 ], shape=(3,), dtype=float32)
CPU times: user 1min 8s, sys: 8.71 s, total: 1min 17s Wall time: 51.5 s
Sanity check logprob and gradients.
In [0]:
for i, (sv, mv) in enumerate(zip(tf.nest.flatten(singleton_vals),
tf.nest.flatten(multibatch_vals))):
np.testing.assert_allclose(sv, mv, err_msg=i, rtol=1e-5)
In [0]: