[v1] support training with fsdp2 (#9773)

Co-authored-by: frozenleaves <frozen@Mac.local>
Co-authored-by: Yaowei Zheng <hiyouga@buaa.edu.cn>

src/llamafactory/v1/config/arg_parser.py

@@ -20,7 +20,7 @@ from typing import Any
 from omegaconf import OmegaConf
 from transformers import HfArgumentParser
 
-from ...extras.misc import is_env_enabled
+from ..utils.env import is_env_enabled
 from .data_args import DataArguments
 from .model_args import ModelArguments
 from .sample_args import SampleArguments

src/llamafactory/v1/config/training_args.py

@@ -45,6 +45,10 @@ class TrainingArguments:
         default=3,
         metadata={"help": "Number of training epochs."},
     )
+    max_steps: int | None = field(
+        default=None,
+        metadata={"help": "Maximum number of training steps. If set, overrides num_train_epochs."},
+    )
     max_grad_norm: float = field(
         default=1.0,
         metadata={"help": "Maximum gradient norm for training."},