Replace xm.mark_step with torch_xla.sync() wherever possible

[ghpvnist]

Fixes #8862.

[tengyifei]

LGTM modulo one nit on the comment.

I think we have some tests that look for strings (yuck!) and there's a subtle string mismatch.

[tengyifei]

Could you elaborate what this means? By reading the code, I wasn't sure what does "resetting" mean. Does it mean any tracing scope set by the user for the profiler are invalidated? Maybe could you run some simple tests to verify its behavior?

[yaoshiang]

This comment shouldn't hold up this PR, but I want to register that I think that we have an opportunity to rename this again in the future to make it easier for users to understand what this does without reading documentation. The problem with sync is that it is too close to torch.cuda.synchronize, which has a different meaning. So torch_xla.sync is going to be confusing the majority of our users.

What I think we should do is have a

torch_xla.synchronize which ONLY waits for all results from the underlying devices to return, similar to torch.cuda.synchronize.
torch_xla.compile(), which will match the norms of torch developers, who expect a possible graph break at the end of a decorated function.
torch_xla._barrier(), which will become a documented but not recommended usage to indicate that lazytensor should break the graph.

There are unresolved design questions:
What happens if a synchronize is inside a compile()? We'd have to look to how this works in CUDA to try to make this parallel.

Next torch.compile vs torch_xla.compile()

torch.cuda.synchronize(device=None)[SOURCE][SOURCE]
Wait for all kernels in all streams on a CUDA device to complete.

Parameters
device (torch.device or int, optional) – device for which to synchronize. It uses the current device, given by current_device(), if device is None (default).

[bhavya01]

This comment shouldn't hold up this PR, but I want to register that I think that we have an opportunity to rename this again in the future to make it easier for users to understand what this does without reading documentation. The problem with sync is that it is too close to torch.cuda.synchronize, which has a different meaning. So torch_xla.sync is going to be confusing the majority of our users.

What I think we should do is have a

torch_xla.synchronize which ONLY waits for all results from the underlying devices to return, similar to torch.cuda.synchronize. torch_xla.compile(), which will match the norms of torch developers, who expect a possible graph break at the end of a decorated function. torch_xla._barrier(), which will become a documented but not recommended usage to indicate that lazytensor should break the graph.

There are unresolved design questions: What happens if a synchronize is inside a compile()? We'd have to look to how this works in CUDA to try to make this parallel.

Next torch.compile vs torch_xla.compile()

torch.cuda.synchronize(device=None)[SOURCE][SOURCE] Wait for all kernels in all streams on a CUDA device to complete.

Parameters device (torch.device or int, optional) – device for which to synchronize. It uses the current device, given by current_device(), if device is None (default).

Agree that torch_xla.sync() is not the best naming. Should we explicitly try to educate the users by using something like compile_all_lazy_graphs()? torch_xla.barrier is also a good option. My only concern is that a barrier is generally used in the context of threads and processes and we already have APIs like apply_backward_optimization_barrier. Just wanted to avoid the term from being too overloaded.

[tengyifei]

IMO it's very worth having a separate discussion for renaming torch_xla.sync(), but outside of this PR. I agree that sync is also confusing.