[BUG]deepspeed always load the whole model to each gpu, then OOM

i use huggingface trl sfttrainer and peft and deepspeed to train a 6B model.
i have a 4 12GB gpu.

when i use automodel.from_pretrained(device_map="auto"), it works. but the training process is so slow and gpu only has 25% working.

so i try to use deepspeed, but when using deepspeed device_map not work, so i deleted this.

i copy a stage 2 config from huggingface tutorial.
{ "fp16": { "enabled": true, "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 16, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": "auto", "weight_decay": "auto", "torch_adam": true, "adam_w_mode": true } }, "scheduler": { "type": "WarmupDecayLR", "params": { "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto", "total_num_steps": "auto" } }, "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 2e8, "overlap_comm": true, "reduce_scatter": true, "reduce_bucket_size": "auto", "contiguous_gradients": true }, "gradient_accumulation_steps": 1, "gradient_clipping": "auto", "steps_per_print": 2000, "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "wall_clock_breakdown": false }
then try to run with deepspeed.

i suppose that deepspeed should work with shard model with model that can not fit the gpu.

but deepspeed always try to load the whole model on each gpu, it can not shard the model to load. and always return oom.

i suspect that even when you try to train bigger model such as 130B, there is no gpu vram can fit the whole model, you still have to shard the model.

so, it should work , but it does not.

can you tell me why? or is there any thing i need to know

Start with the DeepSpeed ZeRO stage 2 configuration and the model-loading path used with AutoModel.from_pretrained(device_map="auto"). Compare how model placement differs between device_map and DeepSpeed training, using the reported 6B model and four 12GB GPUs as the reproduction setup. Done would require a documented explanation or a supported loading path that avoids replicating the whole model on every GPU.