Train LLaMA 2 with PEFT(LoRA) + Deepspeed Zero3 on v100 * 8, raise assert param.ds_status == ZeroParamStatus.AVAILABLE
Author: ZetangForward
Pull request: no

### System Info

Huggingface Version == 4.31.0

## Environment 

Deepspeed Zero3 Config:


```
{
    "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": 0,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 0,
        "stage3_max_reuse_distance": 0,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "fp16": {
        "enabled": true,
        "auto_cast": false,
        "loss_scale": 0,
        "initial_scale_power": 32,
        "loss_scale_window": 2000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
          "lr": "auto",
          "betas": [
            0.9,
            0.999
          ],
          "eps": 1e-8,
          "weight_decay": 0
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "gradient_accumulation_steps": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```
Launch Config:

```
deepspeed --num_gpus 1 \
    --num_nodes 1 \
    train_vqllama_lora.py \
    --model_name_or_path "/model/llama2" \
    --data_path "/data/data.pkl" \
    --output_dir ${OUTPUT_DIR} \
    --num_train_epochs 100 \
    --model_max_length 1024 \
    --per_device_train_batch_size 68 \
    --per_device_eval_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "steps" \
    --eval_steps 5 \
    --greater_is_better False \
    --save_strategy "steps" \
    --load_best_model_at_end True \
    --save_steps 5 \
    --save_total_limit 10 \
    --learning_rate 3e-5 \
    --warmup_steps 20 \
    --logging_steps 5 \
    --dataloader_num_workers 0 \
    --lr_scheduler_type "cosine" \
    --report_to "tensorboard" \
    --deepspeed configs/deepspeed/stage3_test.json \
    --fp16 True \
    --remove_unused_columns False;
```

LoRA Config:

```
myllama = CustomLLama.from_pretrained(
        model_args.model_name_or_path, 
        config=llamaconfig, 
        cache_dir=training_args.cache_dir
    )

config = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj"],
        bias="none",
        task_type="CAUSAL_LM",
        modules_to_save=["custom_embedding", "custom_head", "wte", "lm_head"]
    )

myllama = get_peft_model(myllama , config)
```

Then, I train `myllama` with Huggingface Trainer


## Errors
I come with this error 

```
Parameter Offload: Total persistent parameters: 8663041 in 198 params
  0%|                                                                             | 0/280000 [00:00<?, ?it/s]Traceback (most recent call last):
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 230, in <module>
    train()
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 224, in train
    trainer.train()
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 1809, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 2654, in training_step
    loss = self.compute_loss(model, inputs)
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 105, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1833, in forward
    loss = self.module(*inputs, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/peft_model.py", line 1073, in forward
    return self.base_model(
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/tuners/tuners_utils.py", line 103, in forward
    return self.model.forward(*args, **kwargs)
  File "/workspace/zecheng/modelzipper/projects/custom_llama/models/vqllama.py", line 84, in forward
    svg_token_embeddings = self.vqvae_embedding(svg_token_ids) # Encode svg tokens
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/utils/other.py", line 219, in forward
    return self.modules_to_save[self.active_adapter](*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1557, in _call_impl
    args_result = hook(self, args)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 392, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 505, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module, forward=prev_grad_state)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 310, in fetch_sub_module
    assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 423, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0, 'shape': (0,), 'ds_shape': (0,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {1038}, 'ds_tensor.shape': torch.Size([0])}
```

Any help for this problem ? THX


### Who can help?

@muellerz @pacman100 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. initiate the llama2 model from HF and add some extra modules for training, such as LoRA, one extra embedding / LM Head
2. use peft to wrap the model above
3. apply Deepspeed Zero3 (with my config) and HF Trainer to start training. 

### Expected behavior

Seek for help and may solve some potential bugs

Comments:
Hey, would you mind updating to the latest version of transformers, test again and also  share the output of `transformers-cli env`

> ransformers-cli env

Hi， sorry for the late reply. Sure, I update the transformers version, and here is my transformers version:

```
Python 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import transformers
>>> transformers.__version__
'4.36.2'
>>> 
```

I test again and meet the same issues
```
arameter Offload: Total persistent parameters: 177793 in 57 params
  0%|                                                                | 0/248900 [00:00<?, ?it/s]/opt/conda/envs/llama/lib/python3.10/site-packages/torch/utils/data/dataloader.py:557: UserWarning: This DataLoader will create 32 worker processes in total. Our suggested max number of worker in current system is 28, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.
  warnings.warn(_create_warning_msg(
Traceback (most recent call last):
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 229, in <module>
    train()
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 223, in train
    trainer.train()
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 1556, in train
    # number of training epochs: num_train_epochs
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 1838, in _inner_training_loop
    # Skip past any already trained steps if resuming training
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/transformers/trainer.py", line 2693, in training_step
    A helper wrapper to group together context managers.
  File "/workspace/zecheng/modelzipper/projects/custom_llama/train_vqllama_lora.py", line 105, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1833, in forward
    loss = self.module(*inputs, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/peft_model.py", line 1073, in forward
    return self.base_model(
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/tuners/tuners_utils.py", line 103, in forward
    return self.model.forward(*args, **kwargs)
  File "/workspace/zecheng/modelzipper/projects/custom_llama/models/vqllama.py", line 84, in forward
    svg_token_embeddings = self.vqvae_embedding(svg_token_ids) # Encode svg tokens
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/peft/utils/other.py", line 219, in forward
    return self.modules_to_save[self.active_adapter](*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1557, in _call_impl
    args_result = hook(self, args)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 392, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 505, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module, forward=prev_grad_state)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/llama/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 310, in fetch_sub_module
    assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 347, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0, 'shape': (0,), 'ds_shape': (0,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {846}, 'ds_tensor.shape': torch.Size([0])}

```

here is the transformers-cli env

```
╰─➤  transformers-cli env

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.36.2
- Platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.27
- Python version: 3.10.13
- Huggingface_hub version: 0.20.2
- Safetensors version: 0.4.1
- Accelerate version: 0.26.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```


same issue

@ArthurZucker any solutions?

To me it looks like deep speed 0 is not available. 
Pinging @pacman100 as he might have an idea of a quick fix! 

Ok, this is a strange issue, and I also find similar questions in many other issues... @pacman100 

> To me it looks like deep speed 0 is not available. Pinging @pacman100 as he might have an idea of a quick fix!



Hello @ZetangForward, I think this is due to  `modules_to_save` not being compatible with DeepSpeed. A workaround in the meantime would be to target the modules (add them in `target_modules`) mentioned in `modules_to_save`. If you are targeting embedding layers, the lora modules along with the embedding layers are saved when calling `save_pretrained` to support adding new tokens to the embedding layers. 

> Hello @ZetangForward, I think this is due to `modules_to_save` not being compatible with DeepSpeed. A workaround in the meantime would be to target the modules (add them in `target_modules`) mentioned in `modules_to_save`. If you are targeting embedding layers, the lora modules along with the embedding layers are saved when calling `save_pretrained` to support adding new tokens to the embedding layers.

Sorry for the late reply. Actually, I already put them in the ``target_modules`` at the beginning. Below is my code:

```

svgllama = VQSVGLlama.from_pretrained(
        model_args.model_name_or_path, 
        config=llamaconfig, 
        codebook_size=vqvae_config.vqvae.l_bins,
        cache_dir=training_args.cache_dir
    )


 config = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj"],
        bias="none",
        task_type="CAUSAL_LM",
        modules_to_save=["vqvae_embedding", "vqvae_head", "up_adapter", "down_adapter"]
    )
    svgllama = get_peft_model(svgllama, config)
```

and this is my model:

```
class VQSVGLlama(LlamaForCausalLM):  
    def __init__(self, config, vq_loss_weight=2.0, convert_token_weight=1.5, tokenizer=None, svg_begin_token_id=None, vqvae=None, codebook_size=8192):  
        super(VQSVGLlama, self).__init__(config)
        self.config = config
        self.tokenizer = tokenizer
        self.svg_begin_token_id = svg_begin_token_id
        self.vq_loss_weight = vq_loss_weight
        self.convert_token_weight = convert_token_weight
        self.codebook_size = codebook_size + 1  # add one for svg end token
        self.svg_end_token_id = codebook_size
        self.vqvae = vqvae
        self.up_adapter = nn.Linear(config.hidden_size, config.hidden_size)
        self.down_adapter = nn.Linear(config.hidden_size, config.hidden_size)
        self.vqvae_embedding = nn.Embedding(self.codebook_size, config.hidden_size)
        self.vqvae_head = nn.Linear(config.hidden_size, self.codebook_size)

        self.post_init()
        if config.frozen_llm: 
            print_c("Attention! LLM is freezed!")
            self.base_model.requires_grad_ = False 
            self.lm_head.requires_grad_ = False
            self.base_model.embed_tokens.requires_grad_ = True 

```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@ZetangForward @invisifire Have you solved this issue? I met the same situation when finetuing llama2 with lora and deepspeed zero3.