Crash when using save_state with deepspeed: `model.state_dict` functions incompatible with new deepspeed.
- Dominant language
- Python
- Stars
- 4.8k
- Forks
- 487
- PR merge metrics
- No merged PRs in 30d
Description
### 🐛 Describe the bug
I've recently been using the code provided in https://github.com/tlc4418/llm_optimization, which in turn uses trlx.
In doing so I encountered a bug causing trlx to crash when trying to save, caused by a recent change in deepspeeed.
To reproduce, use this https://gist.github.com/JohannesAck/feb31ee5c491ca30771335296ec8b295 and start it with deepspeed by using `accelerate launch` with a config that enables deepspeed:
```
Traceback (most recent call last):
File "/workspaces/llm_optimization/crash_example.py", line 111, in
main(hparams)
File "/workspaces/llm_optimization/crash_example.py", line 101, in main
trlx.train(
File "/usr/local/lib/python3.10/dist-packages/trlx/trlx.py", line 142, in train
trainer.learn()
File "/usr/local/lib/python3.10/dist-packages/trlx/trainer/accelerate_base_trainer.py", line 598, in learn
self.save(directory)
File "/usr/local/lib/python3.10/dist-packages/trlx/trainer/accelerate_base_trainer.py", line 312, in save
self.accelerator.save_state(dst_dir, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py", line 2944, in save_state
model.save_checkpoint(output_dir, ckpt_id, **save_model_func_kwargs)
File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 3105, in save_checkpoint
self._save_checkpoint(save_dir,
File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 3299, in _save_checkpoint
module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters)
File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/engine.py", line 2540, in module_state_dict
sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
File "/usr/local/lib/python3.10/dist-packages/trlx/models/modeling_ppo.py", line 460, in state_dict
state_dict = self.v_head.state_dict(*args, **dict(prefix="v_head.", **kwargs))
TypeError: dict() got multiple values for keyword argument 'prefix'```
```
This is caused by this change in deepspeed https://github.com/microsoft/DeepSpeed/pull/5408, that changes the call to state_dict to use a keyword instead of positional argument:
```
--- sd = self.module.state_dict(destination, prefix, keep_vars)
+++ sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
```
TRLX however assumes that the argument will be passed
https://github.com/CarperAI/trlx/blob/3340c2f3a56d1d14fdd5f13ad575121fa26b6d92/trlx/models/modeling_ppo.py#L354-L359
In L359: `dict(prefix="v_head.", **kwargs)` becomes `dict(prefix="v_head.", prefix="")` and thus has two values for `prefix` and crashes.
## Workaround:
Downgrade deepspeed to a version < 0.14.1:
```
pip install 'deepspeed<0.14.1'
```
I'm not sure what the proper solution here would be, just ignoring the `prefix` argument doesn't sound great either. One option might be to just ignore it if it's an empty string and raise an exception otherwise.
Hope this helps somebody!
### Which trlX version are you using?
trlx=0.7.0
### Additional system and package information
deepspeed=0.14.4
Contributor guide
Research direction
Start with trlx/models/modeling_ppo.py around the state_dict implementation and reproduce the failure using the linked gist with accelerate launch and a DeepSpeed-enabled config. Done means save_state completes with DeepSpeed 0.14.4 without the duplicate-prefix crash while preserving state-dict prefix handling.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100