Support of device ordering (#95)

* Trying out new version of redistribute_local_tensor

Taken from pytorch/pytorch#160266, but I'm hitting an assertion for now

* update ordered sharding

* relocate redistribute tensor function (jax way map tensor dim to mesh dim)

* fix loss curve mismatch

* imporve ordering logic and bring back _optimize_same_nd_sharding_as_1d

* lint

* fix small bug

* adress review feedback

* fix CI

---------

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

Author: zpcore

.github/workflows/test_torchtitan.yml

@@ -47,4 +47,4 @@ jobs:
           --model.name llama3_auto_parallel \
           --parallelism.tensor_parallel_degree 4 \
           --training.dataset c4 \
-          --training.compile
+          --compile.enable

autoparallel/apply_sharding.py

@@ -51,23 +51,28 @@ def run_node(self, n):
         self._curr_node = n
         return super().run_node(n)
 
-    def redistribute_tensor(self, arg, curr_spec, tgt_spec, src_tgt_nodes=None):
+    def redistribute_tensor(self, arg, curr_spec, tgt_spec, node=None):
         tgt_placements = tuple(
             p if not p.is_partial() else Replicate() for p in tgt_spec.placements
         )
         x = arg
+        if node in self.param_placement_order and self.param_placement_order[node][1]:
+            assert curr_spec.placements != tgt_spec.placements
         if curr_spec.placements != tgt_spec.placements:
             tgt_spec_c = DTensorSpec(
                 tgt_spec.mesh, tgt_placements, tensor_meta=tgt_spec.tensor_meta
             )
-            placement_order = None
-            if (
-                src_tgt_nodes is not None
-                and src_tgt_nodes in self.param_placement_order
-            ):
-                placement_order = self.param_placement_order[src_tgt_nodes]
+            origin_order = None
+            tgt_order = None
+            if node in self.param_placement_order:
+                tgt_order, do_reorder = self.param_placement_order[node]
+                origin_order = tgt_order[::-1] if do_reorder else tgt_order
             x = ordered_redistribute_local_tensor(
-                arg, curr_spec, tgt_spec_c, placement_order
+                arg,
+                curr_spec,
+                tgt_spec_c,
+                src_placement_order=origin_order,
+                tgt_placement_order=tgt_order,
             )
         return x
 
@@ -109,7 +114,7 @@ def redistribute_args(self, args):
         for n, arg, curr_spec, tgt_spec in zip(
             all_input_nodes, flat_args_t, curr_specs, tgt_specs
         ):
-            x = self.redistribute_tensor(arg, curr_spec, tgt_spec, (node, n))
+            x = self.redistribute_tensor(arg, curr_spec, tgt_spec, node)
             new_flat_args_t.append(x)
             self.tgt_spec = tgt_spec
         new_flat_args = []