[BUG] Pipeline Parallel: Only few trainable parameters, Only GPU0 has parameters whose requires_grad is "True". Cause ValueError: optimizer got an empty parameter list when deepspeed.initialize

Describe the bug
I want to use Pipeline Parallel to train a MLLM like MiniGPT4, only training the projection layer. Most parameters of the model are frozen, so I set the partition_method to 'uniform'.

model = PipelineModule(layers=get_model(model), num_stages=2, partition_method='uniform')

However, only a nn.Linear is trainable, lead to the fact that only GPU0 has trainable parameters. GPU1, GPU2, ...., GPUn don't have parameters whose requires_grad are True.

So raise the error when run deepspeed.initialize like following:

engine, _, _, _ = deepspeed.initialize(model=model, config=OmegaConf.to_container(ds_cfg), model_parameters=[p for p in model.parameters() if p.requires_grad])

ValueError: optimizer got an empty parameter list

AttributeError: 'DeepSpeedCPUAdam' object has no attribute 'ds_opt_adam'

To Reproduce
Steps to reproduce the behavior:

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Expected behavior
A clear and concise description of what you expected to happen.

ds_report output
Please run ds_report to give us details about your setup.

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [e.g. two machines with x8 A100s each]
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version
• Any other relevant info about your setup

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

### Tasks

Start with the PipelineModule configuration using partition_method='uniform' and the deepspeed.initialize call that filters parameters by requires_grad. Reproduce the case where only one nn.Linear is trainable, then determine the expected behavior for pipeline stages with no trainable parameters and verify that initialization no longer fails with an empty parameter list.