[v1] support training with fsdp2

examples/accelerate/hf2dcp.py

@@ -0,0 +1,24 @@
+# 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.accelerate import DistributedPlugin
 from ..utils import logging
 from ..utils.helper import compute_valid_tokens
 from ..utils.types import BatchInput, HFModel, ModelOutput, Tensor, TorchDataset
@@ -67,7 +68,11 @@ 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})
@@ -98,7 +103,10 @@ def _create_batch_generator(self) -> None:
         )
 
     def _shard_model(self) -> None:
-        pass
+        self.model = DistributedPlugin(self.args.dist_config.name)(
+            self.model,
+            self.args.dist_config,
+        )
 
     def _init_optimizer(self) -> None:
         """Init optimizer."""
@@ -162,7 +170,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 +184,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,8 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler
 
-from ...accelerator.interface import DistributedInterface
+from ...accelerator.helper import get_current_device
+from ...accelerator.interface import Dim, DistributedInterface
 from ...config import BatchingStrategy
 from ...utils import logging
 from ...utils.helper import pad_and_truncate
@@ -83,8 +84,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 +126,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,
@@ -142,6 +142,7 @@ def _init_data_provider(self) -> None:
             num_workers=self.batching_workers,
             collate_fn=self.renderer.process_samples,
             pin_memory=self.pin_memory,
+            pin_memory_device=get_current_device().type,
             drop_last=self.drop_last,
         )
         if self.batching_strategy == BatchingStrategy.NORMAL:

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,45 +28,153 @@
     + "-" * 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():
-    command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"
+    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
 
-    if command == "sft":  # train command will fallback to sft command
-        from .trainers.sft_trainer import run_sft
+    logger = get_logger(__name__)
 
-        run_sft()
+    # 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 in _DIST_TRAIN_COMMANDS and (
+        is_env_enabled("FORCE_TORCHRUN") or (get_device_count() > 1 and not use_ray() and not use_kt())
+    ):
+
+        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"
+
+        torchrun_args = [
+            "torchrun",
+            "--nproc-per-node",
+            nproc_per_node,
+        ]
+        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}"
+
+            torchrun_args.extend(
+                [
+                    "--nnodes",
+                    rdzv_nnodes,
+                    "--rdzv-id",
+                    rdzv_id,
+                    "--rdzv-backend",
+                    "c10d",
+                    "--rdzv-endpoint",
+                    f"{master_addr}:{master_port}",
+                    "--max-restarts",
+                    max_restarts,
+                ]
+            )
+        else:
+            # NOTE: DO NOT USE shell=True to avoid security risk
+            torchrun_args.extend(
+                [
+                    "--nnodes",
+                    nnodes,
+                    "--node_rank",
+                    node_rank,
+                    "--master_addr",
+                    master_addr,
+                    "--master_port",
+                    master_port,
+                ]
+            )
+
+        script_args = [__file__, command] + sys.argv[1:]
+        process = subprocess.run(
+            torchrun_args + script_args,
+            env=env,
+            check=True,
+        )
+
+        sys.exit(process.returncode)
 
     elif command == "chat":
         from .samplers.cli_sampler import run_chat
 
         run_chat()
 
     elif command == "env":
-        print_env()
+        raise NotImplementedError("Environment information is not implemented yet.")
 
     elif command == "version":
-        print(WELCOME)
+        raise NotImplementedError("Version information is not implemented yet.")
 
     elif command == "help":
         print(USAGE)
 
+    elif command in _DIST_TRAIN_COMMANDS:
+        # Single GPU training without torchrun
+        if command in ("train", "sft"):
+            from llamafactory.v1.trainers.sft_trainer import run_sft
+
+            run_sft()
+        elif command == "dpo":
+            raise NotImplementedError("DPO trainer is not implemented yet.")
+        elif command == "rm":
+            raise NotImplementedError("RM trainer is not implemented yet.")
+
     else:
         print(f"Unknown command: {command}.\n{USAGE}")
 
 
+def main():
+    # 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.
+    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()

src/llamafactory/v1/plugins/trainer_plugins/distributed/accelerate.py

@@ -0,0 +1,36 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from ....config.arg_utils import PluginConfig
+from ....utils.plugin import BasePlugin
+from ....utils.types import HFModel
+
+
+class DistributedPlugin(BasePlugin):
+    def __call__(self, model: HFModel, dist_config: PluginConfig, **kwargs) -> HFModel:
+        return super().__call__(model, dist_config, **kwargs)
+
+
+@DistributedPlugin("fsdp2").register()
+def shard_model_fsdp2(model: HFModel, dist_config: PluginConfig) -> HFModel:
+    from .fsdp2 import FSDP2Engine
+
+    return FSDP2Engine(dist_config).shard_model(model)
+
+
+@DistributedPlugin("deepspeed").register()
+def shard_model_deepspeed(model: HFModel, dist_config: PluginConfig) -> HFModel:
+    return model