add script to train with ft

Summary:
the script adds configuration options to run training locally with ft enabled

Author: tushar00jain

run_train_ft.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -ex
+
+export CONFIG_FILE=${CONFIG_FILE:-"./torchtitan/models/llama3/train_configs/debug_model_ft.toml"}
+
+export FT_REPLICA_ID="${FT_REPLICA_ID:-0}"
+export FT_GROUP_SIZE="${FT_GROUP_SIZE:-1}"
+
+./run_train.sh \
+    --fault_tolerance.group_size="${FT_GROUP_SIZE}" \
+    --fault_tolerance.replica_id="${FT_REPLICA_ID}"

torchtitan/models/llama3/train_configs/debug_model_ft.toml

@@ -0,0 +1,98 @@
+[job]
+dump_folder = "./outputs"
+description = "Llama 3 debug training"
+print_args = false
+
+[profiling]
+enable_profiling = true
+save_traces_folder = "profile_trace"
+profile_freq = 10
+profiler_active = 10
+profiler_warmup = 0
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 1
+disable_color_printing = false
+enable_tensorboard = false
+save_tb_folder = "tb"
+enable_wandb = false
+
+[model]
+name = "llama3"
+flavor = "debugmodel"
+# test folder with tokenizer.json, for debug purpose only
+hf_assets_path = "./tests/assets/tokenizer"
+# converters = ["float8"]
+
+[optimizer]
+name = "AdamW"
+lr = 8e-4
+eps = 1e-8
+
+[lr_scheduler]
+warmup_steps = 2  # lr scheduler warm up, normally 20% of the train steps
+decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
+decay_type = "linear"
+min_lr_factor = 0.0
+
+[training]
+local_batch_size = 8
+seq_len = 2048
+max_norm = 1.0  # grad norm clipping
+steps = 100
+dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = -1
+fsdp_reshard_after_forward = "default"  # default / never / always
+tensor_parallel_degree = 1
+enable_async_tensor_parallel = false
+pipeline_parallel_degree = 1
+context_parallel_degree = 1
+
+[checkpoint]
+enable = false
+folder = "checkpoint"
+interval = 10
+last_save_model_only = false
+export_dtype = "float32"
+async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem"]
+
+[activation_checkpoint]
+mode = "selective"  # ["none", "selective", "full"]
+selective_ac_option = (
+    "2"  # 'int' = ac every positive int layer or 'op', ac based on ops policy
+)
+
+[compile]
+enable = false
+components = ["model", "loss"]
+
+[quantize.linear.float8]
+enable_fsdp_float8_all_gather = false
+precompute_float8_dynamic_scale_for_fsdp = false
+filter_fqns = ["output"]
+
+[validation]
+enable = false
+dataset = "c4_validation"
+freq = 5
+steps = 10
+
+[comm]
+train_timeout_seconds = 15
+
+[fault_tolerance]
+enable = true
+sync_steps = 10
+num_fragments = 2
+semi_sync_method = "diloco"
+data_parallel_shard_degree = 1
+process_group = "nccl"
+process_group_timeout_ms = 10000
+
+[experimental]
+custom_args_module = "torchtitan.components.ft.config"