# Deepspeed stage 3 partition

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fdeepspeed-stage-3-partition%2F60629%3E%5C
**Category:** DeepSpeed\
**Created:** [October 31, 2023, 11:07pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fdeepspeed-stage-3-partition%2F60629 "2023-10-31T23:07:57Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![lzy337](https://avatars.discourse-cdn.com/v4/letter/l/d2c977/32.png) [@lzy337](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Flzy337)%5C
**Post date:** [October 31, 2023, 11:07pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fdeepspeed-stage-3-partition%2F60629%2F1 "2023-10-31T23:07:57Z")

</div>

I tried to run the code from HF deepspeed integration page.

```auto

from transformers import AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM
from transformers.integrations import HfDeepSpeedConfig
import deepspeed
import os
import torch
import torch.distributed as dist

def which_process(idx) -> bool:
    """Check if the current process is the main process."""
    return not dist.is_initialized() or dist.get_rank() == idx

os.environ["TOKENIZERS_PARALLELISM"] = "false" # To avoid warnings about parallelism in tokenizers

# distributed setup
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
torch.cuda.set_device(local_rank)
deepspeed.init_distributed()

model_name = "bigscience/T0_3B"

config = AutoConfig.from_pretrained(model_name)
model_hidden_size = config.d_model

# batch size has to be divisible by world_size, but can be bigger than world_size
train_batch_size = 1 * world_size

# ds_config notes
#
# - enable bf16 if you use Ampere or higher GPU - this will run in mixed precision and will be
# faster.
#
# - for older GPUs you can enable fp16, but it'll only work for non-bf16 pretrained models - e.g.
# all official t5 models are bf16-pretrained
#
# - set offload_param.device to "none" or completely remove the `offload_param` section if you don't
# - want CPU offload
#
# - if using `offload_param` you can manually finetune stage3_param_persistence_threshold to control
# - which params should remain on gpus - the larger the value the smaller the offload size
#
# For indepth info on Deepspeed config see
# https://huggingface.co/docs/transformers/main/main_classes/deepspeed

# keeping the same format as json for consistency, except it uses lower case for true/false
# fmt: off
ds_config = {
    "fp16": {
        "enabled": False
    },
    "bf16": {
        "enabled": False
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": True,
        "contiguous_gradients": True,
        "reduce_bucket_size": model_hidden_size * model_hidden_size,
        "stage3_prefetch_bucket_size": 0.9 * model_hidden_size * model_hidden_size,
        "stage3_param_persistence_threshold": 10 * model_hidden_size
    },
    "steps_per_print": 2000,
    "train_batch_size": train_batch_size,
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": False
}
# fmt: on

# next line instructs transformers to partition the model directly over multiple gpus using
# deepspeed.zero.Init when model's `from_pretrained` method is called.
#
# **it has to be run before loading the model AutoModelForSeq2SeqLM.from_pretrained(model_name)**
#
# otherwise the model will first be loaded normally and only partitioned at forward time which is
# less efficient and when there is little CPU RAM may fail
dschf = HfDeepSpeedConfig(ds_config) # keep this object alive

# now a model can be loaded.
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# initialise Deepspeed ZeRO and store only the engine object
ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval() # inference

dist.barrier()
if which_process(0):
    for name,param in ds_engine.named_parameters():
        print(name,param.device)
        print(name,param)

    for name,param in ds_engine.module.named_parameters():
        print(name,param.device)
        print(name,param)

    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)

dist.barrier()
if which_process(1):
    for name,param in ds_engine.named_parameters():
        print(name,param.device)
        print(name,param)

    for name,param in ds_engine.module.named_parameters():
        print(name,param.device)
        print(name,param)

    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)
    print("*"*10)

```

I set debugger after .from\_pretrained and print out the parameters. As you can see they are all empty tensor.

 ![image](https://us1.discourse-cdn.com/hellohellohello/original/3X/6/a/6a7362ca41ea2f7f023e3dc8773c3b9ab1075082.png)

So my question is what itmes that actually occupy the VRAM spaces as shown by the figure below.

 ![image](https://us1.discourse-cdn.com/hellohellohello/original/3X/6/7/67eaae793a2545eb3c9ff26214752bd1d0e6f071.png)

Why they are 7k and 22k, which in total is 30k which is larger than the size of T0-3B(3 \* 4 = 12k)…  
That’s weird…

Also How could I ensure that they are actually partitioned across gpus?