Update on "Reordered TP parallel plan to follow execution order"

- Llama uses pre-norm (norm before attention and before FFN), so we can move these up.
- The root norm is before output, so we can swap this order too.



[ghstack-poisoned]

torchtitan/parallelisms/parallelize_llama.py

@@ -332,6 +332,7 @@ def apply_tp(model, world_mesh, parallel_dims, job_config: JobConfig):
     """
     Apply tensor parallelism.
     """
+
     tp_mesh = world_mesh["tp"]
     (
         row_parallel_strategy,