linkedin / linkedin/Liger-Kernel
LigerORPOTrainer doesn't work when not using FSDP
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
### 🐛 Describe the bug
Hi all,
I tried using LigerORPOTrainer but it doesn't seem to work when running on a single GPU.
I tried running the orpo trainer example but it fails on `assert isinstance(wrapper_module, FullyShardedDataParallel)` .
As far as I can tell from the code, this is because _FSDPForwardRedirection is used without checking if the model is FullyShardedDataParallel, as can be seen [here](https://github.com/linkedin/Liger-Kernel/blame/30b8486a2bd48dff97f22c5f0c88520b8825cb35/src/liger_kernel/transformers/trainer/orpo_trainer.py#L77).
Please let me know if I understood the issue correctly, perhaps I can submit a fix 😃
### Reproduce
Link to reproduction code - [here](https://gist.github.com/Ola-Vish/d93567a8f4a9896fa80e70d9e2da3cec)
Tried running this code (which is basically the same as the this [example](https://github.com/linkedin/Liger-Kernel/blob/main/examples/alignment/run_orpo.py) code by running "python liger_orpo_trainer.py", and what I get is the following error:
` 0%| | 0/100 [00:00
main()
File "/home/ubuntu/training_code/liger_orpo_trainer.py", line 39, in main
trainer.train()
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/transformers/trainer.py", line 1424, in train
return inner_training_loop(
^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/transformers/trainer.py", line 1506, in _inner_training_loop
self._run_epoch(
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/transformers/trainer.py", line 1734, in _run_epoch
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/transformers/trainer.py", line 1906, in training_step
loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/trl/trainer/orpo_trainer.py", line 873, in compute_loss
loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/liger_kernel/transformers/trainer/orpo_trainer.py", line 101, in get_batch_loss_metrics
loss, aux_outputs = self.concatenated_forward(model, batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/liger_kernel/transformers/trainer/orpo_trainer.py", line 77, in concatenated_forward
orpo_loss, aux_outputs = _FSDPForwardRedirection()(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/training_code/sota_venv/lib/python3.12/site-packages/liger_kernel/transformers/fsdp.py", line 40, in __call__
assert isinstance(wrapper_module, FullyShardedDataParallel)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
0%| | 0/100 [00:00
Contributor guide
Research direction
Start with src/liger_kernel/transformers/trainer/orpo_trainer.py at the concatenated_forward call and inspect the _FSDPForwardRedirection implementation in src/liger_kernel/transformers/fsdp.py. Reproduce the failure with examples/alignment/run_orpo.py or the linked reproduction, then verify that the ORPO trainer works on a single GPU without the assertion while preserving FSDP behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 65/100