deepspeed stage 3 memory cost

The environment is: transformers==4.33.1, torch==2.0.1, deepspeed==0.10.1

I am running the following code:

from transformers.deepspeed import HfDeepSpeedConfig
import deepspeed
from transformers import AutoModelForCausalLM, AutoConfig
import torch
import time
import os
import json
from torch.optim import AdamW

world_size = int(os.getenv("WORLD_SIZE", '1'))
with open("deepspeed_config.json", 'r', encoding='utf-8') as f:
    deepspeed_config = json.load(f)
if deepspeed_config["zero_optimization"]["stage"] == 3:
    deepspeed_config["zero_optimization"]['mics_shard_size'] = world_size
    if world_size > 8:
        deepspeed_config["zero_optimization"]['mics_hierarchical_params_gather'] = True

if deepspeed_config["zero_optimization"]["stage"] == 3:
    dschf = HfDeepSpeedConfig(deepspeed_config)
    with deepspeed.zero.MiCS_Init(config_dict_or_path=deepspeed_config):
        model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.float16)

time.sleep(300)

The deepspeed_config.json file is as follows:

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 4,
   "wall_clock_breakdown": false,
  "gradient_clipping": 1.0,
  "steps_per_print": 100,
  "fp16": {
    "enabled": true,
    "auto_cast": "auto",
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "amp": {
    "enabled": false,
    "opt_level": "O2"
  },
  "bfloat16": {
    "enabled": false
  },
 "zero_optimization": {
        "stage": 3,
       "allgather_partitions": false,
        "allgather_bucket_size": 5e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true,
        "sub_group_size": "auto",
        "stage3_prefetch_bucket_size": 5e8
        "stage3_param_persistence_threshold": 1e6,
        "stage3_max_live_parameters": 5e8,
        "stage3_max_reuse_distance": 5e8,
        "mics_shard_size": 8,
        "mics_hierarchical_params_gather": false,
        "stage3_gather_fp16_weights_on_model_save": true,
        "offload_optimizer": {
          "device": "none"
        },
        "offload_param":{
          "device": "none"
        }
    },
  "zero_allow_untested_optimizer": true,
  "data_efficiency": {
    "enabled": true,
    "seed": 42
  },
  "data_sampling": {
    "enabled": true,
    "num_workers": 8
  }
}

And I found that the memory cost is much larger than the model size (14GB / 4 = 3.5G):

What is the extra 26G comes from (communication should not be that costly)? and how can I reduce it?

I change the communication memory (e.g., stage3_max_live_parameters, stage3_prefetch_bucket_size) from 5e8 to 3e7, and the actual cost does not change.

This results in OOM when calling deepspeed.initialize() [16 40G 33B-model]

deepspeed.initialize() code:

def getOptimizerGroup(model, weight_decay=0.0):
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [
                p for n, p in model.named_parameters()
                if (not any(nd in n
                            for nd in no_decay) and p.requires_grad)
            ],
            "weight_decay":
            weight_decay,
        },
        {
            "params": [
                p for n, p in model.named_parameters()
                if (any(nd in n
                        for nd in no_decay) and p.requires_grad)
            ],
            "weight_decay":
            0.0,
        },
    ]
    
    return optimizer_grouped_parameters

optimizer_parameters = getOptimizerGroup(model)
optimizer = AdamW(optimizer_parameters, lr=2e-5, betas=[0.9, 0.95])

model, optimizer, _, _ = deepspeed.initlialize(
    model=model,
    optimizer=optimizer,
    config_params=deepspeed_config,
    dist_init_required=True,
)

Is my entire process for training setup correct when using stage 3?

Reproduce the reported setup from deepspeed_config.json with transformers==4.33.1, torch==2.0.1, and deepspeed==0.10.1, then inspect memory behavior around HfDeepSpeedConfig, deepspeed.zero.MiCS_Init, and deepspeed.initialize. Compare allocations with the configured stage-3 and MiCS settings; done means identifying the source of the extra memory and whether the training setup is valid or needs a documented configuration change.