Revert "Modularize zero step function and make it customizable"

[JackCaoG]

Reverts #7233

break CI

======================================================================
ERROR: test_zero1 (__mp_main__.XlaZeRO1Test)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/__w/xla/xla/pytorch/xla/test/test_zero1.py", line 59, in test_zero1
    opt2.step()
  File "/usr/local/lib/python3.10/site-packages/torch/optim/optimizer.py", line 484, in wrapper
    out = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch_xla/distributed/zero_redundancy_optimizer.py", line 433, in step
    self._update_parameters(**kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch_xla/distributed/zero_redundancy_optimizer.py", line 410, in _update_parameters
    self.allgather_weights_and_update_full_parameter()
  File "/usr/local/lib/python3.10/site-packages/torch_xla/distributed/zero_redundancy_optimizer.py", line 438, in allgather_weights_and_update_full_parameter
    sharding_scheme = self._get_sharding_scheme(**kwargs)
NameError: name 'kwargs' is not defined

----------------------------------------------------------------------

Please create from branch directly so CI can be run. I can give write access if needed.

[ManfeiBai]

Thanks, LGTM

[JackCaoG]

cc @jeffhataws @karakusc