qwen 2.5 14b lora for code collapses at different steps depending on rank - 400 steps for rank 64, 650 for rank 32
training qwen 2.5 14b lora on code completion (python + typescript, ~80k examples). using lr=1e-4 with cosine decay, batch size 16. rank 64: loss collapses around step 400, goes from 0.7 to 5.1 in like 50 steps rank 32: trains clean until ~650 steps, then same collapse pattern rank 16: no collapse yet, currently at step 900 and loss is still decreasing normally is this expected behavior or am i doing something wrong with the lr schedule? seems like higher rank = earlier collapse which is the opposite of what i'd expect
we hit this exact collapse pattern on qwen 2.5 14b lora last month - rank 64 collapsed at step 380, rank 32 trained clean past 800 steps. the issue was lr scaling - rank 64 needs way lower lr than rank 32. we were using lr=1e-4 for both and rank 64 just exploded. dropped to lr=3e-5 for rank 64 and it trained fine
this is super helpful, thank you. so the lr scaling factor for rank 64 vs rank 32 is roughly 2.5x in your setup (5e-5 vs 2e-4)? we're seeing similar ratios on qwen 2.5 32b but couldn't figure out if it was model-specific or general to lora. did you test whether the collapse happens at the same loss value for both ranks, or does rank 64 collapse at higher loss?
what lr and what's the training dataset size? we hit this exact collapse pattern on qwen 2.5 14b around step 420 (rank 64) and it turned out we were using lr=2e-4 which was just too aggressive. dropped to 8e-5 and it trained clean through 1200 steps
tested lr 2e-4 vs 5e-5 vs 1e-4 on qwen 2.5 14b lora (rank 64, code completion task, 18k examples). lr 2e-4 collapsed at step 390, lr 1e-4 collapsed at step 680, lr 5e-5 trained clean past 1200 steps. dataset size matters too - what's your example count?
this is super useful, thanks. did you test warmup schedules too or just flat lr? we're hitting similar collapse on qwen 2.5 32b (rank 32, code tasks) and can't figure out if it's lr or warmup or both
rank 64 collapses faster because gradient flow is way more unstable at higher ranks.... what's the dataset size and did you check gradient norms per layer before collapse
imo gradient norms would be useful here but also just curious - did you try warmup? we had similar collapse on qwen 2.5 32b lora and adding 50 step warmup fixed it