when using huggingface pretrained model with multi-gpu, model parameters were duplicate for every gpu in ram

[linyubupa]

Bug description

when using huggingface pretrained model with multi-gpu, model parameters were duplicate for every gpu in ram

How to reproduce the bug

trainer = Trainer(
        max_epochs=1,
        devices=args.num_devices,
        precision=16,
        strategy="deepspeed_stage_3",
        accelerator='gpu',
        num_nodes=args.num_nodes,
    
    )
from transformers import (
    AdamW,
    GPTNeoForCausalLM,
    GPT2Tokenizer,
    AutoTokenizer,
    AutoModelForCausalLM,
    get_linear_schedule_with_warmup,
)
class AlpsModule(LightningModule):
    def __init__(
        self,
        model_name_or_path: str = "EleutherAI/gpt-j-6B",
        cache_dir: str ="/mntnlp/yumu/gpt-neo-x/" ,
        num_labels: int = 2,
        learning_rate: float = 5e-6,
        adam_epsilon: float = 3e-8,
        warmup_steps: int = 30,
        weight_decay: float = 0.01,
        **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters()
    

        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path
                                                        ,pad_token_id=self.tokenizer.pad_token_id
                                                        ,bos_token_id=self.tokenizer.bos_token_id
                                                        ,eos_token_id=self.tokenizer.eos_token_id
                                                        , cache_dir=cache_dir
                                                         # ,low_cpu_mem_usage=True
                                                        ).half()

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 @awaelchli

[linyubupa]

if you got multi gpu, the cost cpu memory = 2 * model_size* gpu_numbers

[awaelchli]

Hey @linyubupa

This is expected in the way you are initializing the model. I can see from the code snippet that you create the model in __init__. This isn't wrong, but for large models like yours it is inefficient. I recommend moving the initialization into this special Lightning hook:

def configure_sharded_model(self):
      self.model = AutoModelForCausalLM.from_pretrained(...)

Here is the documentation for working with deepspeed models (and also documentation for configure_sharded_model): https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#shard-model-instantly-to-reduce-initialization-time-memory

[awaelchli]

Please let me know if that helps :)

[cgd-bot]

Please let me know if that helps :)

I had the same problem, but this method didn't solve it

Hey @linyubupa

This is expected in the way you are initializing the model. I can see from the code snippet that you create the model in __init__. This isn't wrong, but for large models like yours it is inefficient. I recommend moving the initialization into this special Lightning hook:

def configure_sharded_model(self):
      self.model = AutoModelForCausalLM.from_pretrained(...)

Here is the documentation for working with deepspeed models (and also documentation for configure_sharded_model): https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#shard-model-instantly-to-reduce-initialization-time-memory

I had the same problem, but this method didn't solve it

[imraviagrawal]

Hey @linyubupa

This is expected in the way you are initializing the model. I can see from the code snippet that you create the model in __init__. This isn't wrong, but for large models like yours it is inefficient. I recommend moving the initialization into this special Lightning hook:

def configure_sharded_model(self):
      self.model = AutoModelForCausalLM.from_pretrained(...)

Here is the documentation for working with deepspeed models (and also documentation for configure_sharded_model): https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#shard-model-instantly-to-reduce-initialization-time-memory

Yeah had same issue, and above does not solve it

[linyubupa]

def configure_sharded_model(self):

sorry for late reply，I build up model in configure_sharded_model , but the cpu memory still cost amountly

[KzZheng]

Same issue. After I put the model initialization into the configure_sharded_model, I return a new error that shows the loaded parameters are trying to assign to empty tensors.