[BUG] AutoTP training runs into missing gradient error

I am running into the following error with AutoTP across all TP ranks for the first step

[rank2]: Traceback (most recent call last):
[rank2]:   File "/work1/mzhang/khoadp2/experiments/train.py", line 379, in <module>
[rank2]:     main()
[rank2]:   File "/work1/mzhang/khoadp2/experiments/train.py", line 359, in main
[rank2]:     engine.step()
[rank2]:   File "/home1/khoadp2/.local/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2414, in step
[rank2]:     self._take_model_step(lr_kwargs)
[rank2]:   File "/home1/khoadp2/.local/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2314, in _take_model_step
[rank2]:     self.optimizer.step()
[rank2]:   File "/share/sw/ai/pytorch/2.7.1/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:   File "/home1/khoadp2/.local/lib/python3.9/site-packages/deepspeed/runtime/bf16_optimizer.py", line 279, in step
[rank2]:     non_expert_groups_norm = get_global_norm_of_tensors(input_tensors=non_expert_grads_for_norm,
[rank2]:   File "/home1/khoadp2/.local/lib/python3.9/site-packages/deepspeed/runtime/utils.py", line 892, in get_global_norm_of_tensors
[rank2]:     device_total_norm = compute_buffer[0].float().detach()
[rank2]: IndexError: list index out of range

I am doing an AutoTP trainging with TP=4, DP=1, ZeRO stage 0. Heres a snippet of my training script. I am using the default clip grad value, enabling gradient checkpointing, and not using any MoE setup.

    with deepspeed.zero.Init(enabled=zero_stage == 3):
        model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16)
        if args.deepspeed_activation_checkpointing:
            model.gradient_checkpointing_enable()
        # Initialize DeepSpeed engine
        engine, _, dataloader, _ = deepspeed.initialize(
            model=model,
            training_data=train_dataset,
            collate_fn=data_collator,
            config=args.deepspeed_config
        )
    for batch in dataloader:
        loss_value = None
        try:
            gpu_batch = {}
            for k, v in batch.items():
                gpu_batch[k] = v.to(engine.device)
            outputs = engine(**gpu_batch)
            loss = outputs.loss
            loss_value = loss.item()
            engine.backward(loss)
            engine.step()
            del loss, outputs, gpu_batch

        except Exception as e:
            print(f"Error in training step {engine.global_steps}: {e}")
            raise

        if engine.global_rank == 0:
            print(f"[step {engine.global_steps}] loss = {loss_value:.4f}")

Start by reproducing the TP=4, DP=1, ZeRO stage 0 training case with gradient checkpointing and inspect deepspeed/runtime/bf16_optimizer.py at step(), then follow get_global_norm_of_tensors in deepspeed/runtime/utils.py. Done means the first training step completes without the missing-gradient IndexError under the reported AutoTP configuration.