Optimizer error when resuming from checkpoint with adafactor
Author: rumourscape
Pull request: no
Labels: name: trainer

name: bug

name: optimization

### System Info

- `transformers` version: 4.45.0.dev0
- Platform: Linux-5.15.0-1062-nvidia-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.24.7
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 16
        - machine_rank: 0
        - num_machines: 2
        - main_process_ip: 
        - main_process_port: 10000
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - deepspeed_config: {'deepspeed_multinode_launcher': 'standard', 'gradient_accumulation_steps': 2, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': False, 'zero3_save_16bit_model': False, 'zero_stage': 2}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.5.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Yes
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-40GB

### Who can help?

I receive the following error when resuming from checkpoint for continued pretraining of Llama 3.2 3B model:
```bash
rank11]:   File "/home/advaitj/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/optimization.py", line 803, in _get_options
[rank11]:     use_first_moment = param_group["beta1"] is not None
[rank11]: KeyError: 'beta1'
```

I am using the adafactor optimizer and continuing pretraining a model with pytorch 2.5, previously pretrained for one epoch with pytorch 2.4.

Tags: @ArthurZucker, @muellerzr 

### 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

```python
model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    device_map="cuda",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    local_files_only=True,
    trust_remote_code=True
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    data_collator=collator,
    args=training_arguments,
)

trainer.train(resume_from_checkpoint=base_model_path)

if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(output_model_path)
```

### Expected behavior

The model should continue training from previous checkpoint without any errors.

Comments:
Oh nevermind, I was using Zero 2 above while my model was previously trained with Zero 3. Unfortunately, I am running into https://github.com/microsoft/DeepSpeed/issues/3377 now.

Thanks for updating us! 🤗