# 11B model gets OOM after using deepspeed zero 3 setting with 8 32G V100 **URL:** \ **Category:** 🤗Accelerate\ **Created:** [April 8, 2024, 5:30pm UTC](https://discuss.huggingface.co/t/11b-model-gets-oom-after-using-deepspeed-zero-3-setting-with-8-32g-v100/80748 "2024-04-08T17:30:02Z")\ **Posts on this page:** 3\ **Page:** 1
**Author:** ![Charlie911](https://avatars.discourse-cdn.com/v4/letter/c/7cd45c/32.png) [@Charlie911](https://discuss.huggingface.co/u/Charlie911)\ **Post date:** [April 8, 2024, 5:30pm UTC](https://discuss.huggingface.co/t/11b-model-gets-oom-after-using-deepspeed-zero-3-setting-with-8-32g-v100/80748/1 "2024-04-08T17:30:02Z")
Can anyone help me with one-node distributed training? The idea of my custom model is to merge some layers of 2 7B llama models like MOE models do. So that yields an 11B model. I’m loading the model in `torch.float16` except for the lm\_head weight which is in `torch.float32`. I’m assuming this would require at least 11\*2\*2 = 44 GB of GPU RAM for training. I used `estimate_zero3_model_states_mem_needs_all_live` to check out the zero3 setting recommended and created an accelerate config file accordingly. ![B894BABD-FC66-40A0-8CE8-32A14D48B344](https://us1.discourse-cdn.com/hellohellohello/original/3X/5/b/5b2f85f3057c96acd32ff42831cadea9bb4bc203.jpeg) The model can be successfully loaded (and I’m guessing into 8 GPUs evenly cause their used RAMs are all somewhere around 6000 MiB). However, once the forward call starts, the code fails due to OOM error. BTW, 8 GPUs seem a bit too extravagant for a model of this size, but I got no luck with 2 & 4 GPUs either. This is my `config.yaml` `In which compute environment are you running? This machine ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------Which type of machine are you using? multi-GPU How many different machines will you use (use more than 1 for multi-node training)? [1]: 1 Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: NO Do you wish to optimize your script with torch dynamo?[yes/NO]:NO Do you want to use DeepSpeed? [yes/NO]: yes Do you want to specify a json file to a DeepSpeed config? [yes/NO]: NO ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------What should be your DeepSpeed's ZeRO optimization stage? 3 ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------Where to offload optimizer states? none ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------Where to offload parameters? none How many gradient accumulation steps you're passing in your script? [1]: 1 Do you want to use gradient clipping? [yes/NO]: NO Do you want to save 16-bit model weights when using ZeRO Stage-3? [yes/NO]: yes Do you want to enable `deepspeed.zero.Init`when using ZeRO Stage-3 for constructing massive models? [yes/NO]: NO How many GPU(s) should be used for distributed training? [1]:8 ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------Do you wish to use FP16 or BF16 (mixed precision)? fp16 ` This is my training script. I implemented my custom model by overwriting and inheriting classes in transformers library. The 2 models to be merged would be loaded using `LlamaForCausalLM.from_pretrained()` in `MoeLlamaForCausalLM` class. I would then copy the weights into my model layer by layer. ```auto model = MoeLlamaForCausalLM(config, trainable=True) for name, param in model.named_parameters(): if param.ndim == 1: # cast the small parameters (e.g. layernorm) to fp32 for stability param.data = param.data.to(torch.float32) if "lora" in name: param.requires_grad = True else: param.requires_grad = False tokenizer = AutoTokenizer.from_pretrained(...) tokenizer.pad_token = tokenizer.eos_token dataset = load_dataset(...) training_config = { "learning_rate": 5e-5, "gradient_accumulation_steps": 1, "batch_size": 2, "num_epochs": 1, "save_step": 100, "save_dir": "./accelerate_output", "lr_scheduler_type": "linear", "num_warmup_steps": 0 } accelerator = Accelerator(mixed_precision='fp16') if accelerator.is_main_process: datasets.utils.logging.set_verbosity_warning() transformers.utils.logging.set_verbosity_info() else: datasets.utils.logging.set_verbosity_error() transformers.utils.logging.set_verbosity_error() # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer optimizer_cls = ( torch.optim.AdamW if accelerator.state.deepspeed_plugin is None or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config else DummyOptim ) optimizer = optimizer_cls(model.parameters(), lr=training_config["learning_rate"]) train_dataloader = DataLoader(dataset, shuffle=True, batch_size=training_config["batch_size"]) num_training_steps = training_config["num_epochs"] * len(train_dataloader) # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler if ( accelerator.state.deepspeed_plugin is None or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config ): lr_scheduler = get_scheduler( training_config["lr_scheduler_type"], optimizer=optimizer, num_warmup_steps=training_config["num_warmup_steps"], num_training_steps=num_training_steps, ) else: lr_scheduler = DummyScheduler( optimizer, total_num_steps=num_training_steps, warmup_num_steps=training_config["num_warmup_steps"] ) train_dataloader, model, optimizer, lr_scheduler = accelerator.prepare( train_dataloader, model, optimizer, lr_scheduler ) progress_bar = tqdm(range(training_config["num_epochs"] * len(train_dataloader)), disable=not accelerator.is_main_process) step = 0 model.train() for epoch in range(training_config["num_epochs"]): for batch in train_dataloader: optimizer.zero_grad() outputs = model(**batch) loss = outputs.loss accelerator.backward(loss) optimizer.step() lr_scheduler.step() accelerator.step() progress_bar.update(1) step += 1 if accelerator.is_main_process and step % training_config["save_step"] == 0: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) unwrapped_model.save_pretrained( training_config["save_dir"], save_function=accelerator.save, state_dict=accelerator.get_state_dict(model), is_main_process=accelerator.is_main_process, ) ``` Thanks in advance for any advice! ---
**Author:** ![lixiqqq123](https://avatars.discourse-cdn.com/v4/letter/l/e480ec/32.png) [@lixiqqq123](https://discuss.huggingface.co/u/lixiqqq123)\ **Post date:** [April 26, 2025, 12:35pm UTC](https://discuss.huggingface.co/t/11b-model-gets-oom-after-using-deepspeed-zero-3-setting-with-8-32g-v100/80748/2 "2025-04-26T12:35:21Z")
Have you solved the problem? I have encountered a similar problem recently. I have an 8B model here. On 8xA100, no matter stage1, 2, 3, it will oom ---
**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\ **Post date:** [April 26, 2025, 1:12pm UTC](https://discuss.huggingface.co/t/11b-model-gets-oom-after-using-deepspeed-zero-3-setting-with-8-32g-v100/80748/3 "2025-04-26T13:12:33Z")
Recent similar issues also… > [@OutOfMemoryError: CUDA out of memory(LLM) Tuning](https://discuss.huggingface.co/t/outofmemoryerror-cuda-out-of-memory-llm-tuning/151672): > > I am currently working on large language model (LLM) and frequently encounter the error: “OutOfMemoryError: CUDA out of memory.” Despite trying various solutions from multiple sources, I have been unable to resolve the issue. Could you please anyone can guide me to resolve this? TIA > [@2B Model Fill Up Memory Usage on 4xA100s](https://discuss.huggingface.co/t/2b-model-fill-up-memory-usage-on-4xa100s/149680): > > With 4xA100s, I am training a 2B model (Gemma-2-2b-it) to perform SequenceClassification on a custom dataset. My code is quite boilerplate, but what I don’t understand is why would model training fill up the GPU memory? As I see it, With a model like gemma-2b, the total number of parameters is 2B. If we assume full precision training, i.e. float32, then the total GPU memory required to hold the model is 2B\*(32/8=4 bytes/parameter) = 8GB of memory. When the model undergoes training, there is also…