On Memory Consumption for Prefix Tuning

I'm trying to prefix-tune kakaobrain/kogpt 6B model for a topic-specific text generation task. Meanwhile, I was quite curious on how to compute runtime requirements for gRAMs on this prefix tuning task.

kogpt model consists of 6B parameters which means 6B * 2(half-precision) / 1024**3 ~= 11.17GB.
According to description in ZeRO paper, Adam optimizer typically requires 6 * model_params for optimizer states. This results in 11.17 * 6 = 67.02GB for this specific model. Thus finetuning KoGPT model requires 11.17 + 11.17(grads) + 67.02 = 89.36GB gRAM.

However, to prefix-tune this model, only the prefix encoder needs training. In my case, 956575744 parameters ~= 1.78GB needs training. Is it correct that I need 1.78 * 8 + 11.17 = 25.96 GB for training this specific model?

I used Microsoft/DeepSpeed to benefit from cpu offload and ZeRO optimization. My Implementation is roughly

model_engine, optimizer, _, scheduler = deepspeed.initialize(
        args=cmd_args,
        model=model,
        model_parameters=model.prefix_encoder.parameters(),
        config=cmd_args.deepspeed_config,
    )
# model is composed of only two parts : kogpt and prefix_encoder

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001,
            "betas": [
                0.8,
                0.999
            ],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 0.001,
            "warmup_num_steps": 1000
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": 5e8,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_fp16_weights_on_model_save": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": 1,
    "steps_per_print": 2000,
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": false
}

Though I passed only encoder parameters to the initializer, I got CUDA runtime Error on my 2 RTX 3090 gpus (each 24GB gRAM) and 64GB cpu RAM. Since I set zero 3 optimizer, I guess models and batches are split across two gpus and each gpu takes 13GB. But I got OOM error at the time of initialization. I'm having hard time finding out why. If you can give me slight hints that would be much appreciated. Thanks in advance!

Here is my full error trace in case you need it.

Start at deepspeed.initialize with model_parameters=model.prefix_encoder.parameters() and the shown ZeRO-3 configuration; reproduce initialization on two RTX 3090 GPUs using the supplied model and error trace. Determine why initialization exceeds available GPU memory, then document the correct memory accounting and configuration outcome.