fixed some issues with setting leaves and made partitioning deterministic (#880)

Author: Chillee

functorch/_src/aot_autograd.py

@@ -147,6 +147,10 @@ class CompiledFunction(torch.autograd.Function):
         def forward(ctx, *flat_tensor_args):
             nonlocal compiled_fw, compiled_bw, num_outs
             if compiled_fw is None:
+                # Set input tensors that require grad to leaves
+                flat_tensor_args = pytree.tree_map(
+                    lambda x: x.detach().requires_grad_(x.requires_grad), flat_tensor_args
+                )
                 with torch.set_grad_enabled(grad_state):
                     out = flat_fn(*flat_tensor_args)
                 out = pytree.tree_map(

functorch/_src/partitioners.py

@@ -369,6 +369,9 @@ def get_node_weight(node):
         node_name = node_in[:-3]
         cut_nodes.add(node_name)
 
+    # To make this stuff deterministic
+    node_idx = {node: idx for idx, node in enumerate(joint_module.graph.nodes)}
+    saved_values = sorted((name_to_node[node] for node in cut_nodes), key=lambda x: node_idx[x])
     saved_values = [name_to_node[node] for node in cut_nodes]
 
     return _extract_fwd_bwd_modules(joint_module, saved_values)