[usability] Accelerate Support

.gitignore

@@ -18,6 +18,7 @@ log/
 regression_test/*/new_output_models
 regression_test/*/new_log
 output_dir/
+tests_out
 
 # data files
 data/

README.md

@@ -69,7 +69,6 @@ An extensible, convenient, and efficient toolbox for finetuning large machine le
 - [LMFlow](#lmflow)
   - [Latest News](#latest-news)
   - [Table of Contents](#table-of-contents)
-  - [Supported Models](#supported-models)
   - [Quick Start](#quick-start)
     - [Setup](#setup)
     - [Prepare Dataset](#prepare-dataset)
@@ -85,21 +84,6 @@ An extensible, convenient, and efficient toolbox for finetuning large machine le
   - [License](#license)
   - [Citation](#citation)
 
-## Supported Models
-
-See all conversation template details [here](https://optimalscale.github.io/LMFlow/examples/supported_conversation_template.html).
-
-|  Model  | Conversation Template |
-|  :---:  | :-------------------: |
-| DeepSeek | `deepseek` <br> `deepseek_v2` <br> `deepseek_r1` <br> `deepseek_r1_distill` <br> `deepseek_v3` |
-| Gemma | `gemma` |
-| Hymba | `hymba` |
-| InternLM2 | `internlm2` |
-| LLaMA | `llama2` <br> `llama3` <br> `llama3_for_tool`|
-| Phi | `phi3` |
-| Qwen | `qwen2` <br> `qwen2_for_tool` <br> `qwen2_5` <br> `qwen2_5_1m` <br> `qwen2_5_math` <br> `qwen_qwq` |
-| Yi | `yi` <br> `yi1_5` |
-| Zephyr | `zephyr` |
 
 ## Quick Start
 
@@ -162,6 +146,16 @@ Please refer to our [doc](https://optimalscale.github.io/LMFlow/examples/DATASET
 
 ### Finetuning
 
+#### Estimated Hardware Requirement
+
+| Method                 | 0.5B |  3B  |  7B  |  14B  |  30B  |  70B  |  `x`B   |
+| ---------------------- | ---- | ---- | ---- | ----- | ----- | ----- | ------- |
+| Full `bf16`/`fp16`     |  9GB | 55GB |120GB | 240GB | 600GB | 1200GB| `18x`GB |
+| LoRA                   |  1GB | 6GB  | 16GB |  32GB |  64GB | 160GB |  `2x`GB |
+| QLoRA `quant_bit=8`    | 0.7GB| 3GB  | 10GB |  20GB |  40GB |   80GB|  `x`GB  |
+| QLoRA `quant_bit=4`    | 0.4GB| 1.5GB|  6GB |  12GB |  24GB |   48GB| `x/2`GB |
+
+
 #### Full Finetuning
 
 Full training updates all the parameters to finetune a language model.

configs/accelerate_fsdp_config.yaml

@@ -0,0 +1,29 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_min_num_params: 1000000
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_forward_prefetch: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: true
+
+downcast_bf16: true
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8 # NOTE: distributed_type should be `NO` if you're training on a single GPU
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
+main_process_port: 1204

configs/accelerate_singlegpu_config.yaml

@@ -0,0 +1,29 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: 'NO'
+
+fsdp_config:
+  fsdp_auto_wrap_policy: SIZE_BASED_WRAP
+  fsdp_min_num_params: 1000000
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_forward_prefetch: false
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: 'NO_SHARD'
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: true
+
+downcast_bf16: true
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
+main_process_port: 1204

configs/iterative_dpo.yaml

@@ -17,7 +17,6 @@ preprocessing_num_workers: 16
 output_dir: ./output_models/iterative_dpo
 run_name: iterative_dpo
 random_seed: 42
-use_accelerator: True
 enable_distributed_inference: True
 distributed_inference_num_instances: 8
 initial_iter_idx: 0 # 0 refers to the first dataset in dataset_path_list

contrib/rlhflow/run_reward_modeling.sh

@@ -30,7 +30,7 @@ deepspeed ${deepspeed_args} \
     --block_size 512 \
     --per_device_train_batch_size 1 \
     --per_device_eval_batch_size 1\
-    --deepspeed configs/ds_config_zero2.json \
+    --deepspeed configs/archive/ds_config_zero2.json \
     --bf16 \
     --run_name rm_test \
     --validation_split_percentage 10 \

contrib/tool-finetune/run_function_call_finetune.sh

@@ -65,7 +65,7 @@ deepspeed ${deepspeed_args} \
     --disable_group_texts 1 \
     --block_size 1024 \
     --per_device_train_batch_size 1 \
-    --deepspeed configs/ds_config_zero3.json \
+    --deepspeed configs/archive/ds_config_zero3.json \
     --fp16 \
     --run_name finetune \
     --validation_split_percentage 0 \

docs/dev_notes/finetuning.mmd

@@ -0,0 +1,62 @@
+sequenceDiagram
+    participant User
+    participant Finetuner as LMFlow Finetuner
+    participant Model as LMFlow Model
+    participant Dataset as LMFlow Dataset
+    participant Trainer as Trainer
+
+    User->>Finetuner: tune(model, dataset)
+
+    %% Tokenization
+    Finetuner->>Model: tokenize(dataset)
+    Model->>Dataset: Apply tokenization to dataset
+
+    alt if not disable_group_texts
+        Finetuner->>Finetuner: group_text(tokenized_dataset, model_max_length)
+    end
+
+    %% Prepare for training
+    Finetuner->>Finetuner: Prepare dataset for trainer
+
+    %% Create appropriate trainer based on configuration
+    alt if model_args.use_lora
+        Finetuner->>Finetuner: Initialize PeftTrainer
+    else
+        Finetuner->>Finetuner: Initialize standard Trainer
+    end
+
+    alt if training_args.use_customized_optim
+        Finetuner->>Finetuner: create_customized_optimizer()
+    end
+
+    alt if training_args.use_lisa
+        Finetuner->>Finetuner: Create DynamicLayerActivationCallback
+    end
+
+    %% Start training
+    Finetuner->>Trainer: train(resume_from_checkpoint)
+
+    %% Training loop (simplified)
+    loop Training iterations (Trainer._inner_training_loop simplified)
+        Trainer->>Model: Forward pass
+        Model-->>Trainer: Return predictions
+        Trainer->>Trainer: Compute loss
+        Trainer->>Model: Backward pass
+        Model->>Model: Compute Gradient  
+        Trainer->>Trainer: Optimizer step
+    end
+
+    %% Save the model
+    alt if not model_args.use_lora
+        Trainer->>Trainer: save_model()
+    else
+        alt if model_args.save_aggregated_lora
+            Finetuner->>Model: merge_lora_weights()
+        end
+        Finetuner->>Model: save(output_dir, save_aggregated_lora)
+    end
+
+    %% Finish and return
+    Trainer-->>Finetuner: Return train result
+    Finetuner->>Finetuner: Log metrics
+    Finetuner-->>User: Return fine-tuned model

examples/benchmarking.py

@@ -214,7 +214,7 @@ def main():
     dataset_name = benchmarking_args.dataset_name
     # metric = pipeline_args.metric
     if is_lmflow_local_benchmarking(dataset_name):   # TODO (@Jipeng)
-        model = AutoModel.get_model(model_args, tune_strategy='none', ds_config=ds_config)
+        model = AutoModel.get_model(model_args, do_train=False, ds_config=ds_config)
         run_lmflow_local_benchmarking(dataset_name,pipeline_name,model_args,pipeline_args,model)  # Pass args TODO (@Jipeng)
     elif is_lm_evaluation_benchmarking(dataset_name):
         model = model_args.model_name_or_path

examples/chatbot.py

@@ -64,10 +64,9 @@ def main():
 
     model = AutoModel.get_model(
         model_args,
-        tune_strategy='none',
+        do_train=False,
         ds_config=ds_config,
         device=pipeline_args.device,
-        use_accelerator=True,
     )
 
     # We don't need input data, we will read interactively from stdin

examples/chatbot_gradio.py

@@ -110,7 +110,7 @@ class ChatbotArguments:
 
 model = AutoModel.get_model(
     model_args,
-    tune_strategy='none',
+    do_train=False,
     ds_config=ds_config,
     device=pipeline_args.device,
     torch_dtype=torch.float16

examples/detail_memory.py

@@ -13,7 +13,7 @@
 
 LISA = True if sys.argv[3] == "1" else False
 LORA = True if sys.argv[4] == "1" else False
-lora_rank = int(sys.argv[5])
+lora_r = int(sys.argv[5])
 # Check if the model name is provided as a command-line argument
 if len(sys.argv) < 6:
     print("Usage: python script_name.py <model_name>")
@@ -26,7 +26,7 @@
 print("token_length     : ", sys.argv[2])
 print("LISA             : ", LISA)
 print("LORA             : ", LORA)
-print("lora_rank        : ", lora_rank)
+print("lora_r        : ", lora_r)
 # Model initialization
 model_name = sys.argv[1]
 token_length = sys.argv[2]
@@ -48,7 +48,7 @@
     peft_config = LoraConfig(
         task_type=TaskType.CAUSAL_LM,
         inference_mode=False,
-        r=lora_rank,
+        r=lora_r,
         lora_alpha=32,
         lora_dropout=0.1,
         target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head"],

examples/evaluation.py

@@ -36,9 +36,8 @@
 
 model = AutoModel.get_model(
     model_args, 
-    tune_strategy='none', 
+    do_train=False, 
     ds_config=ds_config, 
-    use_accelerator=pipeline_args.use_accelerator_for_evaluator
 )
 dataset = Dataset(data_args)
 

examples/finetune_multi_modal.py

@@ -59,7 +59,7 @@ def main():
     # do not resiger deepspeed in the model.
     # with_deepspeed flag may be removed
     # by modifying the tune strategy in the future.
-    model = AutoModel.get_model(model_args, tune_strategy='none',
+    model = AutoModel.get_model(model_args, do_train=True,
                                 ds_config=pipeline_args.deepspeed,
                                 custom_model=True,
                                 with_deepspeed=False,