Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
曾经的那一瞬间
Models
提交
5842f323
M
Models
项目概览
曾经的那一瞬间
/
Models
大约 1 年 前同步成功
通知
1
Star
0
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
M
Models
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
5842f323
编写于
12月 01, 2022
作者:
A
A. Unique TensorFlower
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Internal change
PiperOrigin-RevId: 492385091
上级
bdf4a7bb
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
110 addition
and
133 deletion
+110
-133
official/nlp/configs/encoders.py
official/nlp/configs/encoders.py
+53
-2
official/nlp/modeling/layers/moe.py
official/nlp/modeling/layers/moe.py
+37
-91
official/nlp/modeling/layers/moe_test.py
official/nlp/modeling/layers/moe_test.py
+6
-21
official/nlp/modeling/networks/sparse_mixer.py
official/nlp/modeling/networks/sparse_mixer.py
+14
-19
未找到文件。
official/nlp/configs/encoders.py
浏览文件 @
5842f323
...
...
@@ -258,13 +258,37 @@ class FNetEncoderConfig(hyperparams.Config):
initializer_range
:
float
=
0.02
embedding_width
:
Optional
[
int
]
=
None
output_range
:
Optional
[
int
]
=
None
return_all_encoder_outputs
:
bool
=
False
# Pre/Post-LN Transformer
norm_first
:
bool
=
False
use_fft
:
bool
=
False
attention_layers
:
Sequence
[
int
]
=
()
@
dataclasses
.
dataclass
class
SparseMixerEncoderConfig
(
hyperparams
.
Config
):
"""SparseMixer encoder configuration."""
vocab_size
:
int
=
30522
hidden_size
:
int
=
768
num_layers
:
int
=
14
moe_layers
:
Sequence
[
int
]
=
(
5
,
6
,
7
,
8
)
attention_layers
:
Sequence
[
int
]
=
(
10
,
11
,
12
,
13
)
num_experts
:
int
=
16
train_capacity_factor
:
float
=
1.
eval_capacity_factor
:
float
=
1.
examples_per_group
:
float
=
1.
use_fft
:
bool
=
False
num_attention_heads
:
int
=
8
max_sequence_length
:
int
=
512
type_vocab_size
:
int
=
2
inner_dim
:
int
=
3072
inner_activation
:
str
=
"gelu"
output_dropout
:
float
=
0.1
attention_dropout
:
float
=
0.1
initializer_range
:
float
=
0.02
output_range
:
Optional
[
int
]
=
None
embedding_width
:
Optional
[
int
]
=
None
norm_first
:
bool
=
False
@
dataclasses
.
dataclass
class
EncoderConfig
(
hyperparams
.
OneOfConfig
):
"""Encoder configuration."""
...
...
@@ -279,6 +303,7 @@ class EncoderConfig(hyperparams.OneOfConfig):
xlnet
:
XLNetEncoderConfig
=
XLNetEncoderConfig
()
query_bert
:
QueryBertConfig
=
QueryBertConfig
()
fnet
:
FNetEncoderConfig
=
FNetEncoderConfig
()
sparse_mixer
:
SparseMixerEncoderConfig
=
SparseMixerEncoderConfig
()
# If `any` is used, the encoder building relies on any.BUILDER.
any
:
hyperparams
.
Config
=
hyperparams
.
Config
()
...
...
@@ -607,6 +632,32 @@ def build_encoder(config: EncoderConfig,
use_fft
=
encoder_cfg
.
use_fft
,
attention_layers
=
encoder_cfg
.
attention_layers
)
if
encoder_type
==
"sparse_mixer"
:
return
networks
.
SparseMixer
(
vocab_size
=
encoder_cfg
.
vocab_size
,
hidden_size
=
encoder_cfg
.
hidden_size
,
num_layers
=
encoder_cfg
.
num_layers
,
moe_layers
=
encoder_cfg
.
moe_layers
,
attention_layers
=
encoder_cfg
.
attention_layers
,
num_experts
=
encoder_cfg
.
num_experts
,
train_capacity_factor
=
encoder_cfg
.
train_capacity_factor
,
eval_capacity_factor
=
encoder_cfg
.
eval_capacity_factor
,
examples_per_group
=
encoder_cfg
.
examples_per_group
,
use_fft
=
encoder_cfg
.
use_fft
,
num_attention_heads
=
encoder_cfg
.
num_attention_heads
,
max_sequence_length
=
encoder_cfg
.
max_sequence_length
,
type_vocab_size
=
encoder_cfg
.
type_vocab_size
,
inner_dim
=
encoder_cfg
.
inner_dim
,
inner_activation
=
tf_utils
.
get_activation
(
encoder_cfg
.
inner_activation
),
output_dropout
=
encoder_cfg
.
output_dropout
,
attention_dropout
=
encoder_cfg
.
attention_dropout
,
initializer
=
tf
.
keras
.
initializers
.
TruncatedNormal
(
stddev
=
encoder_cfg
.
initializer_range
),
output_range
=
encoder_cfg
.
output_range
,
embedding_width
=
encoder_cfg
.
embedding_width
,
norm_first
=
encoder_cfg
.
norm_first
,
embedding_layer
=
embedding_layer
)
bert_encoder_cls
=
networks
.
BertEncoder
if
encoder_type
==
"bert_v2"
:
bert_encoder_cls
=
networks
.
BertEncoderV2
...
...
official/nlp/modeling/layers/moe.py
浏览文件 @
5842f323
...
...
@@ -17,8 +17,6 @@
import
dataclasses
from
typing
import
Any
,
Callable
,
Optional
,
Tuple
from
absl
import
logging
import
numpy
as
np
import
tensorflow
as
tf
from
official.modeling
import
tf_utils
...
...
@@ -48,11 +46,13 @@ def _router_z_loss(router_logits: tf.Tensor) -> float:
Returns:
Scalar router z-loss <float32>.
"""
num_groups
,
tokens_per_group
,
_
=
router_logits
.
shape
num_groups
=
tf
.
shape
(
router_logits
)[
0
]
tokens_per_group
=
router_logits
.
shape
[
1
]
log_z
=
tf
.
math
.
reduce_logsumexp
(
router_logits
,
axis
=-
1
)
z_loss
=
log_z
**
2
return
tf
.
math
.
reduce_sum
(
z_loss
)
/
(
num_groups
*
tokens_per_group
)
return
tf
.
math
.
reduce_sum
(
z_loss
)
/
tf
.
cast
(
num_groups
*
tokens_per_group
,
tf
.
float32
)
@
dataclasses
.
dataclass
...
...
@@ -187,7 +187,7 @@ class Router(tf.keras.layers.Layer):
"""
if
apply_jitter
and
self
.
jitter_noise
>
0
:
inputs
*=
tf
.
random
.
uniform
(
inputs
.
shape
,
tf
.
shape
(
inputs
)
,
minval
=
1.0
-
self
.
jitter_noise
,
maxval
=
1.0
+
self
.
jitter_noise
,
dtype
=
inputs
.
dtype
)
...
...
@@ -259,7 +259,9 @@ class ExpertsChooseMaskedRouter(MaskedRouter):
Returns:
Dispatch and combine arrays for routing with masked matmuls.
"""
num_groups
,
tokens_per_group
,
_
=
router_probs
.
shape
num_groups
=
tf
.
shape
(
router_probs
)[
0
]
tokens_per_group
=
router_probs
.
shape
[
1
]
router_probs_t
=
tf
.
transpose
(
router_probs
,
perm
=
[
0
,
2
,
1
])
# router_probs_t: <float32>[num_groups, num_experts, tokens_per_group]
...
...
@@ -296,8 +298,10 @@ class ExpertsChooseMaskedRouter(MaskedRouter):
num_tokens
=
num_groups
*
tokens_per_group
num_tokens_dispatched_somewhere
=
tf
.
math
.
reduce_sum
(
tf
.
math
.
reduce_max
(
dispatch_mask
,
axis
=
(
-
1
,
-
2
)))
fraction_tokens_left_behind
=
1.0
-
num_tokens_dispatched_somewhere
/
float
(
num_tokens
)
fraction_tokens_left_behind
=
1.0
-
tf
.
cast
(
num_tokens_dispatched_somewhere
,
tf
.
float32
)
/
tf
.
cast
(
num_tokens
,
tf
.
float32
)
# Total number of tokens that were dispatched (one token could be
# dispatched to multiple experts).
num_tokens_dispatched
=
tf
.
math
.
reduce_sum
(
dispatch_mask
)
...
...
@@ -513,9 +517,7 @@ class MoeLayer(tf.keras.layers.Layer):
*
,
train_capacity_factor
:
float
=
1.0
,
eval_capacity_factor
:
float
=
1.0
,
min_expert_capacity
:
int
=
4
,
max_group_size
:
int
=
4096
,
strict_group_size
:
bool
=
False
,
examples_per_group
:
float
=
1.0
,
name
:
str
=
"moe"
,
**
kwargs
):
"""Init.
...
...
@@ -537,22 +539,18 @@ class MoeLayer(tf.keras.layers.Layer):
tokens that an expert will process AND will indirectly increase the
number of experts that a given token is routed to.
eval_capacity_factor: As above, but used during evaluation.
min_expert_capacity: Minimum token processing capacity for each expert.
max_group_size: The total number of tokens on each device is subdivided
into groups of this size. Router computations are then performed on a
per-group basis. A larger group size will result in slower but more
accurate top-k and sorting computations, whereas a smaller group size
will result in faster but more approximate (and potentially less stable)
routing choices. Note that actual group size may be smaller than
max_group_size for consistency with the number of experts and tokens;
see also `strict_group_size` attribute. In practice,
we find that imperfect routing choices are tolerable and recommend
choosing a group size on the order of 4096 tokens, although this number
will vary based on model configuration and size.
strict_group_size: If True, fail if unable to set the token group size
equal to max_group_size. If False (default), the actual group size may
be smaller than max_group_size for consistency with the number of
experts and tokens.
examples_per_group: Number of examples to form a group. Router then
performs top_k token selection for each expert on a per group basis.
E.g. when `examples_per_group=4.0`, tokens are assigned to experts in
groups formed from 4 examples. When `examples_per_group=0.5`,
each example is split into 2 groups.
`examples_per_group` must divide the local batch size.
A larger group size will result in slower but more accurate top-k and
sorting computations, whereas a smaller group size will result in faster
but more approximate (and potentially less stable) routing choices.
In practice, we find that imperfect routing choices are tolerable and
recommend choosing a group size on the order of 4096 tokens, although
this number will vary based on model configuration and size.
name: Layer name.
**kwargs: Forwarded to super.
"""
...
...
@@ -565,9 +563,7 @@ class MoeLayer(tf.keras.layers.Layer):
self
.
_train_capacity_factor
=
train_capacity_factor
self
.
_eval_capacity_factor
=
eval_capacity_factor
self
.
_max_group_size
=
max_group_size
self
.
_min_expert_capacity
=
min_expert_capacity
self
.
_strict_group_size
=
strict_group_size
self
.
_examples_per_group
=
examples_per_group
def
call
(
self
,
inputs
:
tf
.
Tensor
,
...
...
@@ -592,10 +588,15 @@ class MoeLayer(tf.keras.layers.Layer):
training
=
tf
.
keras
.
backend
.
learning_phase
()
# inputs shape [batch_size, seq_length, hidden_dim]
per_device_batch_size
,
seq_length
,
hidden_dim
=
inputs
.
shape
num_tokens
=
per_device_batch_size
*
seq_length
num_groups
=
self
.
_num_groups
(
num_tokens
,
self
.
_max_group_size
)
tokens_per_group
=
num_tokens
//
num_groups
batch_size
,
seq_length
,
hidden_dim
=
inputs
.
shape
if
batch_size
is
not
None
:
if
self
.
_examples_per_group
>
batch_size
:
raise
ValueError
(
f
"examples_per_group=
{
self
.
_examples_per_group
}
is larger than the "
"number of examples available in the local (per-device) batch_size="
f
"
{
batch_size
}
. Either decrease examples_per_group or increase the "
"batch_size."
)
tokens_per_group
=
int
(
seq_length
*
self
.
_examples_per_group
)
if
training
:
capacity_factor
=
self
.
_train_capacity_factor
...
...
@@ -604,71 +605,16 @@ class MoeLayer(tf.keras.layers.Layer):
# Each group will send expert_capacity tokens to each expert.
expert_capacity
=
int
(
round
(
capacity_factor
*
tokens_per_group
/
self
.
num_experts
))
expert_capacity
=
max
(
expert_capacity
,
self
.
_min_expert_capacity
)
logging
.
info
(
"Selected expert_capacity=%d for num_experts=%d and training=%r."
,
expert_capacity
,
self
.
num_experts
,
training
)
# Reshape batch and sequence/token dimensions for expert routing.
x
=
tf
.
reshape
(
inputs
,
(
num_groups
,
tokens_per_group
,
hidden_dim
))
x
=
tf
.
reshape
(
inputs
,
(
-
1
,
tokens_per_group
,
hidden_dim
))
x
=
self
.
_mask_and_dispatch_to_experts
(
x
,
expert_capacity
,
training
)
# Return to original input shape.
x
=
tf
.
reshape
(
x
,
(
per_device_batch_size
,
seq_length
,
hidden_dim
))
x
=
tf
.
reshape
(
x
,
(
-
1
,
seq_length
,
hidden_dim
))
return
x
def
_num_groups
(
self
,
num_tokens
:
int
,
max_group_size
:
int
)
->
int
:
"""Returns the number of token routing groups.
Note that the quantities are local to the device.
We select the smallest num_groups such that:
- num_groups >= num_tokens / max_group_size (ensuring the group size is no
larger than max_group_size),
- num_tokens % num_groups = 0 (ensuring that the group size evenly divides
into the num_tokens),
Args:
num_tokens: Number of tokens from input batch.
max_group_size: Maximum size of each token routing group. Actual group
size may end up being smaller unless strict_group_size==True.
Returns:
Number of token routing groups.
Raises:
ValueError if we cannot find a group_size satisfying the above
requirements.
"""
# Increase the number of groups (and decrease the group size) until we have
# a viable number of groups.
min_num_groups
=
int
(
np
.
ceil
(
num_tokens
/
max_group_size
))
num_groups
=
min_num_groups
while
num_groups
<
num_tokens
and
num_tokens
%
num_groups
!=
0
:
num_groups
+=
1
group_size
=
num_tokens
//
num_groups
logging
.
info
(
"Selected group_size=%d and num_groups=%d for input num_tokens=%d, "
"max_group_size=%d, num_experts=%d."
,
group_size
,
num_groups
,
num_tokens
,
max_group_size
,
self
.
num_experts
)
if
group_size
<
self
.
_min_expert_capacity
:
raise
ValueError
(
f
"Local (per-device) group_size
{
group_size
}
is smaller than "
f
"min_expert_capacity
{
self
.
_min_expert_capacity
}
, which is probably "
"not intended. Please increase max_group_size {max_group_size} to"
" seq_length or increase batch_size or decrease min_expert_capacity."
)
if
self
.
_strict_group_size
and
group_size
!=
self
.
_max_group_size
:
raise
ValueError
(
f
"Selected group_size=
{
group_size
}
is less than the "
f
"max_group_size=
{
max_group_size
}
. Exiting because strict mode is "
"active (strict_group_size=True)"
)
return
num_groups
def
_mask_and_dispatch_to_experts
(
self
,
inputs
:
tf
.
Tensor
,
expert_capacity
:
int
,
training
:
bool
)
->
tf
.
Tensor
:
...
...
official/nlp/modeling/layers/moe_test.py
浏览文件 @
5842f323
...
...
@@ -32,14 +32,13 @@ def small_config():
config
[
'jitter_noise'
]
=
0.1
config
[
'train_capacity_factor'
]
=
1.0
config
[
'eval_capacity_factor'
]
=
1.0
config
[
'min_expert_capacity'
]
=
1
config
[
'max_group_size'
]
=
9
config
[
'examples_per_group'
]
=
2.0
config
[
'backbone_d_ff'
]
=
13
return
config
def
make_input_ones
(
batch_size
:
int
=
2
,
def
make_input_ones
(
batch_size
:
int
=
4
,
seq_length
:
int
=
10
,
hidden_dim
:
int
=
7
)
->
tf
.
Tensor
:
return
tf
.
ones
((
batch_size
,
seq_length
,
hidden_dim
),
dtype
=
tf
.
float32
)
...
...
@@ -64,11 +63,7 @@ class MoeTest(tf.test.TestCase):
y
=
moe
.
_router_z_loss
(
x
)
expected
=
(
5
+
np
.
log
(
np
.
exp
(
5
)
+
1
))
**
2
self
.
assertAllClose
(
expected
,
y
,
atol
=
1e-7
)
x
=
tf
.
constant
([[[
10.0
,
5.0
]]],
dtype
=
tf
.
bfloat16
)
y
=
moe
.
_router_z_loss
(
x
)
expected
=
100.0
self
.
assertAllClose
(
expected
,
y
,
atol
=
1e-7
)
self
.
assertDTypeEqual
(
y
,
tf
.
float32
)
def
test_router_z_loss_shape
(
self
):
x
=
make_input_ones
(
2
,
5
,
7
)
...
...
@@ -199,21 +194,12 @@ class MoeTest(tf.test.TestCase):
router
,
train_capacity_factor
=
config
[
'train_capacity_factor'
],
eval_capacity_factor
=
config
[
'eval_capacity_factor'
],
max_group_size
=
config
[
'max_group_size'
],
min_expert_capacity
=
config
[
'min_expert_capacity'
])
examples_per_group
=
config
[
'examples_per_group'
])
inputs
=
make_input_ones
()
with
self
.
assertLogs
(
'absl'
,
level
=
'INFO'
)
as
cm
:
outputs
=
moe_layer
(
inputs
,
training
=
True
)
outputs
=
moe_layer
(
inputs
,
training
=
True
)
self
.
assertAllEqual
(
tf
.
shape
(
inputs
),
tf
.
shape
(
outputs
))
self
.
assertEqual
(
cm
.
output
,
[(
'INFO:absl:Selected group_size=5 and num_groups=4 for input '
'num_tokens=20, max_group_size=9, num_experts=2.'
),
(
'INFO:absl:Selected expert_capacity=2 for num_experts=2 and '
'training=True.'
)])
var_names
=
sorted
([
v
.
name
for
v
in
moe_layer
.
trainable_variables
])
self
.
assertAllEqual
([
'moe/experts/intermediate/bias:0'
,
'moe/experts/intermediate/kernel:0'
,
...
...
@@ -241,8 +227,7 @@ class MoeTest(tf.test.TestCase):
router
,
train_capacity_factor
=
config
[
'train_capacity_factor'
],
eval_capacity_factor
=
config
[
'eval_capacity_factor'
],
max_group_size
=
config
[
'max_group_size'
],
min_expert_capacity
=
config
[
'min_expert_capacity'
])
examples_per_group
=
config
[
'examples_per_group'
])
layer
=
moe
.
MoeLayerWithBackbone
(
moe_layer
,
config
[
'backbone_d_ff'
])
inputs
=
make_input_ones
()
...
...
official/nlp/modeling/networks/sparse_mixer.py
浏览文件 @
5842f323
...
...
@@ -123,7 +123,7 @@ class SparseMixer(tf.keras.layers.Layer):
num_experts
:
int
=
16
,
train_capacity_factor
:
float
=
1.
,
eval_capacity_factor
:
float
=
1.
,
max_group_size
:
int
=
4096
,
examples_per_group
:
float
=
1.
,
mixing_mechanism
:
layers
.
MixingMechanism
=
layers
.
MixingMechanism
.
LINEAR
,
use_fft
:
bool
=
False
,
num_attention_heads
:
int
=
8
,
...
...
@@ -157,7 +157,7 @@ class SparseMixer(tf.keras.layers.Layer):
'num_experts'
:
num_experts
,
'train_capacity_factor'
:
train_capacity_factor
,
'eval_capacity_factor'
:
eval_capacity_factor
,
'
max_group_size'
:
max_group_size
,
'
examples_per_group'
:
examples_per_group
,
'mixing_mechanism'
:
mixing_mechanism
,
'use_fft'
:
use_fft
,
'attention_layers'
:
attention_layers
,
...
...
@@ -243,7 +243,7 @@ class SparseMixer(tf.keras.layers.Layer):
name
=
'router'
),
train_capacity_factor
=
train_capacity_factor
,
eval_capacity_factor
=
eval_capacity_factor
,
max_group_size
=
max_group_size
,
examples_per_group
=
examples_per_group
,
name
=
'moe'
)
else
:
feedforward_layer
=
None
# Fallback to default (dense) MLP class
...
...
@@ -273,18 +273,15 @@ class SparseMixer(tf.keras.layers.Layer):
if
with_dense_inputs
:
self
.
inputs
=
dict
(
input_word_ids
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,),
dtype
=
tf
.
int32
),
input_mask
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,),
dtype
=
tf
.
int32
),
input_type_ids
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,),
dtype
=
tf
.
int32
),
# The total length of token ids and dense inputs still has to be
# max_sequence_length. It is checked in call().
input_word_ids
=
tf
.
keras
.
Input
(
shape
=
(
None
,),
dtype
=
tf
.
int32
),
input_mask
=
tf
.
keras
.
Input
(
shape
=
(
None
,),
dtype
=
tf
.
int32
),
input_type_ids
=
tf
.
keras
.
Input
(
shape
=
(
None
,),
dtype
=
tf
.
int32
),
dense_inputs
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,
embedding_width
),
dtype
=
tf
.
float32
),
dense_mask
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,),
dtype
=
tf
.
int32
),
dense_type_ids
=
tf
.
keras
.
Input
(
shape
=
(
max_sequence_length
,),
dtype
=
tf
.
int32
),
shape
=
(
None
,
embedding_width
),
dtype
=
tf
.
float32
),
dense_mask
=
tf
.
keras
.
Input
(
shape
=
(
None
,),
dtype
=
tf
.
int32
),
dense_type_ids
=
tf
.
keras
.
Input
(
shape
=
(
None
,),
dtype
=
tf
.
int32
),
)
else
:
self
.
inputs
=
dict
(
...
...
@@ -320,11 +317,9 @@ class SparseMixer(tf.keras.layers.Layer):
type_ids
=
tf
.
concat
([
type_ids
,
dense_type_ids
],
axis
=
1
)
mask
=
tf
.
concat
([
mask
,
dense_mask
],
axis
=
1
)
seq_length
=
word_embeddings
.
shape
[
1
]
if
seq_length
!=
self
.
_max_sequence_length
:
raise
ValueError
(
'Sparse Mixer: Sequence length must be the same as '
'`max_sequence_length` ({}), but it is {}.'
.
format
(
self
.
_max_sequence_length
,
seq_length
))
# SparseMixer: Sequence length must be the same as `max_sequence_length`.
word_embeddings
=
tf
.
ensure_shape
(
word_embeddings
,
[
None
,
self
.
_max_sequence_length
,
None
])
# Absolute position embeddings.
position_embeddings
=
self
.
_position_embedding_layer
(
word_embeddings
)
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录