DoRA fixes

[ebsmothers]

Co-authored by: @mirceamironenco

Fixes #2129

This PR fixes our broken DoRA distributed training. It also makes some quality-of-life improvements to how we initialize DoRALinear and (in a small way) LoRALinear.

We do the following:

1. Define a to_empty method on both DoRALinear and LoRALinear. This consolidates the move of all the potentially-still-on-meta-device params onto a real device in a single place (lora_a.weight and lora_b.weight for LoRA, and additionally magnitude for DoRA) and matches the contract of nn.Module. For DoRALinear's magnitude we create a new parameter with torch.empty_like and manually set requires_grad, then use the torch.utils.swap_tensors API to ensure self.magnitude winds up on a real device.

2. For DoRA, we modify the initialize_dora_magnitude method to (a) check that base weight and LoRA weights are not on meta device, (b) calculate ||W + (alpha/rank)BA||, and (c) copy this value into self.magnitude. Since this should always be called after to_empty, we now know that self.magnitude

BC note: From the perspective of someone importing from our library to their own custom recipe, the changes to LoRALinear's contract with respect to to_empty are backwards compatible: calling lora_linear.to_empty(device) (new version) vs lora_linear.lora_a.to_empty(device); lora_linear.lora_b.to_empty(device) (old version) are identical. For DoRALinear, the recipe will be broken without the update to the new version.. but the recipe was broken anyways, so I claim this is OK.

Other changes:

• Deprecate (but don't yet delete) load_dora_magnitudes, which is redundant.
• Add unit tests for proper calculation of DoRA magnitude on single-device and distributed w/ meta device init

Test plan

Unit tests

pytest tests/torchtune/modules/peft/test_dora.py
...
======== 12 passed in 11.73s ========

E2E tests

Single-device recipe

Compare the following three cases:

1. This PR
2. Resume intermediate checkpoint on this PR
3. Baseline (run on main)

Loss curves for all three are the same:

Command for (1):

tune run lora_finetune_single_device --config llama3/8B_dora_single_device max_steps_per_epoch=100 epochs=2 \
metric_logger=torchtune.training.metric_logging.WandBLogger metric_logger.project=test-dora-fixes \
metric_logger.name=dora-single-new

Commands for (2):

# First delete final checkpoint to automatically resume from intermediate checkpoint
rm -r /tmp/torchtune/llama3_8B/dora_single_device/epoch_1

tune run lora_finetune_single_device --config llama3/8B_dora_single_device max_steps_per_epoch=100 epochs=2 \
metric_logger=torchtune.training.metric_logging.WandBLogger metric_logger.project=test-dora-fixes \
metric_logger.name=dora-single-new-resumed resume_from_checkpoint=True

Distributed recipe

1. This PR
2. Resume intermediate checkpoint on this PR

(Note that (3) from the single-device recipe is not applicable, as this is broken on main). Loss curves are the same after resuming from checkpoint:

Command for (1):

tune run --nproc_per_node 2 lora_finetune_distributed --config llama3/8B_dora max_steps_per_epoch=100 epochs=2 \
metric_logger=torchtune.training.metric_logging.WandBLogger metric_logger.project=test-dora-fixes \
metric_logger.name=dora-distributed-new

Commands for (2):

# First delete final checkpoint to automatically resume from intermediate checkpoint
rm -r /tmp/torchtune/llama3_8B/dora/epoch_1

tune run --nproc_per_node 2 lora_finetune_distributed --config llama3/8B_dora max_steps_per_epoch=100 epochs=2 \
metric_logger=torchtune.training.metric_logging.WandBLogger metric_logger.project=test-dora-fixes \
metric_logger.name=dora-distributed-new-resume resume_from_checkpoint=True

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2139

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 031d89b with merge base 06a8379 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[SalmanMohammadi]

When would this be False?

[ebsmothers]

Offhand I can’t think of a case but imo better to just be explicit here

[SalmanMohammadi]

not 100% sure here, but will the original torch.empty(out_dim) still be floating around after this? This probably isn't an issue?

[ebsmothers]

Yeah I was wondering about this too, need to check myself. Offhand not sure about the implications for refcounting with this utility. But magnitude is just embed_dim number of elements so even if still hanging around it’s pretty small

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 88.18182% with 13 lines in your changes missing coverage. Please review.

Project coverage is 67.75%. Comparing base (06a8379) to head (c7e937f).
Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
recipes/lora_dpo_distributed.py	0.00%	4 Missing ⚠️
recipes/knowledge_distillation_single_device.py	0.00%	3 Missing ⚠️
torchtune/modules/peft/lora.py	50.00%	2 Missing ⚠️
recipes/knowledge_distillation_distributed.py	0.00%	1 Missing ⚠️
recipes/lora_finetune_distributed.py	0.00%	1 Missing ⚠️
recipes/qat_lora_finetune_distributed.py	0.00%	1 Missing ⚠️
tests/recipes/test_lora_finetune_distributed.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##            main    #2139       +/-   ##
==========================================
+ Coverage   9.33%   67.75%   +58.41%     
==========================================
  Files        289      334       +45     
  Lines      16956    19281     +2325     
==========================================
+ Hits        1583    13064    +11481     
+ Misses     15373     6217     -9156     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.