Source code for pytorch_lightning.plugins.precision.colossalai

# Copyright The PyTorch Lightning 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Callable, Optional, Union

from lightning_utilities.core.rank_zero import WarningCache
from torch import Tensor
from torch.optim import Optimizer

import pytorch_lightning as pl
from lightning_lite.utilities.types import Steppable
from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
from pytorch_lightning.utilities.enums import PrecisionType

warning_cache = WarningCache()

[docs]class ColossalAIPrecisionPlugin(PrecisionPlugin): """Precision plugin for ColossalAI integration. Args: precision: Half precision (16). Raises: ValueError: If precison is not 16. """ def __init__(self, precision: Union[str, int] = 16) -> None: if not (precision == PrecisionType.HALF): raise ValueError( f"`Trainer(strategy='colossalai', precision={precision!r})` is not supported." " Consider setting `precision=16`." ) super().__init__() self.precision = precision
[docs] def backward( # type: ignore[override] self, tensor: Tensor, model: "pl.LightningModule", optimizer: Optional[Steppable], optimizer_idx: Optional[int], *args: Any, **kwargs: Any, ) -> None: assert optimizer is not None optimizer.backward(tensor)
[docs] def clip_grad_by_norm(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None: optimizer.clip_grad_norm(None, clip_val)
[docs] def clip_grad_by_value(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None: raise NotImplementedError("`clip_grad_by_value` is not supported by `ColossalAI`")
[docs] def optimizer_step( # type: ignore[override] self, optimizer: Steppable, model: "pl.LightningModule", optimizer_idx: int, closure: Callable[[], Any], **kwargs: Any, ) -> Any: closure_result = closure() self._after_closure(model, optimizer, optimizer_idx) skipped_backward = closure_result is None if isinstance(model, pl.LightningModule) and model.automatic_optimization and skipped_backward: raise ValueError( "Skipping backward by returning `None` from your `training_step` is not supported by `ColossalAI`." ) optimizer.step()
def _track_grad_norm(self, trainer: "pl.Trainer") -> None: if trainer.track_grad_norm == -1: return # the gradients are not available in the model due to gradient partitioning in zero stage >= 2 warning_cache.warn( f"You set `Trainer(track_grad_norm={trainer.track_grad_norm!r})' but this is not supported for ColossalAI." " The setting will be ignored." )

© Copyright Copyright (c) 2018-2023, Lightning AI et al...

Built with Sphinx using a theme provided by Read the Docs.