LoRA FSDP2 recipe migration (#1517)

Author: ebsmothers

docs/source/tutorials/memory_optimizations.rst

@@ -321,8 +321,3 @@ to see how to configure.
 
 .. All our ``_distributed`` recipes use `FSDP <https://pytorch.org/docs/stable/fsdp.html>`.
 .. .. _glossary_fsdp2:
-
-.. (Experimental) Fully Sharded Data Parallel 2 (FSDP2)
-.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-.. This directory contains distributed training recipes for LoRA and QLoRA using `FSDP2 <https://github.com/pytorch/pytorch/issues/114299>`_.
-.. Currently FSDP2 is only available in PyTorch nightly releases.

recipes/configs/dev/llama2/13B_lora_fsdp2.yaml

@@ -1,16 +1,17 @@
-# Config for multi-device QLoRA in lora_finetune_fsdp2.py
+# Config for multi-device QLoRA in lora_finetune_distributed.py
 # using a Llama2 70B model
 #
-# This config requires PyTorch nightlies to run.
-# See https://github.com/pytorch/torchtune/blob/main/recipes/dev/fsdp2_recipes.md
-# for setup instructions.
-#
 # This config assumes that you've run the following command before launching
 # this run:
 #   tune download meta-llama/Llama-2-70b-hf --output-dir /tmp/Llama-2-70b-hf --hf-token <HF_TOKEN>
 #
 # This config needs 8 GPUs to run
-#   # tune run --nproc_per_node 8 lora_finetune_fsdp2 --config llama2/70B_qlora
+#   # tune run --nproc_per_node 8 lora_finetune_distributed --config llama2/70B_qlora
+#
+# You can add specific overrides through the command line. For example
+# to override the checkpointer directory while launching training
+# you can run:
+#   tune run --nproc_per_node 8 lora_finetune_distributed --config llama2/70B_qlora checkpointer.checkpoint_dir=<YOUR_CHECKPOINT_DIR>
 #
 
 # Model Arguments

recipes/configs/dev/llama2/70B_lora_fsdp2.yaml

@@ -1,21 +1,17 @@
-# Config for single device QLoRA with lora_finetune_fsdp2.py
+# Config for single device QLoRA with lora_finetune_distributed.py
 # using a Llama2 7B model
 #
-# This config requires PyTorch nightlies to run.
-# See https://github.com/pytorch/torchtune/blob/main/recipes/dev/fsdp2_recipes.md
-# for setup instructions.
-#
 # This config assumes that you've run the following command before launching
 # this run:
 #   tune download meta-llama/Llama-2-7b-hf --output-dir /tmp/Llama-2-7b-hf --hf-token <HF_TOKEN>
 #
 # To launch on 2 devices, run the following command from root:
-#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_fsdp2 --config llama2/7B_qlora
+#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_qlora
 #
 # You can add specific overrides through the command line. For example
 # to override the checkpointer directory while launching training
 # you can run:
-#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_fsdp2 --config llama2/7B_qlora checkpointer.checkpoint_dir=<YOUR_CHECKPOINT_DIR>
+#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_qlora checkpointer.checkpoint_dir=<YOUR_CHECKPOINT_DIR>
 #
 # This config works best when the model is being fine-tuned on 2+ GPUs.
 # For single device LoRA finetuning please use 7B_lora_single_device.yaml

recipes/configs/dev/llama2/7B_lora_fsdp2.yaml

@@ -1,17 +1,17 @@
-# Config for multi-device DoRA finetuning in lora_finetune_fsdp2.py
+# Config for multi-device DoRA finetuning in lora_finetune_distributed.py
 # using a Llama3 8B Instruct model
 #
 # This config assumes that you've run the following command before launching
 # this run:
 #   tune download meta-llama/Meta-Llama-3-8B-Instruct --output-dir /tmp/Meta-Llama-3-8B-Instruct --hf-token <HF_TOKEN>
 #
 # To launch on 2 devices, run the following command from root:
-#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_fsdp2 --config llama3/8B_dora
+#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama3/8B_dora
 #
 # You can add specific overrides through the command line. For example
 # to override the checkpointer directory while launching training
 # you can run:
-#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_fsdp2 --config llama3/8B_dora checkpointer.checkpoint_dir=<YOUR_CHECKPOINT_DIR>
+#   tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama3/8B_dora checkpointer.checkpoint_dir=<YOUR_CHECKPOINT_DIR>
 
 
 # Model Arguments
@@ -30,7 +30,7 @@ tokenizer:
   path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model
 
 checkpointer:
-  _component_: torchtune.utils.FullModelMetaCheckpointer
+  _component_: torchtune.training.FullModelMetaCheckpointer
   checkpoint_dir: /tmp/Meta-Llama-3-8B-Instruct/original/
   checkpoint_files: [
     consolidated.00.pth