Input batch sharding strategy BATCH #884
Conversation
apoorvtintin
force-pushed
the
mainline_upstream_input_sharding
branch
2 times, most recently
from
79a8c21 to
03052e4
Compare
apoorvtintin
force-pushed
the
mainline_upstream_input_sharding
branch
from
03052e4 to
056536d
Compare
| @@ -607,6 +610,7 @@ def host_to_global_device_array( | |||
| host_arrays: Nested[Union[np.ndarray, Tensor]], | |||
| *, | |||
| partition: DataPartitionType = DataPartitionType.FULL, | |||
There was a problem hiding this comment.
I think @markblee plans to remove the DataPartitionType enum and rely on https://jax.readthedocs.io/en/latest/_autosummary/jax.make_array_from_process_local_data.html to support flexible partition specs.
There was a problem hiding this comment.
Thanks that sounds promising.
Hello @markblee, let me know if this PR is needed till you make your changes, or if you have your design in mind I can reshape the PR to make it compatible with your design.
There was a problem hiding this comment.
I think make_array_from_process_local_data is already merged: #781
|
@apoorvtintin With this new sharding strategy, how do you plan to support model parallelism across hosts (i.e. per host bs < 1)? If this is not a supported use case, then this change LGTM. I could imagine the support for model parallelism across hosts is not needed in the short term for |
|
Sorry for the late review, I should've checked this earlier. axlearn/axlearn/common/input_dispatch.py Line 24 in 2d1fb29 |
| @@ -423,6 +423,7 @@ def get_trainer_kwargs( | |||
| raise NotImplementedError(f"Unknown model size {model_size}.") | |||
| model_kwargs = trainer_kwargs.pop("model_kwargs") | |||
| model_kwargs.setdefault("vocab_size", vocab_size) | |||
| trainer_kwargs["input_partition_type"] = None if backend != "neuron" else DataPartitionType.BATCH | |||
There was a problem hiding this comment.
Please go with a configmodifier instead of hardcode in the code, otherwise it becomes hard for people to debug
| ), | ||
| out_shardings=dict( | ||
| replicated=None, | ||
| per_example=utils.input_partition_spec(), | ||
| per_example=utils.data_partition_type_to_spec( partition=self.config.input_partition_type, batch_axis_names=self.config.batch_axis_names), |
There was a problem hiding this comment.
Did you run through pylint? this line seems quite long.
| """Returns a PartitionSpec for the given partition type.""" | ||
| if partition == DataPartitionType.FULL: | ||
| return input_partition_spec() | ||
| elif partition == DataPartitionType.REPLICATED: | ||
| return None | ||
| elif partition == DataPartitionType.BATCH: | ||
| return PartitionSpec(batch_axis_names) |
There was a problem hiding this comment.
Instead of directly assigning PartitionSpec, maybe extend input_partition_spec with an argument batch_axis_names, with default None?
| @@ -37,16 +36,15 @@ def is_supported( | |||
| ) | |||
| ) | |||
|
|
|||
|
|
|||
There was a problem hiding this comment.
Are you using a different pylint style? Those lines should not be removed.
I think the primary purpose of this PR (as stated in the description) is to avoid the resharding cost from physical bs to logical bs when there's no need to shard data along model axis. This could remove the collective for inputs re-sharding and thus reduce step time. @apoorvtintin May I know how much time do we save by removing this collective? If the benefit is not significant, we can stick to the old code. |
But input collective is only triggered within host, which should be insignificant? |
Axlearn currently supports DataPartitionType strategies FULL and REPLICATED for input batches. This PR adds support for a third sharding strategy DataPartitionType.BATCH.
This sharding strategy shards the input batch on sharding axis inferred as "batch_axis". This ensures batches are replicated on the "model" dimension and thus avoids unnecessary collectives to reshard input batches when Tensor parallelism is used.