make dtensor shared test util more generic

[vkuzo]

Summary:

Refactoring to make adding MX easier in next PR, no logic change. Specifically:

• move _test_lowp_mlp_tensor_parallelism_base to shared util file
• move config generation out of this function to the callsite
• instead of gating lowp allgather by rowwise flag, gate it by allgather_in_lowp flag

Test Plan:

./test/float8/test_dtensor.sh

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• enable tensor parallelism for MXLinear #2434
• enable to_mxfp8 cast for DTensor #2420
• -> make dtensor shared test util more generic #2416

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2416

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.