[Question] Padding in ZeRO optimizer

[e-sha]

Why do you use padding here, here in ZeRO optimizer?
Moreover, later you split tensors as if it was not padded. For example, here.
From my perspective padding is not necessary at all.

[tjruwase]

We pad the tensors to align to the data parallel size to ensure load-balancing across the ranks.

The tensor splitting here is on the padded tensors.

[e-sha]

According to this line the optimizer can split tensor that is not aligned. Thus, padding seems to be an unnecessary feature that just makes code more sophisticated and harder to maintain.

[tjruwase]

Apologies for the confusion, but that line is old and redundant since the input to the function is the padded tensor. In addition to load-balancing, padding helps to handle cases where the tensor.numel() is smaller than GPU count, which although is a rare case but one that we have seen and led to this unit test.

[tjruwase]

Closing for lack of activity. Please reopen as needed.