refactor ray integration, support save ckpt

.env.local

@@ -12,6 +12,7 @@ FORCE_CHECK_IMPORTS=
 LLAMAFACTORY_VERBOSITY=
 USE_MODELSCOPE_HUB=
 USE_OPENMIND_HUB=
+USE_RAY=
 RECORD_VRAM=
 # torchrun
 FORCE_TORCHRUN=

examples/README.md

@@ -95,6 +95,12 @@ FORCE_TORCHRUN=1 NNODES=2 NODE_RANK=1 MASTER_ADDR=192.168.0.1 MASTER_PORT=29500
 FORCE_TORCHRUN=1 llamafactory-cli train examples/train_lora/llama3_lora_sft_ds3.yaml
 ```
 
+#### Supervised Fine-Tuning with Ray on 4 GPUs
+
+```bash
+USE_RAY=1 llamafactory-cli train examples/train_full/llama3_lora_sft_ray.yaml
+```
+
 ### QLoRA Fine-Tuning
 
 #### Supervised Fine-Tuning with 4/8-bit Bitsandbytes/HQQ/EETQ Quantization (Recommended)

examples/README_zh.md

@@ -95,6 +95,12 @@ FORCE_TORCHRUN=1 NNODES=2 NODE_RANK=1 MASTER_ADDR=192.168.0.1 MASTER_PORT=29500
 FORCE_TORCHRUN=1 llamafactory-cli train examples/train_lora/llama3_lora_sft_ds3.yaml
 ```
 
+#### 使用 Ray 在 4 张 GPU 上微调
+
+```bash
+USE_RAY=1 llamafactory-cli train examples/train_full/llama3_lora_sft_ray.yaml
+```
+
 ### QLoRA 微调
 
 #### 基于 4/8 比特 Bitsandbytes/HQQ/EETQ 量化进行指令监督微调（推荐）

examples/train_lora/llama3_lora_sft.yaml

@@ -9,7 +9,6 @@ finetuning_type: lora
 lora_target: all
 
 ### dataset
-dataset_dir: /home/ray/default/LLaMA-Factory/data/
 dataset: identity,alpaca_en_demo
 template: llama3
 cutoff_len: 2048
@@ -39,10 +38,3 @@ val_size: 0.1
 per_device_eval_batch_size: 1
 eval_strategy: steps
 eval_steps: 500
-
-
-### ray setup
-resources_per_worker:
-  GPU: 1
-num_workers: 4
-# placement_strategy: ...

examples/train_lora/llama3_lora_sft_ray.yaml

@@ -0,0 +1,48 @@
+### model
+model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct  # or use local absolute path
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_target: all
+
+### dataset
+dataset: identity,alpaca_en_demo
+dataset_dir: REMOTE:llamafactory/demo_data  # or use local absolute path
+template: llama3
+cutoff_len: 2048
+max_samples: 1000
+overwrite_cache: true
+preprocessing_num_workers: 16
+
+### output
+output_dir: tmp_dir
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+learning_rate: 1.0e-4
+num_train_epochs: 3.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+
+### eval
+val_size: 0.1
+per_device_eval_batch_size: 1
+eval_strategy: steps
+eval_steps: 500
+
+### ray
+ray_run_name: llama3_8b_sft_lora
+ray_num_workers: 4  # number of GPUs to use
+resources_per_worker:
+  GPU: 1
+placement_strategy: PACK

src/llamafactory/cli.py

@@ -24,8 +24,7 @@
 from .eval.evaluator import run_eval
 from .extras import logging
 from .extras.env import VERSION, print_env
-from .extras.misc import get_device_count
-from .integrations.ray.ray_utils import should_use_ray
+from .extras.misc import get_device_count, use_ray
 from .train.tuner import export_model, run_exp
 from .webui.interface import run_web_demo, run_web_ui
 
@@ -88,8 +87,7 @@ def main():
         export_model()
     elif command == Command.TRAIN:
         force_torchrun = os.getenv("FORCE_TORCHRUN", "0").lower() in ["true", "1"]
-        use_ray = should_use_ray()
-        if force_torchrun or (get_device_count() > 1 and not use_ray):
+        if force_torchrun or (get_device_count() > 1 and not use_ray()):
             master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
             master_port = os.getenv("MASTER_PORT", str(random.randint(20001, 29999)))
             logger.info_rank0(f"Initializing distributed tasks at: {master_addr}:{master_port}")

src/llamafactory/extras/misc.py

@@ -229,7 +229,7 @@ def skip_check_imports() -> None:
     r"""
     Avoids flash attention import error in custom model files.
     """
-    if os.environ.get("FORCE_CHECK_IMPORTS", "0").lower() not in ["true", "1"]:
+    if os.getenv("FORCE_CHECK_IMPORTS", "0").lower() not in ["true", "1"]:
         transformers.dynamic_module_utils.check_imports = get_relative_imports
 
 
@@ -275,8 +275,12 @@ def try_download_model_from_other_hub(model_args: "ModelArguments") -> str:
 
 
 def use_modelscope() -> bool:
-    return os.environ.get("USE_MODELSCOPE_HUB", "0").lower() in ["true", "1"]
+    return os.getenv("USE_MODELSCOPE_HUB", "0").lower() in ["true", "1"]
 
 
 def use_openmind() -> bool:
-    return os.environ.get("USE_OPENMIND_HUB", "0").lower() in ["true", "1"]
+    return os.getenv("USE_OPENMIND_HUB", "0").lower() in ["true", "1"]
+
+
+def use_ray() -> bool:
+    return os.getenv("USE_RAY", "0").lower() in ["true", "1"]

src/llamafactory/extras/packages.py

@@ -62,6 +62,10 @@ def is_pillow_available():
     return _is_package_available("PIL")
 
 
+def is_ray_available():
+    return _is_package_available("ray")
+
+
 def is_requests_available():
     return _is_package_available("requests")
 

src/llamafactory/hparams/__init__.py

@@ -17,7 +17,8 @@
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
 from .model_args import ModelArguments
-from .parser import get_eval_args, get_infer_args, get_train_args
+from .parser import get_eval_args, get_infer_args, get_ray_args, get_train_args, read_args
+from .training_args import RayArguments, TrainingArguments
 
 
 __all__ = [
@@ -26,7 +27,11 @@
     "FinetuningArguments",
     "GeneratingArguments",
     "ModelArguments",
+    "RayArguments",
+    "TrainingArguments",
     "get_eval_args",
     "get_infer_args",
+    "get_ray_args",
     "get_train_args",
+    "read_args",
 ]

src/llamafactory/hparams/parser.py

@@ -19,12 +19,12 @@
 import os
 import sys
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import transformers
 import yaml
-from transformers import HfArgumentParser, Seq2SeqTrainingArguments
+from transformers import HfArgumentParser
 from transformers.integrations import is_deepspeed_zero3_enabled
 from transformers.trainer_utils import get_last_checkpoint
 from transformers.training_args import ParallelMode
@@ -34,73 +34,54 @@
 from ..extras import logging
 from ..extras.constants import CHECKPOINT_NAMES
 from ..extras.misc import check_dependencies, get_current_device
-from ..integrations.ray.ray_train_args import RayTrainArguments
 from .data_args import DataArguments
 from .evaluation_args import EvaluationArguments
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
 from .model_args import ModelArguments
+from .training_args import RayArguments, TrainingArguments
 
 
 logger = logging.get_logger(__name__)
 
 check_dependencies()
 
 
-_TRAIN_ARGS = [
-    ModelArguments,
-    DataArguments,
-    Seq2SeqTrainingArguments,
-    FinetuningArguments,
-    GeneratingArguments,
-    RayTrainArguments,
-]
-_TRAIN_CLS = Tuple[
-    ModelArguments,
-    DataArguments,
-    Seq2SeqTrainingArguments,
-    FinetuningArguments,
-    GeneratingArguments,
-    RayTrainArguments,
-]
+_TRAIN_ARGS = [ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
+_TRAIN_CLS = Tuple[ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
 _INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
 _INFER_CLS = Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
 _EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
 _EVAL_CLS = Tuple[ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
 
 
-def _read_args(args: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+def read_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> Union[Dict[str, Any], List[str]]:
     if args is not None:
         return args
 
     if len(sys.argv) == 2 and (sys.argv[1].endswith(".yaml") or sys.argv[1].endswith(".yml")):
-        # read yaml file
         return yaml.safe_load(Path(sys.argv[1]).absolute().read_text())
     elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        # read json file
         return json.loads(Path(sys.argv[1]).absolute().read_text())
     else:
-        return {}
+        return sys.argv[1:]
 
 
 def _parse_args(
-    parser: "HfArgumentParser", args: Optional[Dict[str, Any]] = None, allow_extra_keys: bool = False
+    parser: "HfArgumentParser", args: Optional[Union[Dict[str, Any], List[str]]] = None, allow_extra_keys: bool = False
 ) -> Tuple[Any]:
-    args_dict = _read_args(args)
+    args = read_args(args)
+    if isinstance(args, dict):
+        return parser.parse_dict(args, allow_extra_keys=allow_extra_keys)
 
-    if args_dict:
-        return parser.parse_dict(args_dict, allow_extra_keys=allow_extra_keys)
-    else:
-        (*parsed_args, unknown_args) = parser.parse_args_into_dataclasses(
-            args=args_dict, return_remaining_strings=True
-        )
+    (*parsed_args, unknown_args) = parser.parse_args_into_dataclasses(args=args, return_remaining_strings=True)
 
-        if unknown_args:
-            print(parser.format_help())
-            print(f"Got unknown args, potentially deprecated arguments: {unknown_args}")
-            raise ValueError(f"Some specified arguments are not used by the HfArgumentParser: {unknown_args}")
+    if unknown_args:
+        print(parser.format_help())
+        print(f"Got unknown args, potentially deprecated arguments: {unknown_args}")
+        raise ValueError(f"Some specified arguments are not used by the HfArgumentParser: {unknown_args}")
 
-        return (*parsed_args,)
+    return (*parsed_args,)
 
 
 def _set_transformers_logging() -> None:
@@ -141,7 +122,7 @@ def _verify_model_args(
 def _check_extra_dependencies(
     model_args: "ModelArguments",
     finetuning_args: "FinetuningArguments",
-    training_args: Optional["Seq2SeqTrainingArguments"] = None,
+    training_args: Optional["TrainingArguments"] = None,
 ) -> None:
     if os.getenv("DISABLE_VERSION_CHECK", "0").lower() in ["true", "1"]:
         logger.warning_once("Version checking has been disabled, may lead to unexpected behaviors.")
@@ -177,31 +158,29 @@ def _check_extra_dependencies(
         require_version("rouge_chinese", "To fix: pip install rouge-chinese")
 
 
-def _parse_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
+def _parse_train_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _TRAIN_CLS:
     parser = HfArgumentParser(_TRAIN_ARGS)
     return _parse_args(parser, args)
 
 
-def _parse_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
+def _parse_infer_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _INFER_CLS:
     parser = HfArgumentParser(_INFER_ARGS)
     return _parse_args(parser, args)
 
 
-def _parse_eval_args(args: Optional[Dict[str, Any]] = None) -> _EVAL_CLS:
+def _parse_eval_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _EVAL_CLS:
     parser = HfArgumentParser(_EVAL_ARGS)
     return _parse_args(parser, args)
 
 
-def _parse_ray_args(args: Optional[Dict[str, Any]] = None) -> RayTrainArguments:
-    parser = HfArgumentParser(RayTrainArguments)
-    ray_args = _parse_args(parser, args, allow_extra_keys=True)[0]
-    if ray_args.use_ray:
-        require_version("ray", "To fix: pip install ray")
+def get_ray_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> RayArguments:
+    parser = HfArgumentParser(RayArguments)
+    (ray_args,) = _parse_args(parser, args, allow_extra_keys=True)
     return ray_args
 
 
-def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
-    model_args, data_args, training_args, finetuning_args, generating_args, _ = _parse_train_args(args)
+def get_train_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _TRAIN_CLS:
+    model_args, data_args, training_args, finetuning_args, generating_args = _parse_train_args(args)
 
     # Setup logging
     if training_args.should_log:
@@ -410,7 +389,7 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
     return model_args, data_args, training_args, finetuning_args, generating_args
 
 
-def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
+def get_infer_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _INFER_CLS:
     model_args, data_args, finetuning_args, generating_args = _parse_infer_args(args)
 
     _set_transformers_logging()
@@ -443,7 +422,7 @@ def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
     return model_args, data_args, finetuning_args, generating_args
 
 
-def get_eval_args(args: Optional[Dict[str, Any]] = None) -> _EVAL_CLS:
+def get_eval_args(args: Optional[Union[Dict[str, Any], List[str]]] = None) -> _EVAL_CLS:
     model_args, data_args, eval_args, finetuning_args = _parse_eval_args(args)
 
     _set_transformers_logging()

src/llamafactory/hparams/training_args.py

@@ -0,0 +1,48 @@
+import json
+from dataclasses import dataclass, field
+from typing import Literal, Optional, Union
+
+from transformers import Seq2SeqTrainingArguments
+from transformers.training_args import _convert_str_dict
+
+from ..extras.misc import use_ray
+
+
+@dataclass
+class RayArguments:
+    r"""
+    Arguments pertaining to the Ray training.
+    """
+
+    ray_run_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The training results will be saved at `saves/ray_run_name`."},
+    )
+    ray_num_workers: int = field(
+        default=1,
+        metadata={"help": "The number of workers for Ray training. Default is 1 worker."},
+    )
+    resources_per_worker: Union[dict, str] = field(
+        default_factory=lambda: {"GPU": 1},
+        metadata={"help": "The resources per worker for Ray training. Default is to use 1 GPU per worker."},
+    )
+    placement_strategy: Literal["SPREAD", "PACK", "STRICT_SPREAD", "STRICT_PACK"] = field(
+        default="PACK",
+        metadata={"help": "The placement strategy for Ray training. Default is PACK."},
+    )
+
+    def __post_init__(self):
+        self.use_ray = use_ray()
+        if isinstance(self.resources_per_worker, str) and self.resources_per_worker.startswith("{"):
+            self.resources_per_worker = _convert_str_dict(json.loads(self.resources_per_worker))
+
+
+@dataclass
+class TrainingArguments(RayArguments, Seq2SeqTrainingArguments):
+    r"""
+    Arguments pertaining to the trainer.
+    """
+
+    def __post_init__(self):
+        Seq2SeqTrainingArguments.__post_init__(self)
+        RayArguments.__post_init__(self)