Zero Stage-2 Frozen Layers[BUG]

Describe the bug

I saw this issue has been resolved: https://github.com/microsoft/DeepSpeed/issues/2615 and tried freezing some layers of the model. Same training script works fine without the following.

for name, param in model.named_parameters():
    if any(ln in name for ln in ["embed", "lm_head"]):
        param.requires_grad = True
    else:
        param.requires_grad = False

Traceback (most recent call last):
  File "pretrain_deepspeed.py", line 947, in <module>
    main()
  File "pretrain_deepspeed.py", line 456, in main
    model, optimizer, train_dl, valid_dl, lr_scheduler = accelerator.prepare(
  File "/usr/local/lib/python3.8/dist-packages/accelerate/accelerator.py", line 1198, in prepare
    result = self._prepare_deepspeed(*args)
  File "/usr/local/lib/python3.8/dist-packages/accelerate/accelerator.py", line 1537, in _prepare_deepspeed
    engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/__init__.py", line 171, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 310, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 1209, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 1444, in _configure_zero_optimizer
    optimizer = DeepSpeedZeroOptimizer(
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 312, in __init__
    self.flatten_dense_tensors_aligned(
  File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 834, in flatten_dense_tensors_aligned
    return self.flatten(align_dense_tensors(tensor_list, alignment))
  File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 451, in _flatten_dense_tensors
    return torch._C._nn.flatten_dense_tensors(tensors)
RuntimeError: torch.cat(): expected a non-empty list of Tensors

{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "none",
            "pin_memory": true
        },
        "offload_param": {
            "device": "none",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 21,
    "gradient_clipping": 1.0,
    "steps_per_print": 1000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

This is a work around but not favorable because of unnecessary resource usage:

    if args.freeze_all_but_embed:        
        no_decay = ["bias", "norm.weight"]
        freeze_but = ["embed", "lm_head"]
        optimizer_grouped_parameters = [
            # embedding and lm_head layers.
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in freeze_but)],
                "weight_decay": args.weight_decay,
                "lr": args.learning_rate
            },
            # all layers excluding embedding, lm_head and layer/rmsnorm.
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay+freeze_but)],
                "weight_decay": 0.0, # args.weight_decay
                "lr": 0.0
            },
            # layer/rmsnorm.
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
                "lr": 0.0
            }
        ]