PATCH: add back n-dim device-mesh + fix tp hook registration

[S1ro1]

1. Fixes ndim check on device_mesh - This was merged in previously with Allow device_mesh have multiple dim  #38949 but was by mistake reverted by Add ep #39501, we need this for upcoming accelerate/axolotl release.
2. Changes TP plans to make it work as before, we need to replicate (potentiall seqparallel) layer norms and other modules, as else we'd get mixed tensor+dtensor in optimizer.step()
3. Adds merging of basemodel tp_plan with the specific model, it wasn't applied before and i.e. lm_heads were not sharded.

TODO: Add proper plans for other TP models
TODO: Add correct plan merging even for distribute_and_shard

cc @ArthurZucker @SunMarc

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: llama