Repository: hiyouga/LLaMA-Factory
File: src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py (Python)

Code before:
    def shard_model(self, model: torch.nn.Module) -> "AccelerateDeepSpeedEngine":
        """No-op shard — actual model wrapping happens in prepare().

        Returns self so the caller gets the engine instance via the hub interface.
        """
        return self

    def prepare(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        lr_scheduler: Optional[Any] = None,
    ) -> tuple[torch.nn.Module, torch.optim.Optimizer, Any]:
        """Prepare model, optimizer, and lr_scheduler using accelerate.

        Internally calls deepspeed.initialize() and wraps the returned objects.
        """
        if lr_scheduler is not None:
            model, optimizer, lr_scheduler = self.accelerator.prepare(model, optimizer, lr_scheduler)
        else:
            model, optimizer = self.accelerator.prepare(model, optimizer)

        logger.info_rank0("Model, optimizer, and lr_scheduler prepared via accelerate")
        return model, optimizer, lr_scheduler

    def backward(self, loss: torch.Tensor) -> None:
        """Backward pass using accelerate.

        Delegates to DeepSpeedEngineWrapper.backward() which respects
        sync_gradients to control gradient accumulation boundaries.
        When sync_gradients=True: engine.backward(loss) + engine.step()
        When sync_gradients=False: engine.backward(loss) only
        """
        self.accelerator.backward(loss)

    def get_grad_norm(self) -> float:
        """Get the global gradient norm from the DeepSpeed engine."""
        engine_wrapper = getattr(self.accelerator, "deepspeed_engine_wrapped", None)
        if engine_wrapper is not None:
            return engine_wrapper.engine.get_global_grad_norm() or 0.0
        return 0.0

    def save_model(self, trainer) -> None:
        """Save model using accelerate's built-in ZeRO-aware utilities.

        Handles ZeRO-3 parameter gathering automatically via
        accelerator.get_state_dict().
        """
        output_dir = trainer.args.output_dir
        unwrapped_model = self.accelerator.unwrap_model(trainer.model)
        state_dict = self.accelerator.get_state_dict(trainer.model)

Review comment (suggestion):
same as others, even if you use accelerate to do deepspeed strategy. It's better to focus on `deepspeed` 

Code after:

    def shard_model(self, model: HFModel) -> "DeepSpeedEngine":
        """No-op shard — actual model wrapping happens in prepare().

        Returns self so the caller gets the engine instance via the hub interface.
        """
        return self

    def prepare(
        self,
        model: HFModel,
        optimizer: torch.optim.Optimizer,
        lr_scheduler: Optional[Any] = None,
    ) -> tuple[HFModel, torch.optim.Optimizer, Any]:
        """Prepare model, optimizer, and lr_scheduler using accelerate.

        Internally calls deepspeed.initialize() and wraps the returned objects.
        """
        if lr_scheduler is not None:
            model, optimizer, lr_scheduler = self.accelerator.prepare(model, optimizer, lr_scheduler)
        else:
            model, optimizer = self.accelerator.prepare(model, optimizer)

        model._accelerator = self.accelerator  # type: ignore[assignment]

        logger.info_rank0("Model, optimizer, and lr_scheduler prepared via accelerate")
        return model, optimizer, lr_scheduler

    def backward(self, loss: torch.Tensor) -> None:
        """Backward pass using accelerate.

        Delegates to DeepSpeedEngineWrapper.backward() which respects
        sync_gradients to control gradient accumulation boundaries.
        When sync_gradients=True: engine.backward(loss) + engine.step()
        When sync_gradients=False: engine.backward(loss) only
        """
        self.accelerator.backward(loss)

    def get_grad_norm(self) -> float:
        """Get the global gradient norm from the DeepSpeed engine."""
        engine_wrapper = getattr(self.accelerator, "deepspeed_engine_wrapped", None)
        if engine_wrapper is not None:
            return engine_wrapper.engine.get_global_grad_norm() or 0.0
        return 0.0


def save_model(model: HFModel, output_dir: str, processor: Processor) -> None:
    """Save model using accelerate's built-in ZeRO-aware utilities.

    Expects model._accelerator to be set during prepare().
    Handles ZeRO-3 parameter gathering automatically via