Repository: huggingface/accelerate

Issue:
AttributeErrorAttributeError: : 'DeepSpeedEngineWrapper' object has no attribute '_zero3_consolidated_fp16_state_dict'
  Accelerator with Deepspeed stage 3 when saving the model giving following error
 Code I used:

      import torch
      import torch.nn.functional as F
      from datasets import load_dataset
      from accelerate import Accelerator, DeepSpeedPlugin
      from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
      from data import CustomDataset
      import pandas as pd
      from tqdm import tqdm
      from transformers import AdamW
      set_seed(42)
      
      
      BATCH_SIZE = 4
      
      deepspeed_plugin = DeepSpeedPlugin(zero_stage=3, gradient_accumulation_steps=4, offload_optimizer_device='cpu')
      accelerator = Accelerator(fp16=True, deepspeed_plugin=deepspeed_plugin)
      #accelerator = Accelerator()
      
      model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m") #, torch_dtype=torch.float16)
      tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m", use_fast=False)
      
      
      df = pd.read_excel("dataset.xlsx")
      dataset = CustomDataset(df, tokenizer, max_len=512)
      data = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True,  num_workers=5)
      
      optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001) #, correct_bias=True)
      
      model, optimizer, data = accelerator.prepare(model, optimizer, data)
      device = accelerator.device
      accelerator.print("Device: ", device)
      
      gradient_accum_steps = 4
      
      model.train()
      for epoch in range(10):
          
          running_loss = 0
          for batch_idx, batch in enumerate(data):
              
              input_ids = batch["Input_ids"].to(device)
              attention_mask = batch["attentions"].to(device)
              labels = batch["Input_ids"].to(device)
              
              #print("Input_ids shape", input_ids.shape)
              
              optimizer.zero_grad()
      
              output = model(input_ids,
                              attention_mask=attention_mask,
                              labels=labels
                            )
              
              #wandb.log({"Loss", output['loss'].item()})
              
              accelerator.print(f"Loss - {round(output['loss'].item(), 3)}, Batch_id - {batch_idx}, Epoch - {epoch}")
              
              
              accelerator.backward(output['loss']/gradient_accum_steps)
              
              if (batch_idx+1) % gradient_accum_steps == 0:
                  optimizer.step()
              
              torch.cuda.empty_cache()
              
              running_loss += output['loss'].item()
              
          avg_loss = running_loss / ((batch_idx+1) * BATCH_SIZE)
              
          accelerator.print("Avg_loss:", avg_loss)
          if avg_loss < 100 :
              accelerator.wait_for_everyone()
              unwrapped_model = accelerator.unwrap_model(model)
              save_dir = f"MODEL/{epoch}/"
              print(f"Saving Model to {save_dir}")
              unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
              
        
The above code is giving following error:


             Traceback (most recent call last):
              File "untitled.py", line 79, in <module>
            Traceback (most recent call last):
            Traceback (most recent call last):
              File "untitled.py", line 79, in <module>
              File "untitled.py", line 79, in <module>
                unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
              File "/opt/conda/envs/t5/lib/python3.8/site-packages/accelerate/accelerator.py", line 884, in get_state_dict
                    unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
          
            File "/opt/conda/envs/t5/lib/python3.8/site-packages/accelerate/accelerator.py", line 884, in get_state_dict
            File "/opt/conda/envs/t5/lib/python3.8/site-packages/accelerate/accelerator.py", line 884, in get_state_dict
              state_dict = model._zero3_consolidated_fp16_state_dict()
            File "/opt/conda/envs/t5/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1177, in __getattr__
              state_dict = model._zero3_consolidated_fp16_state_dict()    
          state_dict = model._zero3_consolidated_fp16_state_dict()  File "/opt/conda/envs/t5/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1177, in __getattr__
          
            File "/opt/conda/envs/t5/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1177, in __getattr__
              raise AttributeError("'{}' object has no attribute '{}'".format(
          AttributeError: 'DeepSpeedEngineWrapper' object has no attribute '_zero3_consolidated_fp16_state_dict'
                  raise AttributeError("'{}' object has no attribute '{}'".format(raise AttributeError("'{}' object has no attribute '{}'".format(
          
          AttributeErrorAttributeError: : 'DeepSpeedEngineWrapper' object has no attribute '_zero3_consolidated_fp16_state_dict'


Can someone help?


Discussion:

Resolving patch:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -881,14 +881,15 @@ def get_state_dict(self, model):
                 is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
 
         if is_zero_3:
-            state_dict = model._zero3_consolidated_fp16_state_dict()
+            state_dict = model._zero3_consolidated_16bit_state_dict()
         else:
             model = self.unwrap_model(model)
             state_dict = model.state_dict()
 
-        for k in state_dict:
-            if state_dict[k].dtype == torch.float16:
-                state_dict[k] = state_dict[k].float()
+        if state_dict is not None:
+            for k in state_dict:
+                if state_dict[k].dtype == torch.float16:
+                    state_dict[k] = state_dict[k].float()
 
         return state_dict
 


Test patch: