src/llamafactory/train/hyper_parallel/trainer.py

# 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.

"""HyperParallel distributed trainer for LlamaFactory."""

import logging
import os
import types
from contextlib import nullcontext
from functools import partial
from typing import Any, Optional

import torch
from hyper_parallel.integration.llamafactory import (
    HSDPModule,
    HyperParallelArguments,
    export_to_hf_format,
    fsdp2_prepare_model,
    hsdp_sync_stream,
    load_hsdp_model,
    load_hsdp_optimizer_and_scheduler,
    save_hsdp_checkpoint,
    wrap_optimizer_with_skip_dtensor_dispatch,
)
from hyper_parallel.integration.llamafactory import (
    clip_grad_norm_ as hp_clip_grad_norm_,
)
from hyper_parallel.integration.llamafactory.context_parallel import (
    cp_prepare_model,
    get_cp_rank,
    get_dp_rank,
    shard_inputs_for_cp,
)
from hyper_parallel.platform import get_platform
from torch import nn

from ..sft.trainer import CustomSeq2SeqTrainer


logger = logging.getLogger(__name__)


class _CPBatchRepeatedBatchSampler(torch.utils.data.BatchSampler):
    """Repeat logical batches so Accelerate shards CP peers onto the same samples."""

    def __init__(self, sampler, batch_size: int, drop_last: bool, repeat_factor: int, logical_group_size: int):
        super().__init__(sampler, batch_size, drop_last)
        self.repeat_factor = repeat_factor
        self.logical_group_size = logical_group_size

    def __len__(self):
        logical_length = super().__len__()
        if not self.drop_last and logical_length > 0:
            logical_length = _ceil_div(logical_length, self.logical_group_size) * self.logical_group_size
        return logical_length * self.repeat_factor

    def __iter__(self):
        initial_data = []
        logical_count = 0
        pad_cursor = 0
        max_initial_data = self.batch_size * self.logical_group_size

        def collect_initial_data(batch):
            if len(initial_data) < max_initial_data:
                initial_data.extend(batch[: max_initial_data - len(initial_data)])

        def get_padding_item():
            nonlocal pad_cursor
            item = initial_data[pad_cursor % len(initial_data)]
            pad_cursor += 1
            return item

        def pad_batch(batch):
            batch = list(batch)
            if self.drop_last or len(batch) == self.batch_size:
                return batch

            while len(batch) < self.batch_size:
                batch.append(get_padding_item())
            return batch

        def make_padding_batch():
            return [get_padding_item() for _ in range(self.batch_size)]

        def repeat_batch(batch):
            for _ in range(self.repeat_factor):
                yield list(batch)

        for batch in super().__iter__():
            collect_initial_data(batch)
            batch = pad_batch(batch)
            logical_count += 1
            yield from repeat_batch(batch)

        if self.drop_last or logical_count == 0:
            return

        while logical_count % self.logical_group_size != 0:
            logical_count += 1
            yield from repeat_batch(make_padding_batch())


class _CPDataLoaderLengthProxy:
    """Keep baseline logical dataloader length while yielding CP-repeated batches."""

    def __init__(self, dataloader, logical_length: int):
        self._dataloader = dataloader
        self._logical_length = logical_length

    def __iter__(self):
        return iter(self._dataloader)

    def __len__(self):
        return self._logical_length

    def __getattr__(self, name):
        return getattr(self._dataloader, name)


def _ceil_div(numerator: int, denominator: int) -> int:
    return (numerator + denominator - 1) // denominator


class HyperParallelTrainer(CustomSeq2SeqTrainer):
    """Trainer that replaces Accelerate FSDP2 with HyperParallel fully_shard.

    Inherits CustomSeq2SeqTrainer for training algorithm logic (loss, metrics,
    prediction, sampler, etc.) and only overrides HSDP-specific behavior.
    """

    def __init__(
        self,
        hp_args: HyperParallelArguments,
        finetuning_args=None,
        processor=None,
        ref_model: Optional[nn.Module] = None,
        **kwargs,
    ):
        self._hp_args = hp_args

        # Let CustomSeq2SeqTrainer handle everything except ref_model —
        # Custom would prepare it with accelerate's fsdp2_prepare_model,
        # but we need HP's version instead.
        super().__init__(
            finetuning_args=finetuning_args,
            processor=processor,
            ref_model=None,
            **kwargs,
        )

        if not getattr(self.accelerator, "is_fsdp2", False):
            raise ValueError("HyperParallel trainer requires Accelerate FSDP2 mode to be enabled.")

        self._cp_size = hp_args.cp_size
        self._cp_rank = get_cp_rank(hp_args) if self._cp_size > 1 else 0
        self._dp_rank = get_dp_rank(hp_args) if self._cp_size > 1 else get_platform().get_rank()

        # Prepare ref_model with the same CP + HSDP path as the train model.
        self.ref_model = ref_model
        if self.ref_model is not None:
            self.ref_model = self._prepare_model_for_hyper_parallel(self.ref_model)

        self._orig_accelerator_clip_grad_norm = self.accelerator.clip_grad_norm_
        self._orig_fsdp2_prepare_model = None
        self._accelerator_patches_active = False

    def _prepare_model_for_hyper_parallel(self, model: nn.Module) -> nn.Module:
        """Apply CP runtime hooks before delegating to HyperParallel FSDP2 preparation."""
        if self._cp_size > 1:
            model = cp_prepare_model(model, self.accelerator, self._hp_args)
        return fsdp2_prepare_model(self.accelerator, model, self._hp_args)

    def _activate_accelerator_patches(self) -> None:
        """Patch Accelerate to use HyperParallel fsdp2_prepare_model and clip_grad_norm_."""
        if self._accelerator_patches_active:
            return

        import accelerate.accelerator as acc_module  # pylint: disable=C0415

        self._orig_fsdp2_prepare_model = acc_module.fsdp2_prepare_model

        def _hp_fsdp2_prepare_model(accelerator, model):
            return self._prepare_model_for_hyper_parallel(model)

        acc_module.fsdp2_prepare_model = _hp_fsdp2_prepare_model

        def _hp_clip_grad_norm(accelerator, parameters, max_norm, norm_type=2):
            if getattr(accelerator, "is_fsdp2", False):
                accelerator.unscale_gradients()
                parameter_list = list(parameters)
                parameter_ids = {id(param) for param in parameter_list}
                for model in accelerator._models:  # pylint: disable=protected-access
                    if not isinstance(model, HSDPModule):
                        continue
                    model_param_ids = {id(param) for param in model.parameters()}
                    if parameter_ids and parameter_ids.issubset(model_param_ids):
                        return hp_clip_grad_norm_(parameter_list, max_norm, norm_type=norm_type)
            return self._orig_accelerator_clip_grad_norm(parameters, max_norm, norm_type=norm_type)

        self.accelerator.clip_grad_norm_ = types.MethodType(_hp_clip_grad_norm, self.accelerator)
        self._accelerator_patches_active = True

    def _restore_accelerator_patches(self) -> None:
        """Restore original Accelerate methods."""
        if not self._accelerator_patches_active:
            return

        import accelerate.accelerator as acc_module  # pylint: disable=C0415

        if self._orig_fsdp2_prepare_model is not None:
            acc_module.fsdp2_prepare_model = self._orig_fsdp2_prepare_model
        self.accelerator.clip_grad_norm_ = self._orig_accelerator_clip_grad_norm
        self._accelerator_patches_active = False

    def _wrap_model(self, model: nn.Module, training: bool = True, dataloader=None) -> nn.Module:
        """Let Accelerate own FSDP2/HSDP wrapping so optimizer remapping stays correct."""
        del dataloader
        if isinstance(model, HSDPModule):
            return model
        if training and getattr(self.accelerator, "is_fsdp2", False):
            return model
        return super()._wrap_model(model, training=training)

    def _get_train_sampler(self, train_dataset=None):
        """Match the no-CP baseline sampler semantics before CP repeats whole logical batches."""
        if train_dataset is None:
            train_dataset = self.train_dataset
        if getattr(self.finetuning_args, "disable_shuffling", False):
            return torch.utils.data.SequentialSampler(train_dataset)
        return super()._get_train_sampler(train_dataset)

    def _build_cp_batch_sampler(self, dataset, shuffle: bool, batch_size: int, drop_last: bool):
        """Repeat complete logical batches so CP groups consume the same baseline batch."""
        sampler = self._get_train_sampler(dataset) if shuffle else torch.utils.data.SequentialSampler(dataset)
        return _CPBatchRepeatedBatchSampler(
            sampler,
            batch_size=batch_size,
            drop_last=drop_last,
            repeat_factor=self._cp_size,
            logical_group_size=max(1, get_platform().get_world_size() // self._cp_size),
        )

    def _get_cp_dataloader(self, dataset, batch_size: int, shuffle: bool):
        """Create a train dataloader whose logical batches are shared within each CP group."""
        if isinstance(dataset, torch.utils.data.IterableDataset):
            raise NotImplementedError(
                "HyperParallel CP training requires a map-style dataset because iterable datasets cannot "
                "repeat logical batches across CP ranks."
            )

        try:
            import datasets  # pylint: disable=C0415
        except ImportError:  # pragma: no cover
            datasets = None

        if datasets is not None and isinstance(dataset, datasets.Dataset):
            dataset = self._remove_unused_columns(dataset, description="Training")
            data_collator = self.data_collator
        else:
            data_collator = self._get_collator_with_removed_columns(self.data_collator, description="Training")

        batch_sampler = self._build_cp_batch_sampler(
            dataset,
            shuffle=shuffle,
            batch_size=batch_size,
            drop_last=self.args.dataloader_drop_last,
        )
        logical_batches = len(batch_sampler) // self._cp_size
        dp_size = max(1, get_platform().get_world_size() // self._cp_size)
        logical_length = logical_batches // dp_size if self.args.dataloader_drop_last else _ceil_div(logical_batches, dp_size)

        dataloader_params = {
            "batch_sampler": batch_sampler,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers
            if self.args.dataloader_num_workers > 0
            else False,
        }
        if self.args.dataloader_num_workers > 0:
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        from transformers.trainer import seed_worker  # pylint: disable=C0415

        dataloader_params["worker_init_fn"] = partial(
            seed_worker,
            num_workers=self.args.dataloader_num_workers,
            rank=self.args.process_index,
        )

        dataloader = self.accelerator.prepare(torch.utils.data.DataLoader(dataset, **dataloader_params))
        return _CPDataLoaderLengthProxy(dataloader, logical_length)

    def get_train_dataloader(self):
        """Keep the no-CP logical batch stream, then repeat each whole batch across CP peers."""
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")
        if self._cp_size <= 1:
            return super().get_train_dataloader()

        shuffle = not getattr(self.finetuning_args, "disable_shuffling", False)
        return self._get_cp_dataloader(
            dataset=self.train_dataset,
            batch_size=self._train_batch_size,
            shuffle=shuffle,
        )

    def _move_model_to_device(self, model: nn.Module, device: Optional[torch.device] = None):
        """Skip redundant device moves for HSDP-wrapped models."""
        if isinstance(model, HSDPModule):
            return model
        if device is None:
            return model
        return model.to(device)

    def train(self, *args, **kwargs):
        """Activate HP patches during training and restore afterwards."""
        self._activate_accelerator_patches()
        try:
            return super().train(*args, **kwargs)
        finally:
            self._restore_accelerator_patches()

    def training_step(
        self,
        model: nn.Module,
        inputs: dict[str, Any],
        num_items_in_batch: Optional[int] = None,
    ) -> torch.Tensor:
        """Standard training step with HSDP sync plus optional CP input sharding."""
        model.train()
        inputs = self._prepare_inputs(inputs)

        if self._cp_size > 1:
            inputs = shard_inputs_for_cp(inputs, self._cp_rank, self._cp_size)

        sync_gradients = getattr(self.accelerator, "sync_gradients", True)
        if isinstance(model, HSDPModule):
            model.set_is_last_backward(sync_gradients)
            model.set_requires_gradient_sync(sync_gradients)

        compute_loss_context_manager = getattr(self, "compute_loss_context_manager", nullcontext)
        with compute_loss_context_manager():
            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)

        if self.args.n_gpu > 1:
            loss = loss.mean()

        if not getattr(self, "model_accepts_loss_kwargs", False) and getattr(self, "compute_loss_func", None) is None:
            loss = loss / self.args.gradient_accumulation_steps

        self.accelerator.backward(loss)

        if isinstance(model, HSDPModule) and sync_gradients:
            hsdp_sync_stream()

        return loss.detach()

    def create_optimizer(self):
        """Create optimizer and wrap step with SkipDTensorDispatch."""
        optimizer = super().create_optimizer()
        wrap_optimizer_with_skip_dtensor_dispatch(optimizer)
        return optimizer

    def _save_optimizer_and_scheduler(self, output_dir: str) -> None:
        """Save model/optimizer shards per-rank and scheduler."""
        save_hsdp_checkpoint(
            model=self.model,
            optimizer=self.optimizer,
            lr_scheduler=self.lr_scheduler,
            output_dir=output_dir,
            should_save_scheduler=self.args.should_save and self.lr_scheduler is not None,
        )

    def _load_from_checkpoint(self, resume_from_checkpoint: str, model: Optional[nn.Module] = None) -> None:
        """Load model from HSDP sharded checkpoint."""
        target = model if model is not None else self.model
        loaded = load_hsdp_model(target, resume_from_checkpoint)
        if not loaded:
            return super()._load_from_checkpoint(resume_from_checkpoint, model=model)
        self._pending_hsdp_checkpoint = resume_from_checkpoint
        return None

    def _load_optimizer_and_scheduler(self, checkpoint: Optional[str] = None) -> None:
        """Load optimizer/scheduler from per-rank checkpoint files."""
        ckpt_dir = getattr(self, "_pending_hsdp_checkpoint", None) or checkpoint
        if ckpt_dir is None:
            return
        load_hsdp_optimizer_and_scheduler(self.optimizer, self.lr_scheduler, ckpt_dir)

    def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
        """Save model weights in HuggingFace-compatible format."""
        save_dir = output_dir or self.args.output_dir
        os.makedirs(save_dir, exist_ok=True)
        export_to_hf_format(self.model, getattr(self, "processing_class", None), save_dir)