AttributeError: 'LlamaForCausalLM' object has no attribute 'save_checkpoint'
Author: dittops
Pull request: no
Labels: name: solved

### System Info

I'm training a llama model. When I try using Deepspeed along with the trainer, I get the below error

<img width="695" alt="image" src="https://github.com/huggingface/transformers/assets/12937285/3cd5a3c1-34c9-46b3-93e0-e5dd454b8a3b">

Training script

```
tokenizer = LlamaTokenizer.from_pretrained(model_args.base_model)
    tokenizer.pad_token_id = 0

    model = LlamaForCausalLM.from_pretrained(
        model_args.base_model,
        device_map="auto"
    )

    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("total params: ", total_params)


    if data_args.data_path.endswith(".json") or data_args.data_path.endswith(".jsonl"):
        data = load_dataset("json", data_files=data_args.data_path)
    else:
        data = load_dataset(data_args.data_path)
    
    dataset =  load_data(tokenizer, data['train'], data_args.max_seq_length)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)


    trainer = Trainer(
        model=model, 
        tokenizer=tokenizer, 
        args=training_args, 
        data_collator=data_collator,
        **dataset
    )

    model.config.use_cache = False
    trainer.train()
    trainer.save_model()
```

I'm using deepspeed stage 2

```
{
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "zero_allow_untested_optimizer": true,
    "fp16": {
      "enabled": "auto",
      "loss_scale": 0,
      "initial_scale_power": 16,
      "loss_scale_window": 1000,
      "hysteresis": 2,
      "min_loss_scale": 1
    },  
    "zero_optimization": {
      "stage": 2,
      "allgather_partitions": true,
      "allgather_bucket_size": 5e8,
      "reduce_scatter": true,
      "reduce_bucket_size": 5e8,
      "overlap_comm": false,
      "contiguous_gradients": true
    }
  }

```

If I use FSDP instead of deepspeed, the checkpoints are getting saved.

transformers==4.34.0.dev0 (I tried with transformers==4.33.0 as well)
deepspeed==0.10.3

### Who can help?

@pacman100 

### Information

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

### Tasks

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

### Reproduction

Command used for launching the script

`CUDA_VISIBLE_DEVICES=0,3 python train.py --data_path /path/to/data.json --output_dir output  --per_device_train_batch_size 4 --num_train_epochs 2 --learning_rate 2e-5 --fp16 True --logging_steps 10 --save_strategy steps --save_steps 20 --save_total_limit 5 --adam_beta2 0.95 --deepspeed ds_config.json`
``

### Expected behavior

The model gets saved without any errors.

Comments:
The model is saving fine while changing to stage 3

> CUDA_VISIBLE_DEVICES=0,3 python train.py --data_path /path/to/data.json --output_dir output  --per_device_train_batch_size 4 --num_train_epochs 2 --learning_rate 2e-5 --fp16 True --logging_steps 10 --save_strategy steps --save_steps 20 --save_total_limit 5 --adam_beta2 0.95 --deepspeed ds_config.json


The command used for launching is incorrect. Please use `torchrun`, `deepspeed` or `accelerate launch` when running in distributed setup. Please go through the docs here: https://huggingface.co/docs/transformers/main_classes/deepspeed#deployment-with-multiple-gpus

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.

Hi, I encounter the same problem. I cannot use ZERO3 because mpi4py cannot be installed. I cannot use torchrun, deepspeed or accelerate launch because these command will cause OOM error (they seems to fit one model in one GPU). 
Is there any method to use deepspeed ZERO 2 and save the model smoothly?

Hi @jyhong836, could you open a new issue? This will help us track what is and isn't resolved and your issue seems somewhat different, specifically the issue with mpi4py. 

I'd suggest asking in our [forums](https://discuss.huggingface.co/). As github issues are primarily for feature requests and bug reports for transformers