mirror of https://github.com/InternLM/InternLM
1256 lines
59 KiB
Python
1256 lines
59 KiB
Python
# Copyright (c) InternLM. All rights reserved.
|
||
import math
|
||
from typing import Optional
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from einops import rearrange
|
||
from torch import nn
|
||
|
||
from internlm.core.context import IS_SEQUENCE_PARALLEL, IS_TENSOR_PARALLEL, ParallelMode
|
||
from internlm.core.context.parallel_context import global_context as gpc
|
||
from internlm.core.naive_amp import set_fp32_attr_to_module
|
||
from internlm.initialize.initialize_tensor import (
|
||
normal_,
|
||
scaled_init_method_normal,
|
||
scaled_init_method_uniform,
|
||
uniform_,
|
||
)
|
||
from internlm.model.embedding import Embedding1D, RotaryEmbedding
|
||
from internlm.model.linear import (
|
||
ColumnParallelLinearTorch,
|
||
FeedForward,
|
||
RewardModelLinear,
|
||
RowParallelLinearTorch,
|
||
ScaleColumnParallelLinear,
|
||
)
|
||
from internlm.model.moe import MoE
|
||
from internlm.model.utils import gather_forward_split_backward, try_import_RMSNorm
|
||
from internlm.solver.pipeline_utils import partition_uniform
|
||
from internlm.utils.checkpoint import activation_checkpoint
|
||
from internlm.utils.common import filter_kwargs
|
||
from internlm.utils.logger import get_logger
|
||
from internlm.utils.registry import MODEL_INITIALIZER
|
||
|
||
try:
|
||
from flash_attn import flash_attn_varlen_kvpacked_func
|
||
from flash_attn.flash_attn_interface import FlashAttnVarlenKVPackedFunc
|
||
from flash_attn.modules.embedding import ParallelGPT2Embeddings
|
||
from flash_attn.modules.mha import (
|
||
CrossAttention,
|
||
FlashCrossAttention,
|
||
FlashSelfAttention,
|
||
SelfAttention,
|
||
_update_kv_cache,
|
||
)
|
||
from flash_attn.modules.mlp import ParallelFusedMLP
|
||
from flash_attn.ops.layer_norm import dropout_add_layer_norm
|
||
except ImportError:
|
||
pass
|
||
|
||
MODEL_TYPE = "LLAMA_MoE"
|
||
|
||
logger = get_logger(__file__)
|
||
RMSNorm = try_import_RMSNorm()
|
||
|
||
|
||
class MHA(nn.Module):
|
||
"""
|
||
Multi-head self-attention and cross-attention.
|
||
|
||
Args:
|
||
embed_dim (int): The dimention of hidden state.
|
||
num_heads (int): The number of attention heads.
|
||
num_kv_heads (int): The number of kv attention heads.
|
||
process_group (torch.distributed.ProcessGroup): The group of the current device for `parallel_mode`.
|
||
bias (boolean): Whether the bias is needed for linears. Will be used when initializing QKV matrix and
|
||
output projection. True by default.
|
||
dropout (float): The dropout rate for cross attention and self attention. 0.0 by default.
|
||
softmax_scale (float): The temperature to use for the softmax attention.
|
||
causal (boolean): Whether to apply causal attention mask. False by default.
|
||
layer_idx (int): The index of current layer. None by default.
|
||
rope_base (int): The value of `base` for rotary position embeddings. 10000 by default.
|
||
rotary_emb_dim (int): The dimention of Rotary Embedding. 0 by default.
|
||
rotary_emb_scale_base (int): The scaling factor of Rotary Embedding. If scale_base > 0, this implements
|
||
XPos(Sun et al., https://arxiv.org/abs/2212.10554). 0 by default.
|
||
use_flash_attn (boolean): Whether to use flash attention or not.If False, vanilla attention module will be used.
|
||
True by default.
|
||
device (Optional[Union[str, torch.device]]): The device will be used.
|
||
dtype (Optional[torch.dtype]): The type of data.
|
||
rot_embed_HF_impl: rotary embedding hf implementation. False by default.
|
||
|
||
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
embed_dim: int,
|
||
num_heads: int,
|
||
num_kv_heads: int,
|
||
process_group: Optional[torch.distributed.ProcessGroup],
|
||
bias: bool = True,
|
||
dropout: float = 0.0,
|
||
softmax_scale: float = None,
|
||
causal: bool = False,
|
||
layer_idx: int = None,
|
||
rope_base: int = 10000,
|
||
rotary_emb_dim: int = 0,
|
||
rotary_emb_scale_base: int = 0,
|
||
use_flash_attn: bool = True,
|
||
device: Optional[torch.device] = None,
|
||
dtype: Optional[torch.dtype] = None,
|
||
rot_embed_HF_impl: Optional[bool] = False,
|
||
) -> None:
|
||
factory_kwargs = {"device": device, "dtype": dtype}
|
||
super().__init__()
|
||
self.embed_dim = embed_dim
|
||
self.num_heads = num_heads
|
||
assert self.embed_dim % num_heads == 0, "embedding dim must be divisible by num_heads"
|
||
|
||
self.head_dim = self.embed_dim // num_heads
|
||
self.num_kv_heads = num_kv_heads
|
||
self.kv_dim = self.head_dim * num_kv_heads
|
||
self.causal = causal
|
||
self.layer_idx = layer_idx
|
||
self.rotary_emb_dim = rotary_emb_dim
|
||
self.use_flash_attn = use_flash_attn
|
||
self.dtype = dtype
|
||
|
||
self.rot_embed_HF_impl = rot_embed_HF_impl
|
||
sequence_parallel = gpc.config.parallel.get("sequence_parallel", False)
|
||
|
||
if self.rotary_emb_dim > 0:
|
||
self.rotary_emb = RotaryEmbedding(
|
||
self.rotary_emb_dim, base=rope_base, scale_base=rotary_emb_scale_base, device=device
|
||
)
|
||
|
||
# notice here should change bias=True
|
||
self.wq = ColumnParallelLinearTorch(
|
||
embed_dim,
|
||
embed_dim,
|
||
process_group,
|
||
bias=bias,
|
||
sequence_parallel=sequence_parallel,
|
||
**factory_kwargs,
|
||
)
|
||
self.wk = ColumnParallelLinearTorch(
|
||
embed_dim,
|
||
self.kv_dim,
|
||
process_group,
|
||
bias=bias,
|
||
sequence_parallel=sequence_parallel,
|
||
**factory_kwargs,
|
||
)
|
||
self.wv = ColumnParallelLinearTorch(
|
||
embed_dim,
|
||
self.kv_dim,
|
||
process_group,
|
||
bias=bias,
|
||
sequence_parallel=sequence_parallel,
|
||
**factory_kwargs,
|
||
)
|
||
|
||
inner_attn_cls = FlashSelfAttention if use_flash_attn else SelfAttention
|
||
inner_cross_attn_cls = FlashCrossAttention if use_flash_attn else CrossAttention
|
||
self.inner_attn = inner_attn_cls(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
|
||
self.inner_cross_attn = inner_cross_attn_cls(
|
||
causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout
|
||
)
|
||
|
||
self.inner_cross_attn_causal = causal
|
||
self.inner_cross_attn_softmax_scale = softmax_scale
|
||
self.inner_cross_attn_dropout = dropout
|
||
|
||
# output projection always have the bias (for now)
|
||
self.wo = RowParallelLinearTorch(
|
||
embed_dim,
|
||
embed_dim,
|
||
process_group,
|
||
bias=bias,
|
||
sequence_parallel=sequence_parallel,
|
||
**factory_kwargs,
|
||
)
|
||
# need to assign tp attribute so that internlm know it is tensor parallel module
|
||
if gpc.get_world_size(ParallelMode.TENSOR) > 1:
|
||
for name in ["wo", "wq", "wk", "wv"]:
|
||
for param in getattr(self, name).parameters():
|
||
setattr(param, IS_TENSOR_PARALLEL, True)
|
||
|
||
def forward(self, x, seqlen=None, inference_params=None, **kwargs):
|
||
if kwargs.get("indexes", None) is not None:
|
||
return self._packed_forward(x=x, inference_params=inference_params, **kwargs)
|
||
else:
|
||
return self._forward(x=x, seqlen=seqlen, inference_params=inference_params, **kwargs)
|
||
|
||
def _forward(self, x, seqlen=None, inference_params=None, **kwargs): # pylint: disable=W0613
|
||
"""
|
||
Arguments:
|
||
x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) if seqlen=None.
|
||
If seqlen is not None, x is (batch * seqlen, hidden_dim). This is so that when we
|
||
split x during sequence parallel, we split the batch * seqlen dimension
|
||
(in case batch is small).
|
||
"""
|
||
bsz, _, _ = x.shape
|
||
q, k, v = self.wq(x), self.wk(x), self.wv(x)
|
||
if seqlen is None:
|
||
q = rearrange(q, "b s (h d) -> b s h d", d=self.head_dim)
|
||
k = rearrange(k, "b s (h d) -> b s h d", d=self.head_dim)
|
||
v = rearrange(v, "b s (h d) -> b s h d", d=self.head_dim)
|
||
else:
|
||
q = rearrange(q, "(b s) (h d) -> b s h d", s=seqlen, d=self.head_dim)
|
||
k = rearrange(k, "(b s) (h d) -> b s h d", s=seqlen, d=self.head_dim)
|
||
v = rearrange(v, "(b s) (h d) -> b s h d", s=seqlen, d=self.head_dim)
|
||
|
||
if not self.rot_embed_HF_impl:
|
||
q = torch.cat([q[..., ::2], q[..., 1::2]], dim=-1)
|
||
k = torch.cat([k[..., ::2], k[..., 1::2]], dim=-1)
|
||
if inference_params is None:
|
||
if self.rotary_emb_dim > 0:
|
||
q = self.rotary_emb._single_eval_forward(q)
|
||
k = self.rotary_emb._single_eval_forward(k)
|
||
kv = torch.concat([k.unsqueeze(2), v.unsqueeze(2)], dim=2)
|
||
if self.dtype is torch.float32 and self.use_flash_attn:
|
||
if q.dtype not in [torch.float16, torch.bfloat16]:
|
||
q = q.to(torch.bfloat16)
|
||
if kv.dtype not in [torch.float16, torch.bfloat16]:
|
||
kv = kv.to(torch.bfloat16)
|
||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||
context = self.inner_cross_attn(q, kv).to(self.dtype)
|
||
else:
|
||
context = self.inner_cross_attn(q, kv)
|
||
|
||
else:
|
||
assert self.rotary_emb_dim > 0
|
||
if hasattr(inference_params, "attention_mask") and inference_params.attention_mask is not None:
|
||
empties = inference_params.attention_mask[..., -1].sum(dim=-1)
|
||
moved_q = q.clone()
|
||
moved_k = k.clone()
|
||
if inference_params.sequence_len_offset == 0:
|
||
for i in range(len(empties)):
|
||
if empties[i] != 0:
|
||
moved_q[i][: -empties[i]] = q[i][empties[i] :]
|
||
moved_k[i][: -empties[i]] = k[i][empties[i] :]
|
||
moved_q = self.rotary_emb._single_eval_forward(
|
||
moved_q, seqlen_offset=inference_params.sequence_len_offset
|
||
)
|
||
moved_k = self.rotary_emb._single_eval_forward(
|
||
moved_k, seqlen_offset=inference_params.sequence_len_offset
|
||
)
|
||
for i in range(len(empties)):
|
||
if empties[i] != 0:
|
||
q[i][empties[i] :] = moved_q[i][: -empties[i]]
|
||
k[i][empties[i] :] = moved_k[i][: -empties[i]]
|
||
else:
|
||
q[i] = moved_q[i]
|
||
k[i] = moved_k[i]
|
||
else:
|
||
q = q.squeeze(1)
|
||
k = k.squeeze(1)
|
||
q = self.rotary_emb._single_forward(
|
||
q,
|
||
inference_params.sequence_len_offset * torch.ones(q.size(0), dtype=torch.int, device=q.device)
|
||
- empties,
|
||
).unsqueeze(1)
|
||
k = self.rotary_emb._single_forward(
|
||
k,
|
||
inference_params.sequence_len_offset * torch.ones(k.size(0), dtype=torch.int, device=k.device)
|
||
- empties,
|
||
).unsqueeze(1)
|
||
else:
|
||
raise NotImplementedError(
|
||
"You should make sure you are aware that you are changing the method of generating."
|
||
"According to your generation function instead of inference/seq_generator_module.py, "
|
||
"You may implement here for normal running."
|
||
)
|
||
|
||
kv = torch.stack([k, v], dim=2)
|
||
|
||
assert self.layer_idx is not None, "Generation requires layer_idx in the constructor"
|
||
if hasattr(inference_params, "window_size") and inference_params.window_size is not None:
|
||
if inference_params.window_size <= inference_params.sequence_len_offset:
|
||
assert kv.size(1) == 1, "update kv lenth more than 1"
|
||
inference_params.key_value_memory_dict[self.layer_idx][
|
||
:, inference_params.keep_first : inference_params.window_size - 1, ...
|
||
] = inference_params.key_value_memory_dict[self.layer_idx][
|
||
:, -(inference_params.window_size - 1 - inference_params.keep_first) :, ...
|
||
].clone()
|
||
inference_params.real_sequence_len_offset = inference_params.sequence_len_offset
|
||
inference_params.sequence_len_offset = inference_params.window_size - 1
|
||
|
||
kv = _update_kv_cache(kv, inference_params, self.layer_idx)
|
||
|
||
inference_params.sequence_len_offset = inference_params.real_sequence_len_offset
|
||
else:
|
||
kv = _update_kv_cache(kv, inference_params, self.layer_idx)
|
||
else:
|
||
kv = _update_kv_cache(kv, inference_params, self.layer_idx)
|
||
|
||
# When using FP16, there is a high probability of NAN in the KV.
|
||
# Since NAN cannot be removed by multiplying with and 0, it needs
|
||
# to be removed manually here.
|
||
kv = torch.where(torch.isnan(kv), 0, kv)
|
||
|
||
if hasattr(inference_params, "attention_mask") and inference_params.attention_mask is not None:
|
||
assert self.use_flash_attn is True
|
||
if inference_params.sequence_len_offset == 0: # First entrance, attnmask (bs*seqlen*seqlen)
|
||
attn_mask = inference_params.attention_mask[:, None, ...]
|
||
attn_mask = torch.logical_or(
|
||
torch.ones_like(attn_mask, dtype=torch.bool).triu(diagonal=1), attn_mask
|
||
)
|
||
attn_mask4flsh = ~attn_mask[:, :, -1, :].view(bsz, -1)
|
||
cu_seqlens = torch.concat(
|
||
[
|
||
torch.tensor([0], dtype=torch.int32, device=attn_mask4flsh.device),
|
||
attn_mask4flsh.sum(dim=-1).to(dtype=torch.int32),
|
||
],
|
||
dim=0,
|
||
)
|
||
cu_seqlens = cu_seqlens.cumsum(dim=0, dtype=torch.int32)
|
||
max_seqlen_q = attn_mask4flsh.shape[-1]
|
||
max_seqlen_k = attn_mask4flsh.shape[-1]
|
||
total_q = q.masked_select(attn_mask4flsh.view(bsz, -1, 1, 1)).view(-1, q.shape[-2], q.shape[-1])
|
||
total_kv = kv.masked_select(attn_mask4flsh.view(bsz, -1, 1, 1, 1)).view(
|
||
-1, kv.shape[-3], kv.shape[-2], kv.shape[-1]
|
||
)
|
||
if self.dtype is torch.float32:
|
||
if total_q.dtype not in [torch.float16, torch.bfloat16]:
|
||
total_q = total_q.to(torch.bfloat16)
|
||
if total_kv.dtype not in [torch.float16, torch.bfloat16]:
|
||
total_kv = total_kv.to(torch.bfloat16)
|
||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||
output = FlashAttnVarlenKVPackedFunc.apply(
|
||
total_q,
|
||
total_kv,
|
||
cu_seqlens,
|
||
cu_seqlens,
|
||
max_seqlen_q,
|
||
max_seqlen_k,
|
||
0.0,
|
||
None,
|
||
True,
|
||
False,
|
||
).to(self.dtype)
|
||
else:
|
||
output = FlashAttnVarlenKVPackedFunc.apply(
|
||
total_q,
|
||
total_kv,
|
||
cu_seqlens,
|
||
cu_seqlens,
|
||
max_seqlen_q,
|
||
max_seqlen_k,
|
||
0.0,
|
||
None,
|
||
True,
|
||
False,
|
||
)
|
||
|
||
context = torch.zeros_like(q)
|
||
context = context.masked_scatter_(attn_mask4flsh.view(bsz, -1, 1, 1), output)
|
||
|
||
else:
|
||
attn_mask = inference_params.attention_mask[:, -1, :].view(bsz, 1, 1, -1)
|
||
if hasattr(inference_params, "window_size") and inference_params.window_size is not None:
|
||
if inference_params.window_size <= inference_params.sequence_len_offset:
|
||
attn_mask = torch.concat(
|
||
[
|
||
attn_mask[..., : inference_params.keep_first],
|
||
attn_mask[..., -(inference_params.window_size - inference_params.keep_first) :],
|
||
],
|
||
dim=-1,
|
||
)
|
||
|
||
k, v = torch.chunk(kv, 2, dim=2)
|
||
k = k.squeeze(2)
|
||
v = v.squeeze(2)
|
||
sp = k.shape
|
||
expansion = q.size(2) // k.size(2)
|
||
scores = torch.einsum(
|
||
"blhd,bnhd->bhln",
|
||
q,
|
||
k.unsqueeze(3).expand(-1, -1, -1, expansion, -1).reshape(sp[0], sp[1], q.size(2), sp[3]),
|
||
) / math.sqrt(q.size(-1))
|
||
scores = scores.masked_fill(attn_mask, -65000.0)
|
||
scores = F.softmax(scores, dim=-1) # bsz x h x L x L
|
||
context = torch.einsum(
|
||
"bhmn,bnhd->bmhd",
|
||
scores,
|
||
v.unsqueeze(3).expand(-1, -1, -1, expansion, -1).reshape(sp[0], sp[1], q.size(2), sp[3]),
|
||
)
|
||
else:
|
||
if self.dtype is torch.float32 and self.use_flash_attn:
|
||
if q.dtype not in [torch.float16, torch.bfloat16]:
|
||
q = q.to(torch.bfloat16)
|
||
if kv.dtype not in [torch.float16, torch.bfloat16]:
|
||
kv = kv.to(torch.bfloat16)
|
||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||
context = self.inner_cross_attn(q, kv, causal=True).to(self.dtype)
|
||
else:
|
||
context = self.inner_cross_attn(q, kv, causal=True)
|
||
if seqlen is None:
|
||
context = rearrange(context, "b s h d -> b s (h d)")
|
||
else:
|
||
context = rearrange(context, "b s h d -> (b s) (h d)")
|
||
out = self.wo(context)
|
||
return out
|
||
|
||
def _packed_forward(self, x, inference_params=None, **kwargs):
|
||
"""
|
||
we delete seqlen=None for lint check, cause this arg is not used.
|
||
|
||
Arguments:
|
||
x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) if seqlen=None.
|
||
If seqlen is not None, x is (batch * seqlen, hidden_dim). This is so that when we
|
||
split x during sequence parallel, we split the batch * seqlen dimension
|
||
(in case batch is small).
|
||
"""
|
||
assert self.use_flash_attn is True
|
||
q, k, v = self.wq(x), self.wk(x), self.wv(x)
|
||
q = rearrange(q, "t (h d) -> t h d", d=self.head_dim)
|
||
k = rearrange(k, "t (h d) -> t h d", d=self.head_dim)
|
||
v = rearrange(v, "t (h d) -> t h d", d=self.head_dim)
|
||
|
||
# qkv shift
|
||
# the rotary embedding in flash attention module in performed by separating the front and back parts, while
|
||
# most of others are done by odd-even methods.
|
||
if not self.rot_embed_HF_impl:
|
||
q = torch.cat([q[..., ::2], q[..., 1::2]], dim=-1)
|
||
k = torch.cat([k[..., ::2], k[..., 1::2]], dim=-1)
|
||
|
||
indexes = kwargs.pop("indexes")
|
||
q = self.rotary_emb._single_forward(q, indexes=indexes)
|
||
k = self.rotary_emb._single_forward(k, indexes=indexes)
|
||
|
||
if inference_params is None:
|
||
kv = torch.concat([k.unsqueeze(1), v.unsqueeze(1)], dim=1)
|
||
if self.dtype is torch.float32:
|
||
if q.dtype not in [torch.float16, torch.bfloat16]:
|
||
q = q.to(torch.bfloat16)
|
||
if kv.dtype not in [torch.float16, torch.bfloat16]:
|
||
kv = kv.to(torch.bfloat16)
|
||
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
||
context = flash_attn_varlen_kvpacked_func(
|
||
q=q,
|
||
kv=kv,
|
||
cu_seqlens_q=kwargs["cu_seqlens"],
|
||
cu_seqlens_k=kwargs["cu_seqlens"],
|
||
max_seqlen_q=kwargs["max_seqlen"],
|
||
max_seqlen_k=kwargs["max_seqlen"],
|
||
dropout_p=self.inner_cross_attn_dropout,
|
||
softmax_scale=self.inner_cross_attn_softmax_scale,
|
||
causal=self.inner_cross_attn_causal,
|
||
).to(self.dtype)
|
||
else:
|
||
context = flash_attn_varlen_kvpacked_func(
|
||
q=q,
|
||
kv=kv,
|
||
cu_seqlens_q=kwargs["cu_seqlens"],
|
||
cu_seqlens_k=kwargs["cu_seqlens"],
|
||
max_seqlen_q=kwargs["max_seqlen"],
|
||
max_seqlen_k=kwargs["max_seqlen"],
|
||
dropout_p=self.inner_cross_attn_dropout,
|
||
softmax_scale=self.inner_cross_attn_softmax_scale,
|
||
causal=self.inner_cross_attn_causal,
|
||
)
|
||
else:
|
||
raise RuntimeError("Not support this right now")
|
||
|
||
context = rearrange(context, "b h d -> b (h d)") # recover shape
|
||
out = self.wo(context)
|
||
return out
|
||
|
||
|
||
class PackedFlashLlamaLayer1D(nn.Module):
|
||
"""
|
||
1D Packed Flash Llama Layer.
|
||
|
||
Args:
|
||
hidden_size (int): The hidden size of model. 768 by default.
|
||
num_attention_heads (int): The number of attention heads. 12 by default.
|
||
num_kv_attention_heads (int): The number of kv attention heads. 12 by default.
|
||
mlp_ratio (int): The ratio of MLP layers. 4 by default.
|
||
attn_drop_rate (float): The dropout rate of attention module. 0 by default.
|
||
drop_rate (float): The dropout rate of the input hidden state. 0.0 by default.
|
||
dtype (torch.dtype): Type of data. torch.float by default.
|
||
layer_norm_epsilon (float): A value added to the denominator for numerical stability. 1e-5 by default.
|
||
checkpoint (bool): Whether to use checkpointing to save VRAM. True by default.
|
||
layer_idx (int): The index of current layer. 0 by default.
|
||
residual_in_fp32 (bool): Whether to use residual in fp32. False by default.
|
||
device (Optional[Union[str, torch.device]]): The device will be used.
|
||
apply_post_layer_norm (bool): Whether use post layer norm. False by default.
|
||
fused_dropout_add_ln (bool): Whether use fused dropout add ln. True by default.
|
||
no_bias (bool): Whether remove bias. False by default.
|
||
norm_type (str): Use RMS norm or layernorm."rmsnorm" by default.
|
||
adapt_hf (bool): Whether adapt hf. False by default.
|
||
dropout_selective_checkpoint (bool): Whether use dropout selective checkpoint. True by default.
|
||
use_scaled_init (bool): Whether use scaled init. True by default.
|
||
use_swiglu (bool): Whether use swiglu. True by default.
|
||
use_flash_attn (bool): Whether use flash-attn. True by default.
|
||
attn_wqkv_init_std (float): std used to init attn_wqkv weight. 0.02 by default,
|
||
attn_other_init_std (float): std used to init attn_other weight. 0.02 by default,
|
||
ffn_uplayer_init_std (float): std used to init w1, w2 weight in ffn when using glu
|
||
otherwise init fc1 weight in ffn. 0.02 by default,
|
||
ffn_other_init_std (float): std used to init ffn_other weight. 0.02 by default,
|
||
init_type (str): Initialization type. Use uniform or normal. "normal" by default,
|
||
rope_base (int): The value of `base` for rotary position embeddings. 10000 by default.
|
||
num_experts (int): The number of experts. <=1 means dense, >1 means MoE. 1 by default.
|
||
moe_gate_k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.
|
||
moe_capacity_factor (float, optional): default=1.0, the capacity of the expert at training time.
|
||
moe_eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time.
|
||
moe_min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor.
|
||
moe_noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample'.
|
||
moe_drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent
|
||
to infinite capacity).
|
||
moe_use_rts (bool, optional): default=True, whether to use Random Token Selection.
|
||
moe_use_residual (bool, optional): default=False, make this MoE layer a Residual MoE
|
||
(https://arxiv.org/abs/2201.05596) layer.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
hidden_size: int = 768,
|
||
num_attention_heads: int = 12,
|
||
num_kv_attention_heads: int = 12,
|
||
mlp_ratio: int = 4,
|
||
attn_drop_rate: float = 0,
|
||
drop_rate: float = 0.0,
|
||
dtype: torch.dtype = torch.float,
|
||
layer_norm_epsilon: float = 1e-6,
|
||
checkpoint: bool = False,
|
||
layer_idx: int = 0,
|
||
residual_in_fp32: bool = False,
|
||
device: Optional[torch.device] = None,
|
||
apply_post_layer_norm: bool = False,
|
||
fused_dropout_add_ln: bool = True,
|
||
no_bias: bool = False,
|
||
norm_type: str = "rmsnorm",
|
||
adapt_hf: bool = False,
|
||
dropout_selective_checkpoint: bool = True,
|
||
use_scaled_init: bool = True,
|
||
use_swiglu: bool = True,
|
||
use_flash_attn: bool = True,
|
||
attn_wqkv_init_std: float = 0.02,
|
||
attn_other_init_std: float = 0.02,
|
||
ffn_uplayer_init_std: float = 0.02,
|
||
ffn_other_init_std: float = 0.02,
|
||
init_type: str = "normal",
|
||
rope_base: int = 10000,
|
||
num_experts: int = 1,
|
||
moe_gate_k: int = 1,
|
||
moe_capacity_factor: float = 1.0,
|
||
moe_eval_capacity_factor: float = 1.0,
|
||
moe_min_capacity: int = 4,
|
||
moe_noisy_gate_policy: str = None,
|
||
moe_drop_tokens: bool = True,
|
||
moe_use_rts: bool = True,
|
||
moe_use_residual: bool = False,
|
||
):
|
||
super().__init__()
|
||
self.checkpoint = checkpoint
|
||
# dropout selective checkpoint can only be enabled when checkpoint is disabled.
|
||
self.dropout_selective_checkpoint = dropout_selective_checkpoint is True and checkpoint is False
|
||
self.layer_idx = layer_idx
|
||
self.use_flash_attn = use_flash_attn
|
||
self.prenorm = not apply_post_layer_norm
|
||
assert not fused_dropout_add_ln, "dropout_add_layer_norm can not be used here"
|
||
self.fused_dropout_add_ln = fused_dropout_add_ln
|
||
self.attn_wqkv_init_std = attn_wqkv_init_std
|
||
self.attn_other_init_std = attn_other_init_std
|
||
self.ffn_uplayer_init_std = ffn_uplayer_init_std
|
||
self.ffn_other_init_std = ffn_other_init_std
|
||
|
||
head_dim = hidden_size // num_attention_heads
|
||
self.attention = MHA(
|
||
embed_dim=hidden_size,
|
||
num_heads=num_attention_heads,
|
||
num_kv_heads=num_kv_attention_heads,
|
||
process_group=gpc.get_group(ParallelMode.TENSOR),
|
||
dropout=attn_drop_rate,
|
||
softmax_scale=1 / math.sqrt(head_dim),
|
||
causal=True,
|
||
layer_idx=layer_idx,
|
||
rotary_emb_dim=head_dim,
|
||
rotary_emb_scale_base=0,
|
||
use_flash_attn=use_flash_attn,
|
||
device=device,
|
||
dtype=dtype,
|
||
rot_embed_HF_impl=adapt_hf,
|
||
bias=not no_bias,
|
||
rope_base=rope_base,
|
||
)
|
||
|
||
self.dropout1 = nn.Dropout(drop_rate)
|
||
if norm_type == "rmsnorm":
|
||
self.attention_norm = RMSNorm(hidden_size, eps=layer_norm_epsilon)
|
||
self.ffn_norm = RMSNorm(hidden_size, eps=layer_norm_epsilon)
|
||
else:
|
||
self.attention_norm = nn.LayerNorm(hidden_size, eps=layer_norm_epsilon)
|
||
self.ffn_norm = nn.LayerNorm(hidden_size, eps=layer_norm_epsilon)
|
||
set_fp32_attr_to_module(self.attention_norm)
|
||
set_fp32_attr_to_module(self.ffn_norm)
|
||
if self.fused_dropout_add_ln:
|
||
assert dropout_add_layer_norm is not None, "dropout_add_ln is not installed"
|
||
assert isinstance(self.attention_norm, nn.LayerNorm) and isinstance(self.dropout1, nn.Dropout)
|
||
|
||
sequence_parallel = gpc.config.parallel.get("sequence_parallel", False)
|
||
self.num_experts = num_experts
|
||
self.moe_gate_k = moe_gate_k
|
||
self.moe_capacity_factor = moe_capacity_factor
|
||
self.moe_eval_capacity_factor = moe_eval_capacity_factor
|
||
self.moe_min_capacity = moe_min_capacity
|
||
self.moe_noisy_gate_policy = moe_noisy_gate_policy
|
||
self.moe_drop_tokens = moe_drop_tokens
|
||
self.moe_use_rts = moe_use_rts
|
||
self.moe_use_residual = moe_use_residual
|
||
ep_size = gpc.get_world_size(ParallelMode.EXPERT)
|
||
if num_experts <= 1: # dense, not MoE
|
||
if use_swiglu:
|
||
self.feed_forward = FeedForward(
|
||
hidden_size,
|
||
int(hidden_size * mlp_ratio),
|
||
out_features=hidden_size,
|
||
process_group=gpc.get_group(ParallelMode.TENSOR),
|
||
bias=False,
|
||
device=device,
|
||
dtype=dtype,
|
||
)
|
||
else:
|
||
self.feed_forward = ParallelFusedMLP(
|
||
hidden_size,
|
||
int(hidden_size * mlp_ratio),
|
||
out_features=hidden_size,
|
||
activation="gelu_approx",
|
||
process_group=gpc.get_group(ParallelMode.TENSOR),
|
||
bias1=False,
|
||
bias2=False,
|
||
sequence_parallel=sequence_parallel,
|
||
checkpoint_lvl=0,
|
||
heuristic="auto",
|
||
device=device,
|
||
dtype=dtype,
|
||
)
|
||
for _, param in self.feed_forward.named_parameters():
|
||
if gpc.get_world_size(ParallelMode.TENSOR) > 1:
|
||
setattr(param, IS_TENSOR_PARALLEL, True)
|
||
else:
|
||
# replace mlp by MoE module. The expert in MoE is a FeedForward module.
|
||
self.feed_forward = MoE(
|
||
hidden_size=hidden_size,
|
||
num_experts=num_experts,
|
||
ep_size=ep_size,
|
||
k=moe_gate_k,
|
||
capacity_factor=moe_capacity_factor,
|
||
eval_capacity_factor=moe_eval_capacity_factor,
|
||
min_capacity=moe_min_capacity,
|
||
noisy_gate_policy=moe_noisy_gate_policy,
|
||
drop_tokens=moe_drop_tokens,
|
||
use_rts=moe_use_rts,
|
||
use_residual=moe_use_residual,
|
||
device=device,
|
||
dtype=dtype,
|
||
)
|
||
for _, param in self.feed_forward.moe_layer.experts.named_parameters():
|
||
if gpc.get_world_size(ParallelMode.TENSOR) > 1:
|
||
setattr(param, IS_TENSOR_PARALLEL, True)
|
||
set_fp32_attr_to_module(self.feed_forward.moe_layer.gate)
|
||
|
||
for param in self.attention_norm.parameters():
|
||
if gpc.config.parallel.sequence_parallel is True:
|
||
setattr(param, IS_SEQUENCE_PARALLEL, True)
|
||
for param in self.ffn_norm.parameters():
|
||
if gpc.config.parallel.sequence_parallel is True:
|
||
setattr(param, IS_SEQUENCE_PARALLEL, True)
|
||
|
||
self.dropout2 = nn.Dropout(drop_rate)
|
||
self.use_swiglu = use_swiglu
|
||
self.use_scaled_init = use_scaled_init
|
||
self.residual_in_fp32 = residual_in_fp32 # only make sense when using prenorm
|
||
self.return_residual = False
|
||
|
||
if init_type == "normal":
|
||
self.init_func = normal_
|
||
self.scaled_init_func = scaled_init_method_normal
|
||
else:
|
||
self.init_func = uniform_
|
||
self.scaled_init_func = scaled_init_method_uniform
|
||
|
||
self.reset_parameters()
|
||
|
||
def reset_parameters(self):
|
||
with torch.no_grad():
|
||
for name, param in self.attention.named_parameters():
|
||
if param.ndim == 1:
|
||
param.data.zero_()
|
||
elif "wq" in name or "wk" in name or "wv" in name:
|
||
self.init_func(std=self.attn_wqkv_init_std)(param.data)
|
||
elif self.use_scaled_init: # wo
|
||
self.scaled_init_func(sigma=self.attn_other_init_std, num_layers=self.layer_idx + 1)(param.data)
|
||
else:
|
||
self.init_func(std=self.attn_other_init_std)(param.data)
|
||
|
||
for name, param in self.feed_forward.named_parameters():
|
||
if self.use_swiglu:
|
||
if self.use_scaled_init and "w2" in name:
|
||
self.scaled_init_func(sigma=self.ffn_other_init_std, num_layers=self.layer_idx + 1)(param.data)
|
||
else:
|
||
self.init_func(
|
||
std=self.ffn_uplayer_init_std if "w1" in name or "w3" in name else self.ffn_other_init_std
|
||
)(param.data)
|
||
else:
|
||
if self.use_scaled_init and "fc1" not in name:
|
||
self.scaled_init_func(sigma=self.ffn_other_init_std, num_layers=self.layer_idx + 1)(param.data)
|
||
else:
|
||
self.init_func(std=self.ffn_uplayer_init_std if "fc1" in name else self.ffn_other_init_std)(
|
||
param.data
|
||
)
|
||
|
||
def forward(
|
||
self, hidden_states, residual=None, cu_seqlens=None, indexes=None, inference_params=None, max_seqlen=None
|
||
):
|
||
if self.checkpoint and self.training:
|
||
return activation_checkpoint(
|
||
self._forward, False, hidden_states, residual, cu_seqlens, indexes, inference_params, max_seqlen
|
||
)
|
||
else:
|
||
return self._forward(hidden_states, residual, cu_seqlens, indexes, inference_params, max_seqlen)
|
||
|
||
def _forward(
|
||
self, hidden_states=None, residual=None, cu_seqlens=None, indexes=None, inference_params=None, max_seqlen=None
|
||
):
|
||
r"""Pass the input through the encoder layer.
|
||
|
||
Args:
|
||
hidden_states: the sequence to the encoder layer (required).
|
||
residual: hidden_states = Attn/MLP(LN(residual))
|
||
cu_seqlens: 1d LongTensor, len(cu_seqlens) = hidden_states + 1
|
||
indexes: the length of index is same as hidden states, which stand for the current position
|
||
"""
|
||
if self.prenorm:
|
||
|
||
def _dropout_and_norm_attn(_residual, _hidden_states):
|
||
_dropped = self.dropout1(_hidden_states)
|
||
_residual = (_dropped + _residual) if _residual is not None else _dropped
|
||
_hidden_states = self.attention_norm(_residual.to(dtype=self.attention_norm.weight.dtype))
|
||
|
||
return _residual, _hidden_states
|
||
|
||
if self.dropout_selective_checkpoint:
|
||
residual, hidden_states = activation_checkpoint(_dropout_and_norm_attn, False, residual, hidden_states)
|
||
else:
|
||
residual, hidden_states = _dropout_and_norm_attn(residual, hidden_states)
|
||
|
||
if self.residual_in_fp32:
|
||
residual = residual.to(torch.float32)
|
||
mixer_kwargs = {
|
||
"cu_seqlens": cu_seqlens,
|
||
"max_seqlen": max_seqlen,
|
||
"indexes": indexes,
|
||
"inference_params": inference_params,
|
||
}
|
||
hidden_states = self.attention(hidden_states, **mixer_kwargs)
|
||
|
||
if not isinstance(self.feed_forward, nn.Identity):
|
||
if not self.fused_dropout_add_ln:
|
||
|
||
def _dropout_and_norm_ffn(_residual, _hidden_states):
|
||
_dropped = self.dropout2(_hidden_states)
|
||
_residual = (_dropped + _residual) if _residual is not None else _dropped
|
||
_hidden_states = self.ffn_norm(_residual.to(torch.float32))
|
||
|
||
return _residual, _hidden_states
|
||
|
||
if self.dropout_selective_checkpoint:
|
||
residual, hidden_states = activation_checkpoint(
|
||
_dropout_and_norm_ffn, False, residual, hidden_states
|
||
)
|
||
else:
|
||
residual, hidden_states = _dropout_and_norm_ffn(residual, hidden_states)
|
||
|
||
if self.residual_in_fp32:
|
||
residual = residual.to(torch.float32)
|
||
|
||
# MLP.
|
||
moe_loss = torch.tensor(0.0, device=hidden_states.device, dtype=hidden_states.dtype)
|
||
if self.num_experts <= 1: # dense mlp output
|
||
hidden_states = self.feed_forward(hidden_states)
|
||
else: # MoE output
|
||
hidden_states, moe_loss, _ = self.feed_forward(hidden_states)
|
||
|
||
return hidden_states + residual, moe_loss
|
||
else:
|
||
assert residual is None
|
||
mixer_kwargs = {
|
||
"cu_seqlens": cu_seqlens,
|
||
"max_seqlen": max_seqlen,
|
||
"indexes": indexes,
|
||
"inference_params": inference_params,
|
||
}
|
||
mixer_out = self.attention(hidden_states, **mixer_kwargs)
|
||
if self.return_residual: # mixer out is actually a pair here
|
||
mixer_out, hidden_states = mixer_out
|
||
hidden_states = self.attention_norm(self.dropout1(mixer_out) + hidden_states).to(
|
||
dtype=self.attention_norm.weight.dtype
|
||
)
|
||
if not isinstance(self.feed_forward, nn.Identity):
|
||
# MLP.
|
||
moe_loss = torch.tensor(0.0, device=hidden_states.device, dtype=hidden_states.dtype)
|
||
if self.num_experts <= 1: # dense mlp output
|
||
mlp_out = self.feed_forward(hidden_states)
|
||
else: # MoE output
|
||
mlp_out, moe_loss, _ = self.feed_forward(hidden_states)
|
||
if self.return_residual: # mlp out is actually a pair here
|
||
# NOTE: should not be here
|
||
mlp_out, hidden_states = mlp_out
|
||
hidden_states = self.ffn_norm((self.dropout2(mlp_out)) + hidden_states).to(
|
||
dtype=self.ffn_norm.weight.dtype
|
||
)
|
||
return hidden_states, moe_loss
|
||
|
||
|
||
class PackedFlashLlama1D(nn.Module):
|
||
"""
|
||
1D Packed Flash Llama.
|
||
|
||
Args:
|
||
num_layers (int): The number of layer. 12 by default.
|
||
hidden_size (int): The size of hidden state. 768 by default.
|
||
num_attention_heads (int): The number of attention head. 12 by default.
|
||
num_kv_attention_heads (int): The number of kv attention head. 12 by default.
|
||
vocab_size (int): The size of vocabulary. 50304 by default.
|
||
mlp_ratio (int): The ratio of MLP layers. 4 by default.
|
||
attn_drop_rate (float): The dropout rate of attention module. 0.0 by default.
|
||
drop_rate (float): The dropout rate of input hidden state. 0.0 by default.
|
||
dtype (torch.dtype): The type of data. torch.float by default.
|
||
checkpoint (bool): Whether to use checkpointing to save VRAM. True by default.
|
||
checkpoint_fraction (float): The proportion of layers that need to be checkpointed compared to the total number
|
||
of layers. 1.0 by default.
|
||
layer_norm_epsilon (float): A value added to the denominator for numerical stability. 1e-6 by default.
|
||
first (bool): Whether input embedding layer or not. False by default.
|
||
last (bool): Whether output embedding layer or not. False by default.
|
||
embed_split_hidden (bool): Split the embedding layer in the hidden state dimention or vocabulary dimention.
|
||
True by default.
|
||
embed_grad_scale (float): Refer to GLM-130B, for training stability. 0.1 by default.
|
||
parallel_output (bool): If it is necessary to collect the output of parallel computing. True by default.
|
||
start_layer_idx (int): The index of start layer in the pipeline. 0 by default.
|
||
device (Optional[Union[str, torch.device]]): The device will be used. None by default.
|
||
apply_post_layer_norm (bool): Whether use post layer norm. False by default.
|
||
no_bias (bool): Whether remove bias. False by default.
|
||
residual_in_fp32 (bool): Whether to use residual in fp32. False by default.
|
||
norm_type (str): Normalization type. Use RMSNorm or LayerNorm. "rmsnorm" by default.
|
||
adapt_hf (bool): Whether adapt hf. False by default.
|
||
is_reward (bool): Whether use is_reward. False by default.
|
||
dropout_selective_checkpoint (bool): Whether dropout selective checkpoint. True by default.
|
||
use_scaled_init (bool): Whether use scaled init. True by default.
|
||
use_swiglu (bool): Whether use swiglu. True by default.
|
||
use_flash_attn (bool): Whether to use flash-attn. True by default.
|
||
embedding_init_std (float): std used to init embedding weight. 0.02 by default,
|
||
attn_wqkv_init_std (float): std used to init attn_wqkv weight. 0.02 by default,
|
||
attn_other_init_std (float): std used to init attn_other weight. 0.02 by default,
|
||
ffn_uplayer_init_std (float): std used to init w1, w2 weight in ffn when using glu
|
||
otherwise init fc1 weight in ffn. 0.02 by default,
|
||
ffn_other_init_std (float): std used to init ffn_other weight. 0.02 by default,
|
||
out_head_init_std (float): std used to init output lmhead weight. 0.02 by default,
|
||
init_type (str): Initialization type. Use uniform or normal. "normal" by default,
|
||
rope_base (int): The value of `base` for rotary position embeddings. 10000 by default.
|
||
num_experts (int): The number of experts. <=1 means dense, >1 means MoE. 1 by default.
|
||
moe_gate_k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.
|
||
moe_capacity_factor (float, optional): default=1.0, the capacity of the expert at training time.
|
||
moe_eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time.
|
||
moe_min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor.
|
||
moe_noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample'.
|
||
moe_drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent
|
||
to infinite capacity).
|
||
moe_use_rts (bool, optional): default=True, whether to use Random Token Selection.
|
||
moe_use_residual (bool, optional): default=False, make this MoE layer a Residual MoE
|
||
(https://arxiv.org/abs/2201.05596) layer.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
num_layers: int = 12,
|
||
hidden_size: int = 768,
|
||
num_attention_heads: int = 12,
|
||
num_kv_attention_heads: int = 12,
|
||
vocab_size: int = 50304,
|
||
mlp_ratio: int = 4,
|
||
attn_drop_rate: float = 0.0,
|
||
drop_rate: float = 0.0,
|
||
dtype: torch.dtype = torch.float,
|
||
checkpoint: bool = False,
|
||
checkpoint_fraction: float = 1.0,
|
||
layer_norm_epsilon: float = 1e-5,
|
||
first: bool = False,
|
||
last: bool = False,
|
||
embed_split_hidden: bool = False,
|
||
embed_grad_scale: float = 0.1,
|
||
parallel_output: bool = True,
|
||
start_layer_idx: int = 0,
|
||
device: Optional[torch.device] = None,
|
||
apply_post_layer_norm=False,
|
||
no_bias=False,
|
||
residual_in_fp32: bool = False,
|
||
norm_type: str = "rmsnorm",
|
||
adapt_hf: bool = False,
|
||
is_reward: bool = False,
|
||
dropout_selective_checkpoint: bool = True,
|
||
use_scaled_init: bool = True,
|
||
use_swiglu: bool = True,
|
||
use_flash_attn: bool = True,
|
||
embedding_init_std: float = 0.02,
|
||
attn_wqkv_init_std: float = 0.02,
|
||
attn_other_init_std: float = 0.02,
|
||
ffn_uplayer_init_std: float = 0.02,
|
||
ffn_other_init_std: float = 0.02,
|
||
out_head_init_std: float = 0.02,
|
||
init_type: str = "normal",
|
||
rope_base: int = 10000,
|
||
num_experts: bool = 1,
|
||
moe_gate_k: int = 1,
|
||
moe_capacity_factor: float = 1.0,
|
||
moe_eval_capacity_factor: float = 1.0,
|
||
moe_min_capacity: int = 4,
|
||
moe_noisy_gate_policy: str = None,
|
||
moe_drop_tokens: bool = True,
|
||
moe_use_rts: bool = True,
|
||
moe_use_residual: bool = False,
|
||
):
|
||
super().__init__()
|
||
|
||
self.use_flash_attn = use_flash_attn
|
||
if checkpoint_fraction <= 0:
|
||
checkpoint = False
|
||
if not checkpoint:
|
||
checkpoint_fraction = 0
|
||
checkpoint_layer_num = num_layers * checkpoint_fraction
|
||
sequence_parallel = gpc.config.parallel.get("sequence_parallel", False)
|
||
if is_reward:
|
||
head_cls = RewardModelLinear
|
||
else:
|
||
head_cls = ScaleColumnParallelLinear
|
||
if first:
|
||
if embed_split_hidden:
|
||
self.tok_embeddings = Embedding1D(num_embeddings=vocab_size, embedding_dim=hidden_size)
|
||
else:
|
||
|
||
self.tok_embeddings = ParallelGPT2Embeddings(
|
||
embed_dim=hidden_size,
|
||
vocab_size=vocab_size,
|
||
max_position_embeddings=-1,
|
||
process_group=gpc.get_group(ParallelMode.TENSOR),
|
||
padding_idx=None,
|
||
sequence_parallel=sequence_parallel,
|
||
device=device,
|
||
dtype=dtype,
|
||
)
|
||
for _, param in self.tok_embeddings.named_parameters():
|
||
if init_type == "normal":
|
||
normal_(std=embedding_init_std)(param)
|
||
else:
|
||
uniform_(std=embedding_init_std)(param)
|
||
if gpc.get_world_size(ParallelMode.TENSOR) > 1:
|
||
setattr(param, IS_TENSOR_PARALLEL, True)
|
||
self.embed_grad_scale = embed_grad_scale
|
||
self.layers = nn.ModuleList(
|
||
[
|
||
PackedFlashLlamaLayer1D(
|
||
hidden_size=hidden_size,
|
||
num_attention_heads=num_attention_heads,
|
||
num_kv_attention_heads=num_kv_attention_heads,
|
||
mlp_ratio=mlp_ratio,
|
||
attn_drop_rate=attn_drop_rate,
|
||
drop_rate=drop_rate,
|
||
dtype=dtype,
|
||
layer_norm_epsilon=layer_norm_epsilon,
|
||
checkpoint=lid < checkpoint_layer_num,
|
||
layer_idx=lid + start_layer_idx, # This parameter is used for caching during generation
|
||
residual_in_fp32=residual_in_fp32,
|
||
device=device,
|
||
apply_post_layer_norm=apply_post_layer_norm,
|
||
fused_dropout_add_ln=False,
|
||
no_bias=no_bias,
|
||
norm_type=norm_type,
|
||
dropout_selective_checkpoint=dropout_selective_checkpoint,
|
||
use_scaled_init=use_scaled_init,
|
||
use_swiglu=use_swiglu,
|
||
use_flash_attn=use_flash_attn,
|
||
adapt_hf=adapt_hf,
|
||
attn_wqkv_init_std=attn_wqkv_init_std,
|
||
attn_other_init_std=attn_other_init_std,
|
||
ffn_uplayer_init_std=ffn_uplayer_init_std,
|
||
ffn_other_init_std=ffn_other_init_std,
|
||
init_type=init_type,
|
||
rope_base=rope_base,
|
||
num_experts=num_experts,
|
||
moe_gate_k=moe_gate_k,
|
||
moe_capacity_factor=moe_capacity_factor,
|
||
moe_eval_capacity_factor=moe_eval_capacity_factor,
|
||
moe_min_capacity=moe_min_capacity,
|
||
moe_noisy_gate_policy=moe_noisy_gate_policy,
|
||
moe_drop_tokens=moe_drop_tokens,
|
||
moe_use_rts=moe_use_rts,
|
||
moe_use_residual=moe_use_residual,
|
||
)
|
||
for lid in range(num_layers)
|
||
]
|
||
)
|
||
|
||
if last:
|
||
if not apply_post_layer_norm:
|
||
if norm_type == "rmsnorm":
|
||
self.norm = RMSNorm(hidden_size, eps=layer_norm_epsilon)
|
||
else:
|
||
self.norm = nn.LayerNorm(hidden_size, eps=layer_norm_epsilon)
|
||
for param in self.norm.parameters():
|
||
if gpc.config.parallel.sequence_parallel is True:
|
||
setattr(param, IS_SEQUENCE_PARALLEL, True)
|
||
|
||
self.output = head_cls(
|
||
in_features=hidden_size,
|
||
out_features=gpc.get_world_size(ParallelMode.TENSOR) if is_reward else vocab_size,
|
||
process_group=gpc.get_group(ParallelMode.TENSOR),
|
||
bias=False,
|
||
device=device,
|
||
dtype=dtype,
|
||
weight_scale=embed_grad_scale,
|
||
)
|
||
|
||
for _, param in self.output.named_parameters():
|
||
if init_type == "normal":
|
||
normal_(std=out_head_init_std)(param)
|
||
else:
|
||
uniform_(std=out_head_init_std)(param)
|
||
if gpc.get_world_size(ParallelMode.TENSOR) > 1:
|
||
setattr(param, IS_TENSOR_PARALLEL, True)
|
||
|
||
self.parallel_output = parallel_output
|
||
|
||
def forward(self, hidden_states=None, cu_seqlens=None, input_ids=None, indexes=None, inference_params=None):
|
||
# attention_mask: compute attention on the places where the value is 1
|
||
if hasattr(self, "tok_embeddings"):
|
||
hidden_states = self.tok_embeddings(input_ids)
|
||
if self.embed_grad_scale != 1:
|
||
hidden_states = (
|
||
self.embed_grad_scale * hidden_states + (1 - self.embed_grad_scale) * hidden_states.detach()
|
||
)
|
||
if isinstance(cu_seqlens, list):
|
||
assert len(cu_seqlens) == 1
|
||
cu_seqlens = cu_seqlens[0].to(hidden_states.device)
|
||
|
||
if cu_seqlens is not None:
|
||
cu_seqlens = cu_seqlens.squeeze(0)
|
||
hidden_states = hidden_states.squeeze(0) # If cu_seqlens is passed in，it indicated a packed state，
|
||
# the batch dimension with a size of 1 should be directly squeezed off.
|
||
|
||
if indexes is not None:
|
||
assert len(indexes) == 1
|
||
# The indexes are used to indicate the actual position IDs of each token in the packed input.
|
||
indexes = indexes[0]
|
||
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() if cu_seqlens is not None else None
|
||
|
||
moe_losses = []
|
||
for _, block in enumerate(self.layers):
|
||
hidden_states, mos_loss = block(
|
||
hidden_states,
|
||
residual=None,
|
||
cu_seqlens=cu_seqlens,
|
||
indexes=indexes,
|
||
inference_params=inference_params,
|
||
max_seqlen=max_seqlen,
|
||
)
|
||
moe_losses.append(mos_loss)
|
||
|
||
if hasattr(self, "norm"):
|
||
hidden_states = self.norm(hidden_states.float())
|
||
|
||
extra_hidden_states_list = None
|
||
if hasattr(self, "output"):
|
||
hidden_states = self.output(hidden_states)
|
||
|
||
if not self.parallel_output:
|
||
hidden_states = gather_forward_split_backward(hidden_states, ParallelMode.TENSOR, dim=-1)
|
||
if extra_hidden_states_list is not None:
|
||
extra_hidden_states_list = [
|
||
gather_forward_split_backward(extra_hidden_states, ParallelMode.TENSOR, dim=-1)
|
||
for extra_hidden_states in extra_hidden_states_list # pylint: disable=E1133
|
||
]
|
||
|
||
if extra_hidden_states_list is not None:
|
||
return (hidden_states, extra_hidden_states_list)
|
||
|
||
return hidden_states, moe_losses
|
||
|
||
|
||
def _build_generic_model_1d(num_layers, num_chunks, device=torch.device("cuda"), **kwargs):
|
||
"""
|
||
build generic model 1d
|
||
|
||
Args:
|
||
num_layers (int): The number of layer.
|
||
num_chunks (int): The number of partitions in pipeline parallel.
|
||
device (Optional[Union[str, torch.device]]): The device will be used. torch.device("cuda") by default.
|
||
|
||
"""
|
||
pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE)
|
||
pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
|
||
|
||
all_parts = partition_uniform(num_layers, pipeline_size, num_chunks)
|
||
parts = all_parts[pipeline_rank]
|
||
if gpc.is_rank_for_log():
|
||
logger.info(f"The layer sharding is {all_parts}.")
|
||
|
||
models = []
|
||
kwargs["checkpoint_fraction"] = 1.0
|
||
start_idx, end_idx = 0, 0
|
||
for start, end in parts:
|
||
start_idx, end_idx = start, end
|
||
kwargs["num_layers"] = end - start
|
||
kwargs["first"] = start == 0
|
||
# If there is no content in the final layer, assign the last layer.
|
||
kwargs["last"] = end == num_layers and len(all_parts[-1]) != 0
|
||
kwargs["device"] = device
|
||
kwargs["start_layer_idx"] = start
|
||
chunk = PackedFlashLlama1D(**filter_kwargs(PackedFlashLlama1D.__init__, kwargs)).to(device)
|
||
|
||
models.append(chunk)
|
||
torch.distributed.barrier()
|
||
if len(models) == 1:
|
||
model = models[0]
|
||
else:
|
||
model = nn.ModuleList(models)
|
||
setattr(model, "first_layer", start_idx)
|
||
setattr(model, "last_layer", end_idx)
|
||
return model
|
||
|
||
|
||
@MODEL_INITIALIZER.register_module(module_name=MODEL_TYPE)
|
||
def build_model_with_moe_cfg(
|
||
num_chunks=1,
|
||
checkpoint=False,
|
||
dtype=torch.float,
|
||
embed_split_hidden=False,
|
||
num_layers=48,
|
||
hidden_size=2048,
|
||
vocab_size=50304,
|
||
embed_grad_scale=1,
|
||
parallel_output=True,
|
||
num_attention_heads=32,
|
||
num_kv_attention_heads=None,
|
||
mlp_ratio=4.0,
|
||
residual_in_fp32=False,
|
||
norm_type="rmsnorm",
|
||
adapt_hf=False,
|
||
drop_rate=0,
|
||
attn_drop_rate=0,
|
||
apply_post_layer_norm=False, # pylint: disable=W0613
|
||
no_bias=False,
|
||
deepnorm=False, # pylint: disable=W0613
|
||
layer_norm_epsilon=1e-5,
|
||
is_reward=False,
|
||
dropout_selective_checkpoint=True,
|
||
use_scaled_init: bool = True,
|
||
use_swiglu: bool = True,
|
||
use_flash_attn: bool = True,
|
||
embedding_init_std: float = 0.02,
|
||
attn_wqkv_init_std: float = 0.02,
|
||
attn_other_init_std: float = 0.02,
|
||
ffn_uplayer_init_std: float = 0.02,
|
||
ffn_other_init_std: float = 0.02,
|
||
out_head_init_std: float = 0.02,
|
||
init_type: str = "normal",
|
||
rope_base: int = 10000,
|
||
num_experts: int = 1,
|
||
moe_gate_k: int = 1,
|
||
moe_capacity_factor: float = 1.0,
|
||
moe_eval_capacity_factor: float = 1.0,
|
||
moe_min_capacity: int = 4,
|
||
moe_noisy_gate_policy: str = None,
|
||
moe_drop_tokens: bool = True,
|
||
moe_use_rts: bool = True,
|
||
moe_use_residual: bool = False,
|
||
):
|
||
"""
|
||
Build model with config.
|
||
|
||
Args:
|
||
num_chunks (int): The number of partitions in pipeline parallel. 1 by default.
|
||
checkpoint (bool): Whether to use checkpointing to save VRAM. False by default.
|
||
dtype (torch.dtype): The type of data. torch.float by default.
|
||
embed_split_hidden (bool): Split the embedding layer in the hidden state dimention or vocabulary dimention.
|
||
False by default.
|
||
num_layers (int): The number of layer. 48 by default.
|
||
hidden_size (int): The size of hidden state. 2048 by default.
|
||
vocab_size (int): The size of vocabulary. 50304 by default.
|
||
embed_grad_scale (float): Refer to GLM-130B, for training stability. 0.1 by default.
|
||
parallel_output (bool): If it is necessary to collect the output of parallel computing. True by default.
|
||
num_attention_heads (int): The number of attention head. 32 by default.
|
||
mlp_ratio (int): The ratio of MLP layers. 4.0 by default.
|
||
residual_in_fp32 (bool): Whether to use residual in fp32. False by default. It cannot be used temporarily
|
||
because this parameter requires inconsistent data types to be passed between pipelines,
|
||
which requires significant modifications to internlm.
|
||
norm_type (str): Normalization type. Use RMSNorm or LayerNorm. "rmsnorm" by default.
|
||
drop_rate (float): The dropout rate of input hidden state. 0 by default.
|
||
attn_drop_rate (float): The dropout rate of attention module. 0 by default.
|
||
apply_post_layer_norm (bool): Whether to apply post layer norm. False by default.
|
||
layer_norm_epsilon (float): A value added to the denominator for numerical stability. 1e-5 by default.
|
||
is_reward (bool): Whether to use reward model. False by default.
|
||
dropout_selective_checkpoint (bool): It can only be enabled when checkpoint is disabled. True by default.
|
||
use_scaled_init (bool): Whether to use scaled init. True by default.
|
||
use_swiglu (bool): Whether to use swiglu. True by default.
|
||
use_flash_attn (bool): Whether to use flash-attn. True by default.
|
||
num_experts (int): The number of experts. <=1 means dense, >1 means MoE. 1 by default.
|
||
moe_gate_k (int, optional): default=1, top-k gating value, only supports k=1 or k=2.
|
||
moe_capacity_factor (float, optional): default=1.0, the capacity of the expert at training time.
|
||
moe_eval_capacity_factor (float, optional): default=1.0, the capacity of the expert at eval time.
|
||
moe_min_capacity (int, optional): default=4, the minimum capacity per expert regardless of the capacity_factor.
|
||
moe_noisy_gate_policy (str, optional): default=None, noisy gate policy, valid options are 'Jitter', 'RSample'.
|
||
moe_drop_tokens (bool, optional): default=True, whether to drop tokens - (setting to False is equivalent
|
||
to infinite capacity).
|
||
moe_use_rts (bool, optional): default=True, whether to use Random Token Selection.
|
||
moe_use_residual (bool, optional): default=False, make this MoE layer a Residual MoE
|
||
(https://arxiv.org/abs/2201.05596) layer.
|
||
"""
|
||
|
||
cfg = dict(
|
||
hidden_size=hidden_size,
|
||
num_attention_heads=num_attention_heads,
|
||
num_kv_attention_heads=num_kv_attention_heads if num_kv_attention_heads else num_attention_heads,
|
||
checkpoint=checkpoint,
|
||
dtype=dtype,
|
||
embed_split_hidden=embed_split_hidden,
|
||
vocab_size=vocab_size,
|
||
embed_grad_scale=embed_grad_scale,
|
||
parallel_output=parallel_output,
|
||
mlp_ratio=mlp_ratio,
|
||
apply_post_layer_norm=apply_post_layer_norm,
|
||
no_bias=no_bias,
|
||
residual_in_fp32=residual_in_fp32,
|
||
norm_type=norm_type,
|
||
adapt_hf=adapt_hf,
|
||
drop_rate=drop_rate,
|
||
attn_drop_rate=attn_drop_rate,
|
||
layer_norm_epsilon=layer_norm_epsilon,
|
||
is_reward=is_reward,
|
||
dropout_selective_checkpoint=dropout_selective_checkpoint,
|
||
use_scaled_init=use_scaled_init,
|
||
use_swiglu=use_swiglu,
|
||
use_flash_attn=use_flash_attn,
|
||
embedding_init_std=embedding_init_std,
|
||
attn_wqkv_init_std=attn_wqkv_init_std,
|
||
attn_other_init_std=attn_other_init_std,
|
||
ffn_uplayer_init_std=ffn_uplayer_init_std,
|
||
ffn_other_init_std=ffn_other_init_std,
|
||
out_head_init_std=out_head_init_std,
|
||
init_type=init_type,
|
||
rope_base=rope_base,
|
||
num_experts=num_experts,
|
||
moe_gate_k=moe_gate_k,
|
||
moe_capacity_factor=moe_capacity_factor,
|
||
moe_eval_capacity_factor=moe_eval_capacity_factor,
|
||
moe_min_capacity=moe_min_capacity,
|
||
moe_noisy_gate_policy=moe_noisy_gate_policy,
|
||
moe_drop_tokens=moe_drop_tokens,
|
||
moe_use_rts=moe_use_rts,
|
||
moe_use_residual=moe_use_residual,
|
||
)
|
||
|
||
return _build_generic_model_1d(num_layers=num_layers, num_chunks=num_chunks, **cfg)
|