Deepspeed integration

allennlp/training/__init__.py

@@ -9,3 +9,10 @@
     TrainerCallback,
     TrackEpochCallback,
 )
+from allennlp.training.deepspeed import DeepspeedTrainer
+
+# import warnings
+# try:
+#     from allennlp.training.deepspeed import DeepspeedTrainer
+# except ImportError:
+#     warnings.warn('Deepspeed plugin not installed. Ignoring.')

allennlp/training/deepspeed/__init__.py

@@ -0,0 +1,11 @@
+from allennlp.training.deepspeed.trainer import DeepspeedTrainer
+from allennlp.training.deepspeed.optimizers import (
+    FusedAdamOptimizer,
+    DeepspeedCPUAdamOptimizer,
+    FusedLambOptimizer
+)
+
+try:
+    from allennlp.training.deepspeed.sparse_transformer_embedder import SparseTransformerEmbedder
+except ImportError:
+    pass

allennlp/training/deepspeed/checkpointer.py

@@ -0,0 +1,131 @@
+from typing import Union, Dict, Any, List, Tuple, Optional
+
+import logging
+import os
+import re
+import shutil
+import time
+
+from pathlib import Path
+
+import torch
+
+import allennlp
+from allennlp.nn import util as nn_util
+from allennlp.training import util as training_util, Checkpointer
+
+logger = logging.getLogger(__name__)
+_DeepspeedTrainer = "allennlp.training.deepspeed.trainer.DeepspeedTrainer"
+
+
+class DeepspeedCheckpointer(Checkpointer):
+    # def maybe_save_checkpoint(
+    #     self, 
+    #     trainer: _DeepspeedTrainer, 
+    #     epoch: int, 
+    #     batches_this_epoch: int
+    # ) -> None:
+    #     0/0
+
+    def save_checkpoint(
+        self,
+        epoch: Union[int, str],
+        trainer: _DeepspeedTrainer,
+        is_best_so_far: bool = False,
+        save_model_only=False,
+    ) -> None:
+        if self._serialization_dir is None:
+            return
+
+        with trainer.get_checkpoint_state() as state:
+            model_engine, model_state, training_states = state
+
+            checkpoint_id = "deepspeed_epoch_{}".format(epoch)
+            model_path = os.path.join(self._serialization_dir, "model_state_epoch_{}".format(epoch))
+            model_engine.save_checkpoint(self._serialization_dir, checkpoint_id)
+
+            # TODO
+            # Model will need a weight file to load; 
+            # not sure if ZeRO stage 2 will mess this up
+            if not os.path.isfile(model_path):
+                    torch.save(model_state, model_path)
+            if save_model_only:
+                return
+
+            training_path = os.path.join(
+                self._serialization_dir, "training_state_epoch_{}.th".format(epoch)
+            )
+            if not os.path.isfile(training_path):
+                torch.save({**training_states, "epoch": epoch}, training_path)
+
+        # The main checkpointing logic is now done, this is just shuffling files around, to keep
+        # track of best weights, and to remove old checkpoints, if desired.
+        if is_best_so_far:
+            logger.info(
+                "Best validation performance so far. Copying weights to '%s/best.th'.",
+                self._serialization_dir,
+            )
+            shutil.copyfile(model_path, os.path.join(self._serialization_dir, "best.th"))
+
+            engine_dir = os.path.join(self._serialization_dir, "best_deepspeed")
+            shutil.rmtree(engine_dir, ignore_errors=True) # in case no previous checkpoints
+            shutil.copytree(os.path.join(self._serialization_dir, checkpoint_id), engine_dir)
+
+        if (
+            self._num_serialized_models_to_keep is not None
+            and self._num_serialized_models_to_keep >= 0
+        ):
+            self._serialized_paths.append((time.time(), model_path, training_path))
+            if len(self._serialized_paths) > self._num_serialized_models_to_keep:
+                paths_to_remove = self._serialized_paths.pop(0)
+                # Check to see if we should keep this checkpoint, if it has been longer
+                # then self._keep_serialized_model_every_num_seconds since the last
+                # kept checkpoint.
+                remove_path = True
+                if self._keep_serialized_model_every_num_seconds is not None:
+                    save_time = paths_to_remove[0]
+                    time_since_checkpoint_kept = (
+                        save_time - self._last_permanent_saved_checkpoint_time
+                    )
+                    if (
+                        time_since_checkpoint_kept
+                        > self._keep_serialized_model_every_num_seconds
+                    ):
+                        # We want to keep this checkpoint.
+                        remove_path = False
+                        self._last_permanent_saved_checkpoint_time = save_time
+                if remove_path:
+                    for fname in paths_to_remove[1:]:
+                        if os.path.isfile(fname):
+                            os.remove(fname)
+
+    def find_latest_checkpoint(self) -> Optional[Tuple[str, str]]:
+        latest = super().find_latest_checkpoint()
+        if not latest:
+            return None
+
+        model_path, training_state_path = latest
+
+        checkpoints = (self._serialization_dir and Path(self._serialization_dir).glob('deepspeed_epoch_*')) or []
+        checkpoints = sorted(c for c in checkpoints if c.is_dir())
+        if not checkpoints:
+            return None
+
+        engine_path = checkpoints[-1]
+        return engine_path, model_path, training_state_path
+
+    def restore_checkpoint(self) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        latest_checkpoint = self.find_latest_checkpoint()
+
+        if latest_checkpoint is None:
+            # No checkpoint to restore, start at 0
+            return {}, {}, {}
+
+        checkpoint_id, model_path, training_state_path = latest_checkpoint
+
+        model_state = torch.load(model_path, map_location=nn_util.device_mapping(-1))
+        training_state = torch.load(training_state_path, map_location=nn_util.device_mapping(-1))
+        return checkpoint_id, model_state, training_state
+
+    def best_model_state(self) -> Dict[str, Any]:
+        pass

allennlp/training/deepspeed/config.py

@@ -0,0 +1,68 @@
+from typing import Dict, Any
+from enum import IntEnum
+from allennlp.common import FromParams
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class DeepspeedFP16Config(FromParams):
+    enabled: bool = True
+    loss_scale: float = 0.
+    initial_scale_power: int = 32
+    loss_scale_window: int = 1000
+    hysteresis: int = 2
+    min_loss_scale: float = 1.
+
+@dataclass
+class DeepspeedAMPConfig(FromParams):
+    enabled: bool = False
+    opt_level: str = "O1"
+
+@dataclass
+class DeepspeedOptimizerConfig(FromParams):
+    type: str
+    params: Dict[str, Any]
+
+@dataclass
+class DeepspeedLRSchedulerConfig(FromParams):
+    type: str
+    params: Dict[str, Any]
+
+class DeepspeedZeROStage(IntEnum):
+    DISABLED = 0
+    OPTIMIZER = 1
+    GRADIENT = 2
+
+@dataclass
+class DeepspeedZeROConfig(FromParams):
+    stage: DeepspeedZeROStage = DeepspeedZeROStage.GRADIENT
+    allgather_partitions: bool = True
+    allgather_bucket_size: int = 500000000
+    overlap_comm: bool = False
+    reduce_scatter: bool = True
+    reduce_bucket_size: int = 500000000
+    contiguous_gradients: bool = False
+    cpu_offload: bool = False
+
+
+@dataclass
+class DeepspeedConfig(FromParams):
+    zero_optimization: DeepspeedZeROConfig
+    fp16: DeepspeedFP16Config
+    amp: DeepspeedAMPConfig = DeepspeedAMPConfig()
+    optimizer: DeepspeedOptimizerConfig = None
+    scheduler: DeepspeedLRSchedulerConfig = None
+
+    zero_allow_untested_optimizer: bool = True
+    wall_clock_breakdown: bool = False
+
+    def to_dict(self):
+        return asdict(self)
+
+
+@dataclass
+class DeepspeedArgs(FromParams):
+    local_rank: int
+    deepspeed: bool = True
+    deepspeed_mpi: bool = False
+    deepspeed_config: str = None

allennlp/training/deepspeed/optimizers.py

@@ -0,0 +1,87 @@
+from typing import List, Tuple, Dict, Any
+
+import torch
+
+from apex.optimizers.fused_adam import FusedAdam
+from deepspeed.ops.adam import DeepSpeedCPUAdam
+from deepspeed.ops.lamb import FusedLamb
+from deepspeed.runtime.fp16.onebit_adam import OnebitAdam
+
+from allennlp.training.optimizers import Optimizer, make_parameter_groups
+
+@Optimizer.register("fused_adam")
+class FusedAdamOptimizer(Optimizer, FusedAdam):
+    def __init__(
+        self,
+        model_parameters: List[Tuple[str, torch.nn.Parameter]],
+        parameter_groups: List[Tuple[List[str], Dict[str, Any]]] = None,
+        lr: float = 0.001,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-08,
+        weight_decay: float = 0.0,
+        amsgrad: bool = False,
+        bias_correction: bool =True,
+        adam_w_mode: bool = True,
+        set_grad_none: bool = True,
+    ):
+        super().__init__(
+            params=make_parameter_groups(model_parameters, parameter_groups),
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+            bias_correction=bias_correction,
+            adam_w_mode=adam_w_mode,
+            set_grad_none=set_grad_none,
+        )
+
+# This does not currently work
+@Optimizer.register("cpu_adam")
+class DeepspeedCPUAdamOptimizer(Optimizer, DeepSpeedCPUAdam):
+    def __init__(
+        self,
+        model_parameters: List[Tuple[str, torch.nn.Parameter]],
+        parameter_groups: List[Tuple[List[str], Dict[str, Any]]] = None,
+        lr: float = 0.001,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-08,
+        weight_decay: float = 0.0,
+        amsgrad: bool = False,
+    ):
+        super().__init__(
+            model_params=make_parameter_groups(model_parameters, parameter_groups),
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad
+        )
+
+@Optimizer.register("fused_lamb")
+class FusedLambOptimizer(Optimizer, FusedLamb):
+    def __init__(
+        self,
+        model_parameters: List[Tuple[str, torch.nn.Parameter]],
+        parameter_groups: List[Tuple[List[str], Dict[str, Any]]] = None,
+        lr: float = 0.001,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-08,
+        eps_inside_sqrt: bool = False,
+        weight_decay: float = 0.0,
+        amsgrad: bool = False,
+        max_grad_norm: float = 0.,
+        max_coeff: float = 10.0,
+        min_coeff: float = 0.01
+    ):
+        super().__init__(
+            params=make_parameter_groups(model_parameters, parameter_groups),
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+            max_grad_norm=max_grad_norm,
+            max_coeff=max_coeff,
+            min_coeff=min_coeff,
+        )

allennlp/training/deepspeed/sparse_transformer_embedder.py

@@ -0,0 +1,11 @@
+from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
+from allennlp.modules.token_embedders.pretrained_transformer_embedder import PretrainedTransformerEmbedder
+
+from deepspeed.ops.sparse_attention.sparse_attention_utils import SparseAttentionUtils
+
+# Doesn't work yet
+@TokenEmbedder.register('sparse_transformer')
+class SparseTransformerEmbedder(PretrainedTransformerEmbedder):
+    class __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.transformer_model = SparseAttentionUtils.replace_model_self_attention_with_sparse_self_attention(self.transformer_model)