FSDP + lm_head + Liger-Kernel

[gotzmann]

Reminder

• I have read the README and searched the existing issues.

System Info

stage: sft
do_train: true
finetuning_type: lora
lora_target: all
additional_target: embed_tokens,lm_head
lora_rank: 128
lora_alpha: 16
lora_dropout: 0.1
use_rslora: true

Reproduction

[rank3]:   File "/home/git/Liger-Kernel/src/liger_kernel/ops/fused_linear_cross_entropy.py", line 59, in fused_linear_cross_entropy_forward
[rank3]:     logits_chunk = _input_chunk @ weight.t()  # chunk_size x V
[rank3]: RuntimeError: size mismatch, got input (2), mat (2x4096), vec (65667072)

Expected behavior

No response

Others

Please check out discussion on problems with [ lm_head ] fine-tuning:

linkedin/Liger-Kernel#332