295 lines
9.7 KiB
Python
295 lines
9.7 KiB
Python
|
|
|
|
|
|
|
|
from multiprocessing import Process, Manager
|
|
|
|
import numpy as np
|
|
import unittest
|
|
import tempfile
|
|
import shutil
|
|
import logging
|
|
|
|
from hypothesis import given, settings
|
|
import hypothesis.strategies as st
|
|
|
|
from caffe2.python import workspace
|
|
|
|
log = logging.getLogger("parallelize_bmuf_distributed_test")
|
|
log.setLevel(logging.INFO)
|
|
|
|
|
|
def bmuf_process(filestore_dir, process_id, shared_results,
|
|
cpu_device=False, nesterov=False):
|
|
# We need to import caffe2 in every process to initialize CUDA independently.
|
|
from caffe2.python import core, cnn, data_parallel_model, dyndep
|
|
from caffe2.proto import caffe2_pb2
|
|
dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:file_store_handler_ops")
|
|
|
|
if not cpu_device:
|
|
if not workspace.has_gpu_support:
|
|
log.info('No GPU support test is Ignored.')
|
|
return
|
|
if workspace.NumGpuDevices() < 4:
|
|
log.info('Not enough GPU support, test IGNORED')
|
|
return
|
|
|
|
model = cnn.CNNModelHelper(
|
|
order="NHWC",
|
|
name="test"
|
|
)
|
|
if not cpu_device:
|
|
device_type = workspace.GpuDeviceType
|
|
device_prefix = "gpu"
|
|
else:
|
|
device_type = caffe2_pb2.CPU
|
|
device_prefix = "cpu"
|
|
|
|
devices = [0, 1] if process_id == 0 else [2, 3]
|
|
|
|
def _model_build_fun(model, loss_scale):
|
|
fc = model.FC(
|
|
"data", "fc", 16, 1, ("ConstantFill", {}), ("ConstantFill", {})
|
|
)
|
|
fc_fl = model.FlattenToVec(fc, "fc_fl")
|
|
sigm = model.Sigmoid(fc_fl, "sigm")
|
|
sq = model.SquaredL2Distance([sigm, "label"], "sq")
|
|
loss = model.AveragedLoss(sq, "loss")
|
|
loss = model.Scale(loss, scale=loss_scale)
|
|
|
|
# For testing explicit sync
|
|
model.param_init_net.UniformFill([], ["sync_num"], shape=[1])
|
|
return [loss]
|
|
|
|
def _input_builder_fun(model):
|
|
return None
|
|
|
|
def _param_update_fun(model):
|
|
ITER = model.Iter("ITER")
|
|
LR = model.net.LearningRate(
|
|
[ITER],
|
|
"LR",
|
|
base_lr=(-0.1),
|
|
policy="fixed",
|
|
)
|
|
ONE = model.param_init_net.ConstantFill(
|
|
[], "ONE", shape=[1], value=1.0,
|
|
)
|
|
for param in model.GetParams():
|
|
grad = model.param_to_grad[param]
|
|
model.WeightedSum([param, ONE, grad, LR], param)
|
|
|
|
def _generate_data(devices, process_id, device_type, device_prefix):
|
|
np.random.seed(26 + process_id * 10)
|
|
# Each run has same input, independent of number of gpus
|
|
batch_size = 64
|
|
for _ in range(0, 10):
|
|
full_data = np.random.rand(batch_size, 16)
|
|
full_labels = np.round(full_data[:, 0])
|
|
batch_per_device = batch_size // len(devices)
|
|
|
|
for (j, g) in enumerate(devices):
|
|
st = j * batch_per_device
|
|
en = st + batch_per_device
|
|
data = full_data[st:en, :].astype(np.float32)
|
|
labels = full_labels[st:en].astype(np.float32)
|
|
with core.DeviceScope(core.DeviceOption(device_type, g)):
|
|
workspace.FeedBlob("{}_{}/data".format(device_prefix, g), data)
|
|
workspace.FeedBlob("{}_{}/label".format(device_prefix, g), labels)
|
|
|
|
_generate_data(devices, process_id, device_type, device_prefix)
|
|
|
|
workspace.RunOperatorOnce(
|
|
core.CreateOperator(
|
|
"FileStoreHandlerCreate", [], ["store_handler"],
|
|
path=filestore_dir
|
|
)
|
|
)
|
|
rendezvous = dict(
|
|
kv_handler="store_handler",
|
|
shard_id=process_id,
|
|
num_shards=2,
|
|
engine="GLOO",
|
|
exit_nets=None
|
|
)
|
|
|
|
data_parallel_model.Parallelize_BMUF(
|
|
model,
|
|
_input_builder_fun,
|
|
_model_build_fun,
|
|
_param_update_fun,
|
|
devices=devices,
|
|
rendezvous=rendezvous,
|
|
nesterov=nesterov,
|
|
add_blobs_to_sync=["sync_num"],
|
|
cpu_device=cpu_device
|
|
)
|
|
|
|
data_parallel_model.RunInitNet(model)
|
|
|
|
def _device_pid(device, pid):
|
|
if pid == 1:
|
|
return device + 2
|
|
return device
|
|
|
|
np.testing.assert_equal(
|
|
workspace.FetchBlob("{}_{}/fc_w_v".format(
|
|
device_prefix, _device_pid(0, process_id))),
|
|
np.zeros(16).astype(np.float32).reshape(1, 16)
|
|
)
|
|
|
|
# Run the algorithm for one iteration to have non-zero params.
|
|
data_parallel_model.RunNet(model, 1)
|
|
|
|
# Save iteration momentum and post local update params
|
|
results = {}
|
|
v_b_ = workspace.FetchBlob(
|
|
"{}_{}/fc_b_v".format(device_prefix, _device_pid(0, process_id)))
|
|
v_w_ = workspace.FetchBlob(
|
|
"{}_{}/fc_w_v".format(device_prefix, _device_pid(0, process_id)))
|
|
|
|
results['v_b_'] = v_b_
|
|
results['v_w_'] = v_w_
|
|
|
|
workspace.RunNetOnce(model.net)
|
|
|
|
b_0_ = workspace.FetchBlob(
|
|
"{}_{}/fc_b".format(device_prefix, _device_pid(0, process_id)))
|
|
w_0_ = workspace.FetchBlob(
|
|
"{}_{}/fc_w".format(device_prefix, _device_pid(0, process_id)))
|
|
b_1_ = workspace.FetchBlob(
|
|
"{}_{}/fc_b".format(device_prefix, _device_pid(1, process_id)))
|
|
w_1_ = workspace.FetchBlob(
|
|
"{}_{}/fc_w".format(device_prefix, _device_pid(1, process_id)))
|
|
|
|
results['b_0_'] = b_0_
|
|
results['w_0_'] = w_0_
|
|
results['b_1_'] = b_1_
|
|
results['w_1_'] = w_1_
|
|
|
|
# Test sync
|
|
if process_id == 0:
|
|
workspace.FeedBlob(
|
|
device_prefix + "_0/sync_num",
|
|
np.array([2603]).astype(np.float32),
|
|
device_option=core.DeviceOption(device_type, 0))
|
|
|
|
# Compute block gradients.
|
|
b_g_ = workspace.FetchBlob(
|
|
"{}_{}/fc_b_g".format(device_prefix, _device_pid(0, process_id)))
|
|
w_g_ = workspace.FetchBlob(
|
|
"{}_{}/fc_w_g".format(device_prefix, _device_pid(0, process_id)))
|
|
results['b_g_'] = b_g_
|
|
results['w_g_'] = w_g_
|
|
workspace.RunNetOnce(model._global_model_param_updates_net)
|
|
|
|
# g_b = (b_0_ + b_1_) / 2 - b_g_
|
|
# g_w = (w_0_ + w_1_) / 2 - w_g_
|
|
v_b = workspace.FetchBlob(
|
|
"{}_{}/fc_b_v".format(device_prefix, _device_pid(0, process_id)))
|
|
v_w = workspace.FetchBlob(
|
|
"{}_{}/fc_w_v".format(device_prefix, _device_pid(0, process_id)))
|
|
w_g = workspace.FetchBlob(
|
|
"{}_{}/fc_w_g".format(device_prefix, _device_pid(0, process_id)))
|
|
b_g = workspace.FetchBlob(
|
|
"{}_{}/fc_b_g".format(device_prefix, _device_pid(0, process_id)))
|
|
w_0 = workspace.FetchBlob(
|
|
"{}_{}/fc_w".format(device_prefix, _device_pid(0, process_id)))
|
|
b_0 = workspace.FetchBlob(
|
|
"{}_{}/fc_b".format(device_prefix, _device_pid(0, process_id)))
|
|
w_1 = workspace.FetchBlob(
|
|
"{}_{}/fc_w".format(device_prefix, _device_pid(1, process_id)))
|
|
b_1 = workspace.FetchBlob(
|
|
"{}_{}/fc_b".format(device_prefix, _device_pid(1, process_id)))
|
|
results['v_b'] = v_b
|
|
results['v_w'] = v_w
|
|
results['w_g'] = w_g
|
|
results['b_g'] = b_g
|
|
results['w_0'] = w_0
|
|
results['b_0'] = b_0
|
|
results['w_1'] = w_1
|
|
results['b_1'] = b_1
|
|
|
|
# Test add_blobs_to_sync
|
|
for j in devices:
|
|
sync = workspace.FetchBlob(
|
|
device_prefix + "_{}/sync_num".format(j))[0]
|
|
results['sync_{}'.format(j)] = sync
|
|
|
|
shared_results[process_id] = results
|
|
|
|
|
|
class DistributedTest(unittest.TestCase):
|
|
|
|
@given(
|
|
cpu_device=st.booleans(),
|
|
nesterov=st.booleans()
|
|
)
|
|
@settings(deadline=10000)
|
|
def test_bmuf_distributed(self, cpu_device, nesterov):
|
|
if (not cpu_device) and workspace.has_hip_support:
|
|
log.info('Skipping the test on ROCm due to regression in ROCm3.5')
|
|
return
|
|
self._test_bmuf_distributed(cpu_device=cpu_device, nesterov=nesterov)
|
|
|
|
def _test_bmuf_distributed(self, cpu_device=False, nesterov=False):
|
|
processes = []
|
|
filestore_dir = tempfile.mkdtemp()
|
|
results = Manager().dict()
|
|
for idx in range(0, 2):
|
|
process = Process(
|
|
target=bmuf_process,
|
|
args=(filestore_dir, idx, results, cpu_device, nesterov)
|
|
)
|
|
processes.append(process)
|
|
process.start()
|
|
|
|
while len(processes) > 0:
|
|
process = processes.pop()
|
|
process.join()
|
|
shutil.rmtree(filestore_dir)
|
|
|
|
if len(results) == 0:
|
|
return
|
|
|
|
w_0 = results[0]['w_0']
|
|
w_1 = results[0]['w_1']
|
|
b_0 = results[0]['b_0']
|
|
b_1 = results[0]['b_1']
|
|
# Check parameters are in sync.
|
|
np.testing.assert_equal(w_0, w_1)
|
|
np.testing.assert_equal(w_0, results[1]['w_0'])
|
|
np.testing.assert_equal(w_0, results[1]['w_1'])
|
|
np.testing.assert_equal(b_0, b_1)
|
|
np.testing.assert_equal(b_0, results[1]['b_0'])
|
|
np.testing.assert_equal(b_0, results[1]['b_1'])
|
|
|
|
w_g_ = results[0]['w_g_']
|
|
b_g_ = results[0]['b_g_']
|
|
|
|
g_b = (results[0]['b_0_'] + results[1]['b_0_'] + results[0]['b_1_'] +
|
|
results[1]['b_1_']) / 4 - b_g_
|
|
g_w = (results[0]['w_0_'] + results[1]['w_0_'] + results[0]['w_1_'] +
|
|
results[1]['w_1_']) / 4 - w_g_
|
|
v_b_ = results[0]['v_b_']
|
|
v_b = results[0]['v_b']
|
|
v_w_ = results[0]['v_w_']
|
|
v_w = results[0]['v_w']
|
|
|
|
for pid in results.keys():
|
|
for k in results[pid].keys():
|
|
if k.startswith("sync_num"):
|
|
self.assertEqual(2603, results[pid][k])
|
|
|
|
# Check block gradients are correct.
|
|
np.testing.assert_almost_equal(v_b, 0.75 * v_b_ + g_b)
|
|
np.testing.assert_almost_equal(v_w, 0.75 * v_w_ + g_w)
|
|
|
|
# Check params update step
|
|
if nesterov:
|
|
np.testing.assert_equal(w_0, w_g_ + v_w - 0.75 * (v_w - v_w_))
|
|
np.testing.assert_equal(b_0, b_g_ + v_b - 0.75 * (v_b - v_b_))
|
|
else:
|
|
np.testing.assert_equal(w_0, w_g_ + v_w)
|
|
np.testing.assert_equal(b_0, b_g_ + v_b)
|