chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
from .transformers.clip_encoder import DSClipEncoder
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
'''Copyright The Microsoft DeepSpeed Team'''
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ..features.cuda_graph import CUDAGraph
|
||||
|
||||
|
||||
class DSUNet(CUDAGraph, torch.nn.Module):
|
||||
|
||||
def __init__(self, unet, enable_cuda_graph=True):
|
||||
super().__init__(enable_cuda_graph=enable_cuda_graph)
|
||||
self.unet = unet
|
||||
# SD pipeline accesses this attribute
|
||||
self.in_channels = unet.in_channels
|
||||
self.device = self.unet.device
|
||||
self.dtype = self.unet.dtype
|
||||
self.config = self.unet.config
|
||||
self.fwd_count = 0
|
||||
self.unet.requires_grad_(requires_grad=False)
|
||||
self.unet.to(memory_format=torch.channels_last)
|
||||
self.cuda_graph_created = False
|
||||
|
||||
def _graph_replay(self, *inputs, **kwargs):
|
||||
for i in range(len(inputs)):
|
||||
if torch.is_tensor(inputs[i]):
|
||||
self.static_inputs[i].copy_(inputs[i])
|
||||
for k in kwargs:
|
||||
if torch.is_tensor(kwargs[k]):
|
||||
self.static_kwargs[k].copy_(kwargs[k])
|
||||
get_accelerator().replay_graph(self._cuda_graphs)
|
||||
return self.static_output
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
if self.enable_cuda_graph:
|
||||
if self.cuda_graph_created:
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
else:
|
||||
self._create_cuda_graph(*inputs, **kwargs)
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
return outputs
|
||||
else:
|
||||
return self._forward(*inputs, **kwargs)
|
||||
|
||||
def _create_cuda_graph(self, *inputs, **kwargs):
|
||||
# warmup to create the workspace and cublas handle
|
||||
cuda_stream = torch.cuda.Stream()
|
||||
cuda_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(cuda_stream):
|
||||
for i in range(3):
|
||||
ret = self._forward(*inputs, **kwargs)
|
||||
torch.cuda.current_stream().wait_stream(cuda_stream)
|
||||
|
||||
# create cuda_graph and assign static_inputs and static_outputs
|
||||
self._cuda_graphs = get_accelerator().create_graph()
|
||||
self.static_inputs = inputs
|
||||
self.static_kwargs = kwargs
|
||||
|
||||
with get_accelerator().capture_to_graph(self._cuda_graphs):
|
||||
self.static_output = self._forward(*self.static_inputs, **self.static_kwargs)
|
||||
|
||||
self.cuda_graph_created = True
|
||||
|
||||
def _forward(self,
|
||||
sample,
|
||||
timestamp,
|
||||
encoder_hidden_states,
|
||||
return_dict=True,
|
||||
cross_attention_kwargs=None,
|
||||
timestep_cond=None,
|
||||
added_cond_kwargs=None):
|
||||
if cross_attention_kwargs:
|
||||
return self.unet(sample,
|
||||
timestamp,
|
||||
encoder_hidden_states,
|
||||
return_dict,
|
||||
cross_attention_kwargs=cross_attention_kwargs)
|
||||
else:
|
||||
return self.unet(sample, timestamp, encoder_hidden_states, return_dict)
|
||||
@@ -0,0 +1,151 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ..features.cuda_graph import CUDAGraph
|
||||
|
||||
|
||||
class DSVAE(CUDAGraph, torch.nn.Module):
|
||||
|
||||
def __init__(self, vae, enable_cuda_graph=True):
|
||||
super().__init__(enable_cuda_graph=enable_cuda_graph)
|
||||
self.vae = vae
|
||||
self.config = vae.config
|
||||
self.device = self.vae.device
|
||||
self.dtype = self.vae.dtype
|
||||
self.vae.requires_grad_(requires_grad=False)
|
||||
self.decoder_cuda_graph_created = False
|
||||
self.encoder_cuda_graph_created = False
|
||||
self.all_cuda_graph_created = False
|
||||
|
||||
def _graph_replay_decoder(self, *inputs, **kwargs):
|
||||
for i in range(len(inputs)):
|
||||
if torch.is_tensor(inputs[i]):
|
||||
self.static_decoder_inputs[i].copy_(inputs[i])
|
||||
for k in kwargs:
|
||||
if torch.is_tensor(kwargs[k]):
|
||||
self.static_decoder_kwargs[k].copy_(kwargs[k])
|
||||
get_accelerator().replay_graph(self._decoder_cuda_graph)
|
||||
return self.static_decoder_output
|
||||
|
||||
def _decode(self, x, return_dict=True, generator=None):
|
||||
return self.vae.decode(x, return_dict=return_dict)
|
||||
|
||||
def _create_cuda_graph_decoder(self, *inputs, **kwargs):
|
||||
# warmup to create the workspace and cublas handle
|
||||
cuda_stream = torch.cuda.Stream()
|
||||
cuda_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(cuda_stream):
|
||||
for i in range(3):
|
||||
ret = self._decode(*inputs, **kwargs)
|
||||
torch.cuda.current_stream().wait_stream(cuda_stream)
|
||||
|
||||
# create cuda_graph and assign static_inputs and static_outputs
|
||||
self._decoder_cuda_graph = get_accelerator().create_graph()
|
||||
self.static_decoder_inputs = inputs
|
||||
self.static_decoder_kwargs = kwargs
|
||||
|
||||
with get_accelerator().capture_to_graph(self._decoder_cuda_graph):
|
||||
self.static_decoder_output = self._decode(*self.static_decoder_inputs, **self.static_decoder_kwargs)
|
||||
|
||||
self.decoder_cuda_graph_created = True
|
||||
|
||||
def decode(self, *inputs, **kwargs):
|
||||
if self.enable_cuda_graph:
|
||||
if self.decoder_cuda_graph_created:
|
||||
outputs = self._graph_replay_decoder(*inputs, **kwargs)
|
||||
else:
|
||||
self._create_cuda_graph_decoder(*inputs, **kwargs)
|
||||
outputs = self._graph_replay_decoder(*inputs, **kwargs)
|
||||
return outputs
|
||||
else:
|
||||
return self._decode(*inputs, **kwargs)
|
||||
|
||||
def _graph_replay_encoder(self, *inputs, **kwargs):
|
||||
for i in range(len(inputs)):
|
||||
if torch.is_tensor(inputs[i]):
|
||||
self.static_encoder_inputs[i].copy_(inputs[i])
|
||||
for k in kwargs:
|
||||
if torch.is_tensor(kwargs[k]):
|
||||
self.static_encoder_kwargs[k].copy_(kwargs[k])
|
||||
get_accelerator().replay_graph(self._encoder_cuda_graph)
|
||||
return self.static_encoder_output
|
||||
|
||||
def _encode(self, x, return_dict=True):
|
||||
return self.vae.encode(x, return_dict=return_dict)
|
||||
|
||||
def _create_cuda_graph_encoder(self, *inputs, **kwargs):
|
||||
# warmup to create the workspace and cublas handle
|
||||
cuda_stream = torch.cuda.Stream()
|
||||
cuda_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(cuda_stream):
|
||||
for i in range(3):
|
||||
ret = self._encode(*inputs, **kwargs)
|
||||
torch.cuda.current_stream().wait_stream(cuda_stream)
|
||||
|
||||
# create cuda_graph and assign static_inputs and static_outputs
|
||||
self._encoder_cuda_graph = get_accelerator().create_graph()
|
||||
self.static_encoder_inputs = inputs
|
||||
self.static_encoder_kwargs = kwargs
|
||||
|
||||
with get_accelerator().capture_to_graph(self._encoder_cuda_graph):
|
||||
self.static_encoder_output = self._encode(*self.static_encoder_inputs, **self.static_encoder_kwargs)
|
||||
|
||||
self.encoder_cuda_graph_created = True
|
||||
|
||||
def encode(self, *inputs, **kwargs):
|
||||
if self.enable_cuda_graph:
|
||||
if self.encoder_cuda_graph_created:
|
||||
outputs = self._graph_replay_encoder(*inputs, **kwargs)
|
||||
else:
|
||||
self._create_cuda_graph_encoder(*inputs, **kwargs)
|
||||
outputs = self._graph_replay_encoder(*inputs, **kwargs)
|
||||
return outputs
|
||||
else:
|
||||
return self._encode(*inputs, **kwargs)
|
||||
|
||||
def _graph_replay(self, *inputs, **kwargs):
|
||||
for i in range(len(inputs)):
|
||||
if torch.is_tensor(inputs[i]):
|
||||
self.static_inputs[i].copy_(inputs[i])
|
||||
for k in kwargs:
|
||||
if torch.is_tensor(kwargs[k]):
|
||||
self.static_kwargs[k].copy_(kwargs[k])
|
||||
get_accelerator().replay_graph(self._all_cuda_graph)
|
||||
return self.static_output
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
if self.enable_cuda_graph:
|
||||
if self.cuda_graph_created:
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
else:
|
||||
self._create_cuda_graph(*inputs, **kwargs)
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
return outputs
|
||||
else:
|
||||
return self._forward(*inputs, **kwargs)
|
||||
|
||||
def _create_cuda_graph(self, *inputs, **kwargs):
|
||||
# warmup to create the workspace and cublas handle
|
||||
cuda_stream = torch.cuda.Stream()
|
||||
cuda_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(cuda_stream):
|
||||
for i in range(3):
|
||||
ret = self._forward(*inputs, **kwargs)
|
||||
torch.cuda.current_stream().wait_stream(cuda_stream)
|
||||
|
||||
# create cuda_graph and assign static_inputs and static_outputs
|
||||
self._all_cuda_graph = get_accelerator().create_graph()
|
||||
self.static_inputs = inputs
|
||||
self.static_kwargs = kwargs
|
||||
|
||||
with get_accelerator().capture_to_graph(self._all_cuda_graph):
|
||||
self.static_output = self._forward(*self.static_inputs, **self.static_kwargs)
|
||||
|
||||
self.all_cuda_graph_created = True
|
||||
|
||||
def _forward(self, sample, timestamp, encoder_hidden_states, return_dict=True):
|
||||
return self.vae(sample, timestamp, encoder_hidden_states, return_dict)
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
'''Copyright The Microsoft DeepSpeed Team'''
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class CUDAGraph(ABC):
|
||||
|
||||
def __init__(self, enable_cuda_graph=False):
|
||||
super().__init__()
|
||||
self.enable_cuda_graph = enable_cuda_graph
|
||||
|
||||
@abstractmethod
|
||||
def _create_cuda_graph(self):
|
||||
"""
|
||||
Create CUDA graph(s)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _graph_replay(self):
|
||||
"""
|
||||
Replay CUDA graph(s)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
'''Copyright The Microsoft DeepSpeed Team'''
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ..features.cuda_graph import CUDAGraph
|
||||
|
||||
|
||||
class DSClipEncoder(CUDAGraph, torch.nn.Module):
|
||||
|
||||
def __init__(self, enc, enable_cuda_graph=False):
|
||||
super().__init__(enable_cuda_graph=enable_cuda_graph)
|
||||
enc.text_model._build_causal_attention_mask = self._build_causal_attention_mask
|
||||
self.enc = enc
|
||||
self.device = self.enc.device
|
||||
self.dtype = self.enc.dtype
|
||||
self.cuda_graph_created = [False, False]
|
||||
self.static_inputs = [None, None]
|
||||
self.static_kwargs = [None, None]
|
||||
self.static_output = [None, None]
|
||||
self._cuda_graphs = [None, None]
|
||||
self.iter = 0
|
||||
self.config = self.enc.config
|
||||
|
||||
def _build_causal_attention_mask(self, bsz, seq_len, dtype):
|
||||
mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype, device=get_accelerator().current_device_name())
|
||||
mask.fill_(torch.tensor(torch.finfo(dtype).min))
|
||||
mask.triu_(1)
|
||||
mask = mask.unsqueeze(1)
|
||||
return mask
|
||||
|
||||
def _graph_replay(self, *inputs, **kwargs):
|
||||
for i in range(len(inputs)):
|
||||
if torch.is_tensor(inputs[i]):
|
||||
self.static_inputs[self.iter][i].copy_(inputs[i])
|
||||
for k in kwargs:
|
||||
if torch.is_tensor(kwargs[k]):
|
||||
self.static_kwargs[self.iter][k].copy_(kwargs[k])
|
||||
get_accelerator().replay_graph(self._cuda_graphs[self.iter])
|
||||
return self.static_output[self.iter]
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
if self.enable_cuda_graph:
|
||||
if self.cuda_graph_created[self.iter]:
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
else:
|
||||
self._create_cuda_graph(*inputs, **kwargs)
|
||||
outputs = self._graph_replay(*inputs, **kwargs)
|
||||
self.iter = (self.iter + 1) % 2
|
||||
return outputs
|
||||
else:
|
||||
return self.enc(*inputs, **kwargs)
|
||||
|
||||
def _create_cuda_graph(self, *inputs, **kwargs):
|
||||
# warmup to create the workspace and cublas handle
|
||||
cuda_stream = torch.cuda.Stream()
|
||||
cuda_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(cuda_stream):
|
||||
for i in range(3):
|
||||
ret = self._forward(*inputs, **kwargs)
|
||||
torch.cuda.current_stream().wait_stream(cuda_stream)
|
||||
|
||||
# create cuda_graph and assign static_inputs and static_outputs
|
||||
self._cuda_graphs[self.iter] = get_accelerator().create_graph()
|
||||
self.static_inputs[self.iter] = inputs
|
||||
self.static_kwargs[self.iter] = kwargs
|
||||
|
||||
with get_accelerator().capture_to_graph(self._cuda_graphs[self.iter]):
|
||||
self.static_output[self.iter] = self._forward(*self.static_inputs[self.iter],
|
||||
**self.static_kwargs[self.iter])
|
||||
|
||||
self.cuda_graph_created[self.iter] = True
|
||||
|
||||
def _forward(self, *inputs, **kwargs):
|
||||
return self.enc(*inputs, **kwargs)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class DeepSpeedTransformerBase(nn.module):
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
# this would be the new clean base class that will replace DeepSpeedTransformerInference.
|
||||
# we currently don't know how this will look like but keeping it here as a placeholder.
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedBERTInference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed BERT Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedBloomInference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed Bloom Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedGPTInference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed GPT Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedLlama2Inference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed OPT Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
|
||||
input = args[0]
|
||||
input_mask = None
|
||||
get_present = True
|
||||
|
||||
self.allocate_workspace(input.size())
|
||||
|
||||
# We set the prev key/value to None when there is a prompt
|
||||
if input.shape[1] > 1:
|
||||
self.layer_past = None
|
||||
layer_past = self.layer_past
|
||||
|
||||
input_type = input.dtype
|
||||
|
||||
if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \
|
||||
and input.dtype == torch.float:
|
||||
target_dtype = torch.half if self.dtype == torch.int8 else self.dtype
|
||||
input = input.to(target_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
attention_output, key, value, context_outputtn_ctx, inp_norm = \
|
||||
self.attention(input,
|
||||
input_mask,
|
||||
None,
|
||||
layer_past,
|
||||
get_present,
|
||||
None, None, None,
|
||||
self.norm_w,
|
||||
self.norm_b,
|
||||
None)
|
||||
self.layer_past = (key, value)
|
||||
output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob)
|
||||
|
||||
output = output.to(input_type)
|
||||
return output
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedMegatronGPTInference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed Megatron GPT Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
|
||||
|
||||
|
||||
class DeepSpeedOPTInference(DeepSpeedTransformerInference):
|
||||
"""Initialize the DeepSpeed OPT Transformer Layer.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super().__init__(config, mp_group, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping)
|
||||
@@ -0,0 +1,198 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from deepspeed import comm as dist
|
||||
from deepspeed.ops.transformer.inference.op_binding.layer_norm import LayerNormOp
|
||||
from deepspeed.utils.logging import log_dist
|
||||
|
||||
from deepspeed.ops.transformer.inference.ds_mlp import DeepSpeedMLP
|
||||
from deepspeed.ops.transformer.inference.ds_attention import DeepSpeedSelfAttention, BloomSelfAttention
|
||||
from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
import deepspeed
|
||||
# Import the triton kernels whenever triton is installed. Previously this was also
|
||||
# gated on is_triton_supported(), which reads the GPU compute capability at import
|
||||
# time and thereby creates a CUDA context, breaking fork()-based multiprocessing
|
||||
# (issue #7918). Triton use is gated at runtime via self.config.use_triton below.
|
||||
if deepspeed.HAS_TRITON:
|
||||
from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP
|
||||
from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention
|
||||
|
||||
|
||||
class DeepSpeedTransformerInference(nn.Module):
|
||||
"""Initialize the DeepSpeed Transformer Layer.
|
||||
Arguments:
|
||||
layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers,
|
||||
layer_id will be 0,1,2...23 when each layer object is instantiated
|
||||
config: An object of DeepSpeedInferenceConfig
|
||||
mp_group: Model parallelism group initialized on the modeling side.
|
||||
quantize_scales: This argument groups all the layers' scales used for quantization
|
||||
quantize_groups: Number of groups used for quantizing the model
|
||||
merge_count: Shows the number of model-parallel checkpoints merged before running inference.
|
||||
We use this argument to control the quantization scale for the model parameters if a bigger
|
||||
quantize-grouping than 1 is used.
|
||||
mlp_extra_grouping: This flag is used to show a 2x higher number of groups used for the MLP part
|
||||
of a Transformer layer. We use this feature for quantization to reduce the convergence impact
|
||||
for specific downstream tasks.
|
||||
"""
|
||||
layer_id = 0
|
||||
workspace = None
|
||||
|
||||
def __init__(self,
|
||||
config,
|
||||
mp_group=None,
|
||||
quantize_scales=None,
|
||||
quantize_groups=1,
|
||||
merge_count=1,
|
||||
mlp_extra_grouping=False):
|
||||
super(DeepSpeedTransformerInference, self).__init__()
|
||||
|
||||
self.config = config
|
||||
self.config.layer_id = DeepSpeedTransformerInference.layer_id
|
||||
DeepSpeedTransformerInference.layer_id += 1
|
||||
|
||||
data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype
|
||||
|
||||
if DeepSpeedTransformerInference.layer_id == 1:
|
||||
log_dist(f"DeepSpeed-Inference config: {self.config.__dict__}", [0])
|
||||
if deepspeed.HAS_TRITON and self.config.use_triton:
|
||||
log_dist("Injecting Triton kernels ...", [0])
|
||||
|
||||
if self.config.bigscience_bloom:
|
||||
self.attention = BloomSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count)
|
||||
assert not self.config.use_triton
|
||||
else:
|
||||
if deepspeed.HAS_TRITON and self.config.use_triton:
|
||||
self.attention = TritonSelfAttention(self.config)
|
||||
else:
|
||||
self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups,
|
||||
merge_count)
|
||||
|
||||
if deepspeed.HAS_TRITON and self.config.use_triton:
|
||||
self.mlp = TritonMLP(self.config)
|
||||
else:
|
||||
self.mlp = DeepSpeedMLP(self.config, mp_group, quantize_scales, quantize_groups, merge_count,
|
||||
mlp_extra_grouping)
|
||||
|
||||
device = get_accelerator().current_device_name() # if config.bigscience_bloom else 'cpu'
|
||||
if self.config.set_empty_params:
|
||||
self.norm_w = None
|
||||
self.norm_b = None
|
||||
else:
|
||||
self.norm_w = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device),
|
||||
requires_grad=False)
|
||||
self.norm_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device),
|
||||
requires_grad=False)
|
||||
self.layer_past = None
|
||||
self.layer_norm = LayerNormOp()
|
||||
if DeepSpeedTransformerInference.workspace is None:
|
||||
DeepSpeedTransformerInference.workspace = WorkspaceOp(self.config)
|
||||
self._should_allocate_workspace = True
|
||||
|
||||
def allocate_workspace(self, size):
|
||||
# Allocate memory only on first layer forward
|
||||
if self.config.layer_id == 0 and self._should_allocate_workspace:
|
||||
DeepSpeedTransformerInference.workspace.allocate_workspace(
|
||||
self.config.hidden_size, self.config.heads, size[1], size[0], DeepSpeedTransformerInference.layer_id,
|
||||
self.config.mp_size, self.config.bigscience_bloom,
|
||||
dist.get_rank() if dist.is_initialized() else 0, self.config.max_out_tokens,
|
||||
self.config.min_out_tokens)
|
||||
self._should_allocate_workspace = False
|
||||
|
||||
@classmethod
|
||||
def reset_cache(cls):
|
||||
if cls.workspace is not None:
|
||||
cls.workspace.reset_cache()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input=None,
|
||||
input_mask=None,
|
||||
attention_mask=None,
|
||||
attn_mask=None,
|
||||
head_mask=None,
|
||||
layer_past=None,
|
||||
get_key_value=False,
|
||||
get_present=False,
|
||||
encoder_output=None,
|
||||
enc_dec_attn_mask=None,
|
||||
x=None,
|
||||
encoder_hidden_states=None,
|
||||
encoder_attention_mask=None,
|
||||
use_cache=False,
|
||||
alibi=None,
|
||||
output_attentions=False,
|
||||
# TODO(arashb): 'layer_head_mask' and 'past_key_value' are only added to satisfy the OPT models API.
|
||||
# This needs to be redesigned later!
|
||||
layer_head_mask=None,
|
||||
past_key_value=None,
|
||||
**kwargs):
|
||||
|
||||
if x is not None:
|
||||
input = x
|
||||
if "hidden_states" in kwargs:
|
||||
input = kwargs["hidden_states"]
|
||||
|
||||
if layer_past is not None and past_key_value is not None:
|
||||
raise ValueError("Only one of `layer_past` or `past_key_value` can be present.")
|
||||
|
||||
input_mask = (input_mask if attn_mask is None else attn_mask) if attention_mask is None else attention_mask
|
||||
|
||||
self.allocate_workspace(input.size())
|
||||
|
||||
get_present = (get_present or get_key_value or use_cache)
|
||||
input_mask = input_mask if attention_mask is None else attention_mask
|
||||
|
||||
# We set the prev key/value to None when there is a prompt
|
||||
if input.shape[1] > 1:
|
||||
self.layer_past = None
|
||||
_layer_past = layer_past or past_key_value or self.layer_past
|
||||
head_mask = layer_head_mask if layer_head_mask is not None else head_mask
|
||||
|
||||
attn_mask = None
|
||||
if isinstance(input, tuple):
|
||||
attn_mask = input[1]
|
||||
input = input[0]
|
||||
input_type = input.dtype
|
||||
|
||||
if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \
|
||||
and input.dtype == torch.float:
|
||||
target_dtype = torch.half if self.config.dtype == torch.int8 else self.config.dtype
|
||||
input = input.to(target_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
attention_output, key, value, context_outputtn_ctx, inp_norm = \
|
||||
self.attention(input,
|
||||
input_mask,
|
||||
head_mask,
|
||||
_layer_past,
|
||||
get_present,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
output_attentions,
|
||||
self.norm_w,
|
||||
self.norm_b,
|
||||
alibi,
|
||||
**kwargs)
|
||||
|
||||
presents = (key, value)
|
||||
self.layer_past = presents if layer_past is None and past_key_value is None else None
|
||||
output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob)
|
||||
|
||||
if not self.config.pre_layer_norm:
|
||||
output = self.layer_norm(output, self.norm_w, self.norm_b, self.config.epsilon)
|
||||
|
||||
output = output.to(input_type)
|
||||
if get_present:
|
||||
output = (output, presents)
|
||||
|
||||
if self.config.return_single_tuple:
|
||||
return (output, )
|
||||
elif self.config.return_tuple:
|
||||
return output if type(output) is tuple else (output, attn_mask)
|
||||
else:
|
||||
return output
|
||||
Reference in New Issue
Block a user