[MoE training] Assert expert weights are column-major; preserve subclass with transpose

[danielvegamyhre]

Once pytorch/torchtitan#1517 lands we need to:

• remove the temporary workaround of dynamic memory layout transformation from row-major to col-major before every grouped gemm, and just assert weights are col-major.
• preserve subclass through transpose ops, so the weight.transpose(-2, -1) op doesn't lose the subclass and thus no override to fp8 grouped gemm is done

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Pending

As of commit ec07a5c with merge base 7dbc816 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_uneven_shard
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_uneven_shard

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

[danielvegamyhre]

cc @vkuzo @drisspg this is ready for review, the associated titan PR has landed pytorch/torchtitan#1517

[danielvegamyhre]

confirmed test failures are unrelated, and we are looking into these red CI tests to get them fixed