[WIP] enable QLoRA + FSDP2

[weifengpy]

this PR is stacked on

• TorchTune: LoRA + FSDP2 enable LoRA + FSDP2 #855
• TorchAO: NF4Tensor with torch.chunk and ops needed by FSDP [FSDP2][NF4Tensor][2/n] implement torch.chunk and other ops ao#150
• PyTorch: meta init + cpu offloading [FSDP2] support fully_shard(model_on_meta, cpu_offload) pytorch#126305

command:
tune run --nnodes 1 --nproc_per_node 8 lora_finetune_distributed --config recipes/configs/llama2/7B_qlora_single_device.yaml

QLoRA differs from LoRA in config

• model._component_: torchtune.models.llama2.qlora_llama2_7b instead of lora_llama2_7b. LoRALinear(quantize_base=True/False)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 8 New Failures

As of commit b2fd531 with merge base 30c75d4 ():

NEW FAILURES - The following jobs have failed:

• Multi-GPU Recipe Tests / recipe_test_multi_gpu (3.10, stable) (gh)
  tests/recipes/test_lora_finetune_distributed.py::TestLoRAFinetuneDistributedRecipe::test_save_and_load_merged_weights
• Multi-GPU Recipe Tests / recipe_test_multi_gpu (3.11, stable) (gh)
  tests/recipes/test_lora_finetune_distributed.py::TestLoRAFinetuneDistributedRecipe::test_save_and_load_merged_weights
• Multi-GPU Recipe Tests / recipe_test_multi_gpu (3.8, stable) (gh)
  ##[error]The operation was canceled.
• Multi-GPU Recipe Tests / recipe_test_multi_gpu (3.9, stable) (gh)
• Unit Test / unit_tests (3.10) (gh)
  tests/torchtune/utils/test_seed.py::TestSeed::test_deterministic_false
• Unit Test / unit_tests (3.11) (gh)
  tests/torchtune/utils/test_seed.py::TestSeed::test_deterministic_false
• Unit Test / unit_tests (3.8) (gh)
  tests/torchtune/utils/test_seed.py::TestSeed::test_deterministic_false
• Unit Test / unit_tests (3.9) (gh)
  tests/torchtune/utils/test_seed.py::TestSeed::test_deterministic_false

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

[awgu]

nit: 😄 to divide with the all-reduce

Suggested change

	torch.distributed.all_reduce(param.grad)
	torch.distributed.all_reduce(param.grad, op=ReduceOp.AVG)

from torch.distributed.distributed_c10d import ReduceOp

[weifengpy]

neat!

[weifengpy]

updated LoRA PR to use ReduceOp.AVG. This PR will be stacked on it when landing