[PyTorch] Refactor parameter splitting in Linear and LayerNormLinear

[timmoon10]

#533 reports that TransformerLayer doesn't work out-of-the-box with tensor parallelism. The root cause is because the logic for parameter splitting (e.g. for QKV matrices) does not handle tensor parallelism. We've also had another user run into trouble when trying to set parameters_split in Linear because it currently expects the split names to have exactly one underscore at the end (so mysplit and my_split_ would both fail).

I think this is a good opportunity to refactor this logic:

• Adjust parameter split size as needed for tensor parallelism
• Generalize support for split names. To maintain backward compatibility, we now strip all trailing underscores before appending _weight or _bias, resulting in parameter names like q_weight, etc.
• Separate the noop_cat operation so it is independent from the TE modules.

Closes #533.

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[cyanguwa]

Would these changes affect when people try to load existing checkpoints?

[timmoon10]

I don't think so. The resulting param names (q_weight, k_weight, v_weight) are unchanged, and I don't think the actual values in parameters_split are involved in checkpointing.