[v1] support training with fsdp2

examples/v1/train_full/train_full_fsdp2.yaml

@@ -0,0 +1,34 @@
+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-0.6B-dcp
+
+init_config:
+  name: init_on_meta
+
+### data
+train_dataset: data/v1_sft_demo.yaml 
+
+### training
+output_dir: outputs/test_fsdp2
+micro_batch_size: 1
+global_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-4
+bf16: false
+max_steps: 10
+
+### sample
+sample_backend: hf
+max_new_tokens: 128

scripts/hf2dcp.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""Convert a HuggingFace model to DCP checkpoint format.
+
+Usage:
+  python scripts/hf2dcp.py convert --hf_path=/path/to/hf --dcp_path=/path/to/dcp
+
+Arguments:
+  hf_path: Path to the HuggingFace model directory.
+  dcp_path: Output path (directory) for DCP checkpoint.
+"""
+
+import fire
+import torch
+import torch.distributed.checkpoint as dcp
+from transformers import AutoModelForCausalLM
+
+
+def convert(hf_path: str, dcp_path: str) -> None:
+    """Convert HF model weights to DCP.
+
+    Args:
+        hf_path: HuggingFace model directory.
+        dcp_path: Output path (directory) for DCP checkpoint.
+    """
+    if not hf_path or not dcp_path:
+        raise ValueError("Both 'hf_path' and 'dcp_path' are required.")
+
+    print(f"Loading HF model from {hf_path}...")
+    model = AutoModelForCausalLM.from_pretrained(hf_path, device_map="cpu", torch_dtype=torch.bfloat16)
+
+    print(f"Saving to DCP format at {dcp_path}...")
+    dcp.save(model.state_dict(), checkpoint_id=dcp_path)
+    print("Done!")
+
+
+def help() -> None:
+    """Show help message."""
+    print(__doc__)
+
+
+if __name__ == "__main__":
+    fire.Fire({"convert": convert, "help": help, "--convert": convert})

src/llamafactory/v1/accelerator/helper.py

@@ -180,6 +180,16 @@ def operate_tensorlike(fn: Callable[[...], Tensor], data: TensorLike, **kwargs)
         return result.tolist()
 
 
+def get_process_group_backend() -> str:
+    """Get backend for init process group."""
+    if get_current_accelerator().type == DeviceType.NPU:
+        return "hccl"
+    elif get_current_accelerator().type == DeviceType.CUDA:
+        return "nccl"
+    else:
+        return "gloo"
+
+
 def all_gather(tensor: Tensor, group: Optional[ProcessGroup] = None) -> Tensor:
     """Gathers the tensor from all ranks and stacks them at the first dim."""
     world_size = get_world_size()

src/llamafactory/v1/accelerator/interface.py

@@ -145,7 +145,7 @@ def __init__(self, config: DistributedConfig | None = None) -> None:
             timeout = config.get("timeout", 18000)
 
         if self._is_distributed:
-            init_process_group(timeout=timedelta(seconds=timeout))
+            init_process_group(timeout=timedelta(seconds=timeout), backend=helper.get_process_group_backend())
             self.model_device_mesh = init_device_mesh(
                 device_type=self.current_device.type,
                 mesh_shape=self.strategy.model_mesh_shape,

src/llamafactory/v1/config/arg_parser.py

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

src/llamafactory/v1/core/base_trainer.py

@@ -67,7 +67,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 +102,22 @@ def _create_batch_generator(self) -> None:
         )
 
     def _shard_model(self) -> None:
-        pass
+        if self.args.dist_config is None:
+            if DistributedInterface().get_world_size(Dim.DP) > 1:
+                from torch.nn.parallel import DistributedDataParallel as DDP
+
+                logger.warning_rank0(
+                    "dist_config is None but distributed training is enabled; falling back to DistributedDataParallel."
+                )
+                device_ids = None if self.device.type == "cpu" else [self.device.index]
+                self.model = DDP(self.model, device_ids=device_ids)
+        else:
+            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+
+            self.model = DistributedPlugin(self.args.dist_config.name)(
+                self.model,
+                self.args.dist_config,
+            )
 
     def _init_optimizer(self) -> None:
         """Init optimizer."""
@@ -162,7 +181,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()
@@ -172,10 +193,17 @@ def fit(self) -> None:
 
                 step_loss, grad_norm = DistributedInterface().all_reduce([step_loss, grad_norm])
                 DistributedInterface().sync()
-                print(f"Epoch {epoch}, Step {self.global_step}, Loss: {step_loss:.4f}, Grad Norm: {grad_norm:.4f}")
+                if DistributedInterface().get_rank() == 0:
+                    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)
+        model_to_save = self.model.module if hasattr(self.model, "module") else self.model
+        model_to_save.save_pretrained(self.args.output_dir)
         self.renderer.processor.save_pretrained(self.args.output_dir)
         logger.info_rank0(f"Model saved to {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,
@@ -142,6 +141,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=DistributedInterface().current_device.type,
             drop_last=self.drop_last,
         )
         if self.batching_strategy == BatchingStrategy.NORMAL: