[BUG] 1bit-Adam is not compatible with ZeRO

Describe the bug
AssertionError: 1bit-Adam is not compatible with ZeRO

To Reproduce
Steps to reproduce the behavior:
Using this configuration to pass as deepspeed configuration

import evaluate
from transformers import Trainer, TrainingArguments
from transformers import (
GPTJForCausalLM,
AutoTokenizer,
default_data_collator,
)
from transformers.utils.logging import disable_progress_bar, enable_progress_bar
import torch

from ray.air import session

def trainer_init_per_worker(train_dataset, eval_dataset=None, **config):
# Use the actual number of CPUs assigned by Ray
os.environ["OMP_NUM_THREADS"] = str(
session.get_trial_resources().bundles[-1].get("CPU", 1)
)
# Enable tf32 for better performance
torch.backends.cuda.matmul.allow_tf32 = True

batch_size = config.get("batch_size", 4)
epochs = config.get("epochs", 2)
warmup_steps = config.get("warmup_steps", 0)
learning_rate = config.get("learning_rate", 0.00002)
weight_decay = config.get("weight_decay", 0.01)

deepspeed = {
    "fp16": {
        "enabled": "auto",
        "initial_scale_power": 8,
    },
    "bf16": {"enabled": "auto"},
    "optimizer": {
        "type": "OneBitAdam",
            "params": {
              "lr": "auto",
              "betas": "auto",
              "eps": "auto",
              "weight_decay": "auto",
              "freeze_step": 400,
              "cuda_aware": False,
              "comm_backend_name": "nccl"
            }
    },
    "zero_optimization": {
        "stage": 3,
        # "offload_optimizer": {
        #     "device": "cpu",
        #     "pin_memory": False,
        # },
        "overlap_comm": False,
        "contiguous_gradients": False,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "gather_16bit_weights_on_model_save": True,
        "round_robin_gradients": True,
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 10,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": False,
}

print("Preparing training arguments")
training_args = TrainingArguments(
    "output",
    per_device_train_batch_size=batch_size,
    logging_steps=1,
    save_strategy="no",
    per_device_eval_batch_size=batch_size,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    warmup_steps=warmup_steps,
    label_names=["input_ids", "attention_mask"],
    num_train_epochs=epochs,
    push_to_hub=False,
    disable_tqdm=True,  # declutter the output a little
    fp16=True,
    gradient_checkpointing=True,
    deepspeed=deepspeed,
)
disable_progress_bar()

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

print("Loading model")

model = GPTJForCausalLM.from_pretrained(model_name, use_cache=True)
model.resize_token_embeddings(len(tokenizer))

print("Model loaded")

enable_progress_bar()

# metric = evaluate.load("accuracy")
metric = evaluate.load("/domino/datasets/local/ray/metrics/accuracy/accuracy.py")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
    data_collator=default_data_collator,
)
return trainer

Expected behavior

ds_report output

DeepSpeed C++/CUDA extension op report

NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.

JIT compiled ops requires ninja
ninja .................. [OKAY]

op name ................ installed .. compatible

[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-devel package with yum
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]

DeepSpeed general environment info:
torch install path ............... ['/opt/conda/envs/domino-ray/lib/python3.8/site-packages/torch']
torch version .................... 1.13.0
deepspeed install path ........... ['/opt/conda/envs/domino-ray/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.9.2, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.7
deepspeed wheel compiled w. ...... torch 1.13, cuda 11.7

Screenshots
stacktrace

(RayTrainWorker pid=29918) /opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/train/huggingface/transformers/_transformers_utils.py:86: FutureWarning: 'format_type' is deprecated and will be removed in the next major version of datasets. Please use 'formatting=FormattingConfig(format_type=format_type)' instead.
(RayTrainWorker pid=29918)   iterable_dataset = datasets.iterable_dataset.IterableDataset(
(RayTrainWorker pid=29919) Loading model
(RayTrainWorker pid=29918) [2024-01-31 14:59:01,710] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 6.05B parameters
(RayTrainWorker pid=29918) [2024-01-31 14:59:01,710] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 6.05B parameters
(RayTrainWorker pid=29918) [2024-01-31 14:59:01,710] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 6.05B parameters
(RayTrainWorker pid=29919) Model loaded
(RayTrainWorker pid=29918) Using cuda_amp half precision backend
(RayTrainWorker pid=29919) Using cuda_amp half precision backend
(RayTrainWorker pid=29919) Using cuda_amp half precision backend
(RayTrainWorker pid=29918) [2024-01-31 14:59:35,219] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.9.2, git-hash=unknown, git-branch=unknown
(RayTrainWorker pid=29918) [2024-01-31 14:59:35,234] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
2024-01-31 14:59:35,453	ERROR tune_controller.py:873 -- Trial task failed for trial TransformersTrainer_13ddd_00000
Traceback (most recent call last):
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/air/execution/_internal/event_manager.py", line 110, in resolve_future
    result = ray.get(future)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/_private/auto_init_hook.py", line 18, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/_private/worker.py", line 2540, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AssertionError): ray::_Inner.train() (pid=29701, ip=10.68.113.118, actor_id=b334a22e8801a50bf43fcc3901000000, repr=TransformersTrainer)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/tune/trainable/trainable.py", line 389, in train
    raise skipped from exception_cause(skipped)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/train/_internal/utils.py", line 54, in check_for_failure
    ray.get(object_ref)
ray.exceptions.RayTaskError(AssertionError): ray::_RayTrainWorker__execute.get_next() (pid=29919, ip=10.68.113.118, actor_id=e507a0b69b5916b8a7df80e101000000, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x7f79cf3e6880>)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 32, in __execute
    raise skipped from exception_cause(skipped)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/train/_internal/utils.py", line 129, in discard_return_wrapper
    train_func(*args, **kwargs)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/ray/train/huggingface/transformers/transformers_trainer.py", line 482, in _huggingface_train_loop_per_worker
    trainer.train()
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/transformers/trainer.py", line 1543, in train
    return inner_training_loop(
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/transformers/trainer.py", line 1612, in _inner_training_loop
    deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/transformers/deepspeed.py", line 344, in deepspeed_init
    deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/deepspeed/__init__.py", line 165, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 308, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1162, in _configure_optimizer
    basic_optimizer = self._configure_basic_optimizer(model_parameters)
  File "/opt/conda/envs/domino-ray/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1241, in _configure_basic_optimizer
    assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO"
AssertionError: 1bit-Adam is not compatible with ZeRO

System info (please complete the following information):
deepspeed 0.9.2
python 3.8

Start at deepspeed/runtime/engine.py, especially _configure_basic_optimizer, and reproduce the assertion with the reported 1bit-Adam and ZeRO configuration. Determine the intended compatibility behavior and identify the regression coverage needed; the issue does not specify an expected result or test location.