[SPMD] Expose apply_backward_optimization_barrier

[alanwaketan]

Summary:
This PR exposes apply_backward_optimization_barrier to the spmd namespace.

Test Plan:
N.A.

[alanwaketan]

Thanks, Jack!

[bhavya01]

@alanwaketan Should we also merge this change in the 2.4 release? I saw test failure with latest 2.4 wheel

[2024-07-17, 16:47:21 UTC] {logging_mixin.py:150} WARNING - Traceback (most recent call last):
  File "/home/ml-auto-solutions/transformers/examples/pytorch/language-modeling/run_clm.py", line 873, in <module>
    main()
  File "/home/ml-auto-solutions/transformers/examples/pytorch/language-modeling/run_clm.py", line 644, in main
[2024-07-17, 16:47:21 UTC] {logging_mixin.py:150} WARNING -     xs.apply_backward_optimization_barrier(model.model.layers[i])
AttributeError: module 'torch_xla.distributed.spmd' has no attribute 'apply_backward_optimization_barrier'
    ```

[JackCaoG]

which model failed?

[bhavya01]

llama2-train-spmd

Still looking into why this passed in the previous test run

[bhavya01]

I realized that this is failing because of pytorch-tpu/transformers@ccf5b15

I replaced torch_xla.experimental.xla_sharding with torch_xla.distributed.spmd in the test and 2.4 release doesn't expose apply_optimization_barrier through the latter. The test passes locally with the fix.

[alanwaketan]

Sure, can I still backport things? I might have 2-3 PRs need to be back ported still?