Model parameters size is zero when using fabric.sharded_model() context with deepspeed zero-3 strategy

[seraphzl]

Bug description

Hi,

I use fabric with deepspeed zero-3 strategy to shard model among 2 gpus, and get Model params = 0.0 M of model size when using with fabric.sharded_model() context.

import lightning as L

fabric = L.Fabric(accelerator="cuda", strategy='deepspeed_stage_3', precision='bf16-mixed')
fabric.launch()

with fabric.sharded_model():
    net = mymodel()
num_params = sum([param.nelement() for param in net.parameters()])
fabric.print('Model params = %2.1f M' % (num_params / 1000**2))

Without the fabric.sharded_model() context, I get the correct model size as Model params = 13.6 M.
How to solve this issue? Thanks.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @carmocca @justusschock @awaelchli

[awaelchli]

This is the correct and expected behavior. The model gets created with emtpy uninitialized (meta device) parameters. This is what deepspeed does in its Zero3 context manager.

You can access the shapes of the parameters after the module is sharded and materialized in setup():

with fabric.sharded_model():
    net = mymodel()
    
net = fabric.setup(net)  # init and shard model
    
num_params = sum([param.nelement() for param in net.parameters()])
fabric.print('Model params = %2.1f M' % (num_params / 1000**2))

The Fabric.sharded_model has since been renamed to Fabric.init_module to avoid confusion of "where" the sharding happens.

[seraphzl]

This is the correct and expected behavior. The model gets created with emtpy uninitialized (meta device) parameters. This is what deepspeed does in its Zero3 context manager.

You can access the shapes of the parameters after the module is sharded and materialized in setup():

with fabric.sharded_model():
    net = mymodel()
    
net = fabric.setup(net)  # init and shard model
    
num_params = sum([param.nelement() for param in net.parameters()])
fabric.print('Model params = %2.1f M' % (num_params / 1000**2))

The Fabric.sharded_model has since been renamed to Fabric.init_module to avoid confusion of "where" the sharding happens.

@awaelchli
Thank you for reply.
Unfortunately, printing the model params size after fabric.setup() did not work for me, and I still get Model params = 0.0 M.
I tried to remove the Fabric.sharded_model context and test as follows.

net = mymodel()

# first print
num_params = sum([param.nelement() for param in net.parameters()])
fabric.print('Model params = %2.1f M' % (num_params / 1000**2))
fabric.print([param for param in net.parameters()][0])

# init and shard model
net = fabric.setup(net)

# second print
num_params = sum([param.nelement() for param in net.parameters()])
fabric.print('Model params = %2.1f M' % (num_params / 1000**2))
fabric.print([param for param in net.parameters()][0])

The first print show me the correct model size and param matirx as:

Model params = 15.7 M
Parameter containing:
tensor([[[-1.2726e-02, -4.4814e-03,  1.2335e-02,  2.2365e-02,  7.6349e-03,
          -6.1491e-03, -1.9465e-02, -6.7052e-03, -1.7685e-02, -2.7615e-02,
           4.1764e-02,  5.2056e-03, -1.1024e-02,  1.4891e-02,  1.5603e-02,
          -1.4920e-02,  1.5677e-02,  9.0524e-03, -3.1848e-02, -2.9157e-03,
          -2.5880e-02, -1.5773e-02,  1.0548e-02,  4.9336e-03, -8.7119e-03,
          -1.6867e-02,  2.5237e-02, -1.8115e-02,  2.7237e-04,  1.2387e-02,
          -4.9779e-02,  3.2784e-02,  3.3176e-02, -8.6861e-03, -4.0892e-02,
          -1.6123e-02, -7.9835e-03, -6.2689e-03, -2.7889e-02, -9.0203e-03,
           8.4139e-05, -3.1184e-03,  3.7463e-02,  4.0182e-02, -1.0615e-02,
          -1.5290e-02,  1.6402e-02, -3.5148e-02, -1.8098e-02, -4.9479e-02,
           1.9473e-02, -1.7437e-02, -1.4878e-02, -1.4885e-02,  1.1689e-02,
          -1.9244e-02,  3.4017e-02, -1.6559e-02, -1.3335e-02,  4.5822e-02,
           4.1282e-03,  5.0279e-02, -1.8471e-03, -4.9715e-03,  2.8953e-03,
           4.1224e-04,  9.1655e-03, -2.1277e-02,  5.5486e-03, -2.0112e-02,
          -9.0826e-03,  1.0644e-02,  2.2591e-02,  4.2523e-02,  1.4045e-02,
           6.2759e-03,  2.4613e-03,  2.3158e-02, -1.0950e-02,  7.1588e-03,
          -2.0503e-02,  3.4789e-02, -1.5309e-02, -2.3975e-02, -1.0950e-02,
          -2.2816e-03, -2.4227e-02,  1.5146e-02, -3.6122e-03,  3.9893e-03,
           5.5877e-03,  3.9308e-02, -3.3923e-03, -4.9096e-03,  1.4838e-02,
          -1.7904e-02,  3.8857e-03, -4.9425e-03, -3.2249e-02, -3.0654e-02,
          -2.4428e-02, -3.2424e-03, -1.3084e-02, -5.7562e-03,  4.6911e-03,
          -1.7116e-02, -5.3637e-02, -1.4524e-02,  8.2667e-03, -2.7705e-02,
           1.1347e-02,  1.9582e-02,  1.6069e-02, -2.8310e-02, -2.5199e-02,
          -1.8466e-02, -3.7708e-03,  3.1106e-02,  4.5919e-03,  2.3399e-02,
           4.2658e-03,  1.2363e-02,  2.6782e-03,  1.7111e-02,  8.8507e-03,
           2.0644e-02, -1.7001e-02,  7.5520e-03]]], requires_grad=True)

while the second print show me the message as:

Model params = 0.0 M
Parameter containing:
tensor([], device='cuda:0', dtype=torch.bfloat16, requires_grad=True)

Is there any method that I can access the sharded model parameters to ensure the model initialized correctly?