update

aws-samples · Jun 4, 2024 · 952eba3 · 952eba3
1 parent 4ac5496
commit 952eba3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/...rchtune/slurm/tutorials/e2e-llama3-70b-development/configs/full_finetune_distributed.yaml b/...rchtune/slurm/tutorials/e2e-llama3-70b-development/configs/full_finetune_distributed.yaml
@@ -25,10 +25,9 @@ tokenizer:
 
 # Dataset
 dataset:
-  _component_: torchtune.datasets.wiki_text
-  train_on_input: True
+  _component_: torchtune.datasets.wikitext_dataset
 seed: null
-shuffle: True
+shuffle: False
 
 # Model Arguments
 model:
@@ -75,8 +74,8 @@ checkpointer:
 resume_from_checkpoint: False
 
 # Fine-tuning arguments
-batch_size: 2
-epochs: 3
+batch_size: 1
+epochs: 1
 
 optimizer:
   _component_: torch.optim.AdamW
@@ -95,6 +94,7 @@ device: cuda
 # Memory management
 enable_activation_checkpointing: True
 memory_efficient_fsdp_wrap: True
+fsdp_cpu_offload: True
 
 # Reduced precision
 dtype: bf16

diff --git a/...ses/torchtune/slurm/tutorials/e2e-llama3-70b-development/full_finetune_distributed.sbatch b/...ses/torchtune/slurm/tutorials/e2e-llama3-70b-development/full_finetune_distributed.sbatch
@@ -77,7 +77,7 @@ declare -a TORCHRUN_ARGS=(
     --rdzv_endpoint=$(hostname)
 )
 declare -a TRAIN_ARGS=(
-    --config  ${PWD}/tutorials/e2e-llama3-70b-development/configs/lora_finetune_distributed.yaml
+    --config  ${PWD}/tutorials/e2e-llama3-70b-development/configs/full_finetune_distributed.yaml
     tokenizer.path=${MODEL_PATH}/${HF_MODEL}/original/tokenizer.model
     checkpointer.checkpoint_dir=${MODEL_PATH}/${HF_MODEL}
     checkpointer.output_dir=${MODEL_PATH}/${HF_MODEL}-tuned