group get_grads together for torch compile fusions

thunder/tests/test_torch_compile_executor.py

@@ -85,7 +85,7 @@ def test_torch_compile_cat_rope_single_fusion():
 
     backward_execution_trace = thunder.last_backward_traces(jfn)[-1]
     assert len(get_fusions(backward_execution_trace)) == 1
-    assert len(backward_execution_trace.bound_symbols) == 14
+    assert len(backward_execution_trace.bound_symbols) == 17
 
 
 @pytest.mark.skipif(not is_inductor_supported() or platform.system() == "Windows", reason="inductor unsupported")

thunder/transforms/autodiff.py

@@ -302,6 +302,9 @@ def shallow_copy_if_input(p):
     trace, _ = AugmentedForwardProcessor(trace)()
     # run through DCE in case some of the gradients of intermediates are not needed.
     trace = thunder.core.transform_common.dce(trace)
+    # group get_grad symbols together for torch compile fusions
+    # !!! is it preferrable to do this here or in the torch compile fusion pass?
+    _group_get_grad_bsyms(trace)
 
     end_time_ns = time.perf_counter_ns()
     elapsed_time_ns = end_time_ns - start_time_ns
@@ -312,6 +315,18 @@ def shallow_copy_if_input(p):
     return trace
 
 
+def _group_get_grad_bsyms(trace):
+    i = 0
+    while trace.bound_symbols[i].sym != prims.get_grad:
+        i += 1
+    if i == len(trace.bound_symbols):
+        return
+    get_grad_bsyms = list(filter(lambda bsym: bsym.sym == prims.get_grad, trace.bound_symbols))
+    bsyms = list(filter(lambda bsym: bsym.sym != prims.get_grad, trace.bound_symbols))
+    bsyms = bsyms[:i] + list(get_grad_bsyms) + bsyms[i:]
+    trace.bound_symbols = bsyms
+
+
 def split_into_forward_and_backward(joint_trace):
     """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""