delete unused files

colossalai/amp/apex_amp/apex_amp.py

-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-import torch.nn as nn
-try:
-    import apex.amp as apex_amp
-except ImportError:
-    pass
-
-from torch import Tensor
-
-from colossalai.nn.optimizer import ColossalaiOptimizer
-from colossalai.utils import clip_grad_norm_fp32
-
-
-class ApexAMPOptimizer(ColossalaiOptimizer):
-    """ A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm
-    methods
-    """
-
-    def backward(self, loss: Tensor):
-        """Backward pass to get all gradients
-
-        :param loss: Loss computed by a loss function
-        :type loss: torch.Tensor
-        """
-        with apex_amp.scale_loss(loss, self.optim) as scaled_loss:
-            scaled_loss.backward()
-
-    def clip_grad_norm(self, model: nn.Module, max_norm: float):
-        """Clip gradients' norm
-
-        :param model: Your model object
-        :type model: torch.nn.Module
-        :param max_norm: The max norm value for gradient clipping
-        :type max_norm: float
-        """
-        if max_norm > 0:
-            clip_grad_norm_fp32(apex_amp.master_params(self.optim), max_norm)

colossalai/amp/naive_amp/__init__.py

-import torch.nn as nn
-from torch.optim import Optimizer
-from colossalai.utils import is_no_pp_or_last_stage
-
-from .naive_amp import NaiveAMPOptimizer, NaiveAMPModel
-
-
-def convert_to_naive_amp(model: nn.Module,
-                         optimizer: Optimizer,
-                         amp_config):
-    """A helper function to wrap training components with naive AMP modules
-
-    :param model: your model object
-    :type model: :class:`torch.nn.Module`
-    :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
-    :param amp_config: configuration for naive mode amp
-    :type amp_config: :class:`colossalai.context.Config` or dict
-
-    :return: (model, optimizer)
-    :rtype: Tuple
-    """
-    if isinstance(model, nn.ModuleList):
-        # interleaved pipeline
-        module_list = []
-        for chunk, m in enumerate(model):
-            output_to_fp32 = is_no_pp_or_last_stage() and chunk == len(model) - 1
-            module_list.append(NaiveAMPModel(m, output_to_fp32=output_to_fp32))
-        model = nn.ModuleList(module_list)
-    else:
-        output_to_fp32 = is_no_pp_or_last_stage()
-        model = NaiveAMPModel(model, output_to_fp32=output_to_fp32)
-
-    optimizer = NaiveAMPOptimizer(optimizer, **amp_config)
-    return model, optimizer
-
-
-__all__ = ['convert_to_naive_amp', 'NaiveAMPOptimizer']

colossalai/amp/naive_amp/naive_amp.py

-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-import torch
-import torch.nn as nn
-from torch import Tensor
-from typing import Union, List, Any, Dict
-from torch.optim import Optimizer
-import torch.cuda.amp as torch_amp
-
-from colossalai.nn.optimizer import ColossalaiOptimizer
-from ._fp16_optimizer import FP16Optimizer
-
-
-class NaiveAMPOptimizer(ColossalaiOptimizer):
-    """A wrapper class for optimizer to cast all parameters to fp16
-
-    :param optim: A normal optimizer like Adam or SGD
-    :param args: Args used to initialize FP16 optimizer
-    :param kwargs: Kwargs used to initialize FP16 optimizer
-
-    :type optim: torch.optim.Optimizer
-    """
-
-    def __init__(self, optim: Optimizer, *args, **kwargs):
-        optim = FP16Optimizer(optimizer=optim, *args, **kwargs)
-        super().__init__(optim)
-
-    def backward(self, loss: Tensor):
-        """Backward with gradient scaler
-
-        :param loss: loss computed by a loss function
-        :type loss: torch.Tensor
-        """
-        loss = self.optim.scale_loss(loss)
-        loss.backward()
-
-    def step(self):
-        return self.optim.step()
-
-    def clip_grad_norm(self, model: nn.Module, max_norm: float):
-        pass
-
-
-class NaiveAMPModel(nn.Module):
-    """A wrapper class for model to cast the model into fp16 and 
-    automatically cast the input and output
-    """
-
-    def __init__(self,
-                 model: nn.Module,
-                 output_to_fp32: bool = True):
-        super().__init__()
-        self.model = model.half()
-        self._output_to_fp32 = output_to_fp32
-
-    def _convert_to_fp16(self, input_: Any):
-        if isinstance(input_, Tensor) and input_.dtype == torch.float32:
-            input_ = input_.half()
-        return input_
-
-    def _convert_to_fp32(self, input_: Any):
-        if isinstance(input_, Tensor) and input_.dtype == torch.float16:
-            input_ = input_.float()
-        return input_
-
-    def forward(self, *args, **kwargs):
-        if args:
-            args = [self._convert_to_fp16(arg) for arg in args]
-        if kwargs:
-            for k, v in kwargs.items():
-                kwargs[k] = self._convert_to_fp16(v)
-
-        out = self.model(*args, **kwargs)
-
-        if self._output_to_fp32:
-            if isinstance(out, Tensor):
-                out = self._convert_to_fp32(out)
-            elif isinstance(out, (tuple, list)):
-                out = [self._convert_to_fp32(val) for val in out]
-        return out

colossalai/amp/torch_amp/__init__.py

-import torch.nn as nn
-from torch.optim import Optimizer
-from torch.nn.modules.loss import _Loss
-from colossalai.context import Config
-from .torch_amp import TorchAMPOptimizer, TorchAMPModel, TorchAMPLoss
-
-
-def convert_to_torch_amp(model: nn.Module,
-                         optimizer: Optimizer,
-                         criterion: _Loss,
-                         amp_config: Config):
-    """A helper function to wrap training components with Torch AMP modules
-
-    :param model: your model object
-    :type model: :class:`torch.nn.Module`
-    :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
-    :param criterion: your loss function object
-    :type criterion: :class:`torch.nn.modules.loss._Loss`
-    :param amp_config: configuration for different amp modes
-    :type amp_config: :class:`colossalai.context.Config` or dict
-    
-    :return: (model, optimizer, criterion)
-    :rtype: Tuple
-    """
-    model = TorchAMPModel(model)
-    optimizer = TorchAMPOptimizer(optimizer, **amp_config)
-    criterion = TorchAMPLoss(criterion)
-    return model, optimizer, criterion
-
-
-__all__ = ['convert_to_torch_amp', 'TorchAMPModel', 'TorchAMPLoss', 'TorchAMPOptimizer']