[v1] support training with fsdp2

examples/accelerate/hf2dcp.py

@@ -0,0 +1,22 @@
+# convert_hf_to_dcp.py
+import argparse
+import torch
+import torch.distributed.checkpoint as dcp
+from transformers import AutoModelForCausalLM
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--hf_path", type=str, required=True)
+parser.add_argument("--dcp_path", type=str, required=True)
+args = parser.parse_args()
+
+def convert(model_path, save_path):
+    print(f"Loading HF model from {model_path}...")
+    model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cpu", torch_dtype=torch.bfloat16)
+
+    print(f"Saving to DCP format at {save_path}...")
+    dcp.save(model.state_dict(), checkpoint_id=save_path)
+    print("Done!")
+
+
+if __name__ == "__main__":
+    convert(args.hf_path, args.dcp_path)

examples/v1/train_full/train_full_fsdp2.yaml

@@ -0,0 +1,33 @@
+model: Qwen/Qwen3-0.6B
+trust_remote_code: true
+model_class: llm
+
+template: qwen3_nothink
+
+kernel_config:
+    name: auto
+    include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
+
+quant_config: null
+
+dist_config:
+    name: fsdp2
+    dcp_path: null # /mnt/f/pretrain_models/Qwen3-1.7B-dcp
+
+init_config:
+     name: init_on_meta
+
+### data
+train_dataset: data/v1_sft_demo.yaml 
+
+### training
+output_dir: outputs/test_run
+micro_batch_size: 1
+global_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-4
+bf16: false
+
+### sample
+sample_backend: hf
+max_new_tokens: 128

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."},

src/llamafactory/v1/core/base_trainer.py

@@ -35,6 +35,7 @@
 from ..accelerator.helper import ReduceOp
 from ..accelerator.interface import Dim, DistributedInterface
 from ..config import TrainingArguments
+from ..plugins.trainer_plugins.distributed.fsdp2 import FSDP2Plugin
 from ..utils import logging
 from ..utils.helper import compute_valid_tokens
 from ..utils.types import BatchInput, HFModel, ModelOutput, Tensor, TorchDataset
@@ -67,13 +68,17 @@ def __init__(
         self.model_input_names = self.renderer.processor.model_input_names
 
         self._create_batch_generator()
-        self.num_training_steps = self.args.num_train_epochs * len(self.train_batch_generator)
+        # Calculate num_training_steps: max_steps takes priority if set
+        if self.args.max_steps is not None and self.args.max_steps > 0:
+            self.num_training_steps = self.args.max_steps
+        else:
+            self.num_training_steps = self.args.num_train_epochs * len(self.train_batch_generator)
 
         if self.args.enable_activation_checkpointing:
             self.model.gradient_checkpointing_enable({"use_reentrant": False})
 
         if self.args.dist_config is not None:
-            shard_need_optimizer = self.args.dist_config.name == "deepspeed"
+            shard_need_optimizer = self.args.dist_config.get("name", None) == "deepspeed"
         else:
             shard_need_optimizer = False
 
@@ -98,7 +103,17 @@ def _create_batch_generator(self) -> None:
         )
 
     def _shard_model(self) -> None:
-        pass
+        assert self.args.dist_config is not None
+        if self.args.dist_config.get("name", None) == "fsdp2":
+            fsdp_plugin = FSDP2Plugin(self.args.dist_config)
+            dcp_path = self.args.dist_config.get("dcp_path", None)
+            self.model = fsdp_plugin.shard_model(
+                self.model, hf_model_path=self.model.config.name_or_path, dcp_path=dcp_path
+            )
+        elif self.args.dist_config.get("name", None) == "deepspeed":
+            pass
+        else:
+            raise ValueError(f"Unsupported dist config: {self.args.dist_config.get('name')}")
 
     def _init_optimizer(self) -> None:
         """Init optimizer."""
@@ -162,7 +177,9 @@ def fit(self) -> None:
                     step_loss += loss.item()
 
                 grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm).item()
-                if not torch.isfinite(grad_norm):
+
+                # isfinite(): argument 'input' (position 1) must be Tensor, not float
+                if not torch.isfinite(torch.tensor(grad_norm)):  # type: ignore # pyright: ignore [reportUnknownReturnType]
                     logger.warning_rank0(f"Gradient norm is not finite: {grad_norm}")
                 else:
                     self.optimizer.step()
@@ -174,6 +191,11 @@ def fit(self) -> None:
                 DistributedInterface().sync()
                 print(f"Epoch {epoch}, Step {self.global_step}, Loss: {step_loss:.4f}, Grad Norm: {grad_norm:.4f}")
 
+                # Check if max_steps is reached
+                if self.global_step >= self.num_training_steps:
+                    logger.info_rank0(f"Reached max_steps ({self.num_training_steps}), stopping training.")
+                    return
+
     def save_model(self) -> None:
         """Save the model."""
         self.model.save_pretrained(self.args.output_dir)

src/llamafactory/v1/core/utils/batching.py

@@ -30,7 +30,7 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler
 
-from ...accelerator.interface import DistributedInterface
+from ...accelerator.interface import Dim, DistributedInterface
 from ...config import BatchingStrategy
 from ...utils import logging
 from ...utils.helper import pad_and_truncate
@@ -83,8 +83,7 @@ def __init__(
         self.pin_memory = pin_memory
         self.drop_last = drop_last
         # TODO: support length and infinity
-
-        dp_size = DistributedInterface().get_world_size("dp")
+        dp_size = DistributedInterface().get_world_size(Dim.DP)
 
         if self.global_batch_size is None:
             self.global_batch_size = dp_size * micro_batch_size
@@ -126,8 +125,8 @@ def _init_data_provider(self) -> None:
         if len(self.dataset) != -1:
             sampler = StatefulDistributedSampler(
                 self.dataset,
-                num_replicas=DistributedInterface().get_world_size("dp"),
-                rank=DistributedInterface().get_rank("dp"),
+                num_replicas=DistributedInterface().get_world_size(Dim.DP),
+                rank=DistributedInterface().get_rank(Dim.DP),
                 shuffle=True,
                 seed=0,
                 drop_last=self.drop_last,

src/llamafactory/v1/launcher.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
 import sys
-
-from ..extras.env import VERSION, print_env
+from copy import deepcopy
 
 
 USAGE = (
@@ -27,27 +28,112 @@
     + "-" * 70
 )
 
-
-WELCOME = (
-    "-" * 58
-    + "\n"
-    + f"| Welcome to LLaMA Factory, version {VERSION}"
-    + " " * (21 - len(VERSION))
-    + "|\n|"
-    + " " * 56
-    + "|\n"
-    + "| Project page: https://github.com/hiyouga/LLaMA-Factory |\n"
-    + "-" * 58
-)
+_DIST_TRAIN_COMMANDS = ("train", "sft", "dpo", "rm")
 
 
 def launch():
+    from ..extras.env import VERSION, print_env
+    from .accelerator.helper import get_device_count
+    from .utils.env import find_available_port, is_env_enabled, use_kt, use_ray
+    from .utils.logging import get_logger
+
+    logger = get_logger(__name__)
+
+    WELCOME = (
+        "-" * 58
+        + "\n"
+        + f"| Welcome to LLaMA Factory, version {VERSION}"
+        + " " * (21 - len(VERSION))
+        + "|\n|"
+        + " " * 56
+        + "|\n"
+        + "| Project page: https://github.com/hiyouga/LLaMA-Factory |\n"
+        + "-" * 58
+    )
+
+    # NOTE:
+    # `llamafactory-cli <command> ...` enters here first.
+    # We may re-launch via `torchrun` for distributed training. In that case we must
+    # forward `<command>` as argv[1] to the re-executed script, otherwise the script
+    # will misinterpret the first user argument (e.g. yaml config) as the command.
     command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"
 
-    if command == "sft":  # train command will fallback to sft command
-        from .trainers.sft_trainer import run_sft
-
-        run_sft()
+    if command in _DIST_TRAIN_COMMANDS and (
+        is_env_enabled("FORCE_TORCHRUN") or (get_device_count() > 1 and not use_ray() and not use_kt())
+    ):
+        # breakpoint()
+        # launch distributed training
+        nnodes = os.getenv("NNODES", "1")
+        node_rank = os.getenv("NODE_RANK", "0")
+        nproc_per_node = os.getenv("NPROC_PER_NODE", str(get_device_count()))
+        master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
+        master_port = os.getenv("MASTER_PORT", str(find_available_port()))
+        logger.info_rank0(f"Initializing {nproc_per_node} distributed tasks at: {master_addr}:{master_port}")
+        if int(nnodes) > 1:
+            logger.info_rank0(f"Multi-node training enabled: num nodes: {nnodes}, node rank: {node_rank}")
+
+        # elastic launch support
+        max_restarts = os.getenv("MAX_RESTARTS", "0")
+        rdzv_id = os.getenv("RDZV_ID")
+        min_nnodes = os.getenv("MIN_NNODES")
+        max_nnodes = os.getenv("MAX_NNODES")
+
+        env = deepcopy(os.environ)
+        if is_env_enabled("OPTIM_TORCH", "1"):
+            # optimize DDP, see https://zhuanlan.zhihu.com/p/671834539
+            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+            env["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+
+        if rdzv_id is not None:
+            # launch elastic job with fault tolerant support when possible
+            # see also https://docs.pytorch.org/docs/stable/elastic/train_script.html
+            rdzv_nnodes = nnodes
+            # elastic number of nodes if MIN_NNODES and MAX_NNODES are set
+            if min_nnodes is not None and max_nnodes is not None:
+                rdzv_nnodes = f"{min_nnodes}:{max_nnodes}"
+
+            process = subprocess.run(
+                (
+                    "torchrun --nnodes {rdzv_nnodes} --nproc-per-node {nproc_per_node} "
+                    "--rdzv-id {rdzv_id} --rdzv-backend c10d --rdzv-endpoint {master_addr}:{master_port} "
+                    "--max-restarts {max_restarts} {file_name} {args}"
+                )
+                .format(
+                    rdzv_nnodes=rdzv_nnodes,
+                    nproc_per_node=nproc_per_node,
+                    rdzv_id=rdzv_id,
+                    master_addr=master_addr,
+                    master_port=master_port,
+                    max_restarts=max_restarts,
+                    file_name=__file__,
+                    args=" ".join([command] + sys.argv[1:]),
+                )
+                .split(),
+                env=env,
+                check=True,
+            )
+        else:
+            # NOTE: DO NOT USE shell=True to avoid security risk
+            process = subprocess.run(
+                (
+                    "torchrun --nnodes {nnodes} --node_rank {node_rank} --nproc-per-node {nproc_per_node} "
+                    "--master_addr {master_addr} --master_port {master_port} {file_name} {args}"
+                )
+                .format(
+                    nnodes=nnodes,
+                    node_rank=node_rank,
+                    nproc_per_node=nproc_per_node,
+                    master_addr=master_addr,
+                    master_port=master_port,
+                    file_name=__file__,
+                    args=" ".join([command] + sys.argv[1:]),
+                )
+                .split(),
+                env=env,
+                check=True,
+            )
+
+        sys.exit(process.returncode)
 
     elif command == "chat":
         from .samplers.cli_sampler import run_chat
@@ -67,5 +153,35 @@ def launch():
         print(f"Unknown command: {command}.\n{USAGE}")
 
 
+def main():
+    # Use absolute import when script is run directly by torchrun
+    # sys.argv[1] contains the command (sft/dpo/rm/train), sys.argv[2:] contains the rest args
+    command = sys.argv[1] if len(sys.argv) > 1 else "sft"
+
+    # Routing needs the sub-command, but downstream trainers usually expect argv without it.
+    # When launched by `torchrun`, we pass:
+    #   launcher.py <command> <config.yaml> [extra args...]
+    # So remove `<command>` before calling trainer entrypoints.
+    if command in _DIST_TRAIN_COMMANDS:
+        sys.argv.pop(1)
+    else:
+        # Backward-compat: if someone runs `torchrun launcher.py config.yaml`,
+        # treat it as sft by default.
+        if len(sys.argv) > 1 and sys.argv[1].endswith((".yaml", ".yml")):
+            command = "sft"
+    if command in ("train", "sft"):
+        from llamafactory.v1.trainers.sft_trainer import run_sft
+
+        run_sft()
+    elif command == "dpo":
+        # from llamafactory.v1.trainers.dpo_trainer import run_dpo
+        # run_dpo()
+        raise NotImplementedError("DPO trainer is not implemented yet.")
+    elif command == "rm":
+        # from llamafactory.v1.trainers.rm_trainer import run_rm
+        # run_rm()
+        raise NotImplementedError("RM trainer is not implemented yet.")
+
+
 if __name__ == "__main__":
-    pass
+    main()