Fix zenflow_torch_adam.py (#7544)
`_disable_dynamo_if_unsupported` fallback wasn't getting created under
certain conditions. This PR is fixing this. Also removed debug print.
Fixes issue installing deepspeed on torch 2.4 and 2.1 that triggered
this:
```
#42 15.84 Traceback (most recent call last):
#42 15.84 File "<string>", line 2, in <module>
#42 15.84 File "<pip-setuptools-caller>", line 34, in <module>
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/setup.py", line 40, in <module>
#42 15.84 from op_builder import get_default_compute_capabilities, OpBuilder
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/op_builder/__init__.py", line 18, in <module>
#42 15.84 import deepspeed.ops.op_builder # noqa: F401 # type: ignore
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/deepspeed/__init__.py", line 25, in <module>
#42 15.84 from . import ops
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/deepspeed/ops/__init__.py", line 6, in <module>
#42 15.84 from . import adam
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/deepspeed/ops/adam/__init__.py", line 9, in <module>
#42 15.84 from .zenflow_torch_adam import ZenFlowSelectiveAdamW
#42 15.84 File "/tmp/pip-install-qgzd6ybt/deepspeed_b3b4858a062d49c7b8d6ef31332a96cf/deepspeed/ops/adam/zenflow_torch_adam.py", line 685, in <module>
#42 15.84 @_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adamw)
#42 15.84 NameError: name '_disable_dynamo_if_unsupported' is not defined
#42 15.84 [WARNING] ZenFlow disabled: torch internal optimizer symbols could not be imported: cannot import name '_disable_dynamo_if_unsupported' from 'torch.optim.optimizer' (/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py)
```
---------
Signed-off-by: Stas Bekman <stas@stason.org>
