[BUG] partition_balanced return wrong result.

[zjjMaiMai]

Background

In pipeline parallelism, deepspeed uses ds_utils.partition_balanced to balance the partitioning of the model according to the number of parameters or class names.

DeepSpeed/deepspeed/runtime/pipe/module.py

Lines 380 to 395 in 581e44d

	if method == 'uniform':
	num_layers = len(self._layer_specs)
	self.parts = ds_utils.partition_uniform(num_items=num_layers, num_parts=num_stages)
	elif method == 'parameters':
	param_counts = self._count_layer_params()
	self.parts = ds_utils.partition_balanced(weights=param_counts, num_parts=num_stages)
	elif method.startswith('type:'):
	layertype = method.split(':')[1]
	binary_weights = [0] * len(self._layer_specs)
	for idx in self._find_layer_type(layertype):
	binary_weights[idx] = 1
	self.parts = ds_utils.partition_balanced(weights=binary_weights, num_parts=num_stages)
	elif method == 'profile':
	raise NotImplementedError(f'Partitioning method {method} not implemented.')
	else:
	raise NotImplementedError(f'Partitioning method {method} not implemented.')

What wrong?

>>> import deepspeed
>>> deepspeed.__version__
'0.10.3+542dc0d5'
>>> from deepspeed.runtime import utils as ds_utils
>>> ds_utils.partition_balanced([1, 1, 1, 1, 1], 4)
[0, 2, 4, 5, 5]
>>> 

the result [0, 2, 4, 5, 5] means [2, 2, 1, 0] layers for each part, which is not balanced at all. the last part will throw an exception because there are no parameters to training.

i add some unit test for this function, and i will fix it later if anyone need it.

[zjjMaiMai]

@microsoft-github-policy-service agree

[zjjMaiMai]

already fixed! cc @tjruwase

[ShadenSmith]

Thanks for this PR, @zjjMaiMai!

A note on balance: the objective function in the original code minimizes the maximum load per partition. The maximum load on a pipeline stage determines the pipeline throughput, and so the original result is also balanced. Example paper: Fast Optimal Load Balancing Algorithms for 1D Partitioning