New DataPartitionType DATA

[apoorvtintin]

Increases memory efficiency during large scale training, input batches and labels are sharded along the 'data' axis.
Added new input data sharding option DataPartitionType.DATA.

[kelvin-zou]

A high level question, what is the purpose of this change?
I see that we have FULL partition support already, which partitions on axis=0 which is the data axis, how is DATA different from FULL?

[ptoulme-aws]

DATA replicates over the sequence dimension. so the spec is ("data", None) versus ("data", "model") for FULL

[ruomingp]

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

[ptoulme-aws]

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

By replicating the sequence length over TP workers we limit collectives and dynamic-slices introduced by the SPMD partitioner. This lowers overall step time and also allows us to run sequence parallelism over TP workers.

[ruomingp]

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

By replicating the sequence length over TP workers we limit collectives and dynamic-slices introduced by the SPMD partitioner. This lowers overall step time and also allows us to run sequence parallelism over TP workers.

Thanks. Do you have quantitative measurements?

[ptoulme-aws]

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

Increases memory efficiency

Do you have measurements on how DATA improves memory efficiency? Thanks.

By replicating the sequence length over TP workers we limit collectives and dynamic-slices introduced by the SPMD partitioner. This lowers overall step time and also allows us to run sequence parallelism over TP workers.

Thanks. Do you have quantitative measurements?

No, we do not. It is more when we inspect the HLO after SPMD partition pass we see much more optimal sharding. Less all-to-alls and less dynamic-slices on right hand side.