Disable additional transforms for the PyTorch Autograd path

[IvanYashchuk]

Applying transforms on the trace after "forward backward split + transform_for_execution" is not a well-tested path and breaks the careful sorting of operations introduced inside the split_forward_backward function. In the future, we need to revisit the order of transformations and move execution transforms out of the split_forward_backward function. Today this PR is needed to restore distributed communication overlap with computation and unblock fixing TransformerEngine+DataParallel.

[t-vi]

Thank you @IvanYashchuk

[t-vi]

So I generally agree with the patch, but the CI has opinions.

[IvanYashchuk]

Notebooks fail with:

     18 cache_rec, i_, _ = thunder.compile_data(thunder_model).get_computation_and_inputs(x)
---> 19 computation_trace = cache_rec.computation_traces[0]

IndexError: list index out of range

[IvanYashchuk]

Next error in the notebooks:

nbclient.exceptions.CellExecutionError: An error occurred while executing the following cell:
------------------
### DON'T TRY THIS AT HOME
computation_trace.bound_symbols[2].sym = cache_rec.computation_traces[0].bound_symbols[2].subsymbols[0].sym
if cache_rec.computation_traces[0].bound_symbols[3].subsymbols:
    computation_trace.bound_symbols[3] = cache_rec.computation_traces[0].bound_symbols[3].subsymbols[0]
computation_trace.bound_symbols[4].sym = cache_rec.computation_traces[0].bound_symbols[4].subsymbols[0].sym

wrap_as_highlighted_code(computation_trace)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB