[BUG]pipeline/zero1 don't support layers less than 3 in first stage

It seems BF16/Z1/PP doesn't support only one embedding LaySpec in first stage even though the layer is very small. when two decoder layers at least are added both in first and last stages, it could run well.
deepspeed: >=0.13.3
GPU: A100-40G x8
Pipeline Parallelism: the first stage only contains a embedding weight of 152064*8192
model's main configs for debugging:

{
  "hidden_size": 8192,
  "intermediate_size": 24576,
  "num_hidden_layers": 6,
  "vocab_size": 152064
}

When I train my model with Pipeline Parallelism ,bfloat16 and ZERO Stage 1, the utilization of GPU:0 who holds pipeline's first stage and GPU:7 who holds the last stage are always 0% untill watchdog caught collective operation timeout, while other GPUs' utilization are 100%.
NCCL timeout never happened when I change the ZERO Stage to 0, but it requires more GPU memory.
Error info:

 Watchdog caught collective operation timeout: WorkNCCL(OpType=BROADCAST, Timeout(ms)=1800000) ran for 1808499 milliseconds before timing out.

There may be some bugs in bf16's support with zero1 of the pipeline parallelism

Start by reproducing the reported BF16 and ZeRO Stage 1 pipeline-parallel configuration: one embedding-only first stage, six total layers, and the provided model dimensions. Trace the pipeline-parallel and collective handling around the broadcast timeout; done means the configuration runs without the watchdog timeout, while ZeRO Stage 0 remains a comparison case.